public SqlQueryScheduler(
QueryStateMachine queryStateMachine,
LocationFactory locationFactory,
StageExecutionPlan plan,
NodePartitioningManager nodePartitioningManager,
NodeScheduler nodeScheduler,
RemoteTaskFactory remoteTaskFactory,
Session session,
boolean summarizeTaskInfo,
int splitBatchSize,
ExecutorService executor,
OutputBuffers rootOutputBuffers,
NodeTaskMap nodeTaskMap,
ExecutionPolicy executionPolicy) {
this.queryStateMachine = requireNonNull(queryStateMachine, "queryStateMachine is null");
this.executionPolicy = requireNonNull(executionPolicy, "schedulerPolicyFactory is null");
this.summarizeTaskInfo = summarizeTaskInfo;
// todo come up with a better way to build this, or eliminate this map
ImmutableMap.Builder<StageId, StageScheduler> stageSchedulers = ImmutableMap.builder();
ImmutableMap.Builder<StageId, StageLinkage> stageLinkages = ImmutableMap.builder();
// Only fetch a distribution once per query to assure all stages see the same machine
// assignments
Map<PartitioningHandle, NodePartitionMap> partitioningCache = new HashMap<>();
List<SqlStageExecution> stages =
createStages(
Optional.empty(),
new AtomicInteger(),
locationFactory,
plan.withBucketToPartition(Optional.of(new int[1])),
nodeScheduler,
remoteTaskFactory,
session,
splitBatchSize,
partitioningHandle ->
partitioningCache.computeIfAbsent(
partitioningHandle,
handle -> nodePartitioningManager.getNodePartitioningMap(session, handle)),
executor,
nodeTaskMap,
stageSchedulers,
stageLinkages);
SqlStageExecution rootStage = stages.get(0);
rootStage.setOutputBuffers(rootOutputBuffers);
this.rootStageId = rootStage.getStageId();
this.stages = stages.stream().collect(toImmutableMap(SqlStageExecution::getStageId));
this.stageSchedulers = stageSchedulers.build();
this.stageLinkages = stageLinkages.build();
this.executor = executor;
rootStage.addStateChangeListener(
state -> {
if (state == FINISHED) {
queryStateMachine.transitionToFinishing();
} else if (state == CANCELED) {
// output stage was canceled
queryStateMachine.transitionToFailed(
new PrestoException(USER_CANCELED, "Query was canceled"));
}
});
for (SqlStageExecution stage : stages) {
stage.addStateChangeListener(
state -> {
if (queryStateMachine.isDone()) {
return;
}
if (state == FAILED) {
queryStateMachine.transitionToFailed(
stage.getStageInfo().getFailureCause().toException());
} else if (state == ABORTED) {
// this should never happen, since abort can only be triggered in query clean up after
// the query is finished
queryStateMachine.transitionToFailed(
new PrestoException(INTERNAL_ERROR, "Query stage was aborted"));
} else if (queryStateMachine.getQueryState() == QueryState.STARTING) {
// if the stage has at least one task, we are running
if (stage.hasTasks()) {
queryStateMachine.transitionToRunning();
}
}
});
}
}
