@GET public List<BasicQueryInfo> getAllQueryInfo(@QueryParam("state") String stateFilter) { QueryState expectedState = stateFilter == null ? null : QueryState.valueOf(stateFilter.toUpperCase(Locale.ENGLISH)); ImmutableList.Builder<BasicQueryInfo> builder = new ImmutableList.Builder<>(); for (BasicQueryInfo queryInfo : queryManager.getQueries()) { if (stateFilter == null || queryInfo.getState() == expectedState) { builder.add(queryInfo); } } return builder.build(); }
@GET @Produces(MediaType.APPLICATION_JSON) public List<QueryStateInfo> getQueryStateInfos(@QueryParam("user") String user) { List<BasicQueryInfo> queryInfos = queryManager.getQueries(); if (!isNullOrEmpty(user)) { queryInfos = queryInfos.stream() .filter(queryInfo -> Pattern.matches(user, queryInfo.getSession().getUser())) .collect(toImmutableList()); } return queryInfos.stream() .filter(queryInfo -> !queryInfo.getState().isDone()) .map(this::getQueryStateInfo) .collect(toImmutableList()); }
private Set<String> captureQueries(Callable<?> action) throws Exception { Set<QueryId> queryIdsBefore = server.getQueryManager().getQueries().stream() .map(BasicQueryInfo::getQueryId) .collect(toImmutableSet()); action.call(); return server.getQueryManager().getQueries().stream() .filter(queryInfo -> !queryIdsBefore.contains(queryInfo.getQueryId())) .map(BasicQueryInfo::getQuery) .collect(toImmutableSet()); }
public static void waitForQueryCount(DistributedQueryRunner queryRunner, Set<QueryState> countingStates, int expectedCount) throws InterruptedException { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); while (queryManager.getQueries().stream() .filter(q -> countingStates.contains(q.getState())).count() != expectedCount) { MILLISECONDS.sleep(500); } }
private void cancelAllQueries() { QueryManager queryManager = coordinator.getQueryManager(); for (BasicQueryInfo queryInfo : queryManager.getQueries()) { if (!queryInfo.getState().isDone()) { queryManager.cancelQuery(queryInfo.getQueryId()); } } }
private void testNoLeak(@Language("SQL") String query) throws Exception { Map<String, String> properties = ImmutableMap.<String, String>builder() .put("task.verbose-stats", "true") .build(); try (DistributedQueryRunner queryRunner = createQueryRunner(TINY_SESSION, properties)) { executor.submit(() -> queryRunner.execute(query)).get(); for (BasicQueryInfo info : queryRunner.getCoordinator().getQueryManager().getQueries()) { assertEquals(info.getState(), FINISHED); } // Make sure we didn't leak any memory on the workers for (TestingPrestoServer worker : queryRunner.getServers()) { Optional<MemoryPool> reserved = worker.getLocalMemoryManager().getReservedPool(); assertTrue(reserved.isPresent()); assertEquals(reserved.get().getMaxBytes(), reserved.get().getFreeBytes()); MemoryPool general = worker.getLocalMemoryManager().getGeneralPool(); assertEquals(general.getMaxBytes(), general.getFreeBytes()); } } }
List<BasicQueryInfo> queryInfos = queryRunner.getCoordinator().getQueryManager().getQueries(); for (BasicQueryInfo info : queryInfos) { assertEquals(info.getState(), FINISHED);
public static void waitForQueryState(DistributedQueryRunner queryRunner, QueryId queryId, Set<QueryState> expectedQueryStates) throws InterruptedException { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); do { // Heartbeat all the running queries, so they don't die while we're waiting for (BasicQueryInfo queryInfo : queryManager.getQueries()) { if (queryInfo.getState() == RUNNING) { queryManager.recordHeartbeat(queryInfo.getQueryId()); } } MILLISECONDS.sleep(500); } while (!expectedQueryStates.contains(queryManager.getQueryState(queryId))); }
private void waitForQueryToBeKilled(DistributedQueryRunner queryRunner) throws InterruptedException { while (true) { for (BasicQueryInfo info : queryRunner.getCoordinator().getQueryManager().getQueries()) { if (info.getState().isDone()) { assertNotNull(info.getErrorCode()); assertEquals(info.getErrorCode().getCode(), CLUSTER_OUT_OF_MEMORY.toErrorCode().getCode()); return; } } MILLISECONDS.sleep(10); } }
@Test public void testQueryLoggingCount() { QueryManager queryManager = ((DistributedQueryRunner) getQueryRunner()).getCoordinator().getQueryManager(); executeExclusively(() -> { assertUntilTimeout( () -> assertEquals( queryManager.getQueries().stream() .map(BasicQueryInfo::getQueryId) .map(queryManager::getFullQueryInfo) .filter(info -> !info.isFinalQueryInfo()) .collect(toList()), ImmutableList.of()), new Duration(1, MINUTES)); // We cannot simply get the number of completed queries as soon as all the queries are completed, because this counter may not be up-to-date at that point. // The completed queries counter is updated in a final query info listener, which is called eventually. // Therefore, here we wait until the value of this counter gets stable. long beforeCompletedQueriesCount = waitUntilStable(() -> queryManager.getStats().getCompletedQueries().getTotalCount(), new Duration(5, SECONDS)); long beforeSubmittedQueriesCount = queryManager.getStats().getSubmittedQueries().getTotalCount(); assertUpdate("CREATE TABLE test_query_logging_count AS SELECT 1 foo_1, 2 foo_2_4", 1); assertQuery("SELECT foo_1, foo_2_4 FROM test_query_logging_count", "SELECT 1, 2"); assertUpdate("DROP TABLE test_query_logging_count"); assertQueryFails("SELECT * FROM test_query_logging_count", ".*Table .* does not exist"); // TODO: Figure out a better way of synchronization assertUntilTimeout( () -> assertEquals(queryManager.getStats().getCompletedQueries().getTotalCount() - beforeCompletedQueriesCount, 4), new Duration(1, MINUTES)); assertEquals(queryManager.getStats().getSubmittedQueries().getTotalCount() - beforeSubmittedQueriesCount, 4); }); }
List<QueryInfo> queryInfos = queryManager.getQueries().stream() .map(BasicQueryInfo::getQueryId) .map(queryId -> {
List<BasicQueryInfo> currentQueryInfos = queryRunner.getCoordinator().getQueryManager().getQueries(); for (BasicQueryInfo info : currentQueryInfos) { assertFalse(info.getState().isDone()); currentQueryInfos = queryRunner.getCoordinator().getQueryManager().getQueries(); for (BasicQueryInfo info : currentQueryInfos) { assertFalse(info.getState().isDone()); for (BasicQueryInfo info : queryRunner.getCoordinator().getQueryManager().getQueries()) { assertEquals(info.getState(), FINISHED);
long totalCpuTimeSecs = queryManager.getStats().getConsumedCpuTimeSecs().getTotalCount(); for (BasicQueryInfo query : queryManager.getQueries()) { if (query.getState() == QueryState.QUEUED) { queuedQueries++;
private Set<String> captureQueries(Callable<?> action) throws Exception { Set<QueryId> queryIdsBefore = server.getQueryManager().getQueries().stream() .map(BasicQueryInfo::getQueryId) .collect(toImmutableSet()); action.call(); return server.getQueryManager().getQueries().stream() .filter(queryInfo -> !queryIdsBefore.contains(queryInfo.getQueryId())) .map(BasicQueryInfo::getQuery) .collect(toImmutableSet()); }
public static void waitForQueryCount(DistributedQueryRunner queryRunner, Set<QueryState> countingStates, int expectedCount) throws InterruptedException { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); while (queryManager.getQueries().stream() .filter(q -> countingStates.contains(q.getState())).count() != expectedCount) { MILLISECONDS.sleep(500); } }
private void cancelAllQueries() { QueryManager queryManager = coordinator.getQueryManager(); for (BasicQueryInfo queryInfo : queryManager.getQueries()) { if (!queryInfo.getState().isDone()) { queryManager.cancelQuery(queryInfo.getQueryId()); } } }
public static void waitForQueryState(DistributedQueryRunner queryRunner, QueryId queryId, Set<QueryState> expectedQueryStates) throws InterruptedException { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); do { // Heartbeat all the running queries, so they don't die while we're waiting for (BasicQueryInfo queryInfo : queryManager.getQueries()) { if (queryInfo.getState() == RUNNING) { queryManager.recordHeartbeat(queryInfo.getQueryId()); } } MILLISECONDS.sleep(500); } while (!expectedQueryStates.contains(queryManager.getQueryState(queryId))); }
private void testNoLeak(@Language("SQL") String query) throws Exception { Map<String, String> properties = ImmutableMap.<String, String>builder() .put("task.verbose-stats", "true") .build(); try (DistributedQueryRunner queryRunner = createQueryRunner(TINY_SESSION, properties)) { executor.submit(() -> queryRunner.execute(query)).get(); for (BasicQueryInfo info : queryRunner.getCoordinator().getQueryManager().getQueries()) { assertEquals(info.getState(), FINISHED); } // Make sure we didn't leak any memory on the workers for (TestingPrestoServer worker : queryRunner.getServers()) { Optional<MemoryPool> reserved = worker.getLocalMemoryManager().getReservedPool(); assertTrue(reserved.isPresent()); assertEquals(reserved.get().getMaxBytes(), reserved.get().getFreeBytes()); MemoryPool general = worker.getLocalMemoryManager().getGeneralPool(); assertEquals(general.getMaxBytes(), general.getFreeBytes()); } } }
private void waitForQueryToBeKilled(DistributedQueryRunner queryRunner) throws InterruptedException { while (true) { for (BasicQueryInfo info : queryRunner.getCoordinator().getQueryManager().getQueries()) { if (info.getState().isDone()) { assertNotNull(info.getErrorCode()); assertEquals(info.getErrorCode().getCode(), CLUSTER_OUT_OF_MEMORY.toErrorCode().getCode()); return; } } MILLISECONDS.sleep(10); } }
@Test public void testQueryLoggingCount() { QueryManager queryManager = ((DistributedQueryRunner) getQueryRunner()).getCoordinator().getQueryManager(); executeExclusively(() -> { assertUntilTimeout( () -> assertEquals( queryManager.getQueries().stream() .map(BasicQueryInfo::getQueryId) .map(queryManager::getFullQueryInfo) .filter(info -> !info.isFinalQueryInfo()) .collect(toList()), ImmutableList.of()), new Duration(1, MINUTES)); // We cannot simply get the number of completed queries as soon as all the queries are completed, because this counter may not be up-to-date at that point. // The completed queries counter is updated in a final query info listener, which is called eventually. // Therefore, here we wait until the value of this counter gets stable. long beforeCompletedQueriesCount = waitUntilStable(() -> queryManager.getStats().getCompletedQueries().getTotalCount(), new Duration(5, SECONDS)); long beforeSubmittedQueriesCount = queryManager.getStats().getSubmittedQueries().getTotalCount(); assertUpdate("CREATE TABLE test_query_logging_count AS SELECT 1 foo_1, 2 foo_2_4", 1); assertQuery("SELECT foo_1, foo_2_4 FROM test_query_logging_count", "SELECT 1, 2"); assertUpdate("DROP TABLE test_query_logging_count"); assertQueryFails("SELECT * FROM test_query_logging_count", ".*Table .* does not exist"); // TODO: Figure out a better way of synchronization assertUntilTimeout( () -> assertEquals(queryManager.getStats().getCompletedQueries().getTotalCount() - beforeCompletedQueriesCount, 4), new Duration(1, MINUTES)); assertEquals(queryManager.getStats().getSubmittedQueries().getTotalCount() - beforeSubmittedQueriesCount, 4); }); }