@Override public QueryState getQueryState(QueryId queryId) { return queryTracker.getQuery(queryId).getState(); }
/** * Enforce memory limits at the query level */ private void enforceMemoryLimits() { List<QueryExecution> runningQueries = queryTracker.getAllQueries().stream() .filter(query -> query.getState() == RUNNING) .collect(toImmutableList()); memoryManager.process(runningQueries, this::getQueries); }
/** * Enforce memory limits at the query level */ public void enforceMemoryLimits() { memoryManager.process(queries.values().stream() .filter(query -> query.getState() == RUNNING) .collect(toImmutableList())); }
/** * Set up a callback to fire when a query is completed. The callback will be called at most once. */ static void addCompletionCallback(QueryExecution queryExecution, Runnable callback) { AtomicBoolean taskExecuted = new AtomicBoolean(); queryExecution.addStateChangeListener(newValue -> { if (newValue.isDone() && taskExecuted.compareAndSet(false, true)) { callback.run(); } }); // Need to do this check in case the state changed before we added the previous state change listener if (queryExecution.getState().isDone() && taskExecuted.compareAndSet(false, true)) { callback.run(); } } }
@PreDestroy public void stop() { boolean queryCancelled = false; for (QueryExecution queryExecution : queries.values()) { if (queryExecution.getState().isDone()) { continue; } log.info("Server shutting down. Query %s has been cancelled", queryExecution.getQueryId()); queryExecution.fail(new PrestoException(SERVER_SHUTTING_DOWN, "Server is shutting down. Query " + queryExecution.getQueryId() + " has been cancelled")); queryCancelled = true; } if (queryCancelled) { try { TimeUnit.SECONDS.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } } queryManagementExecutor.shutdownNow(); queryExecutor.shutdownNow(); }
/** * Enforce timeout at the query level */ public void enforceQueryMaxRunTimeLimits() { for (QueryExecution query : queries.values()) { if (query.getState().isDone()) { continue; } Duration queryMaxRunTime = SystemSessionProperties.getQueryMaxRunTime(query.getSession()); DateTime executionStartTime = query.getQueryInfo().getQueryStats().getCreateTime(); if (executionStartTime.plus(queryMaxRunTime.toMillis()).isBeforeNow()) { query.fail(new PrestoException(EXCEEDED_TIME_LIMIT, "Query exceeded maximum time limit of " + queryMaxRunTime)); } } }