You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by ja...@apache.org on 2026/03/27 06:51:47 UTC
(iotdb) branch ty/kill-master created (now f5906d86686)
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a change to branch ty/kill-master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
at f5906d86686 [To dev/1.3] Fix kill query doesn't take effect bug (#17358)
This branch includes the following new commits:
new f5906d86686 [To dev/1.3] Fix kill query doesn't take effect bug (#17358)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(iotdb) 01/01: [To dev/1.3] Fix kill query doesn't take effect bug (#17358)
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch ty/kill-master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit f5906d86686a0e515c955d57af554064acf90fec
Author: Jackie Tien <ja...@gmail.com>
AuthorDate: Thu Mar 26 20:51:11 2026 +0800
[To dev/1.3] Fix kill query doesn't take effect bug (#17358)
---
.../protocol/thrift/impl/ClientRPCServiceImpl.java | 27 ++++--
.../db/queryengine/common/MPPQueryContext.java | 14 +--
.../queryengine/execution/QueryStateMachine.java | 4 +-
.../fragment/FragmentInstanceManager.java | 2 +
.../iotdb/db/queryengine/plan/Coordinator.java | 34 ++++++-
.../analyze/schema/ClusterSchemaFetchExecutor.java | 43 +++++----
.../plan/execution/IQueryExecution.java | 25 ++++-
.../queryengine/plan/execution/QueryExecution.java | 101 ++++++++++++++-------
.../plan/execution/config/ConfigExecution.java | 27 ++++--
.../operator/MergeTreeSortOperatorTest.java | 11 ++-
10 files changed, 206 insertions(+), 82 deletions(-)
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java
index d3f6b7d91dd..ed754c418c6 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java
@@ -245,6 +245,7 @@ import static org.apache.iotdb.db.utils.ErrorHandlingUtils.onIoTDBException;
import static org.apache.iotdb.db.utils.ErrorHandlingUtils.onQueryException;
import static org.apache.iotdb.db.utils.QueryDataSetUtils.convertTsBlockByFetchSize;
import static org.apache.iotdb.rpc.RpcUtils.TIME_PRECISION;
+import static org.apache.iotdb.rpc.TSStatusCode.QUERY_WAS_KILLED;
public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
@@ -286,6 +287,9 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
private static final int DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES =
TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes();
+ private static final String NO_QUERY_EXECUTION_ERR_MSG =
+ "Query is not found, it may be killed by others, timeout or some other runtime errors, you can see more details in server log.";
+
@FunctionalInterface
public interface SelectResult {
@@ -1526,15 +1530,18 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
finished = true;
return RpcUtils.getTSFetchResultsResp(getNotLoggedInStatus());
}
- TSFetchResultsResp resp = RpcUtils.getTSFetchResultsResp(TSStatusCode.SUCCESS_STATUS);
queryExecution = COORDINATOR.getQueryExecution(req.queryId);
if (queryExecution == null) {
- resp.setHasResultSet(false);
- resp.setMoreData(false);
- return resp;
+ TSStatus noQueryExecutionStatus = new TSStatus(QUERY_WAS_KILLED.getStatusCode());
+ noQueryExecutionStatus.setMessage(NO_QUERY_EXECUTION_ERR_MSG);
+ return RpcUtils.getTSFetchResultsResp(noQueryExecutionStatus);
}
+
+ TSFetchResultsResp resp = RpcUtils.getTSFetchResultsResp(TSStatusCode.SUCCESS_STATUS);
+
+ queryExecution.updateCurrentRpcStartTime(startTime);
statementType = queryExecution.getStatementType();
try (SetThreadName queryName = new SetThreadName(queryExecution.getQueryId())) {
@@ -2272,16 +2279,16 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
return RpcUtils.getTSFetchResultsResp(getNotLoggedInStatus());
}
- TSFetchResultsResp resp = RpcUtils.getTSFetchResultsResp(TSStatusCode.SUCCESS_STATUS);
-
queryExecution = COORDINATOR.getQueryExecution(req.queryId);
if (queryExecution == null) {
- resp.setHasResultSet(false);
- resp.setMoreData(true);
- return resp;
+ TSStatus noQueryExecutionStatus = new TSStatus(QUERY_WAS_KILLED.getStatusCode());
+ noQueryExecutionStatus.setMessage(NO_QUERY_EXECUTION_ERR_MSG);
+ return RpcUtils.getTSFetchResultsResp(noQueryExecutionStatus);
}
+ queryExecution.updateCurrentRpcStartTime(startTime);
statementType = queryExecution.getStatementType();
+ TSFetchResultsResp resp = RpcUtils.getTSFetchResultsResp(TSStatusCode.SUCCESS_STATUS);
try (SetThreadName queryName = new SetThreadName(queryExecution.getQueryId())) {
Pair<TSQueryDataSet, Boolean> pair =
convertTsBlockByFetchSize(queryExecution, req.fetchSize);
@@ -2291,7 +2298,7 @@ public class ClientRPCServiceImpl implements IClientRPCServiceWithHandler {
resp.setHasResultSet(hasResultSet);
resp.setQueryDataSet(result);
resp.setIsAlign(true);
- resp.setMoreData(finished);
+ resp.setMoreData(!finished);
return resp;
}
} catch (Exception e) {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java
index 0294a14af25..88bd1998f68 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java
@@ -76,7 +76,11 @@ public class MPPQueryContext implements IAuditEntity {
private long localQueryId;
private SessionInfo session;
private QueryType queryType = QueryType.READ;
+
+ /** the max executing time of query in ms. Unit: millisecond */
private long timeOut;
+
+ // time unit is ms
private long startTime;
private TEndPoint localDataBlockEndpoint;
@@ -147,6 +151,7 @@ public class MPPQueryContext implements IAuditEntity {
// Tables in the subquery
private final Map<NodeRef<Query>, List<Identifier>> subQueryTables = new HashMap<>();
+ @TestOnly
public MPPQueryContext(QueryId queryId) {
this.queryId = queryId;
this.endPointBlackList = ConcurrentHashMap.newKeySet();
@@ -161,12 +166,7 @@ public class MPPQueryContext implements IAuditEntity {
SessionInfo session,
TEndPoint localDataBlockEndpoint,
TEndPoint localInternalEndpoint) {
- this(queryId);
- this.sql = sql;
- this.session = session;
- this.localDataBlockEndpoint = localDataBlockEndpoint;
- this.localInternalEndpoint = localInternalEndpoint;
- this.initResultNodeContext();
+ this(sql, queryId, -1, session, localDataBlockEndpoint, localInternalEndpoint);
}
public MPPQueryContext(
@@ -244,10 +244,12 @@ public class MPPQueryContext implements IAuditEntity {
return queryType;
}
+ /** the max executing time of query in ms. Unit: millisecond */
public long getTimeOut() {
return timeOut;
}
+ /** the max executing time of query in ms. Unit: millisecond */
public void setTimeOut(long timeOut) {
this.timeOut = timeOut;
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachine.java
index 146359bc215..19d0fb38749 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachine.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachine.java
@@ -107,10 +107,10 @@ public class QueryStateMachine {
transitionToDoneState(CANCELED);
}
- public void transitionToCanceled(Throwable throwable, TSStatus failureStatus) {
+ public boolean transitionToCanceled(Throwable throwable, TSStatus failureStatus) {
this.failureStatus.compareAndSet(null, failureStatus);
this.failureException.compareAndSet(null, throwable);
- transitionToDoneState(CANCELED);
+ return transitionToDoneState(CANCELED);
}
public void transitionToAborted() {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
index 3f6812943d6..1898cbfe53c 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java
@@ -38,6 +38,7 @@ import org.apache.iotdb.db.queryengine.execution.schedule.DriverScheduler;
import org.apache.iotdb.db.queryengine.execution.schedule.IDriverScheduler;
import org.apache.iotdb.db.queryengine.metric.QueryExecutionMetricSet;
import org.apache.iotdb.db.queryengine.metric.QueryRelatedResourceMetricSet;
+import org.apache.iotdb.db.queryengine.plan.Coordinator;
import org.apache.iotdb.db.queryengine.plan.planner.LocalExecutionPlanner;
import org.apache.iotdb.db.queryengine.plan.planner.PipelineDriverFactory;
import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance;
@@ -435,6 +436,7 @@ public class FragmentInstanceManager {
+ "ms, and now is in flushing state"));
}
});
+ Coordinator.getInstance().cleanUpStaleQueries();
}
public ExecutorService getIntoOperationExecutor() {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
index b0d28f03c02..9342b3a4674 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
@@ -35,6 +35,7 @@ import org.apache.iotdb.commons.memory.MemoryBlockType;
import org.apache.iotdb.db.auth.AuthorityChecker;
import org.apache.iotdb.db.conf.IoTDBConfig;
import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException;
import org.apache.iotdb.db.exception.sql.SemanticException;
import org.apache.iotdb.db.protocol.session.IClientSession;
import org.apache.iotdb.db.protocol.session.PreparedStatementInfo;
@@ -805,7 +806,7 @@ public class Coordinator {
}
public void cleanupQueryExecution(Long queryId, Supplier<String> contentSupplier, Throwable t) {
- IQueryExecution queryExecution = getQueryExecution(queryId);
+ IQueryExecution queryExecution = queryExecutionMap.remove(queryId);
if (queryExecution != null) {
cleanupQueryExecutionInternal(queryId, queryExecution, contentSupplier, t);
}
@@ -813,7 +814,7 @@ public class Coordinator {
public void cleanupQueryExecution(
Long queryId, org.apache.thrift.TBase<?, ?> nativeApiRequest, Throwable t) {
- IQueryExecution queryExecution = getQueryExecution(queryId);
+ IQueryExecution queryExecution = queryExecutionMap.remove(queryId);
if (queryExecution != null) {
Supplier<String> contentSupplier =
new ContentOfQuerySupplier(nativeApiRequest, queryExecution);
@@ -898,6 +899,35 @@ public class Coordinator {
}
}
+ /**
+ * We need to reclaim resources from queries that have exceeded their timeout by more than one
+ * minute. This indicates that the associated clients have failed to perform proper resource
+ * cleanup.
+ */
+ public void cleanUpStaleQueries() {
+ long currentTime = System.currentTimeMillis();
+ queryExecutionMap.forEach(
+ (queryId, queryExecution) -> {
+ if (queryExecution.isActive()) {
+ return;
+ }
+ long timeout = queryExecution.getTimeout();
+ long queryStartTime = queryExecution.getStartExecutionTime();
+ long executeTime = currentTime - queryStartTime;
+ if (timeout > 0 && executeTime - 60_000L > timeout) {
+ LOGGER.warn(
+ "Cleaning up stale query with id {}, which has been running for {} ms, timeout duration is: {}ms",
+ queryId,
+ executeTime,
+ timeout);
+ cleanupQueryExecution(
+ queryId,
+ (org.apache.thrift.TBase<?, ?>) null,
+ new QueryTimeoutRuntimeException(queryStartTime, currentTime, timeout));
+ }
+ });
+ }
+
public void cleanupQueryExecution(Long queryId) {
cleanupQueryExecution(queryId, (org.apache.thrift.TBase<?, ?>) null, null);
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/ClusterSchemaFetchExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/ClusterSchemaFetchExecutor.java
index aad9f50aca0..d575f6420eb 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/ClusterSchemaFetchExecutor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/schema/ClusterSchemaFetchExecutor.java
@@ -36,6 +36,7 @@ import org.apache.iotdb.db.queryengine.exception.MemoryNotEnoughException;
import org.apache.iotdb.db.queryengine.plan.Coordinator;
import org.apache.iotdb.db.queryengine.plan.analyze.ClusterPartitionFetcher;
import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
+import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution;
import org.apache.iotdb.db.queryengine.plan.statement.Statement;
import org.apache.iotdb.db.queryengine.plan.statement.internal.DeviceSchemaFetchStatement;
import org.apache.iotdb.db.queryengine.plan.statement.internal.SeriesSchemaFetchStatement;
@@ -265,30 +266,38 @@ class ClusterSchemaFetchExecutor {
String.format("Fetch Schema failed, because %s", executionResult.status.getMessage()),
executionResult.status.getCode());
}
+ IQueryExecution queryExecution = coordinator.getQueryExecution(queryId);
try (SetThreadName ignored = new SetThreadName(executionResult.queryId.getId())) {
ClusterSchemaTree result = new ClusterSchemaTree();
ClusterSchemaTree.SchemaNodeBatchDeserializer deserializer =
new ClusterSchemaTree.SchemaNodeBatchDeserializer();
Set<String> databaseSet = new HashSet<>();
- 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;
- try {
- tsBlock = coordinator.getQueryExecution(queryId).getBatchResult();
- } catch (IoTDBException e) {
- t = e;
- throw new QuerySchemaFetchFailedException(
- String.format("Fetch Schema failed: %s", e.getMessage()), e.getErrorCode());
- }
- if (!tsBlock.isPresent() || tsBlock.get().isEmpty()) {
- break;
- }
- Column column = tsBlock.get().getColumn(0);
- for (int i = 0; i < column.getPositionCount(); i++) {
- parseFetchedData(column.getBinary(i), result, deserializer, databaseSet, context);
+ if (queryExecution != null) {
+ while (queryExecution.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;
+ try {
+ tsBlock = queryExecution.getBatchResult();
+ } catch (IoTDBException e) {
+ t = e;
+ throw new QuerySchemaFetchFailedException(
+ String.format("Fetch Schema failed: %s", e.getMessage()), e.getErrorCode());
+ }
+ if (!tsBlock.isPresent() || tsBlock.get().isEmpty()) {
+ break;
+ }
+ Column column = tsBlock.get().getColumn(0);
+ for (int i = 0; i < column.getPositionCount(); i++) {
+ parseFetchedData(column.getBinary(i), result, deserializer, databaseSet, context);
+ }
}
+ } else {
+ throw new RuntimeException(
+ String.format("Fetch Schema failed, because queryExecution is null for %s", queryId));
}
+
result.setDatabases(databaseSet);
return result;
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/IQueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/IQueryExecution.java
index f754bc4b10e..9b5a183d98a 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/IQueryExecution.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/IQueryExecution.java
@@ -35,8 +35,6 @@ public interface IQueryExecution {
void stop(Throwable t);
- void stopAndCleanup();
-
void stopAndCleanup(Throwable t);
void cancel();
@@ -61,15 +59,38 @@ public interface IQueryExecution {
String getQueryId();
+ // time unit is ms
long getStartExecutionTime();
+ /**
+ * @param executionTime time unit should be ns
+ */
void recordExecutionTime(long executionTime);
+ /**
+ * update current rpc start time, which is used to calculate rpc execution time and update total
+ * execution time
+ *
+ * @param startTime start time of current rpc, time unit is ns
+ */
+ void updateCurrentRpcStartTime(long startTime);
+
+ /**
+ * Check if there is an active RPC for this query. If {@code startTimeOfCurrentRpc == -1}, it
+ * means there is no active RPC, otherwise there is an active RPC. An active RPC means that the
+ * client is still fetching results and the QueryExecution should not be cleaned up until the RPC
+ * finishes. On the other hand, if there is no active RPC, it means that the client has finished
+ * fetching results or has not started fetching results yet, and the QueryExecution can be safely
+ * cleaned up.
+ */
+ boolean isActive();
+
/**
* @return cost time in ns
*/
long getTotalExecutionTime();
+ /** the max executing time of query in ms. Unit: millisecond */
long getTimeout();
Optional<String> getExecuteSQL();
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java
index 6e940e9816e..7500fbb9c34 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java
@@ -39,6 +39,7 @@ import org.apache.iotdb.db.queryengine.execution.exchange.source.ISourceHandle;
import org.apache.iotdb.db.queryengine.execution.exchange.source.SourceHandle;
import org.apache.iotdb.db.queryengine.metric.QueryExecutionMetricSet;
import org.apache.iotdb.db.queryengine.metric.QueryPlanCostMetricSet;
+import org.apache.iotdb.db.queryengine.plan.Coordinator;
import org.apache.iotdb.db.queryengine.plan.analyze.IAnalysis;
import org.apache.iotdb.db.queryengine.plan.analyze.QueryType;
import org.apache.iotdb.db.queryengine.plan.execution.memory.MemorySourceHandle;
@@ -108,9 +109,14 @@ public class QueryExecution implements IQueryExecution {
private final AtomicBoolean stopped;
- // cost time in ns
+ // cost time in ns of finished rpc
private long totalExecutionTime = 0;
+ // -1 if previous rpc is finished and next client req hasn't come yet, unit is ns
+ // it will be updated in fetchResult rpc
+ // protected by synchronized(this)
+ private volatile long startTimeOfCurrentRpc = System.nanoTime();
+
private static final QueryExecutionMetricSet QUERY_EXECUTION_METRIC_SET =
QueryExecutionMetricSet.getInstance();
private static final QueryPlanCostMetricSet QUERY_PLAN_COST_METRIC_SET =
@@ -133,14 +139,19 @@ public class QueryExecution implements IQueryExecution {
if (!state.isDone()) {
return;
}
+ Throwable cause = null;
if (state == QueryState.FAILED
|| state == QueryState.ABORTED
|| state == QueryState.CANCELED) {
- LOGGER.debug("[ReleaseQueryResource] state is: {}", state);
- Throwable cause = stateMachine.getFailureException();
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("[ReleaseQueryResource] state is: {}", state);
+ }
+ cause = stateMachine.getFailureException();
releaseResource(cause);
+
+ this.cleanUpCoordinatorContextMapIfNeeded(cause);
}
- this.stop(null);
+ this.stop(cause);
}
});
this.stopped = new AtomicBoolean(false);
@@ -321,33 +332,20 @@ public class QueryExecution implements IQueryExecution {
}
}
- // Stop the query and clean up all the resources this query occupied
- @Override
- public void stopAndCleanup() {
- stop(null);
- releaseResource();
- }
-
@Override
public void cancel() {
- stateMachine.transitionToCanceled(
- new KilledByOthersException(),
- new TSStatus(TSStatusCode.QUERY_WAS_KILLED.getStatusCode())
- .setMessage(KilledByOthersException.MESSAGE));
- }
-
- /** Release the resources that current QueryExecution hold. */
- private void releaseResource() {
- // close ResultHandle to unblock client's getResult request
- // Actually, we should not close the ResultHandle when the QueryExecution is Finished.
- // There are only two scenarios where the ResultHandle should be closed:
- // 1. The client fetch all the result and the ResultHandle is finished.
- // 2. The client's connection is closed that all owned QueryExecution should be cleaned up
- // If the QueryExecution's state is abnormal, we should also abort the resultHandle without
- // waiting it to be finished.
- if (resultHandle != null) {
- resultHandle.close();
- cleanUpResultHandle();
+ Throwable cause = new KilledByOthersException();
+ boolean cancelled =
+ stateMachine.transitionToCanceled(
+ cause,
+ new TSStatus(TSStatusCode.QUERY_WAS_KILLED.getStatusCode())
+ .setMessage(KilledByOthersException.MESSAGE));
+ if (!cancelled) {
+ // cancel failed, means this query has already in a done state, we can do nothing to change
+ // the state but clean up the resource if needed
+ // we don't need to do cleanUpCoordinatorContextMapIfNeeded if cancel succeed, because it will
+ // be called in callback logic in QueryStateMachine of this QueryExecution
+ this.cleanUpCoordinatorContextMapIfNeeded(cause);
}
}
@@ -356,7 +354,7 @@ public class QueryExecution implements IQueryExecution {
// We don't need to deal with MemorySourceHandle because it doesn't register to memory pool
// We don't need to deal with LocalSourceHandle because the SharedTsBlockQueue uses the upstream
// FragmentInstanceId to register
- if (resultHandleCleanUp.compareAndSet(false, true) && resultHandle instanceof SourceHandle) {
+ if (resultHandle instanceof SourceHandle) {
TFragmentInstanceId fragmentInstanceId = resultHandle.getLocalFragmentInstanceId();
MPPDataExchangeService.getInstance()
.getMPPDataExchangeManager()
@@ -384,7 +382,7 @@ public class QueryExecution implements IQueryExecution {
// 2. The client's connection is closed that all owned QueryExecution should be cleaned up
// If the QueryExecution's state is abnormal, we should also abort the resultHandle without
// waiting it to be finished.
- if (resultHandle != null) {
+ if (resultHandle != null && resultHandleCleanUp.compareAndSet(false, true)) {
if (t != null) {
resultHandle.abort(t);
} else {
@@ -394,6 +392,32 @@ public class QueryExecution implements IQueryExecution {
}
}
+ /**
+ * Clear up Coordinator.queryExecutionMap by calling Coordinator.cleanupQueryExecution if there is
+ * no RPC in progress for this query (that is, the current RPC has finished and {@code
+ * startTimeOfCurrentRpc == -1}). In cases where an RPC is still active, the finally block in
+ * ClientRPCServiceImpl is responsible for performing the cleanup.
+ */
+ private synchronized void cleanUpCoordinatorContextMapIfNeeded(Throwable t) {
+ if (isActive()) {
+ Coordinator.getInstance()
+ .cleanupQueryExecution(
+ context.getLocalQueryId(), (org.apache.thrift.TBase<?, ?>) null, t);
+ }
+ }
+
+ /**
+ * Check if there is an active RPC for this query. If {@code startTimeOfCurrentRpc == -1}, it
+ * means there is no active RPC, otherwise there is an active RPC. An active RPC means that the
+ * client is still fetching results and the QueryExecution should not be cleaned up until the RPC
+ * finishes. On the other hand, if there is no active RPC, it means that the client has finished
+ * fetching results or has not started fetching results yet, and the QueryExecution can be safely
+ * cleaned up.
+ */
+ public synchronized boolean isActive() {
+ return startTimeOfCurrentRpc == -1;
+ }
+
/**
* This method will be called by the request thread from client connection. This method will block
* until one of these conditions occurs: 1. There is a batch of result 2. There is no more result
@@ -671,13 +695,22 @@ public class QueryExecution implements IQueryExecution {
}
@Override
- public void recordExecutionTime(long executionTime) {
+ public synchronized void recordExecutionTime(long executionTime) {
totalExecutionTime += executionTime;
+ // recordExecutionTime is called after current rpc finished, so we need to set
+ // startTimeOfCurrentRpc to -1
+ this.startTimeOfCurrentRpc = -1;
+ }
+
+ @Override
+ public synchronized void updateCurrentRpcStartTime(long startTime) {
+ this.startTimeOfCurrentRpc = startTime;
}
@Override
- public long getTotalExecutionTime() {
- return totalExecutionTime;
+ public synchronized long getTotalExecutionTime() {
+ return totalExecutionTime
+ + (startTimeOfCurrentRpc == -1 ? 0 : System.nanoTime() - startTimeOfCurrentRpc);
}
@Override
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java
index f431bb2e51f..2b70a667be9 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java
@@ -113,6 +113,12 @@ public class ConfigExecution implements IQueryExecution {
private final StatementType statementType;
private long totalExecutionTime;
+ // -1 if previous rpc is finished and next client req hasn't come yet, unit is ns
+ // it will be updated in fetchResult rpc
+ // currently, ConfigExecution will return result is just one call, so this field is not used. But
+ // we will keep it for future use when ConfigExecution may return result in multiple calls
+ private volatile long startTimeOfCurrentRpc = System.nanoTime();
+
public ConfigExecution(
MPPQueryContext context,
StatementType statementType,
@@ -222,11 +228,6 @@ public class ConfigExecution implements IQueryExecution {
// do nothing
}
- @Override
- public void stopAndCleanup() {
- // do nothing
- }
-
@Override
public void stopAndCleanup(Throwable t) {
// do nothing
@@ -327,11 +328,25 @@ public class ConfigExecution implements IQueryExecution {
@Override
public void recordExecutionTime(long executionTime) {
totalExecutionTime += executionTime;
+ // recordExecutionTime is called after current rpc finished, so we need to set
+ // startTimeOfCurrentRpc to -1
+ this.startTimeOfCurrentRpc = -1;
+ }
+
+ @Override
+ public void updateCurrentRpcStartTime(long startTime) {
+ this.startTimeOfCurrentRpc = startTime;
+ }
+
+ @Override
+ public boolean isActive() {
+ return startTimeOfCurrentRpc == -1;
}
@Override
public long getTotalExecutionTime() {
- return totalExecutionTime;
+ return totalExecutionTime
+ + (startTimeOfCurrentRpc == -1 ? 0 : System.nanoTime() - startTimeOfCurrentRpc);
}
@Override
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/MergeTreeSortOperatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/MergeTreeSortOperatorTest.java
index 1739944be7c..2d0c1bf543c 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/MergeTreeSortOperatorTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/MergeTreeSortOperatorTest.java
@@ -1826,6 +1826,14 @@ public class MergeTreeSortOperatorTest {
@Override
public void recordExecutionTime(long executionTime) {}
+ @Override
+ public void updateCurrentRpcStartTime(long startTime) {}
+
+ @Override
+ public boolean isActive() {
+ return true;
+ }
+
@Override
public long getTotalExecutionTime() {
return 0;
@@ -1857,9 +1865,6 @@ public class MergeTreeSortOperatorTest {
@Override
public void stop(Throwable t) {}
- @Override
- public void stopAndCleanup() {}
-
@Override
public void stopAndCleanup(Throwable t) {}