Examples with QueryManager com.facebook.presto.execution.QueryManager used on opensource projects

Example 11 with QueryManager

use of com.facebook.presto.execution.QueryManager in project presto by prestodb.

the class TestQueues method testBasic.

@Test(timeOut = 240_000)
public void testBasic() throws Exception {
    String dbConfigUrl = getDbConfigUrl();
    H2ResourceGroupsDao dao = getDao(dbConfigUrl);
    try (DistributedQueryRunner queryRunner = createQueryRunner(dbConfigUrl, dao)) {
        QueryManager queryManager = queryRunner.getCoordinator().getQueryManager();
        // submit first "dashboard" query
        QueryId firstDashboardQuery = createQuery(queryRunner, newDashboardSession(), LONG_LASTING_QUERY);
        // wait for the first "dashboard" query to start
        waitForQueryState(queryRunner, firstDashboardQuery, RUNNING);
        assertEquals(queryManager.getStats().getRunningQueries(), 1);
        // submit second "dashboard" query
        QueryId secondDashboardQuery = createQuery(queryRunner, newDashboardSession(), LONG_LASTING_QUERY);
        MILLISECONDS.sleep(2000);
        // wait for the second "dashboard" query to be queued ("dashboard.${USER}" queue strategy only allows one "dashboard" query to be accepted for execution)
        waitForQueryState(queryRunner, secondDashboardQuery, QUEUED);
        assertEquals(queryManager.getStats().getRunningQueries(), 1);
        // Update db to allow for 1 more running query in dashboard resource group
        dao.updateResourceGroup(3, "user-${USER}", "1MB", 3, 4, null, null, null, null, null, 1L);
        dao.updateResourceGroup(5, "dashboard-${USER}", "1MB", 1, 2, null, null, null, null, null, 3L);
        waitForQueryState(queryRunner, secondDashboardQuery, RUNNING);
        QueryId thirdDashboardQuery = createQuery(queryRunner, newDashboardSession(), LONG_LASTING_QUERY);
        waitForQueryState(queryRunner, thirdDashboardQuery, QUEUED);
        assertEquals(queryManager.getStats().getRunningQueries(), 2);
        // submit first non "dashboard" query
        QueryId firstNonDashboardQuery = createQuery(queryRunner, newSession(), LONG_LASTING_QUERY);
        // wait for the first non "dashboard" query to start
        waitForQueryState(queryRunner, firstNonDashboardQuery, RUNNING);
        assertEquals(queryManager.getStats().getRunningQueries(), 3);
        // submit second non "dashboard" query
        QueryId secondNonDashboardQuery = createQuery(queryRunner, newSession(), LONG_LASTING_QUERY);
        // wait for the second non "dashboard" query to start
        waitForQueryState(queryRunner, secondNonDashboardQuery, RUNNING);
        assertEquals(queryManager.getStats().getRunningQueries(), 4);
        // cancel first "dashboard" query, the second "dashboard" query and second non "dashboard" query should start running
        cancelQuery(queryRunner, firstDashboardQuery);
        waitForQueryState(queryRunner, firstDashboardQuery, FAILED);
        waitForQueryState(queryRunner, thirdDashboardQuery, RUNNING);
        assertEquals(queryManager.getStats().getRunningQueries(), 4);
        assertEquals(queryManager.getStats().getCompletedQueries().getTotalCount(), 1);
    }
}

Example 12 with QueryManager

use of com.facebook.presto.execution.QueryManager in project presto by prestodb.

the class LocalQueryProvider method getQuery.

public Query getQuery(QueryId queryId, String slug) {
    Query query = queries.get(queryId);
    if (query != null) {
        if (!query.isSlugValid(slug)) {
            throw notFound("Query not found");
        }
        return query;
    }
    // this is the first time the query has been accessed on this coordinator
    Session session;
    try {
        if (!queryManager.isQuerySlugValid(queryId, slug)) {
            throw notFound("Query not found");
        }
        session = queryManager.getQuerySession(queryId);
    } catch (NoSuchElementException e) {
        throw notFound("Query not found");
    }
    query = queries.computeIfAbsent(queryId, id -> {
        ExchangeClient exchangeClient = exchangeClientSupplier.get(new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), LocalQueryProvider.class.getSimpleName()));
        return Query.create(session, slug, queryManager, transactionManager, exchangeClient, responseExecutor, timeoutExecutor, blockEncodingSerde, retryCircuitBreaker);
    });
    return query;
}

Example 13 with QueryManager

use of com.facebook.presto.execution.QueryManager in project presto by prestodb.

the class TestQueryManager method testQueryCpuLimit.

@Test(timeOut = 60_000L)
public void testQueryCpuLimit() throws Exception {
    try (DistributedQueryRunner queryRunner = builder().setSingleExtraProperty("query.max-cpu-time", "1ms").build()) {
        QueryId queryId = createQuery(queryRunner, TEST_SESSION, "SELECT COUNT(*) FROM lineitem");
        waitForQueryState(queryRunner, queryId, FAILED);
        QueryManager queryManager = queryRunner.getCoordinator().getQueryManager();
        BasicQueryInfo queryInfo = queryManager.getQueryInfo(queryId);
        assertEquals(queryInfo.getState(), FAILED);
        assertEquals(queryInfo.getErrorCode(), EXCEEDED_CPU_LIMIT.toErrorCode());
    }
}

Example 14 with QueryManager

use of com.facebook.presto.execution.QueryManager in project presto by prestodb.

the class TestQueryManager method testFailQueryPrerun.

@Test(timeOut = 60_000L)
public void testFailQueryPrerun() throws Exception {
    DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager();
    QueryManager queryManager = queryRunner.getCoordinator().getQueryManager();
    // Create 3 running queries to guarantee queueing
    createQueries(dispatchManager, 3);
    QueryId queryId = dispatchManager.createQueryId();
    dispatchManager.createQuery(queryId, "slug", 0, new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem").get();
    assertNotEquals(dispatchManager.getStats().getQueuedQueries(), 0L, "Expected 0 queued queries, found: " + dispatchManager.getStats().getQueuedQueries());
    // wait until it's admitted but fail it before it starts
    while (true) {
        QueryState state = dispatchManager.getQueryInfo(queryId).getState();
        if (state.ordinal() >= RUNNING.ordinal()) {
            fail("unexpected query state: " + state);
        }
        if (state.ordinal() >= QUEUED.ordinal()) {
            // cancel query
            dispatchManager.failQuery(queryId, new PrestoException(GENERIC_USER_ERROR, "mock exception"));
            break;
        }
    }
    QueryState state = dispatchManager.getQueryInfo(queryId).getState();
    assertEquals(state, FAILED);
    // Give the stats a time to update
    Thread.sleep(1000);
    assertEquals(queryManager.getStats().getQueuedQueries(), 0);
}

Example 15 with QueryManager

use of com.facebook.presto.execution.QueryManager in project presto by prestodb.

the class TestMemoryManager method testClusterPoolsMultiCoordinator.

@Test(timeOut = 60_000)
public void testClusterPoolsMultiCoordinator() throws Exception {
    Map<String, String> properties = ImmutableMap.<String, String>builder().put("task.verbose-stats", "true").put("resource-manager.memory-pool-fetch-interval", "10ms").put("resource-manager.query-heartbeat-interval", "10ms").put("resource-manager.node-status-timeout", "5s").build();
    try (DistributedQueryRunner queryRunner = createQueryRunner(properties, ImmutableMap.of(), properties, 2)) {
        // Reserve all the memory
        QueryId fakeQueryId = new QueryId("fake");
        for (TestingPrestoServer server : queryRunner.getCoordinatorWorkers()) {
            for (MemoryPool pool : server.getLocalMemoryManager().getPools()) {
                assertTrue(pool.tryReserve(fakeQueryId, "test", pool.getMaxBytes()));
            }
        }
        List<Future<?>> queryFutures = new ArrayList<>();
        for (int i = 0; i < 2; i++) {
            int coordinator = i;
            queryFutures.add(executor.submit(() -> queryRunner.execute(coordinator, "SELECT COUNT(*), clerk FROM orders GROUP BY clerk")));
        }
        ClusterMemoryManager memoryManager = queryRunner.getCoordinator(0).getClusterMemoryManager();
        ClusterMemoryPoolInfo reservedPool;
        while ((reservedPool = memoryManager.getClusterInfo(RESERVED_POOL)) == null) {
            MILLISECONDS.sleep(10);
        }
        ClusterMemoryPoolInfo generalPool = memoryManager.getClusterInfo(GENERAL_POOL);
        assertNotNull(generalPool);
        // Wait for the queries to start running and get assigned to the expected pools
        while (generalPool.getAssignedQueries() != 1 || reservedPool.getAssignedQueries() != 1 || generalPool.getBlockedNodes() != 3 || reservedPool.getBlockedNodes() != 3) {
            generalPool = memoryManager.getClusterInfo(GENERAL_POOL);
            reservedPool = memoryManager.getClusterInfo(RESERVED_POOL);
            MILLISECONDS.sleep(10);
        }
        // Make sure the queries are blocked
        List<BasicQueryInfo> currentQueryInfos;
        do {
            currentQueryInfos = queryRunner.getCoordinators().stream().map(TestingPrestoServer::getQueryManager).map(QueryManager::getQueries).flatMap(Collection::stream).collect(toImmutableList());
            MILLISECONDS.sleep(10);
        } while (currentQueryInfos.size() != 2);
        for (BasicQueryInfo info : currentQueryInfos) {
            assertFalse(info.getState().isDone());
        }
        // Check that the pool information propagated to the query objects
        assertNotEquals(currentQueryInfos.get(0).getMemoryPool(), currentQueryInfos.get(1).getMemoryPool());
        while (!currentQueryInfos.stream().allMatch(TestMemoryManager::isBlockedWaitingForMemory)) {
            MILLISECONDS.sleep(10);
            currentQueryInfos = queryRunner.getCoordinators().stream().map(TestingPrestoServer::getQueryManager).map(QueryManager::getQueries).flatMap(Collection::stream).collect(toImmutableList());
            for (BasicQueryInfo info : currentQueryInfos) {
                assertFalse(info.getState().isDone());
            }
        }
        // Release the memory in the reserved pool
        for (TestingPrestoServer server : queryRunner.getCoordinatorWorkers()) {
            Optional<MemoryPool> reserved = server.getLocalMemoryManager().getReservedPool();
            assertTrue(reserved.isPresent());
            // Free up the entire pool
            reserved.get().free(fakeQueryId, "test", reserved.get().getMaxBytes());
            assertTrue(reserved.get().getFreeBytes() > 0);
        }
        // This also checks that the query in the general pool is successfully moved to the reserved pool.
        for (Future<?> query : queryFutures) {
            query.get();
        }
        queryRunner.getCoordinators().stream().map(TestingPrestoServer::getQueryManager).map(QueryManager::getQueries).flatMap(Collection::stream).forEach(info -> assertEquals(info.getState(), FINISHED));
        // Make sure we didn't leak any memory on the workers
        for (TestingPrestoServer worker : queryRunner.getCoordinatorWorkers()) {
            Optional<MemoryPool> reserved = worker.getLocalMemoryManager().getReservedPool();
            assertTrue(reserved.isPresent());
            assertEquals(reserved.get().getMaxBytes(), reserved.get().getFreeBytes());
            MemoryPool general = worker.getLocalMemoryManager().getGeneralPool();
            // Free up the memory we reserved earlier
            general.free(fakeQueryId, "test", general.getMaxBytes());
            assertEquals(general.getMaxBytes(), general.getFreeBytes());
        }
    }
}