[Query] Fix scheduler stop races and query attempt context leaks

iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java

@@ -20,6 +20,11 @@
 package org.apache.iotdb.db.i18n;
 
 public final class DataNodeQueryMessages {
+  public static final String EXCEPTION_QUERY_ATTEMPT_ARG_IS_CLOSED_12143E89 = "Query attempt %s is closed.";
+  public static final String EXCEPTION_QUERY_ATTEMPT_ARG_ADMISSION_EXPIRED_DEADLINE_ARG_NS_NOW_ARG_NS_16A9846C = "Query attempt %s admission expired: deadline=%d ns, now=%d ns.";
+  public static final String EXCEPTION_INVALID_QUERY_ADMISSION_ISSUER_FOR_ATTEMPT_ARG_9D1C4E9C = "Invalid query admission issuer for attempt %s.";
+
+
   public static final String MESSAGE_FAIL_TO_INSERT_MEASUREMENTS_ARG_CAUSED_BY_ARG_B9858E6B =
       "Fail to insert measurements %s caused by %s";
 

iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java

@@ -20,6 +20,11 @@
 package org.apache.iotdb.db.i18n;
 
 public final class DataNodeQueryMessages {
+  public static final String EXCEPTION_QUERY_ATTEMPT_ARG_IS_CLOSED_12143E89 = "查询执行轮次 %s 已关闭。";
+  public static final String EXCEPTION_QUERY_ATTEMPT_ARG_ADMISSION_EXPIRED_DEADLINE_ARG_NS_NOW_ARG_NS_16A9846C = "查询执行轮次 %s 的接收凭证已过期：截止时间=%d ns，当前时间=%d ns。";
+  public static final String EXCEPTION_INVALID_QUERY_ADMISSION_ISSUER_FOR_ATTEMPT_ARG_9D1C4E9C = "查询执行轮次 %s 的接收凭证签发节点进程不匹配。";
+
+
   public static final String MESSAGE_FAIL_TO_INSERT_MEASUREMENTS_ARG_CAUSED_BY_ARG_B9858E6B =
       "测点 %s 插入失败，原因：%s";
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java

@@ -133,6 +133,7 @@
 import org.apache.iotdb.db.protocol.session.SessionManager;
 import org.apache.iotdb.db.protocol.thrift.OperationType;
 import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
+import org.apache.iotdb.db.queryengine.common.QueryId;
 import org.apache.iotdb.db.queryengine.execution.executor.RegionExecutionResult;
 import org.apache.iotdb.db.queryengine.execution.executor.RegionReadExecutor;
 import org.apache.iotdb.db.queryengine.execution.executor.RegionWriteExecutor;
@@ -323,6 +324,7 @@
 import org.apache.iotdb.mpp.rpc.thrift.TPushTopicMetaResp;
 import org.apache.iotdb.mpp.rpc.thrift.TPushTopicMetaRespExceptionMessage;
 import org.apache.iotdb.mpp.rpc.thrift.TPushTopicOwnerLeaseReq;
+import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken;
 import org.apache.iotdb.mpp.rpc.thrift.TRegionLeaderChangeReq;
 import org.apache.iotdb.mpp.rpc.thrift.TRegionLeaderChangeResp;
 import org.apache.iotdb.mpp.rpc.thrift.TRegionMigrateResult;
@@ -510,6 +512,11 @@ private TSStatus waitForConsensusStarted() {
     }
   }
 
+  @Override
+  public TQueryAdmissionToken fetchQueryAdmissionToken(long validityMs) {
+    return FragmentInstanceManager.getInstance().fetchQueryAdmissionToken(validityMs);
+  }
+
   @Override
   public TSendFragmentInstanceResp sendFragmentInstance(final TSendFragmentInstanceReq req) {
     LOGGER.debug(DataNodeMiscMessages.RECEIVE_FRAGMENT_INSTANCE, req.getConsensusGroupId());
@@ -556,6 +563,9 @@ public TSendFragmentInstanceResp sendFragmentInstance(final TSendFragmentInstanc
                       FragmentInstance.deserializeFrom(
                           fragmentInstanceId, req.fragmentInstance.body));
       fragmentInstance.attachUDFQueryScope(queryScope);
+      if (req.isSetQueryAttemptId() && req.isSetAdmissionToken()) {
+        fragmentInstance.setQueryAttempt(req.getQueryAttemptId(), req.getAdmissionToken());
+      }
     } catch (final Exception e) {
       if (classLoaderManager != null && queryScope != null) {
         classLoaderManager.finalizeUDFQuery(queryScope);
@@ -642,6 +652,14 @@ public TFragmentInstanceInfoResp fetchFragmentInstanceInfo(TFetchFragmentInstanc
   @Override
   public TCancelResp cancelQuery(TCancelQueryReq req) {
     try (SetThreadName threadName = new SetThreadName(req.getQueryId())) {
+      if (req.isSetQueryAttemptId() && req.isSetAdmissionToken()) {
+        FragmentInstanceManager.getInstance()
+            .closeQueryAttempt(
+                new QueryId(req.getQueryId()),
+                req.getQueryAttemptId(),
+                req.getAdmissionToken(),
+                req.hasThrowable);
+      }
       List<FragmentInstanceId> taskIds =
           req.getFragmentInstanceIds().stream()
               .map(FragmentInstanceId::fromThrift)

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

@@ -109,6 +109,8 @@ public enum ExplainType {
   // time unit is ms
   private long startTime;
 
+  private int queryAttemptId;
+
   private TEndPoint localDataBlockEndpoint;
   private TEndPoint localInternalEndpoint;
   private ResultNodeContext resultNodeContext;
@@ -280,6 +282,7 @@ public void releaseMemoryForSchemaTree() {
   }
 
   public void prepareForRetry() {
+    queryAttemptId++;
     if (!isInnerTriggeredQuery()) {
       cleanUpCte();
     }
@@ -303,6 +306,10 @@ public QueryId getQueryId() {
     return queryId;
   }
 
+  public int getQueryAttemptId() {
+    return queryAttemptId;
+  }
+
   public ExternalTsFileQueryResource createExternalTsFileQueryResource(
       String tableName,
       List<String> tsFilePaths,

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

@@ -0,0 +1,59 @@
+/*
+ * 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.common;
+
+/** Identifies one execution of a query, independently of transport-level RPC retries. */
+public final class QueryAttemptId {
+  public static final int LEGACY_ATTEMPT = -1;
+
+  private final QueryId queryId;
+  private final int attemptId;
+
+  public QueryAttemptId(QueryId queryId, int attemptId) {
+    this.queryId = queryId;
+    this.attemptId = attemptId;
+  }
+
+  public boolean isLegacy() {
+    return attemptId == LEGACY_ATTEMPT;
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    if (this == other) {
+      return true;
+    }
+    if (!(other instanceof QueryAttemptId)) {
+      return false;
+    }
+    QueryAttemptId that = (QueryAttemptId) other;
+    return attemptId == that.attemptId && queryId.equals(that.queryId);
+  }
+
+  @Override
+  public int hashCode() {
+    return 31 * queryId.hashCode() + attemptId;
+  }
+
+  @Override
+  public String toString() {
+    return queryId + "/" + attemptId;
+  }
+}

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

@@ -178,7 +178,8 @@ public String getFailureMessage() {
   public Throwable getFailureException() {
     Throwable throwable = failureException.get();
     if (throwable == null) {
-      return new IoTDBException(getFailureStatus());
+      TSStatus status = getFailureStatus();
+      return status == null ? null : new IoTDBException(status);
     } else {
       return throwable;
     }

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

@@ -66,8 +66,6 @@ public class DataNodeQueryContext {
   private static final TableDeviceSchemaCache TABLE_DEVICE_SCHEMA_CACHE =
       TableDeviceSchemaCache.getInstance();
 
-  private final AtomicInteger dataNodeFINum;
-
   // Used for TableModel information table_disk_usage scan
   @GuardedBy("lock")
   private Map<String, List<TTableInfo>> databaseTableInfoMap;
@@ -96,9 +94,8 @@ public TopKRuntimeFilter getTopKRuntimeFilter(String filterId) {
     return runtimeFilters.get(filterId);
   }
 
-  public DataNodeQueryContext(int dataNodeFINum) {
+  public DataNodeQueryContext() {
     this.uncachedPathToSeriesScanInfo = new ConcurrentHashMap<>();
-    this.dataNodeFINum = new AtomicInteger(dataNodeFINum);
     this.deviceCountAndMeasurementValues = new HashMap<>();
   }
 
@@ -184,10 +181,6 @@ public Map<PartialPath, Pair<AtomicInteger, TimeValuePair>> getUncachedPathToSer
     return uncachedPathToSeriesScanInfo;
   }
 
-  public int decreaseDataNodeFINum() {
-    return dataNodeFINum.decrementAndGet();
-  }
-
   public void lock(boolean isDeviceInMultiRegion) {
     // When a device exists in only one region, there will be no intermediate state.
     if (isDeviceInMultiRegion) {

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

@@ -41,7 +41,6 @@
 import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
 import org.apache.iotdb.db.queryengine.common.DeviceContext;
 import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
-import org.apache.iotdb.db.queryengine.common.QueryId;
 import org.apache.iotdb.db.queryengine.metric.DriverSchedulerMetricSet;
 import org.apache.iotdb.db.queryengine.metric.QueryRelatedResourceMetricSet;
 import org.apache.iotdb.db.queryengine.metric.QueryResourceMetricSet;
@@ -183,7 +182,8 @@
   // Outer query deadline (startTime + timeout) for IoTDBLocal UDF
   private long outerQueryDeadlineMs = -1L;
 
-  private final Map<QueryId, DataNodeQueryContext> dataNodeQueryContextMap;
+  private QueryAttemptRegistry.Registration queryAttemptRegistration;
+  private boolean executionInitialized;
   private DataNodeQueryContext dataNodeQueryContext;
 
   // Used for EXPLAIN ANALYZE to cache statistics result when the FI is finished,
@@ -244,7 +244,7 @@
       SessionInfo sessionInfo,
       IDataRegionForQuery dataRegion,
       TimePredicate globalTimePredicate,
-      Map<QueryId, DataNodeQueryContext> dataNodeQueryContextMap,
+      DataNodeQueryContext dataNodeQueryContext,
       long outerQueryDeadlineMs,
       boolean debug,
       boolean isVerbose) {
@@ -255,7 +255,7 @@
             sessionInfo,
             dataRegion,
             globalTimePredicate,
-            dataNodeQueryContextMap,
+            dataNodeQueryContext,
             outerQueryDeadlineMs,
             debug,
             isVerbose);
@@ -359,7 +359,7 @@
       SessionInfo sessionInfo,
       IDataRegionForQuery dataRegion,
       TimePredicate globalTimePredicate,
-      Map<QueryId, DataNodeQueryContext> dataNodeQueryContextMap,
+      DataNodeQueryContext dataNodeQueryContext,
       long outerQueryDeadlineMs,
       boolean debug,
       boolean verbose) {
@@ -375,8 +375,7 @@
             ? null
             : globalTimePredicate.convertPredicateToTimeFilter(
                 sessionInfo.getZoneId(), TimestampPrecisionUtils.currPrecision);
-    this.dataNodeQueryContextMap = dataNodeQueryContextMap;
-    this.dataNodeQueryContext = dataNodeQueryContextMap.get(id.getQueryId());
+    this.dataNodeQueryContext = dataNodeQueryContext;
     this.memoryReservationManager =
         new ThreadSafeMemoryReservationManager(id.getQueryId(), this.getClass().getName());
   }
@@ -392,7 +391,6 @@
     this.stateMachine = stateMachine;
     this.executionEndTime.set(END_TIME_INITIAL_VALUE);
     this.sessionInfo = sessionInfo;
-    this.dataNodeQueryContextMap = null;
     this.dataNodeQueryContext = null;
     this.memoryReservationManager =
         new ThreadSafeMemoryReservationManager(id.getQueryId(), this.getClass().getName());
@@ -410,7 +408,6 @@
     this.stateMachine = stateMachine;
     this.executionEndTime.set(END_TIME_INITIAL_VALUE);
     this.sessionInfo = sessionInfo;
-    this.dataNodeQueryContextMap = null;
     this.dataNodeQueryContext = null;
     this.memoryReservationManager = memoryReservationManager;
   }
@@ -430,7 +427,6 @@
     this.sessionInfo = sessionInfo;
     this.dataRegion = dataRegion;
     this.globalTimeFilter = globalTimeFilter;
-    this.dataNodeQueryContextMap = null;
     this.memoryReservationManager =
         new ThreadSafeMemoryReservationManager(id.getQueryId(), this.getClass().getName());
   }
@@ -452,7 +448,6 @@
     this.queryId = queryId;
     this.id = null;
     this.stateMachine = null;
-    this.dataNodeQueryContextMap = null;
     this.dataNodeQueryContext = null;
     this.dataRegion = dataRegion;
     this.globalTimeFilter = timeFilter;
@@ -547,12 +542,6 @@
       executionEndTime.compareAndSet(END_TIME_INITIAL_VALUE, now);
       endNanos.compareAndSet(0, System.nanoTime());
 
-      // release some query resource in FragmentInstanceContext
-      // why not release them in releaseResourceWhenAllDriversAreClosed() together?
-      // because we may have no chane to run the releaseResourceWhenAllDriversAreClosed which is
-      // called in callback of FragmentInstanceExecution
-      // FragmentInstanceExecution won't be created if we meet some errors like MemoryNotEnough
-      releaseDataNodeQueryContext();
       sourcePaths = null;
     }
   }
@@ -1352,7 +1341,15 @@
    * be decreased.
    */
   public synchronized void releaseResource() {
-    finalizeUDFQuery();
+    try {
+      finalizeUDFQuery();
+      releaseResourceInternal();
+    } finally {
+      releaseDataNodeQueryContext();
+    }
+  }
+
+  private void releaseResourceInternal() {
     try {
       getQueryChunkReaders().close();
     } catch (IOException e) {
@@ -1486,25 +1483,33 @@
         .recordExecutionCost(QUERY_RESOURCE_INIT, getInitQueryDataSourceCost());
   }
 
-  private void releaseDataNodeQueryContext() {
-    if (dataNodeQueryContextMap == null) {
-      // this process is in fetch schema, nothing need to release
-      return;
-    }
+  void setQueryAttemptRegistration(QueryAttemptRegistry.Registration registration) {
+    this.queryAttemptRegistration = registration;
+  }
+
+  void setExecutionInitialized() {
+    executionInitialized = true;
+  }
+
+  boolean isExecutionInitialized() {
+    return executionInitialized;
+  }
 
-    if (dataNodeQueryContext.decreaseDataNodeFINum() == 0) {
-      dataNodeQueryContext = null;
-      dataNodeQueryContextMap.remove(id.getQueryId());
+  synchronized void releaseDataNodeQueryContext() {
+    if (queryAttemptRegistration != null) {
+      queryAttemptRegistration.close();
+      queryAttemptRegistration = null;
     }
+    dataNodeQueryContext = null;
   }
 
   synchronized void adoptUDFQueryScope(UDFClassLoaderManager.QueryScope transferredScope) {
     if (udfQueryScope == null) {
       udfQueryScope = transferredScope;
     }
   }
 
   public synchronized UDFClassLoaderManager.QueryScope getOrCreateUDFQueryScope() {
     UDFClassLoaderManager classLoaderManager = UDFClassLoaderManager.getInstance();
     if (udfQueryScope == null && classLoaderManager != null) {
       udfQueryScope = classLoaderManager.openQueryScope(id.getFullId());

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

@@ -98,6 +98,7 @@ public static FragmentInstanceExecution createFragmentInstanceExecution(
         new FragmentInstanceExecution(
             instanceId, context, drivers, sinkHandle, stateMachine, timeOut, exchangeManager);
     execution.initialize(scheduler, isExplainAnalyze);
+    context.setExecutionInitialized();
     scheduler.submitDrivers(instanceId.getQueryId(), drivers, timeOut, context.getSessionInfo());
     return execution;
   }
@@ -380,6 +381,10 @@ private void initialize(IDriverScheduler scheduler, boolean isExplainAnalyze) {
             try (SetThreadName threadName = new SetThreadName(instanceId.getFullId())) {
               LOGGER.error(DataNodeQueryMessages.ERRORS_FINISHING_FI_PROCESS, t);
             }
+          } finally {
+            if (newState.isDone()) {
+              context.releaseDataNodeQueryContext();
+            }
           }
         });
   }