diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java index 54e607dd27afa0..fdd30dd435fb36 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java +++ b/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 executionTask; protected volatile Optional> 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()) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/JobProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/JobProcessor.java index 6365f26c4f9c2a..50cbdf5a253ced 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/JobProcessor.java +++ b/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(); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/PipelineExecutionTask.java b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/PipelineExecutionTask.java index 6a52e3a6d9f855..59670ccbbf13db 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/PipelineExecutionTask.java +++ b/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> 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); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/AbstractJobProcessorTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/AbstractJobProcessorTest.java new file mode 100644 index 00000000000000..64f57cfec1e111 --- /dev/null +++ b/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) {} + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/PipelineExecutionTaskTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/PipelineExecutionTaskTest.java new file mode 100644 index 00000000000000..762ffdafd18961 --- /dev/null +++ b/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) () -> 0L); + Mockito.when(coordinatorContext.withLock(ArgumentMatchers.>any())) + .thenAnswer(invocation -> invocation.>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"); + } +}