Properly handle aborted dependent stages in PhasedExecutionSchedule

trinodb · Jan 18, 2022 · b265bbf · b265bbf
1 parent 2874b92
commit b265bbf
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 1 deletion.
diff --git a/...trino-main/src/main/java/io/trino/execution/scheduler/policy/PhasedExecutionSchedule.java b/...trino-main/src/main/java/io/trino/execution/scheduler/policy/PhasedExecutionSchedule.java
@@ -179,7 +179,8 @@ Set<StageExecution> getActiveStages()
 
     private Set<PlanFragmentId> removeCompletedStages()
     {
-        Set<StageExecution> completedStages = activeStages.stream()
+        // iterate over all stages, not only active ones; stages which are not yet active could have already been aborted
+        Set<StageExecution> completedStages = stagesByFragmentId.values().stream()
                 .filter(this::isStageCompleted)
                 .collect(toImmutableSet());
         // remove completed stages outside of Java stream to prevent concurrent modification
@@ -192,6 +193,10 @@ private Set<PlanFragmentId> removeCompletedStage(StageExecution stage)
     {
         // start all stages that depend on completed stage
         PlanFragmentId fragmentId = stage.getFragment().getId();
+        if (!fragmentDependency.containsVertex(fragmentId)) {
+            // already gone
+            return ImmutableSet.of();
+        }
         Set<PlanFragmentId> fragmentsToExecute = fragmentDependency.outgoingEdgesOf(fragmentId).stream()
                 .map(FragmentsEdge::getTarget)
                 // filter stages that depend on completed stage only

diff --git a/...o-main/src/test/java/io/trino/execution/scheduler/policy/TestPhasedExecutionSchedule.java b/...o-main/src/test/java/io/trino/execution/scheduler/policy/TestPhasedExecutionSchedule.java
@@ -40,6 +40,7 @@
 import java.util.function.Consumer;
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static io.trino.execution.scheduler.StageExecution.State.ABORTED;
 import static io.trino.execution.scheduler.StageExecution.State.FINISHED;
 import static io.trino.execution.scheduler.StageExecution.State.FLUSHING;
 import static io.trino.execution.scheduler.policy.PlanUtils.createAggregationFragment;
@@ -142,6 +143,40 @@ public void testAggregation()
         assertThat(getActiveFragments(schedule)).containsExactly(buildFragment.getId(), sourceFragment.getId(), aggregationFragment.getId());
     }
 
+    @Test
+    public void testDependentStageAbortedBeforeStarted()
+    {
+        PlanFragment sourceFragment = createTableScanPlanFragment("probe");
+        PlanFragment aggregationFragment = createAggregationFragment("aggregation", sourceFragment);
+        PlanFragment buildFragment = createTableScanPlanFragment("build");
+        PlanFragment joinFragment = createJoinPlanFragment(INNER, REPLICATED, "join", buildFragment, aggregationFragment);
+
+        TestingStageExecution sourceStage = new TestingStageExecution(sourceFragment);
+        TestingStageExecution aggregationStage = new TestingStageExecution(aggregationFragment);
+        TestingStageExecution buildStage = new TestingStageExecution(buildFragment);
+        TestingStageExecution joinStage = new TestingStageExecution(joinFragment);
+
+        PhasedExecutionSchedule schedule = PhasedExecutionSchedule.forStages(ImmutableSet.of(sourceStage, aggregationStage, buildStage, joinStage));
+        assertThat(schedule.getSortedFragments()).containsExactly(buildFragment.getId(), sourceFragment.getId(), aggregationFragment.getId(), joinFragment.getId());
+
+        // aggregation and source stage should start immediately, join stage should wait for build stage to complete
+        DirectedGraph<PlanFragmentId, FragmentsEdge> dependencies = schedule.getFragmentDependency();
+        assertThat(dependencies.edgeSet()).containsExactly(new FragmentsEdge(buildFragment.getId(), joinFragment.getId()));
+        assertThat(getActiveFragments(schedule)).containsExactly(buildFragment.getId(), sourceFragment.getId(), aggregationFragment.getId());
+
+        // abort non-active join stage
+        joinStage.setState(ABORTED);
+
+        // dependencies finish
+        buildStage.setState(FINISHED);
+        aggregationStage.setState(FINISHED);
+        sourceStage.setState(FINISHED);
+
+        // join stage already aborted. Whole schedule should be marked as finished
+        schedule.schedule();
+        assertThat(schedule.isFinished()).isTrue();
+    }
+
     @Test
     public void testStageWithBroadcastAndPartitionedJoin()
     {