Fix the issue that sometimes the FragmentInstance may stuck in FLUSHING states

server/src/main/java/org/apache/iotdb/db/mpp/execution/datatransfer/DataBlockManager.java

@@ -106,10 +106,10 @@ public void onAcknowledgeDataBlockEvent(TAcknowledgeDataBlockEvent e) throws TEx
           e.getEndSequenceId(),
           e.getSourceFragmentInstanceId());
       if (!sinkHandles.containsKey(e.getSourceFragmentInstanceId())) {
-        throw new TException(
-            "Source fragment instance not found. Fragment instance ID: "
-                + e.getSourceFragmentInstanceId()
-                + ".");
+        logger.warn(
+            "received ACK event but target FragmentInstance[{}] is not found.",
+            e.getSourceFragmentInstanceId());
+        return;
       }
       ((SinkHandle) sinkHandles.get(e.getSourceFragmentInstanceId()))
           .acknowledgeTsBlock(e.getStartSequenceId(), e.getEndSequenceId());
@@ -130,10 +130,13 @@ public void onNewDataBlockEvent(TNewDataBlockEvent e) throws TException {
               .get(e.getTargetFragmentInstanceId())
               .get(e.getTargetPlanNodeId())
               .isAborted()) {
-        throw new TException(
-            "Target fragment instance not found. Fragment instance ID: "
-                + e.getTargetFragmentInstanceId()
-                + ".");
+        // In some scenario, when the SourceHandle sends the data block ACK event, its upstream may
+        // have already been stopped. For example, in the query whit LimitOperator, the downstream
+        // FragmentInstance may be finished, although the upstream is still working.
+        logger.warn(
+            "received NewDataBlockEvent but the upstream FragmentInstance[{}] is not found",
+            e.getTargetFragmentInstanceId());
+        return;
       }
 
       SourceHandle sourceHandle =

server/src/main/java/org/apache/iotdb/db/mpp/execution/datatransfer/SourceHandle.java

@@ -320,6 +320,8 @@ public void run() {
             TsBlock tsBlock = serde.deserialize(byteBuffer);
             tsBlocks.add(tsBlock);
           }
+          executorService.submit(
+              new SendAcknowledgeDataBlockEventTask(startSequenceId, endSequenceId));
           synchronized (SourceHandle.this) {
             if (aborted) {
               return;
@@ -331,8 +333,6 @@ public void run() {
               blocked.set(null);
             }
           }
-          executorService.submit(
-              new SendAcknowledgeDataBlockEventTask(startSequenceId, endSequenceId));
           break;
         } catch (Throwable e) {
           logger.error(

server/src/main/java/org/apache/iotdb/db/mpp/plan/analyze/ClusterSchemaFetcher.java

@@ -107,6 +107,8 @@ private SchemaTree executeSchemaFetchQuery(SchemaFetchStatement schemaFetchState
     }
     SchemaTree result = new SchemaTree();
     while (coordinator.getQueryExecution(queryId).hasNextResult()) {
+      // The query will be transited to FINISHED when invoking getBatchResult() at the last time
+      // So we don't need to clean up it manually
       Optional<TsBlock> tsBlock = coordinator.getQueryExecution(queryId).getBatchResult();
       if (!tsBlock.isPresent()) {
         break;
@@ -121,8 +123,6 @@ private SchemaTree executeSchemaFetchQuery(SchemaFetchStatement schemaFetchState
         result.mergeSchemaTree(fetchedSchemaTree);
       }
     }
-    // TODO: (xingtanzjr) need to release this query's resource here. This is a temporary way
-    coordinator.getQueryExecution(queryId).stopAndCleanup();
     return result;
   }
 

server/src/main/java/org/apache/iotdb/db/mpp/plan/execution/QueryExecution.java

@@ -238,7 +238,9 @@ public Optional<TsBlock> getBatchResult() {
       ListenableFuture<Void> blocked = resultHandle.isBlocked();
       blocked.get();
       if (resultHandle.isFinished()) {
-        releaseResource();
+        // Once the resultHandle is finished, we should transit the state of this query to FINISHED.
+        // So that the corresponding cleanup work could be triggered.
+        stateMachine.transitionToFinished();
         return Optional.empty();
       }
       return Optional.of(resultHandle.receive());

server/src/main/java/org/apache/iotdb/db/mpp/plan/scheduler/SimpleQueryTerminator.java

@@ -25,6 +25,7 @@
 import org.apache.iotdb.db.mpp.common.QueryId;
 import org.apache.iotdb.db.mpp.plan.planner.plan.FragmentInstance;
 import org.apache.iotdb.mpp.rpc.thrift.TCancelQueryReq;
+import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId;
 
 import org.apache.thrift.TException;
 import org.slf4j.Logger;
@@ -37,7 +38,8 @@
 import java.util.stream.Collectors;
 
 public class SimpleQueryTerminator implements IQueryTerminator {
-  private static final Logger LOGGER = LoggerFactory.getLogger(SimpleQueryTerminator.class);
+  private static final Logger logger = LoggerFactory.getLogger(SimpleQueryTerminator.class);
+  private static final long TERMINATION_GRACE_PERIOD_IN_MS = 1000L;
   private final ExecutorService executor;
   private final QueryId queryId;
   private final List<FragmentInstance> fragmentInstances;
@@ -58,17 +60,22 @@ public SimpleQueryTerminator(
 
   @Override
   public Future<Boolean> terminate() {
-    List<TEndPoint> relatedHost = getRelatedHost(fragmentInstances);
-
+    List<TEndPoint> relatedHost = getRelatedHost();
+    try {
+      Thread.sleep(TERMINATION_GRACE_PERIOD_IN_MS);
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    }
     return executor.submit(
         () -> {
           for (TEndPoint endPoint : relatedHost) {
             // TODO (jackie tien) change the port
             try (SyncDataNodeInternalServiceClient client =
                 internalServiceClientManager.borrowClient(endPoint)) {
-              client.cancelQuery(new TCancelQueryReq(queryId.getId()));
+              client.cancelQuery(
+                  new TCancelQueryReq(queryId.getId(), getRelatedFragmentInstances(endPoint)));
             } catch (IOException e) {
-              LOGGER.error("can't connect to node {}", endPoint, e);
+              logger.error("can't connect to node {}", endPoint, e);
               return false;
             } catch (TException e) {
               return false;
@@ -78,10 +85,17 @@ public Future<Boolean> terminate() {
         });
   }
 
-  private List<TEndPoint> getRelatedHost(List<FragmentInstance> instances) {
-    return instances.stream()
+  private List<TEndPoint> getRelatedHost() {
+    return fragmentInstances.stream()
         .map(instance -> instance.getHostDataNode().internalEndPoint)
         .distinct()
         .collect(Collectors.toList());
   }
+
+  private List<TFragmentInstanceId> getRelatedFragmentInstances(TEndPoint endPoint) {
+    return fragmentInstances.stream()
+        .filter(instance -> instance.getHostDataNode().internalEndPoint.equals(endPoint))
+        .map(instance -> instance.getId().toThrift())
+        .collect(Collectors.toList());
+  }
 }

server/src/main/java/org/apache/iotdb/db/service/thrift/impl/InternalServiceImpl.java

@@ -72,6 +72,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.stream.Collectors;
 
 public class InternalServiceImpl implements InternalService.Iface {
 
@@ -132,11 +133,14 @@ public TFragmentInstanceStateResp fetchFragmentInstanceState(TFetchFragmentInsta
 
   @Override
   public TCancelResp cancelQuery(TCancelQueryReq req) throws TException {
-
-    // TODO need to be implemented and currently in order not to print NotImplementedException log,
-    // we simply return null
+    List<FragmentInstanceId> taskIds =
+        req.getFragmentInstanceIds().stream()
+            .map(FragmentInstanceId::fromThrift)
+            .collect(Collectors.toList());
+    for (FragmentInstanceId taskId : taskIds) {
+      FragmentInstanceManager.getInstance().cancelTask(taskId);
+    }
     return new TCancelResp(true);
-    //    throw new NotImplementedException();
   }
 
   @Override

thrift/src/main/thrift/mpp.thrift

@@ -105,6 +105,7 @@ struct TFragmentInstanceStateResp {
 
 struct TCancelQueryReq {
   1: required string queryId
+  2: required list<TFragmentInstanceId> fragmentInstanceIds
 }
 
 struct TCancelPlanFragmentReq {