private Response failQuery(QueryId queryId, PrestoException queryException) { requireNonNull(queryId, "queryId is null"); try { QueryState state = queryManager.getQueryState(queryId); // check before killing to provide the proper error code (this is racy) if (state.isDone()) { return Response.status(Status.CONFLICT).build(); } queryManager.failQuery(queryId, queryException); // verify if the query was failed (if not, we lost the race) if (!queryException.getErrorCode().equals(queryManager.getQueryInfo(queryId).getErrorCode())) { return Response.status(Status.CONFLICT).build(); } return Response.status(Status.OK).build(); } catch (NoSuchElementException e) { return Response.status(Status.GONE).build(); } }
if (queryManager.getQueryState(queryId) == QueryState.FAILED) { query.dispose();
private synchronized ListenableFuture<?> getFutureStateChange() { // ensure the query has been submitted submissionFuture.submitQuery(); // if query query submission has not finished, wait for it to finish if (!submissionFuture.isDone()) { return submissionFuture; } // if the exchange client is open, wait for data if (!exchangeClient.isClosed()) { return exchangeClient.isBlocked(); } // otherwise, wait for the query to finish queryManager.recordHeartbeat(queryId); try { return queryDoneFuture(queryManager.getQueryState(queryId)); } catch (NoSuchElementException e) { return immediateFuture(null); } }
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))); }
@Test(timeOut = 60_000L) public void testFailQuery() throws Exception { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); QueryId queryId = queryManager.createQueryId(); queryManager.createQuery( queryId, new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem") .get(); // wait until query starts running while (true) { QueryState state = queryManager.getQueryState(queryId); if (state.isDone()) { fail("unexpected query state: " + state); } if (state == RUNNING) { break; } Thread.sleep(100); } // cancel query queryManager.failQuery(queryId, new PrestoException(GENERIC_INTERNAL_ERROR, "mock exception")); QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId); assertEquals(queryInfo.getState(), FAILED); assertEquals(queryInfo.getErrorCode(), GENERIC_INTERNAL_ERROR.toErrorCode()); assertNotNull(queryInfo.getFailureInfo()); assertEquals(queryInfo.getFailureInfo().getMessage(), "mock exception"); }
@UsedByGeneratedCode public void killQuery(String queryId, String message) { QueryId query = parseQueryId(queryId); try { QueryState state = queryManager.getQueryState(query); // check before killing to provide the proper error message (this is racy) if (state.isDone()) { throw new PrestoException(NOT_SUPPORTED, "Target query is not running: " + queryId); } queryManager.failQuery(query, createKillQueryException(message)); // verify if the query was killed (if not, we lost the race) if (!ADMINISTRATIVELY_KILLED.toErrorCode().equals(queryManager.getQueryInfo(query).getErrorCode())) { throw new PrestoException(NOT_SUPPORTED, "Target query is not running: " + queryId); } } catch (NoSuchElementException e) { throw new PrestoException(NOT_FOUND, "Target query not found: " + queryId); } }
@Override public void run() { try { // Queries are added to the query manager before being recorded in queryIds set. // Therefore, we take a snapshot if queryIds before getting the live queries // from the query manager. Then we remove only the queries in the snapshot and // not live queries set. If we did this in the other order, a query could be // registered between fetching the live queries and inspecting the queryIds set. for (QueryId queryId : ImmutableSet.copyOf(queries.keySet())) { Query query = queries.get(queryId); Optional<QueryState> state = queryManager.getQueryState(queryId); // free up resources if the query completed if (!state.isPresent() || state.get() == QueryState.FAILED) { query.dispose(); } // forget about this query if the query manager is no longer tracking it if (!state.isPresent()) { queries.remove(queryId); } } } catch (Throwable e) { log.warn(e, "Error removing old queries"); } } }
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))); }
@Test(timeOut = 60_000L) public void testFailQuery() throws Exception { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); QueryId queryId = queryManager.createQueryId(); queryManager.createQuery( queryId, new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem") .get(); // wait until query starts running while (true) { QueryState state = queryManager.getQueryState(queryId); if (state.isDone()) { fail("unexpected query state: " + state); } if (state == RUNNING) { break; } Thread.sleep(100); } // cancel query queryManager.failQuery(queryId, new PrestoException(GENERIC_INTERNAL_ERROR, "mock exception")); QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId); assertEquals(queryInfo.getState(), FAILED); assertEquals(queryInfo.getErrorCode(), GENERIC_INTERNAL_ERROR.toErrorCode()); assertNotNull(queryInfo.getFailureInfo()); assertEquals(queryInfo.getFailureInfo().getMessage(), "mock exception"); }