Examples with QueryId io.trino.spi.QueryId used on opensource projects

Example 71 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestQueryContext method testSetMemoryPool.

@Test
public void testSetMemoryPool() {
    try (LocalQueryRunner localQueryRunner = LocalQueryRunner.create(TEST_SESSION)) {
        QueryContext queryContext = new QueryContext(new QueryId("query"), DataSize.ofBytes(10), new MemoryPool(DataSize.ofBytes(10)), new TestingGcMonitor(), localQueryRunner.getExecutor(), localQueryRunner.getScheduler(), DataSize.ofBytes(0), new SpillSpaceTracker(DataSize.ofBytes(0)));
        // Use memory
        queryContext.getQueryMemoryContext().initializeLocalMemoryContexts("test");
        LocalMemoryContext userMemoryContext = queryContext.getQueryMemoryContext().localUserMemoryContext();
        LocalMemoryContext revocableMemoryContext = queryContext.getQueryMemoryContext().localRevocableMemoryContext();
        assertTrue(userMemoryContext.setBytes(3).isDone());
        assertTrue(revocableMemoryContext.setBytes(5).isDone());
        // Free memory
        userMemoryContext.close();
        revocableMemoryContext.close();
    }
}

Example 72 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestTotalReservationLowMemoryKiller method testSkewedQuery.

@Test
public void testSkewedQuery() {
    int memoryPool = 12;
    // q2 is the query with the most total memory reservation, but not the query with the max memory reservation.
    // This also tests the corner case where a node has an empty memory pool
    Map<String, Map<String, Long>> queries = ImmutableMap.<String, Map<String, Long>>builder().put("q_1", ImmutableMap.of("n1", 0L, "n2", 8L, "n3", 0L, "n4", 0L, "n5", 0L)).put("q_2", ImmutableMap.of("n1", 3L, "n2", 5L, "n3", 2L, "n4", 4L, "n5", 0L)).put("q_3", ImmutableMap.of("n1", 0L, "n2", 0L, "n3", 9L, "n4", 0L, "n5", 0L)).buildOrThrow();
    assertEquals(lowMemoryKiller.chooseQueryToKill(toQueryMemoryInfoList(queries), toNodeMemoryInfoList(memoryPool, queries)), Optional.of(KillTarget.wholeQuery(new QueryId("q_2"))));
}

Example 73 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class LowMemoryKillerTestingUtils method toNodeMemoryInfoList.

static List<MemoryInfo> toNodeMemoryInfoList(long memoryPoolMaxBytes, Map<String, Map<String, Long>> queries, Map<String, Map<String, Map<String, Long>>> tasks) {
    Map<InternalNode, NodeReservation> nodeReservations = new HashMap<>();
    for (Map.Entry<String, Map<String, Long>> entry : queries.entrySet()) {
        QueryId queryId = new QueryId(entry.getKey());
        Map<String, Long> reservationByNode = entry.getValue();
        for (Map.Entry<String, Long> nodeEntry : reservationByNode.entrySet()) {
            InternalNode node = new InternalNode(nodeEntry.getKey(), URI.create("http://localhost"), new NodeVersion("version"), false);
            long bytes = nodeEntry.getValue();
            if (bytes == 0) {
                continue;
            }
            nodeReservations.computeIfAbsent(node, ignored -> new NodeReservation()).add(queryId, bytes);
        }
    }
    ImmutableList.Builder<MemoryInfo> result = ImmutableList.builder();
    for (Map.Entry<InternalNode, NodeReservation> entry : nodeReservations.entrySet()) {
        NodeReservation nodeReservation = entry.getValue();
        MemoryPoolInfo memoryPoolInfo = new MemoryPoolInfo(memoryPoolMaxBytes, nodeReservation.getTotalReservedBytes(), 0, nodeReservation.getReservationByQuery(), ImmutableMap.of(), ImmutableMap.of());
        result.add(new MemoryInfo(7, memoryPoolInfo, tasksMemoryInfoForNode(entry.getKey().getNodeIdentifier(), tasks)));
    }
    return result.build();
}

Example 74 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestMemoryPools method setUp.

private void setUp(Supplier<List<Driver>> driversSupplier) {
    checkState(localQueryRunner == null, "Already set up");
    Session session = testSessionBuilder().setCatalog("tpch").setSchema("tiny").setSystemProperty("task_default_concurrency", "1").build();
    localQueryRunner = LocalQueryRunner.builder(session).withInitialTransaction().build();
    // add tpch
    localQueryRunner.createCatalog("tpch", new TpchConnectorFactory(1), ImmutableMap.of());
    userPool = new MemoryPool(TEN_MEGABYTES);
    fakeQueryId = new QueryId("fake");
    SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(DataSize.of(1, GIGABYTE));
    QueryContext queryContext = new QueryContext(new QueryId("query"), TEN_MEGABYTES, userPool, new TestingGcMonitor(), localQueryRunner.getExecutor(), localQueryRunner.getScheduler(), TEN_MEGABYTES, spillSpaceTracker);
    taskContext = createTaskContext(queryContext, localQueryRunner.getExecutor(), localQueryRunner.getDefaultSession());
    drivers = driversSupplier.get();
}

Example 75 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class ClusterMemoryManager method callOomKiller.

private synchronized void callOomKiller(Iterable<QueryExecution> runningQueries) {
    List<QueryMemoryInfo> queryMemoryInfoList = Streams.stream(runningQueries).map(this::createQueryMemoryInfo).collect(toImmutableList());
    List<MemoryInfo> nodeMemoryInfos = nodes.values().stream().map(RemoteNodeMemory::getInfo).filter(Optional::isPresent).map(Optional::get).collect(toImmutableList());
    Optional<KillTarget> killTarget = lowMemoryKiller.chooseQueryToKill(queryMemoryInfoList, nodeMemoryInfos);
    if (killTarget.isPresent()) {
        if (killTarget.get().isWholeQuery()) {
            QueryId queryId = killTarget.get().getQuery();
            log.debug("Low memory killer chose %s", queryId);
            Optional<QueryExecution> chosenQuery = findRunningQuery(runningQueries, killTarget.get().getQuery());
            if (chosenQuery.isPresent()) {
                // See comments in  isQueryGone for why chosenQuery might be absent.
                chosenQuery.get().fail(new TrinoException(CLUSTER_OUT_OF_MEMORY, "Query killed because the cluster is out of memory. Please try again in a few minutes."));
                queriesKilledDueToOutOfMemory.incrementAndGet();
                lastKillTarget = killTarget.get();
                logQueryKill(queryId, nodeMemoryInfos);
            }
        } else {
            Set<TaskId> tasks = killTarget.get().getTasks();
            log.debug("Low memory killer chose %s", tasks);
            ImmutableSet.Builder<TaskId> killedTasksBuilder = ImmutableSet.builder();
            for (TaskId task : tasks) {
                Optional<QueryExecution> runningQuery = findRunningQuery(runningQueries, task.getQueryId());
                if (runningQuery.isPresent()) {
                    runningQuery.get().failTask(task, new TrinoException(CLUSTER_OUT_OF_MEMORY, "Task killed because the cluster is out of memory."));
                    tasksKilledDueToOutOfMemory.incrementAndGet();
                    killedTasksBuilder.add(task);
                }
            }
            // only record tasks actually killed
            ImmutableSet<TaskId> killedTasks = killedTasksBuilder.build();
            if (!killedTasks.isEmpty()) {
                lastKillTarget = KillTarget.selectedTasks(killedTasks);
                logTasksKill(killedTasks, nodeMemoryInfos);
            }
        }
    }
}