[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,8 @@
 package org.apache.iotdb.db.i18n;
 
 public final class DataNodeQueryMessages {
+  public static final String EXCEPTION_QUERY_ATTEMPT_ARG_IS_CLOSED_OR_EXPIRED_F45E3DCF = "Query attempt %s is closed or expired.";
+
   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,8 @@
 package org.apache.iotdb.db.i18n;
 
 public final class DataNodeQueryMessages {
+  public static final String EXCEPTION_QUERY_ATTEMPT_ARG_IS_CLOSED_OR_EXPIRED_F45E3DCF = "查询执行轮次 %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

@@ -132,6 +132,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;
@@ -532,6 +533,10 @@ public TSendFragmentInstanceResp sendFragmentInstance(final TSendFragmentInstanc
     final FragmentInstance fragmentInstance;
     try {
       fragmentInstance = FragmentInstance.deserializeFrom(req.fragmentInstance.body);
+      if (req.isSetQueryAttemptId() && req.isSetQueryDeadline()) {
+        fragmentInstance.setQueryAttempt(req.getQueryAttemptId(), req.getQueryDeadline());
+        fragmentInstance.setAdmissionDeadline(req.getAdmissionDeadline());
+      }
     } catch (final Exception e) {
       LOGGER.warn(DataNodeMiscMessages.DESERIALIZE_FRAGMENT_INSTANCE_FAILED, e);
       TSendFragmentInstanceResp resp = new TSendFragmentInstanceResp(false);
@@ -611,6 +616,14 @@ public TFragmentInstanceInfoResp fetchFragmentInstanceInfo(TFetchFragmentInstanc
   @Override
   public TCancelResp cancelQuery(TCancelQueryReq req) {
     try (SetThreadName threadName = new SetThreadName(req.getQueryId())) {
+      if (req.isSetQueryAttemptId() && req.isSetAdmissionDeadline()) {
+        FragmentInstanceManager.getInstance()
+            .closeQueryAttempt(
+                new QueryId(req.getQueryId()),
+                req.getQueryAttemptId(),
+                req.getAdmissionDeadline(),
+                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

@@ -55,6 +55,7 @@
 import org.apache.iotdb.db.queryengine.statistics.QueryPlanStatistics;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.math.LongMath;
 import org.apache.tsfile.read.filter.basic.Filter;
 import org.apache.tsfile.utils.Pair;
 import org.slf4j.Logger;
@@ -109,6 +110,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 +283,7 @@ public void releaseMemoryForSchemaTree() {
   }
 
   public void prepareForRetry() {
+    queryAttemptId++;
     if (!isInnerTriggeredQuery()) {
       cleanUpCte();
     }
@@ -303,6 +307,14 @@ public QueryId getQueryId() {
     return queryId;
   }
 
+  public int getQueryAttemptId() {
+    return queryAttemptId;
+  }
+
+  public long getQueryDeadline() {
+    return LongMath.saturatedAdd(startTime, timeOut);
+  }
+
   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,58 @@
+/*
+ * 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/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

@@ -40,7 +40,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;
@@ -181,7 +180,8 @@ public class FragmentInstanceContext extends QueryContext {
   // 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,
@@ -242,7 +242,7 @@ public static FragmentInstanceContext createFragmentInstanceContext(
       SessionInfo sessionInfo,
       IDataRegionForQuery dataRegion,
       TimePredicate globalTimePredicate,
-      Map<QueryId, DataNodeQueryContext> dataNodeQueryContextMap,
+      DataNodeQueryContext dataNodeQueryContext,
       long outerQueryDeadlineMs,
       boolean debug,
       boolean isVerbose) {
@@ -253,7 +253,7 @@ public static FragmentInstanceContext createFragmentInstanceContext(
             sessionInfo,
             dataRegion,
             globalTimePredicate,
-            dataNodeQueryContextMap,
+            dataNodeQueryContext,
             outerQueryDeadlineMs,
             debug,
             isVerbose);
@@ -357,7 +357,7 @@ private FragmentInstanceContext(
       SessionInfo sessionInfo,
       IDataRegionForQuery dataRegion,
       TimePredicate globalTimePredicate,
-      Map<QueryId, DataNodeQueryContext> dataNodeQueryContextMap,
+      DataNodeQueryContext dataNodeQueryContext,
       long outerQueryDeadlineMs,
       boolean debug,
       boolean verbose) {
@@ -373,8 +373,7 @@ private FragmentInstanceContext(
             ? 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());
   }
@@ -390,7 +389,6 @@ private FragmentInstanceContext(
     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());
@@ -408,7 +406,6 @@ private FragmentInstanceContext(
     this.stateMachine = stateMachine;
     this.executionEndTime.set(END_TIME_INITIAL_VALUE);
     this.sessionInfo = sessionInfo;
-    this.dataNodeQueryContextMap = null;
     this.dataNodeQueryContext = null;
     this.memoryReservationManager = memoryReservationManager;
   }
@@ -428,7 +425,6 @@ private FragmentInstanceContext(
     this.sessionInfo = sessionInfo;
     this.dataRegion = dataRegion;
     this.globalTimeFilter = globalTimeFilter;
-    this.dataNodeQueryContextMap = null;
     this.memoryReservationManager =
         new ThreadSafeMemoryReservationManager(id.getQueryId(), this.getClass().getName());
   }
@@ -450,7 +446,6 @@ protected FragmentInstanceContext(
     this.queryId = queryId;
     this.id = null;
     this.stateMachine = null;
-    this.dataNodeQueryContextMap = null;
     this.dataNodeQueryContext = null;
     this.dataRegion = dataRegion;
     this.globalTimeFilter = timeFilter;
@@ -545,12 +540,6 @@ private void updateStatsIfDone(FragmentInstanceState newState) {
       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;
     }
   }
@@ -1350,6 +1339,14 @@ private void releaseTVListOwnedByQuery() {
    * be decreased.
    */
   public synchronized void releaseResource() {
+    try {
+      releaseResourceInternal();
+    } finally {
+      releaseDataNodeQueryContext();
+    }
+  }
+
+  private void releaseResourceInternal() {
     try {
       getQueryChunkReaders().close();
     } catch (IOException e) {
@@ -1483,16 +1480,24 @@ public synchronized void releaseResource() {
         .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;
   }
 
   public void setMayHaveTmpFile(boolean mayHaveTmpFile) {

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();
+            }
           }
         });
   }