[Query] Cancel dispatched fragment instances right away when the dispatch fails

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java

@@ -65,6 +65,8 @@ public class ClusterScheduler implements IScheduler {
   private final IFragInstanceDispatcher dispatcher;
   private IFragInstanceStateTracker stateTracker;
   private IQueryTerminator queryTerminator;
+  // Set when the dispatch of the fragment instances failed (see start()).
+  private volatile boolean dispatchFailed = false;
 
   private static final QueryExecutionMetricSet QUERY_EXECUTION_METRICS =
       QueryExecutionMetricSet.getInstance();
@@ -122,6 +124,7 @@ public void start() {
     try {
       FragInstanceDispatchResult result = dispatchResultFuture.get();
       if (!result.isSuccessful()) {
+        dispatchFailed = true;
         if (needRetry(result.getFailureStatus())) {
           stateMachine.transitionToPendingRetry(result.getFailureStatus());
         } else {
@@ -132,6 +135,7 @@ public void start() {
       }
     } catch (InterruptedException | ExecutionException e) {
       // If the dispatch request cannot be sent or TException is caught, we will retry this query.
+      dispatchFailed = true;
       if (e instanceof InterruptedException) {
         Thread.currentThread().interrupt();
       }
@@ -162,7 +166,13 @@ public void stop(Throwable t) {
       stateTracker.abort();
     }
     if (queryTerminator != null) {
-      queryTerminator.terminate(t);
+      if (dispatchFailed) {
+        // The FragmentInstances dispatched before the failure can't finish on their own and hold
+        // operator memory until they are cancelled, so don't wait for the grace period.
+        queryTerminator.terminateImmediately(t);
+      } else {
+        queryTerminator.terminate(t);
+      }
     }
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java

@@ -22,5 +22,13 @@
 import java.util.concurrent.Future;
 
 public interface IQueryTerminator {
+  /** Cancels the unfinished FragmentInstances of the query after a grace period. */
   Future<Boolean> terminate(Throwable t);
+
+  /**
+   * Cancels the unfinished FragmentInstances of the query without a grace period. Used when the
+   * dispatch failed: the dispatched FragmentInstances can't finish on their own and hold their
+   * resources until they are cancelled.
+   */
+  Future<Boolean> terminateImmediately(Throwable t);
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java

@@ -86,15 +86,23 @@ private void calculateParameter(final List<FragmentInstance> instances) {
 
   @Override
   public Future<Boolean> terminate(final Throwable t) {
+    return terminate(t, TERMINATION_GRACE_PERIOD_IN_MS);
+  }
+
+  @Override
+  public Future<Boolean> terminateImmediately(final Throwable t) {
+    return terminate(t, 0L);
+  }
+
+  private Future<Boolean> terminate(final Throwable t, final long delayInMs) {
     // For the failure dispatch, the termination should not be triggered because of connection issue
     this.relatedHost =
         this.relatedHost.stream()
             .filter(endPoint -> !queryContext.getEndPointBlackList().contains(endPoint))
             .collect(Collectors.toList());
     this.hasThrowable = Objects.nonNull(t);
 
-    return scheduledExecutor.schedule(
-        this::syncTerminate, TERMINATION_GRACE_PERIOD_IN_MS, TimeUnit.MILLISECONDS);
+    return scheduledExecutor.schedule(this::syncTerminate, delayInMs, TimeUnit.MILLISECONDS);
   }
 
   public Boolean syncTerminate() {

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java

@@ -0,0 +1,72 @@
+/*
+ * 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.iotdb.db.queryengine.plan.scheduler;
+
+import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
+import org.apache.iotdb.db.queryengine.common.QueryId;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class SimpleQueryTerminatorTest {
+
+  private ScheduledExecutorService scheduledExecutor;
+  private SimpleQueryTerminator terminator;
+
+  @Before
+  public void setUp() {
+    scheduledExecutor = mock(ScheduledExecutorService.class);
+    final MPPQueryContext queryContext = mock(MPPQueryContext.class);
+    when(queryContext.getQueryId()).thenReturn(new QueryId("test_query"));
+    when(queryContext.getEndPointBlackList()).thenReturn(Collections.emptySet());
+    terminator =
+        new SimpleQueryTerminator(
+            scheduledExecutor,
+            queryContext,
+            Collections.emptyList(),
+            null,
+            mock(IFragInstanceStateTracker.class));
+  }
+
+  @Test
+  public void terminateWaitsForTheGracePeriod() {
+    terminator.terminate(new RuntimeException("query failed"));
+
+    verify(scheduledExecutor).schedule(any(Callable.class), eq(1000L), eq(TimeUnit.MILLISECONDS));
+  }
+
+  @Test
+  public void terminateImmediatelySkipsTheGracePeriod() {
+    terminator.terminateImmediately(new RuntimeException("dispatch failed"));
+
+    verify(scheduledExecutor).schedule(any(Callable.class), eq(0L), eq(TimeUnit.MILLISECONDS));
+  }
+}