[fix](fe) Guard fragment cleanup until dispatch completes

fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java

@@ -45,6 +45,9 @@ public abstract class AbstractJobProcessor implements JobProcessor {
     private final Logger logger = LogManager.getLogger(getClass());
 
     protected final AtomicBoolean finished = new AtomicBoolean(false);
+    // FINISHED cleanup must wait until all fragment dispatch RPCs, including phase-two starts, complete.
+    private final AtomicBoolean executionFinished = new AtomicBoolean(false);
+    private final AtomicBoolean fragmentDispatchCompleted = new AtomicBoolean(false);
     protected final CoordinatorContext coordinatorContext;
     protected volatile Optional<PipelineExecutionTask> executionTask;
     protected volatile Optional<Map<BackendFragmentId, SingleFragmentPipelineTask>> backendFragmentTasks;
@@ -74,6 +77,20 @@ protected void afterSetPipelineExecutionTask(PipelineExecutionTask pipelineExecu
 
     @Override
     public void tryFinishSchedule() {
+        executionFinished.set(true);
+        tryBroadcastExecutionFinished();
+    }
+
+    @Override
+    public void markFragmentDispatchCompleted() {
+        fragmentDispatchCompleted.set(true);
+        tryBroadcastExecutionFinished();
+    }
+
+    private void tryBroadcastExecutionFinished() {
+        if (!executionFinished.get() || !fragmentDispatchCompleted.get()) {
+            return;
+        }
         if (finished.compareAndSet(false, true)) {
             this.executionTask.ifPresent(sqlPipelineTask -> {
                 for (MultiFragmentsPipelineTask fragmentsTask : sqlPipelineTask.getChildrenTasks().values()) {

fe/fe-core/src/main/java/org/apache/doris/qe/JobProcessor.java

@@ -29,4 +29,6 @@ public interface JobProcessor {
     void updateFragmentExecStatus(TReportExecStatusParams params);
 
     void tryFinishSchedule();
+
+    void markFragmentDispatchCompleted();
 }

fe/fe-core/src/main/java/org/apache/doris/qe/runtime/PipelineExecutionTask.java

@@ -97,6 +97,7 @@ public void execute() throws Exception {
             if (coordinatorContext.twoPhaseExecution()) {
                 sendAndWaitPhaseTwoRpc();
             }
+            coordinatorContext.getJobProcessor().markFragmentDispatchCompleted();
             return null;
         });
     }
@@ -168,6 +169,9 @@ private Map<TNetworkAddress, List<Long>> waitPipelineRpc(
                         queryOptions.isSetQueryTimeout(), queryOptions.getQueryTimeout(),
                         timeoutDeadline, currentTimeMillis);
             }
+            Status cancelStatus = new Status(TStatusCode.INTERNAL_ERROR, msg);
+            coordinatorContext.updateStatusIfOk(cancelStatus);
+            coordinatorContext.cancelSchedule(cancelStatus);
             throw new UserException(msg);
         }
 

fe/fe-core/src/test/java/org/apache/doris/qe/AbstractJobProcessorTest.java

@@ -0,0 +1,88 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.qe;
+
+import org.apache.doris.common.Status;
+import org.apache.doris.nereids.trees.plans.distribute.worker.BackendWorker;
+import org.apache.doris.qe.runtime.MultiFragmentsPipelineTask;
+import org.apache.doris.qe.runtime.PipelineExecutionTask;
+import org.apache.doris.qe.runtime.SingleFragmentPipelineTask;
+import org.apache.doris.thrift.TReportExecStatusParams;
+
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+import java.util.Collections;
+import java.util.Optional;
+
+class AbstractJobProcessorTest {
+    @Test
+    void finishBeforeFragmentDispatchDoesNotCancelPreparedFragments() {
+        MultiFragmentsPipelineTask fragmentsTask = Mockito.mock(MultiFragmentsPipelineTask.class);
+        TestJobProcessor processor = createProcessor(fragmentsTask);
+
+        processor.tryFinishSchedule();
+        Mockito.verifyNoInteractions(fragmentsTask);
+
+        processor.markFragmentDispatchCompleted();
+        Mockito.verify(fragmentsTask).cancelExecute(Status.FINISHED);
+
+        processor.tryFinishSchedule();
+        processor.markFragmentDispatchCompleted();
+        Mockito.verifyNoMoreInteractions(fragmentsTask);
+    }
+
+    @Test
+    void fragmentDispatchBeforeFinishBroadcastsWhenExecutionFinishes() {
+        MultiFragmentsPipelineTask fragmentsTask = Mockito.mock(MultiFragmentsPipelineTask.class);
+        TestJobProcessor processor = createProcessor(fragmentsTask);
+
+        processor.markFragmentDispatchCompleted();
+        Mockito.verifyNoInteractions(fragmentsTask);
+
+        processor.tryFinishSchedule();
+        Mockito.verify(fragmentsTask).cancelExecute(Status.FINISHED);
+    }
+
+    private static TestJobProcessor createProcessor(MultiFragmentsPipelineTask fragmentsTask) {
+        BackendWorker worker = Mockito.mock(BackendWorker.class);
+        PipelineExecutionTask executionTask = Mockito.mock(PipelineExecutionTask.class);
+        Mockito.when(executionTask.getChildrenTasks()).thenReturn(Collections.singletonMap(worker, fragmentsTask));
+
+        TestJobProcessor processor = new TestJobProcessor(Mockito.mock(CoordinatorContext.class));
+        processor.setExecutionTask(executionTask);
+        return processor;
+    }
+
+    private static class TestJobProcessor extends AbstractJobProcessor {
+        TestJobProcessor(CoordinatorContext coordinatorContext) {
+            super(coordinatorContext);
+        }
+
+        void setExecutionTask(PipelineExecutionTask executionTask) {
+            this.executionTask = Optional.of(executionTask);
+        }
+
+        @Override
+        protected void doProcessReportExecStatus(
+                TReportExecStatusParams params, SingleFragmentPipelineTask fragmentTask) {}
+
+        @Override
+        public void cancel(Status cancelReason) {}
+    }
+}

fe/fe-core/src/test/java/org/apache/doris/qe/runtime/PipelineExecutionTaskTest.java

@@ -0,0 +1,76 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.qe.runtime;
+
+import org.apache.doris.common.Status;
+import org.apache.doris.common.UserException;
+import org.apache.doris.common.jmockit.Deencapsulation;
+import org.apache.doris.nereids.trees.plans.distribute.worker.BackendWorker;
+import org.apache.doris.proto.InternalService.PExecPlanFragmentResult;
+import org.apache.doris.qe.CoordinatorContext;
+import org.apache.doris.rpc.BackendServiceProxy;
+import org.apache.doris.thrift.TQueryOptions;
+import org.apache.doris.thrift.TUniqueId;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentMatchers;
+import org.mockito.Mockito;
+
+import java.util.Collections;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Supplier;
+
+class PipelineExecutionTaskTest {
+    @Test
+    void expiredDeadlineCancelsAlreadySubmittedFragments() throws Exception {
+        CoordinatorContext coordinatorContext = Mockito.mock(CoordinatorContext.class);
+        TQueryOptions queryOptions = new TQueryOptions();
+        queryOptions.setExecutionTimeout(1);
+        queryOptions.setQueryTimeout(1);
+        Deencapsulation.setField(coordinatorContext, "queryOptions", queryOptions);
+        Deencapsulation.setField(coordinatorContext, "queryId", new TUniqueId(1, 2));
+        Deencapsulation.setField(coordinatorContext, "timeoutDeadline", (Supplier<Long>) () -> 0L);
+        Mockito.when(coordinatorContext.withLock(ArgumentMatchers.<Callable<Object>>any()))
+                .thenAnswer(invocation -> invocation.<Callable<Object>>getArgument(0).call());
+        Mockito.when(coordinatorContext.twoPhaseExecution()).thenReturn(false);
+
+        MultiFragmentsPipelineTask fragmentsTask = Mockito.mock(MultiFragmentsPipelineTask.class);
+        Mockito.when(fragmentsTask.getChildrenTasks()).thenReturn(Collections.emptyMap());
+        Mockito.when(fragmentsTask.sendPhaseOneRpc(false))
+                .thenReturn(CompletableFuture.completedFuture(PExecPlanFragmentResult.getDefaultInstance()));
+        PipelineExecutionTask executionTask = new PipelineExecutionTask(
+                coordinatorContext,
+                Mockito.mock(BackendServiceProxy.class),
+                Collections.singletonMap(Mockito.mock(BackendWorker.class), fragmentsTask));
+
+        UserException exception = Assertions.assertThrows(UserException.class, executionTask::execute);
+
+        Assertions.assertTrue(exception.getMessage().contains("timeout before waiting send fragments rpc"));
+        Mockito.verify(fragmentsTask).sendPhaseOneRpc(false);
+        Mockito.verify(coordinatorContext).updateStatusIfOk(ArgumentMatchers.argThat(
+                status -> hasDeadlineTimeoutMessage(status)));
+        Mockito.verify(coordinatorContext).cancelSchedule(ArgumentMatchers.argThat(
+                status -> hasDeadlineTimeoutMessage(status)));
+    }
+
+    private static boolean hasDeadlineTimeoutMessage(Status status) {
+        return status.getErrorMsg().contains("timeout before waiting send fragments rpc");
+    }
+}