diff --git a/server/src/main/java/org/apache/iotdb/db/mpp/execution/datatransfer/DataBlockManager.java b/server/src/main/java/org/apache/iotdb/db/mpp/execution/datatransfer/DataBlockManager.java index c889f18f6d8ab..8cfda3c40b14a 100644 --- a/server/src/main/java/org/apache/iotdb/db/mpp/execution/datatransfer/DataBlockManager.java +++ b/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 = diff --git a/server/src/main/java/org/apache/iotdb/db/mpp/execution/datatransfer/SourceHandle.java b/server/src/main/java/org/apache/iotdb/db/mpp/execution/datatransfer/SourceHandle.java index 1c8bf70880ec2..2102e98bfb515 100644 --- a/server/src/main/java/org/apache/iotdb/db/mpp/execution/datatransfer/SourceHandle.java +++ b/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( diff --git a/server/src/main/java/org/apache/iotdb/db/mpp/plan/analyze/ClusterSchemaFetcher.java b/server/src/main/java/org/apache/iotdb/db/mpp/plan/analyze/ClusterSchemaFetcher.java index 08d82575600d4..275501eaf75db 100644 --- a/server/src/main/java/org/apache/iotdb/db/mpp/plan/analyze/ClusterSchemaFetcher.java +++ b/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 = 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; } diff --git a/server/src/main/java/org/apache/iotdb/db/mpp/plan/execution/QueryExecution.java b/server/src/main/java/org/apache/iotdb/db/mpp/plan/execution/QueryExecution.java index 763ca26a82267..7f6078087fc18 100644 --- a/server/src/main/java/org/apache/iotdb/db/mpp/plan/execution/QueryExecution.java +++ b/server/src/main/java/org/apache/iotdb/db/mpp/plan/execution/QueryExecution.java @@ -238,7 +238,9 @@ public Optional getBatchResult() { ListenableFuture 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()); diff --git a/server/src/main/java/org/apache/iotdb/db/mpp/plan/scheduler/SimpleQueryTerminator.java b/server/src/main/java/org/apache/iotdb/db/mpp/plan/scheduler/SimpleQueryTerminator.java index df6f08f7c206b..357c26fcf5884 100644 --- a/server/src/main/java/org/apache/iotdb/db/mpp/plan/scheduler/SimpleQueryTerminator.java +++ b/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 fragmentInstances; @@ -58,17 +60,22 @@ public SimpleQueryTerminator( @Override public Future terminate() { - List relatedHost = getRelatedHost(fragmentInstances); - + List 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 terminate() { }); } - private List getRelatedHost(List instances) { - return instances.stream() + private List getRelatedHost() { + return fragmentInstances.stream() .map(instance -> instance.getHostDataNode().internalEndPoint) .distinct() .collect(Collectors.toList()); } + + private List getRelatedFragmentInstances(TEndPoint endPoint) { + return fragmentInstances.stream() + .filter(instance -> instance.getHostDataNode().internalEndPoint.equals(endPoint)) + .map(instance -> instance.getId().toThrift()) + .collect(Collectors.toList()); + } } diff --git a/server/src/main/java/org/apache/iotdb/db/service/thrift/impl/InternalServiceImpl.java b/server/src/main/java/org/apache/iotdb/db/service/thrift/impl/InternalServiceImpl.java index f8b5778611c50..e2a1fff0c3eeb 100644 --- a/server/src/main/java/org/apache/iotdb/db/service/thrift/impl/InternalServiceImpl.java +++ b/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 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 diff --git a/thrift/src/main/thrift/mpp.thrift b/thrift/src/main/thrift/mpp.thrift index 0c2acf6cb2785..1d33d58cc4a60 100644 --- a/thrift/src/main/thrift/mpp.thrift +++ b/thrift/src/main/thrift/mpp.thrift @@ -105,6 +105,7 @@ struct TFragmentInstanceStateResp { struct TCancelQueryReq { 1: required string queryId + 2: required list fragmentInstanceIds } struct TCancelPlanFragmentReq {