Examples with QueryState io.prestosql.execution.QueryState used on opensource projects

Example 1 with QueryState

use of io.prestosql.execution.QueryState in project hetu-core by openlookeng.

the class TestQueryManager method testFailQuery.

@Test(timeOut = 60_000L)
public void testFailQuery() throws Exception {
    DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager();
    QueryId queryId = dispatchManager.createQueryId();
    dispatchManager.createQuery(queryId, "slug", new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem").get();
    // wait until query starts running
    while (true) {
        QueryState state = dispatchManager.getQueryInfo(queryId).getState();
        if (state.isDone()) {
            fail("unexpected query state: " + state);
        }
        if (state == RUNNING) {
            break;
        }
        Thread.sleep(100);
    }
    // cancel query
    QueryManager queryManager = queryRunner.getCoordinator().getQueryManager();
    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");
}

Example 2 with QueryState

use of io.prestosql.execution.QueryState in project hetu-core by openlookeng.

the class QueryResource method failQuery.

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();
    }
}

Example 3 with QueryState

use of io.prestosql.execution.QueryState in project hetu-core by openlookeng.

the class TestPrestoDriver method findQueryState.

private Optional<QueryState> findQueryState(String text) throws SQLException {
    String sql = format("SELECT state FROM system.runtime.queries WHERE regexp_like(query, '%s$') /* */", Pattern.quote(text));
    try (Connection connection = createConnection();
        Statement statement = connection.createStatement();
        ResultSet resultSet = statement.executeQuery(sql)) {
        if (!resultSet.next()) {
            return Optional.empty();
        }
        QueryState state = QueryState.valueOf(requireNonNull(resultSet.getString(1)));
        assertFalse(resultSet.next(), "Found multiple queries");
        return Optional.of(state);
    }
}

Example 4 with QueryState

use of io.prestosql.execution.QueryState in project hetu-core by openlookeng.

the class TestPrestoDriver method testUpdateCancelExplicit.

@Test(timeOut = 10000)
public void testUpdateCancelExplicit() throws Exception {
    CountDownLatch queryFinished = new CountDownLatch(1);
    AtomicReference<String> queryId = new AtomicReference<>();
    AtomicReference<Throwable> queryFailure = new AtomicReference<>();
    String queryUuid = "/* " + UUID.randomUUID().toString() + " */";
    try (Connection connection = createConnection("blackhole", "blackhole");
        Statement statement = connection.createStatement()) {
        // execute the slow update on another thread
        executorService.execute(() -> {
            try {
                statement.executeUpdate("CREATE TABLE test_cancel_create AS SELECT * FROM slow_test_table " + queryUuid);
            } catch (SQLException t) {
                queryFailure.set(t);
            } finally {
                queryFinished.countDown();
            }
        });
        // start query and make sure it is not finished
        while (true) {
            Optional<QueryState> state = findQueryState(queryUuid);
            if (state.isPresent()) {
                assertFalse(state.get().isDone());
                break;
            }
            MILLISECONDS.sleep(50);
        }
        // cancel the query from this test thread
        statement.cancel();
        // make sure the query was aborted
        queryFinished.await(10, SECONDS);
        assertNotNull(queryFailure.get());
        assertEquals(findQueryState(queryUuid), Optional.of(FAILED));
    }
}

Example 5 with QueryState

use of io.prestosql.execution.QueryState in project hetu-core by openlookeng.

the class SqlQueryScheduler method schedule.

private void schedule() {
    try (SetThreadName ignored = new SetThreadName("Query-%s", queryStateMachine.getQueryId())) {
        Set<StageId> completedStages = new HashSet<>();
        ExecutionSchedule executionSchedule = executionPolicy.createExecutionSchedule(stages.values());
        while (!executionSchedule.isFinished()) {
            List<ListenableFuture<?>> blockedStages = new ArrayList<>();
            for (SqlStageExecution stage : executionSchedule.getStagesToSchedule()) {
                if (isReuseTableScanEnabled(session) && !SqlStageExecution.getReuseTableScanMappingIdStatus(stage.getStateMachine())) {
                    continue;
                }
                stage.beginScheduling();
                // and then let it schedule 10% of splits.
                if (queryStateMachine.isThrottlingEnabled() && !canScheduleMoreSplits()) {
                    try {
                        SECONDS.sleep(THROTTLE_SLEEP_TIMER[currentTimerLevel]);
                    } catch (InterruptedException e) {
                        throw new PrestoException(GENERIC_INTERNAL_ERROR, "interrupted while sleeping");
                    }
                    currentTimerLevel = Math.min(currentTimerLevel + 1, THROTTLE_SLEEP_TIMER.length - 1);
                    stage.setThrottledSchedule(true);
                } else {
                    stage.setThrottledSchedule(false);
                    currentTimerLevel = 0;
                }
                // perform some scheduling work
                /* Get groupSize specification from the ResourceGroupManager */
                int maxSplitGroupSize = getOptimalSmallSplitGroupSize();
                ScheduleResult result = stageSchedulers.get(stage.getStageId()).schedule(maxSplitGroupSize);
                // modify parent and children based on the results of the scheduling
                if (result.isFinished()) {
                    stage.schedulingComplete();
                } else if (!result.getBlocked().isDone()) {
                    blockedStages.add(result.getBlocked());
                }
                stageLinkages.get(stage.getStageId()).processScheduleResults(stage.getState(), result.getNewTasks());
                schedulerStats.getSplitsScheduledPerIteration().add(result.getSplitsScheduled());
                if (result.getBlockedReason().isPresent()) {
                    switch(result.getBlockedReason().get()) {
                        case WRITER_SCALING:
                            // no-op
                            break;
                        case WAITING_FOR_SOURCE:
                            schedulerStats.getWaitingForSource().update(1);
                            break;
                        case SPLIT_QUEUES_FULL:
                            schedulerStats.getSplitQueuesFull().update(1);
                            break;
                        case MIXED_SPLIT_QUEUES_FULL_AND_WAITING_FOR_SOURCE:
                        case NO_ACTIVE_DRIVER_GROUP:
                            break;
                        default:
                            throw new UnsupportedOperationException("Unknown blocked reason: " + result.getBlockedReason().get());
                    }
                }
            }
            // make sure to update stage linkage at least once per loop to catch async state changes (e.g., partial cancel)
            for (SqlStageExecution stage : stages.values()) {
                if (!completedStages.contains(stage.getStageId()) && stage.getState().isDone()) {
                    stageLinkages.get(stage.getStageId()).processScheduleResults(stage.getState(), ImmutableSet.of());
                    completedStages.add(stage.getStageId());
                }
            }
            // wait for a state change and then schedule again
            if (!blockedStages.isEmpty()) {
                try (TimeStat.BlockTimer timer = schedulerStats.getSleepTime().time()) {
                    tryGetFutureValue(whenAnyComplete(blockedStages), 1, SECONDS);
                }
                for (ListenableFuture<?> blockedStage : blockedStages) {
                    blockedStage.cancel(true);
                }
            }
        }
        for (SqlStageExecution stage : stages.values()) {
            StageState state = stage.getState();
            // Snapshot: if state is resumable_failure, then state of stage and query will change soon again. Don't treat as an error.
            if (state != SCHEDULED && state != RUNNING && !state.isDone() && state != RESUMABLE_FAILURE) {
                throw new PrestoException(GENERIC_INTERNAL_ERROR, format("Scheduling is complete, but stage %s is in state %s", stage.getStageId(), state));
            }
        }
    } catch (Throwable t) {
        queryStateMachine.transitionToFailed(t);
        throw t;
    } finally {
        RuntimeException closeError = new RuntimeException();
        for (StageScheduler scheduler : stageSchedulers.values()) {
            try {
                // Snapshot: when trying to reschedule, then don't close the scheduler (and more importantly, split sources in it)
                QueryState state = queryStateMachine.getQueryState();
                if (state != QueryState.RESCHEDULING && state != QueryState.RESUMING) {
                    scheduler.close();
                }
            } catch (Throwable t) {
                queryStateMachine.transitionToFailed(t);
                // Self-suppression not permitted
                if (closeError != t) {
                    closeError.addSuppressed(t);
                }
            }
        }
        // Snpashot: if resuming, notify the new scheduler so it can start scheduling new stages
        schedulingFuture.set(null);
        if (closeError.getSuppressed().length > 0) {
            throw closeError;
        }
    }
}