Fix kill query doesn't take effect bug (#17358)

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.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 @@
   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 @@
         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 @@
         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 @@
         resp.setHasResultSet(hasResultSet);
         resp.setQueryDataSet(result);
         resp.setIsAlign(true);
-        resp.setMoreData(finished);
+        resp.setMoreData(!finished);
         return resp;
       }
     } catch (Exception e) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java

@@ -76,7 +76,11 @@ public enum ExplainType {
   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 enum ExplainType {
   // 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 MPPQueryContext(
       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 QueryType getQueryType() {
     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;
   }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachine.java

@@ -107,10 +107,10 @@ public void transitionToCanceled() {
     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() {

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.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 @@ private void cancelTimeoutFlushingInstances() {
                             + "ms, and now is in flushing state"));
           }
         });
+    Coordinator.getInstance().cleanUpStaleQueries();
   }
 
   public ExecutorService getIntoOperationExecutor() {

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

@@ -35,6 +35,7 @@
 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,15 +806,15 @@
   }
 
   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);
     }
   }
 
   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 @@
     }
   }
 
+  /**
+   * 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);
   }

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.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;
@@ -254,7 +255,7 @@
     return templateMap;
   }
 
   private ClusterSchemaTree executeSchemaFetchQuery(
       Statement fetchStatement, MPPQueryContext context) {
     long queryId = SessionManager.getInstance().requestQueryId();
     Throwable t = null;
@@ -265,30 +266,38 @@
             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;
       }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/IQueryExecution.java

@@ -35,8 +35,6 @@
 
   void stop(Throwable t);
 
-  void stopAndCleanup();
-
   void stopAndCleanup(Throwable t);
 
   void cancel();
@@ -61,15 +59,38 @@
 
   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();