From 0030ac2c76a3b0d33f0306dea26534f28ac6feba Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Wed, 11 Nov 2026 17:02:31 +0800 Subject: [PATCH 1/4] [Query] Fix scheduler stop races and query attempt context leaks Publish schedulers before starting them and preserve stop requests across creation. Isolate DataNode contexts by query attempt, close partially dispatched attempts, and bound late-dispatch fences with admission deadlines. Add regression coverage for concurrent lifecycle transitions, retry isolation, partial dispatch cleanup, RPC compatibility, and admission expiry. --- .../iotdb/db/i18n/DataNodeQueryMessages.java | 2 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 2 + .../impl/DataNodeInternalRPCServiceImpl.java | 13 + .../queryengine/common/MPPQueryContext.java | 12 + .../db/queryengine/common/QueryAttemptId.java | 58 ++++ .../fragment/DataNodeQueryContext.java | 9 +- .../fragment/FragmentInstanceContext.java | 55 ++-- .../fragment/FragmentInstanceExecution.java | 5 + .../fragment/FragmentInstanceManager.java | 108 ++++--- .../fragment/QueryAttemptRegistry.java | 270 ++++++++++++++++++ .../plan/execution/QueryExecution.java | 56 +++- .../db/queryengine/plan/planner/IPlanner.java | 3 +- .../plan/planner/TreeModelPlanner.java | 3 +- .../plan/planner/plan/FragmentInstance.java | 31 ++ .../relational/planner/TableModelPlanner.java | 3 +- .../plan/scheduler/ClusterScheduler.java | 53 +++- .../FragmentInstanceDispatcherImpl.java | 68 ++++- .../plan/scheduler/IScheduler.java | 1 + .../plan/scheduler/SimpleQueryTerminator.java | 29 +- .../scheduler/load/LoadTsFileScheduler.java | 18 ++ .../fragment/FragmentInstanceManagerTest.java | 46 +++ .../fragment/QueryAttemptRegistryTest.java | 242 ++++++++++++++++ ...nedTreeDeviceViewScanOperatorTreeTest.java | 2 +- .../SinkOperatorDownStreamNodeIdTest.java | 2 +- .../QueryExecutionLifecycleTest.java | 135 +++++++++ .../execution/QueryExecutionRetryTest.java | 25 +- .../plan/planner/LastQueryCacheTest.java | 7 +- .../LocalExecutionPlannerBudgetTest.java | 4 +- .../plan/planner/PipelineBuilderTest.java | 4 +- .../ClusterSchedulerLifecycleTest.java | 116 ++++++++ ...agmentInstanceDispatcherLifecycleTest.java | 133 +++++++++ .../scheduler/QueryAttemptProtocolTest.java | 94 ++++++ .../scheduler/SimpleQueryTerminatorTest.java | 52 +++- .../load/LoadTsFileSchedulerTest.java | 21 +- .../builder/EvaluationDAGBuilderTest.java | 2 +- .../src/main/thrift/datanode.thrift | 8 + 36 files changed, 1572 insertions(+), 120 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/QueryAttemptId.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionLifecycleTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterSchedulerLifecycleTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryAttemptProtocolTest.java diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 355d5a00176b9..7e38eaa85fda6 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/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"; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index bc8b004bbecfb..c6aba75034517 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/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"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java index f03366065cd4c..d4b2d24647389 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java +++ b/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 taskIds = req.getFragmentInstanceIds().stream() .map(FragmentInstanceId::fromThrift) 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 728dd9e7df0b2..9c1f34bc16b8c 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 @@ -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 tsFilePaths, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/QueryAttemptId.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/QueryAttemptId.java new file mode 100644 index 0000000000000..a76cdd7a28153 --- /dev/null +++ b/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; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/DataNodeQueryContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/DataNodeQueryContext.java index a1cd2af1de277..8d353cfa31dfd 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/DataNodeQueryContext.java +++ b/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> 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> 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) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java index 74a1cf529d05c..a9e0596010a4b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java +++ b/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 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 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 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) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecution.java index b2fdeaac4f1f3..cb0ebb5323fc7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecution.java +++ b/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(); + } } }); } 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 f5fb70c94da39..e794e3b4a0dde 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 @@ -31,6 +31,7 @@ import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.QueryAttemptId; import org.apache.iotdb.db.queryengine.common.QueryId; import org.apache.iotdb.db.queryengine.execution.driver.IDriver; import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager; @@ -80,7 +81,7 @@ public class FragmentInstanceManager { private final Map instanceContext; private final Map instanceExecution; - private final Map dataNodeQueryContextMap; + private final QueryAttemptRegistry queryAttempts = new QueryAttemptRegistry(); private final LocalExecutionPlanner planner = LocalExecutionPlanner.getInstance(); private final IDriverScheduler scheduler = DriverScheduler.getInstance(); @@ -105,7 +106,6 @@ public static FragmentInstanceManager getInstance() { private FragmentInstanceManager() { this.instanceContext = new ConcurrentHashMap<>(); this.instanceExecution = new ConcurrentHashMap<>(); - this.dataNodeQueryContextMap = new ConcurrentHashMap<>(); this.instanceManagementExecutor = IoTDBThreadPoolFactory.newScheduledThreadPool( 1, ThreadName.FRAGMENT_INSTANCE_MANAGEMENT.getName()); @@ -152,39 +152,46 @@ public FragmentInstanceInfo execDataQueryFragmentInstance( new FragmentInstanceStateMachine(instanceId, instanceNotificationExecutor); boolean[] contextCreated = new boolean[] {false}; - DataNodeQueryContext[] dataNodeQueryContexts = new DataNodeQueryContext[1]; + QueryAttemptRegistry.Registration[] registrations = + new QueryAttemptRegistry.Registration[1]; FragmentInstanceContext context = instanceContext.computeIfAbsent( instanceId, fragmentInstanceId -> { contextCreated[0] = true; - // Only ensure the DataNodeQueryContext when we actually create the - // FragmentInstanceContext, so the repeated-dispatch path (which rejects - // without creating a context) does not leak a context entry. - dataNodeQueryContexts[0] = - getOrCreateDataNodeQueryContext( - instanceId.getQueryId(), instance.getDataNodeFINum()); - return createFragmentInstanceContext( - fragmentInstanceId, - stateMachine, - instance.getSessionInfo(), - dataRegion, - instance.getGlobalTimePredicate(), - dataNodeQueryContextMap, - DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()), - instance.isDebug(), - instance.isVerbose()); + registrations[0] = queryAttempts.register(instance, stateMachine, true); + try { + return createFragmentInstanceContext( + fragmentInstanceId, + stateMachine, + instance.getSessionInfo(), + dataRegion, + instance.getGlobalTimePredicate(), + registrations[0].getContext(), + DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()), + instance.isDebug(), + instance.isVerbose()); + } catch (Throwable t) { + registrations[0].close(); + throw t; + } }); rejectIfRepeatedDispatch(contextCreated[0], instanceId); + context.setQueryAttemptRegistration(registrations[0]); context.setHighestPriority(instance.isHighestPriority()); try { + if (stateMachine.getState().isDone()) { + failedInstanceInfo.set(context.getInstanceInfo()); + context.releaseDataNodeQueryContext(); + return null; + } List driverFactories = planner.plan( instance.getFragment().getPlanNodeTree(), instance.getFragment().getTypeProvider(), context, - dataNodeQueryContexts[0]); + registrations[0].getContext()); List drivers = new ArrayList<>(); driverFactories.forEach(factory -> drivers.add(factory.createDriver())); @@ -235,7 +242,15 @@ public FragmentInstanceInfo execDataQueryFragmentInstance( // still being created. Capture the failure result from the local context before // returning from computeIfAbsent instead of looking it up from the map later. failedInstanceInfo.set(context.getInstanceInfo()); - clearFIRelatedResources(instanceId); + try { + clearFIRelatedResources(instanceId); + } finally { + // Once execution cleanup is installed, it owns the reference until driver + // teardown finishes, including failures while submitting the drivers. + if (!context.isExecutionInitialized()) { + context.releaseDataNodeQueryContext(); + } + } return null; } }); @@ -295,11 +310,6 @@ private void clearFIRelatedResources(FragmentInstanceId instanceId) { instanceId.getQueryId().getId(), instanceId.getFragmentInstanceId(), false); } - private DataNodeQueryContext getOrCreateDataNodeQueryContext(QueryId queryId, int dataNodeFINum) { - return dataNodeQueryContextMap.computeIfAbsent( - queryId, queryId1 -> new DataNodeQueryContext(dataNodeFINum)); - } - @SuppressWarnings("squid:S1181") public FragmentInstanceInfo execSchemaQueryFragmentInstance( FragmentInstance instance, ISchemaRegion schemaRegion) { @@ -313,23 +323,37 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance( new FragmentInstanceStateMachine(instanceId, instanceNotificationExecutor); boolean[] contextCreated = new boolean[] {false}; + QueryAttemptRegistry.Registration[] registrations = + new QueryAttemptRegistry.Registration[1]; FragmentInstanceContext context = instanceContext.computeIfAbsent( instanceId, fragmentInstanceId -> { contextCreated[0] = true; - return createFragmentInstanceContext( - fragmentInstanceId, - stateMachine, - instance.getSessionInfo(), - DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()), - instance.isDebug(), - instance.isVerbose()); + registrations[0] = queryAttempts.register(instance, stateMachine, false); + try { + return createFragmentInstanceContext( + fragmentInstanceId, + stateMachine, + instance.getSessionInfo(), + DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()), + instance.isDebug(), + instance.isVerbose()); + } catch (Throwable t) { + registrations[0].close(); + throw t; + } }); rejectIfRepeatedDispatch(contextCreated[0], instanceId); + context.setQueryAttemptRegistration(registrations[0]); context.setHighestPriority(instance.isHighestPriority()); try { + if (stateMachine.getState().isDone()) { + failedInstanceInfo.set(context.getInstanceInfo()); + context.releaseDataNodeQueryContext(); + return null; + } List driverFactories = planner.plan( instance.getFragment().getPlanNodeTree(), @@ -371,7 +395,13 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance( // See execDataQueryFragmentInstance for why this result must not be fetched from // instanceContext after computeIfAbsent returns. failedInstanceInfo.set(context.getInstanceInfo()); - clearFIRelatedResources(instanceId); + try { + clearFIRelatedResources(instanceId); + } finally { + if (!context.isExecutionInitialized()) { + context.releaseDataNodeQueryContext(); + } + } return null; } }); @@ -419,6 +449,15 @@ public FragmentInstanceInfo cancelTask(FragmentInstanceId instanceId, boolean ha return null; } + public void closeQueryAttempt( + QueryId queryId, int attemptId, long deadline, boolean hasThrowable) { + queryAttempts.close(new QueryAttemptId(queryId, attemptId), deadline, hasThrowable); + } + + QueryAttemptRegistry getQueryAttempts() { + return queryAttempts; + } + /** * Gets the info for the specified fragment instance. * @@ -460,6 +499,7 @@ public TFetchFragmentInstanceStatisticsResp getFragmentInstanceStatistics( } private void removeOldInstances() { + queryAttempts.removeExpired(); long oldestAllowedInstance = System.currentTimeMillis() - infoCacheTime.toMillis(); instanceContext .entrySet() diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java new file mode 100644 index 0000000000000..1b36d4e05664e --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java @@ -0,0 +1,270 @@ +/* + * 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.execution.fragment; + +import org.apache.iotdb.commons.exception.IoTDBRuntimeException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.QueryAttemptId; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.LongSupplier; + +import static com.google.common.math.LongMath.saturatedAdd; + +/** + * Owns shared DataNode state for one execution attempt. All registration/closing bookkeeping is + * serialized by compute on the attempt key; planning and state transitions run outside that lock. + * Closed entries retain no query data after their last actual user releases it. They only fence + * late dispatch until the request's immutable admission deadline, after which the request itself is + * rejected, even if the entry has been removed. + */ +class QueryAttemptRegistry { + private static final Logger LOGGER = LoggerFactory.getLogger(QueryAttemptRegistry.class); + + private final Map attempts = new ConcurrentHashMap<>(); + private final LongSupplier clock; + // Expiring a fence must remain irreversible even if the wall clock subsequently goes backwards. + private final AtomicLong observedTime = new AtomicLong(); + + QueryAttemptRegistry() { + this(System::currentTimeMillis); + } + + QueryAttemptRegistry(LongSupplier clock) { + this.clock = clock; + } + + private long now() { + return observedTime.accumulateAndGet(clock.getAsLong(), Math::max); + } + + Registration register( + FragmentInstance instance, FragmentInstanceStateMachine stateMachine, boolean dataQuery) { + QueryAttemptId key = + new QueryAttemptId( + instance.getId().getQueryId(), + instance.hasQueryAttempt() + ? instance.getQueryAttemptId() + : QueryAttemptId.LEGACY_ATTEMPT); + long executionDeadline = + instance.hasQueryAttempt() + ? instance.getQueryDeadline() + : saturatedAdd(now(), instance.getTimeOut()); + long admissionDeadline = + instance.hasQueryAttempt() + ? Math.min(executionDeadline, instance.getAdmissionDeadline()) + : executionDeadline; + AtomicReference result = new AtomicReference<>(); + attempts.compute( + key, + (ignored, existing) -> { + long currentTime = now(); + if (admissionDeadline <= currentTime + || (existing != null + && (existing.closed || existing.executionDeadline <= currentTime))) { + throw rejected(key, admissionDeadline <= currentTime); + } + Entry entry = + existing == null + ? new Entry(executionDeadline, Math.max(1, instance.getDataNodeFINum())) + : existing; + // A transport retry must never extend the lifetime of an existing attempt. + entry.executionDeadline = Math.min(entry.executionDeadline, executionDeadline); + entry.fenceDeadline = Math.max(entry.fenceDeadline, admissionDeadline); + if (entry.active.containsKey(instance.getId())) { + throw new IoTDBRuntimeException( + String.format( + DataNodeQueryMessages + .QUERY_EXCEPTION_REPEATED_RPC_CALL_DETECTED_FOR_FRAGMENTINSTANCE_S_REJECT_BF609A26, + instance.getId().getFullId()), + TSStatusCode.REPEATED_RPC_CALL.getStatusCode()); + } + if (dataQuery && entry.context == null) { + entry.context = new DataNodeQueryContext(); + } + Registration registration = + new Registration(key, entry, instance.getId(), stateMachine, entry.context); + entry.active.put(instance.getId(), registration); + result.set(registration); + return entry; + }); + return result.get(); + } + + void close(QueryAttemptId key, long deadline, boolean hasThrowable) { + List toCancel = new ArrayList<>(); + long currentTime = now(); + attempts.compute( + key, + (ignored, existing) -> { + if (existing == null && deadline <= currentTime) { + return null; + } + Entry entry = existing == null ? new Entry(deadline, 0) : existing; + entry.closed = true; + entry.fenceDeadline = Math.max(entry.fenceDeadline, deadline); + if (entry.active != null) { + toCancel.addAll(entry.active.values()); + } + return retireIfUnused(key, entry, currentTime); + }); + terminate(toCancel, hasThrowable); + } + + void removeExpired() { + long currentTime = now(); + List toCancel = new ArrayList<>(); + for (QueryAttemptId key : attempts.keySet()) { + attempts.computeIfPresent( + key, + (ignored, entry) -> { + if ((entry.closed ? entry.fenceDeadline : entry.executionDeadline) > currentTime) { + return entry; + } + entry.closed = true; + if (entry.active != null) { + toCancel.addAll(entry.active.values()); + } + return retireIfUnused(key, entry, currentTime); + }); + } + terminate(toCancel, true); + } + + private static void terminate(List registrations, boolean hasThrowable) { + for (Registration registration : registrations) { + try { + if (hasThrowable) { + registration.stateMachine.cancel(); + } else { + registration.stateMachine.finished(); + } + } catch (RuntimeException e) { + // One failed notification must not prevent cancellation of the other FIs. + LOGGER.warn(DataNodeQueryMessages.ERRORS_FINISHING_FI_PROCESS, e); + } + } + } + + private static Entry retireIfUnused(QueryAttemptId key, Entry entry, long currentTime) { + if ((entry.active == null || entry.active.isEmpty()) + && (entry.closed || entry.remaining == 0)) { + entry.closed = true; + entry.context = null; + entry.active = null; + // Legacy requests have no admission deadline on the wire and may reuse the queryId on + // retry. Preserve their existing behavior instead of fencing a later legacy execution. + if (key.isLegacy() || entry.fenceDeadline <= currentTime) { + return null; + } + } + return entry; + } + + private static IoTDBRuntimeException rejected(QueryAttemptId key, boolean expired) { + return new IoTDBRuntimeException( + String.format( + DataNodeQueryMessages.EXCEPTION_QUERY_ATTEMPT_ARG_IS_CLOSED_OR_EXPIRED_F45E3DCF, key), + (expired ? TSStatusCode.QUERY_TIMEOUT : TSStatusCode.QUERY_WAS_KILLED).getStatusCode()); + } + + int getRetainedEntryCount() { + return attempts.size(); + } + + boolean hasSharedContext(QueryAttemptId key) { + AtomicBoolean present = new AtomicBoolean(); + attempts.computeIfPresent( + key, + (ignored, entry) -> { + present.set(entry.context != null); + return entry; + }); + return present.get(); + } + + private static class Entry { + private long executionDeadline; + private long fenceDeadline; + private int remaining; + private boolean closed; + private DataNodeQueryContext context; + private Map active = new HashMap<>(); + + private Entry(long deadline, int remaining) { + this.executionDeadline = deadline; + this.remaining = remaining; + } + } + + final class Registration implements AutoCloseable { + private final QueryAttemptId key; + private final Entry entry; + private final FragmentInstanceId instanceId; + private final FragmentInstanceStateMachine stateMachine; + private final DataNodeQueryContext context; + private final AtomicBoolean released = new AtomicBoolean(); + + private Registration( + QueryAttemptId key, + Entry entry, + FragmentInstanceId instanceId, + FragmentInstanceStateMachine stateMachine, + DataNodeQueryContext context) { + this.key = key; + this.entry = entry; + this.instanceId = instanceId; + this.stateMachine = stateMachine; + this.context = context; + } + + DataNodeQueryContext getContext() { + return context; + } + + @Override + public void close() { + if (!released.compareAndSet(false, true)) { + return; + } + attempts.computeIfPresent( + key, + (ignored, current) -> { + if (current != entry || !entry.active.remove(instanceId, this)) { + return current; + } + entry.remaining--; + return retireIfUnused(key, entry, now()); + }); + } + } +} 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 99ae864c465e0..1dbffd48ff2f8 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 @@ -94,7 +94,7 @@ public class QueryExecution implements IQueryExecution { private static final long RETRY_INTERVAL_IN_MS = 2000; private int retryCount = 0; private final MPPQueryContext context; - private IScheduler scheduler; + private volatile SchedulingAttempt schedulingAttempt = new SchedulingAttempt(); private final QueryStateMachine stateMachine; private final IPlanner planner; @@ -110,8 +110,6 @@ public class QueryExecution implements IQueryExecution { // used for cleaning resultHandle up exactly once private final AtomicBoolean resultHandleCleanUp; - private final AtomicBoolean stopped; - // cost time in ns of finished rpc private long totalExecutionTime = 0; @@ -168,7 +166,6 @@ public QueryExecution(IPlanner planner, MPPQueryContext context, ExecutorService this.stop(cause); } }); - this.stopped = new AtomicBoolean(false); this.resultHandleCleanUp = new AtomicBoolean(false); } @@ -296,7 +293,7 @@ private ExecutionResult retry() { // clear runtime variables in MPPQueryContext context.prepareForRetry(); // re-stop - this.stopped.compareAndSet(true, false); + this.schedulingAttempt = new SchedulingAttempt(); this.resultHandleCleanUp.compareAndSet(true, false); // re-analyze the query planner.beginAnalysisAttempt(); @@ -339,7 +336,13 @@ private IAnalysis analyze(MPPQueryContext context) { private void schedule() { final long startTime = System.nanoTime(); - this.scheduler = planner.doSchedule(analysis, distributedPlan, context, stateMachine); + SchedulingAttempt attempt = schedulingAttempt; + IScheduler scheduler = + planner.createScheduler(analysis, distributedPlan, context, stateMachine); + if (stateMachine.getState().isDone()) { + attempt.stop(stateMachine.getFailureException()); + } + attempt.start(scheduler); PERFORMANCE_OVERVIEW_METRICS.recordScheduleCost(System.nanoTime() - startTime); } @@ -389,9 +392,44 @@ private String printFragmentInstances(List instances) { // Stop the workers for this query @Override public void stop(Throwable t) { - // only stop once - if (stopped.compareAndSet(false, true) && this.scheduler != null) { - this.scheduler.stop(t); + schedulingAttempt.stop(t); + } + + /** One-shot handoff between scheduler creation and concurrent stop requests. */ + static class SchedulingAttempt { + private IScheduler scheduler; + private boolean stopped; + private Throwable stopCause; + + void start(IScheduler candidate) { + boolean stopBeforeStart; + Throwable cause; + synchronized (this) { + scheduler = candidate; + stopBeforeStart = stopped; + cause = stopCause; + } + if (stopBeforeStart) { + candidate.stop(cause); + } else { + // The scheduler also supports stop-before-start: stop may win after publication above. + candidate.start(); + } + } + + void stop(Throwable cause) { + IScheduler toStop; + synchronized (this) { + if (stopped) { + return; + } + stopped = true; + stopCause = cause; + toStop = scheduler; + } + if (toStop != null) { + toStop.stop(cause); + } } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/IPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/IPlanner.java index 99cd988bbb7b6..9ffd88be05731 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/IPlanner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/IPlanner.java @@ -45,7 +45,8 @@ default void commitAnalysisAttempt() {} DistributedQueryPlan doDistributionPlan( IAnalysis analysis, LogicalQueryPlan logicalPlan, MPPQueryContext context); - IScheduler doSchedule( + /** Creates an unstarted scheduler. QueryExecution publishes it before calling start. */ + IScheduler createScheduler( IAnalysis analysis, DistributedQueryPlan distributedPlan, MPPQueryContext context, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/TreeModelPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/TreeModelPlanner.java index 8a614809559ec..464f6082ecd51 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/TreeModelPlanner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/TreeModelPlanner.java @@ -136,7 +136,7 @@ public DistributedQueryPlan doDistributionPlan( } @Override - public IScheduler doSchedule( + public IScheduler createScheduler( IAnalysis analysis, DistributedQueryPlan distributedPlan, MPPQueryContext context, @@ -168,7 +168,6 @@ public IScheduler doSchedule( asyncInternalServiceClientManager); } - scheduler.start(); return scheduler; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java index 463aeb5413196..295ff71dd1cb9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java @@ -71,6 +71,37 @@ public class FragmentInstance implements IConsensusRequest { // The num of all FI on the dispatched DataNode in this query private int dataNodeFINum; + // Admission metadata travels in the optional RPC envelope, not the binary FI body, so older + // DataNodes can still deserialize the plan. Local dispatch sets the same fields directly. + private int queryAttemptId = -1; + private long queryDeadline; + private long admissionDeadline; + + public void setQueryAttempt(int queryAttemptId, long queryDeadline) { + this.queryAttemptId = queryAttemptId; + this.queryDeadline = queryDeadline; + } + + public boolean hasQueryAttempt() { + return queryAttemptId >= 0 && queryDeadline > 0; + } + + public int getQueryAttemptId() { + return queryAttemptId; + } + + public long getQueryDeadline() { + return queryDeadline; + } + + public void setAdmissionDeadline(long admissionDeadline) { + this.admissionDeadline = admissionDeadline; + } + + public long getAdmissionDeadline() { + return admissionDeadline; + } + private boolean isHighestPriority; // indicate which index we are retrying diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlanner.java index cca616e0a3b17..f2710e28be4ef 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlanner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlanner.java @@ -180,7 +180,7 @@ public DistributedQueryPlan doDistributionPlan( } @Override - public IScheduler doSchedule( + public IScheduler createScheduler( final IAnalysis analysis, final DistributedQueryPlan distributedPlan, final MPPQueryContext context, @@ -210,7 +210,6 @@ public IScheduler doSchedule( syncInternalServiceClientManager, asyncInternalServiceClientManager); } - scheduler.start(); return scheduler; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java index 03fa0b666a071..c15de059af9b9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java @@ -67,6 +67,9 @@ public class ClusterScheduler implements IScheduler { private IQueryTerminator queryTerminator; // Set when the dispatch of the fragment instances failed (see start()). private volatile boolean dispatchFailed = false; + private boolean started; + private boolean stopped; + private boolean dispatchComplete; private static final QueryExecutionMetricSet QUERY_EXECUTION_METRICS = QueryExecutionMetricSet.getInstance(); @@ -84,12 +87,13 @@ public ClusterScheduler( this.rootSubPlan = distributedQueryPlan.getRootSubPlan(); this.instances = distributedQueryPlan.getInstances(); this.queryType = queryType; - this.dispatcher = + FragmentInstanceDispatcherImpl fragmentDispatcher = new FragmentInstanceDispatcherImpl( queryType, queryContext, syncInternalServiceClientManager, asyncInternalServiceClientManager); + this.dispatcher = fragmentDispatcher; if (isQuery()) { this.stateTracker = new FixedRateFragInsStateTracker( @@ -100,10 +104,25 @@ public ClusterScheduler( queryContext, instances, syncInternalServiceClientManager, - stateTracker); + stateTracker, + fragmentDispatcher::getLastAdmissionDeadline); } } + ClusterScheduler( + QueryStateMachine stateMachine, + IFragInstanceDispatcher dispatcher, + IFragInstanceStateTracker stateTracker, + IQueryTerminator queryTerminator) { + this.stateMachine = stateMachine; + this.queryType = QueryType.READ; + this.rootSubPlan = null; + this.instances = java.util.Collections.emptyList(); + this.dispatcher = dispatcher; + this.stateTracker = stateTracker; + this.queryTerminator = queryTerminator; + } + private boolean needRetry(TSStatus failureStatus) { return failureStatus != null && isQuery() @@ -114,17 +133,29 @@ && isQuery() @Override public void start() { + synchronized (this) { + if (started || stopped || stateMachine.getState().isDone()) { + return; + } + started = true; + } stateMachine.transitionToDispatching(); long startTime = System.nanoTime(); - Future dispatchResultFuture = - dispatcher.dispatch(rootSubPlan, instances); // NOTICE: the FragmentInstance may be dispatched to another Host due to consensus redirect. // So we need to start the state fetcher after the dispatching stage. try { + Future dispatchResultFuture = + dispatcher.dispatch(rootSubPlan, instances); FragInstanceDispatchResult result = dispatchResultFuture.get(); + synchronized (this) { + if (stopped) { + return; + } + dispatchComplete = true; + dispatchFailed = !result.isSuccessful(); + } if (!result.isSuccessful()) { - dispatchFailed = true; if (needRetry(result.getFailureStatus())) { stateMachine.transitionToPendingRetry(result.getFailureStatus()); } else { @@ -133,7 +164,7 @@ public void start() { } return; } - } catch (InterruptedException | ExecutionException e) { + } catch (InterruptedException | ExecutionException | RuntimeException e) { // If the dispatch request cannot be sent or TException is caught, we will retry this query. dispatchFailed = true; if (e instanceof InterruptedException) { @@ -161,12 +192,20 @@ public void start() { @Override public void stop(Throwable t) { + boolean terminateImmediately; + synchronized (this) { + if (stopped) { + return; + } + stopped = true; + terminateImmediately = dispatchFailed || !dispatchComplete; + } dispatcher.abort(); if (stateTracker != null) { stateTracker.abort(); } if (queryTerminator != null) { - if (dispatchFailed) { + if (terminateImmediately) { // The FragmentInstances dispatched before the failure can't finish on their own and hold // operator memory until they are cancelled, so don't wait for the grace period. queryTerminator.terminateImmediately(t); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java index 9af76e1b40121..76f2542f7f295 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java @@ -39,6 +39,7 @@ import org.apache.iotdb.db.audit.DataNodeUserDataTransferAuditor; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.exception.mpp.FragmentInstanceDispatchException; +import org.apache.iotdb.db.exception.query.KilledByOthersException; import org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; @@ -62,6 +63,7 @@ import org.apache.iotdb.rpc.RpcUtils; import org.apache.iotdb.rpc.TSStatusCode; +import com.google.common.math.LongMath; import org.apache.thrift.TException; import org.apache.thrift.transport.TTransportException; import org.apache.tsfile.external.commons.lang3.exception.ExceptionUtils; @@ -93,6 +95,12 @@ public class FragmentInstanceDispatcherImpl implements IFragInstanceDispatcher { private static final CommonConfig COMMON_CONFIG = CommonDescriptor.getInstance().getConfig(); private final QueryType type; private final MPPQueryContext queryContext; + private final int queryAttemptId; + private final long queryDeadline; + private final long admissionWindowMs; + private long lastAdmissionDeadline; + private volatile boolean aborted; + private volatile LinkedBlockingQueue> dispatchQueue; private final String localhostIpAddr; private final int localhostInternalPort; private final IClientManager @@ -120,6 +128,13 @@ public FragmentInstanceDispatcherImpl( asyncInternalServiceClientManager) { this.type = type; this.queryContext = queryContext; + this.queryAttemptId = queryContext.getQueryAttemptId(); + this.queryDeadline = queryContext.getQueryDeadline(); + // A zero socket timeout means unlimited waiting, but admission fences must remain finite. + this.admissionWindowMs = + COMMON_CONFIG.getDnConnectionTimeoutInMS() > 0 + ? COMMON_CONFIG.getDnConnectionTimeoutInMS() + : TimeUnit.MINUTES.toMillis(1); this.syncInternalServiceClientManager = syncInternalServiceClientManager; this.asyncInternalServiceClientManager = asyncInternalServiceClientManager; this.localhostIpAddr = IoTDBDescriptor.getInstance().getConfig().getInternalAddress(); @@ -134,6 +149,9 @@ public FragmentInstanceDispatcherImpl( public Future dispatch( SubPlan root, List instances) { if (isQuery()) { + for (FragmentInstance instance : instances) { + instance.setQueryAttempt(queryAttemptId, queryDeadline); + } return instances.size() == 1 || root == null ? dispatchRead(instances) : topologicalParallelDispatchRead(root, instances); @@ -145,11 +163,12 @@ public Future dispatch( private Future topologicalParallelDispatchRead( SubPlan root, List instances) { long startTime = System.nanoTime(); - LinkedBlockingQueue> queue = new LinkedBlockingQueue<>(instances.size()); + LinkedBlockingQueue> queue = new LinkedBlockingQueue<>(); + dispatchQueue = queue; List> futures = new ArrayList<>(instances.size()); queue.add(new Pair<>(root, true)); try { - while (futures.size() < instances.size()) { + while (!aborted && futures.size() < instances.size()) { Pair pair = queue.take(); SubPlan next = pair.getLeft(); boolean previousSuccess = pair.getRight(); @@ -188,6 +207,7 @@ private Future topologicalParallelDispatchRead( RpcUtils.getStatus( TSStatusCode.INTERNAL_SERVER_ERROR, UNEXPECTED_ERRORS + t.getMessage()))); } finally { + dispatchQueue = null; long queryDispatchReadTime = System.nanoTime() - startTime; QUERY_EXECUTION_METRIC_SET.recordExecutionCost(DISPATCH_READ, queryDispatchReadTime); queryContext.recordDispatchCost(queryDispatchReadTime); @@ -218,7 +238,7 @@ private Future asyncDispatchOneInstance( // If the failure is not notified and the child FI is not sent, the outer loop // won't be able to exit. queue.add(new Pair<>(null, false)); - } else { + } else if (!aborted) { for (SubPlan child : plan.getChildren()) { queue.add(new Pair<>(child, true)); } @@ -463,8 +483,10 @@ private long getRetrySleepTime(int retryTimes) { private void dispatchOneInstance(FragmentInstance instance) throws FragmentInstanceDispatchException { + checkNotAborted(); TEndPoint endPoint = instance.getHostDataNode().getInternalEndPoint(); if (isDispatchedToLocal(endPoint)) { + prepareAdmission(instance); dispatchLocally(instance); } else { dispatchRemote(instance, endPoint); @@ -493,6 +515,13 @@ private void dispatchRemoteHelper(final FragmentInstance instance, final TEndPoi case READ: final TSendFragmentInstanceReq sendFragmentInstanceReq = new TSendFragmentInstanceReq(new TFragmentInstance(instance.serializeToByteBuffer())); + prepareAdmission(instance); + if (instance.hasQueryAttempt()) { + sendFragmentInstanceReq + .setQueryAttemptId(instance.getQueryAttemptId()) + .setQueryDeadline(instance.getQueryDeadline()) + .setAdmissionDeadline(instance.getAdmissionDeadline()); + } if (instance.getExecutorType().isStorageExecutor()) { sendFragmentInstanceReq.setConsensusGroupId( instance.getRegionReplicaSet().getRegionId()); @@ -662,6 +691,7 @@ private void dispatchRemote(FragmentInstance instance, TEndPoint endPoint) currentTime))); } // we just retry once to clear stale connection for a restart node. + checkNotAborted(); try { dispatchRemoteHelper(instance, endPoint); } catch (ClientManagerException @@ -767,7 +797,37 @@ private void dispatchLocally(final FragmentInstance instance) } @Override - public void abort() {} + public synchronized void abort() { + aborted = true; + LinkedBlockingQueue> queue = dispatchQueue; + if (queue != null) { + queue.offer(new Pair<>(null, false)); + } + } + + private synchronized void prepareAdmission(FragmentInstance instance) + throws FragmentInstanceDispatchException { + checkNotAborted(); + long deadline = instance.getAdmissionDeadline(); + if (deadline == 0) { + deadline = + Math.min( + queryDeadline, LongMath.saturatedAdd(System.currentTimeMillis(), admissionWindowMs)); + instance.setAdmissionDeadline(deadline); + } + lastAdmissionDeadline = Math.max(lastAdmissionDeadline, deadline); + } + + public synchronized long getLastAdmissionDeadline() { + return lastAdmissionDeadline; + } + + private void checkNotAborted() throws FragmentInstanceDispatchException { + if (aborted) { + throw new FragmentInstanceDispatchException( + RpcUtils.getStatus(TSStatusCode.QUERY_WAS_KILLED, KilledByOthersException.MESSAGE)); + } + } private boolean isQuery() { return type == QueryType.READ || type == QueryType.READ_WRITE; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IScheduler.java index 4cc6a5c7eaf5e..d023fe4edc18e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IScheduler.java @@ -25,6 +25,7 @@ public interface IScheduler { + /** Starts once; a scheduler stopped before this call must not dispatch work. */ void start(); void stop(Throwable t); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java index 8af9117fe0546..d0b6da402e266 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java @@ -43,6 +43,7 @@ import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.function.LongSupplier; import java.util.stream.Collectors; public class SimpleQueryTerminator implements IQueryTerminator { @@ -50,6 +51,9 @@ public class SimpleQueryTerminator implements IQueryTerminator { private static final long TERMINATION_GRACE_PERIOD_IN_MS = 1000L; protected ScheduledExecutorService scheduledExecutor; private final QueryId queryId; + private final int queryAttemptId; + private final LongSupplier admissionDeadlineSupplier; + private long admissionDeadline; private final MPPQueryContext queryContext; private final IFragInstanceStateTracker stateTracker; @@ -67,9 +71,12 @@ public SimpleQueryTerminator( final List fragmentInstances, final IClientManager internalServiceClientManager, - final IFragInstanceStateTracker stateTracker) { + final IFragInstanceStateTracker stateTracker, + final LongSupplier admissionDeadlineSupplier) { this.scheduledExecutor = scheduledExecutor; this.queryId = queryContext.getQueryId(); + this.queryAttemptId = queryContext.getQueryAttemptId(); + this.admissionDeadlineSupplier = admissionDeadlineSupplier; this.queryContext = queryContext; this.internalServiceClientManager = internalServiceClientManager; this.stateTracker = stateTracker; @@ -101,6 +108,9 @@ private Future terminate(final Throwable t, final long delayInMs) { .filter(endPoint -> !queryContext.getEndPointBlackList().contains(endPoint)) .collect(Collectors.toList()); this.hasThrowable = Objects.nonNull(t); + // ClusterScheduler aborts the dispatcher before termination: no later request can obtain an + // admission deadline beyond this cutoff. Transport retries reuse their original deadline. + this.admissionDeadline = admissionDeadlineSupplier.getAsLong(); return scheduledExecutor.schedule(this::syncTerminate, delayInMs, TimeUnit.MILLISECONDS); } @@ -108,10 +118,11 @@ private Future terminate(final Throwable t, final long delayInMs) { public Boolean syncTerminate() { boolean succeed = true; for (final TEndPoint endPoint : relatedHost) { - // We only send cancel query request if there is remaining unfinished FI in that node + // Attempt closure is required even when every received FI finished: planned but unsent FIs + // can still keep the shared DataNode context alive. Legacy receivers only use the FI list. final List unfinishedFIs = stateTracker.filterUnFinishedFIs(ownedFragmentInstance.get(endPoint)); - if (unfinishedFIs.isEmpty()) { + if (unfinishedFIs.isEmpty() && admissionDeadline <= 0) { continue; } @@ -119,6 +130,10 @@ public Boolean syncTerminate() { final int internalPort = IoTDBDescriptor.getInstance().getConfig().getInternalPort(); if (internalAddress.equalsIgnoreCase(endPoint.getIp()) && internalPort == endPoint.getPort()) { + if (admissionDeadline > 0) { + FragmentInstanceManager.getInstance() + .closeQueryAttempt(queryId, queryAttemptId, admissionDeadline, hasThrowable); + } for (final FragmentInstanceId insId : unfinishedFIs) { FragmentInstanceManager.getInstance().cancelTask(insId, hasThrowable); } @@ -127,13 +142,17 @@ public Boolean syncTerminate() { try (final SyncDataNodeInternalServiceClient client = internalServiceClientManager.borrowClient(endPoint)) { - client.cancelQuery( + TCancelQueryReq request = new TCancelQueryReq( queryId.getId(), unfinishedFIs.stream() .map(FragmentInstanceId::toThrift) .collect(Collectors.toList()), - hasThrowable)); + hasThrowable); + if (admissionDeadline > 0) { + request.setQueryAttemptId(queryAttemptId).setAdmissionDeadline(admissionDeadline); + } + client.cancelQuery(request); } catch (final ClientManagerException e) { logger.warn(DataNodeQueryMessages.CAN_T_CONNECT_TO_NODE, endPoint, e); // we shouldn't return here and need to cancel queryTasks in other nodes diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java index 18678a043eeea..24c78410dd5a5 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java @@ -49,6 +49,7 @@ import java.util.Objects; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; /** * LOAD scheduler: {@link LoadTsFileScheduler} is the coordinator of a batch of {@link @@ -269,6 +270,8 @@ * */ public class LoadTsFileScheduler implements IScheduler { + private final AtomicBoolean started = new AtomicBoolean(); + private volatile boolean stopped; private static final Logger LOGGER = LoggerFactory.getLogger(LoadTsFileScheduler.class); @@ -350,12 +353,21 @@ public LoadTsFileScheduler( /** Starts the batch loading orchestration across all input TsFile nodes. */ @Override public void start() { + if (!started.compareAndSet(false, true)) { + return; + } try { + if (stopped || stateMachine.getState().isDone()) { + return; + } stateMachine.transitionToRunning(); boolean allFilesSucceeded = true; final int totalFiles = tsFileNodeList.size(); for (int i = 0; i < totalFiles; ++i) { + if (stopped) { + return; + } final LoadSingleTsFileNode node = tsFileNodeList.get(i); final String filePath = node.getTsFileResource().getTsFilePath(); final String userName = queryContext.getSession().getUserName(); @@ -586,9 +598,15 @@ static String getPartitionQueryDatabase( /** Aborts active dispatchers and transitions the query state machine to failed. */ @Override public void stop(final Throwable t) { + stopped = true; try { dispatcher.abort(); } finally { + // The factory can now publish this scheduler before start. Release constructor-owned + // memory when stop wins that race; otherwise start's finally owns teardown. + if (started.compareAndSet(false, true)) { + teardown(); + } if (t != null) { stateMachine.transitionToFailed(t); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java index d04fb263a31d6..6fe3c2232231b 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java @@ -66,6 +66,52 @@ public void testSchemaQueryFailureInfoAfterConcurrentContextRemoval() throws Exc assertFailureInfoAfterConcurrentContextRemoval(false); } + @Test + public void testPartialDispatchReleasesAttemptContextAndRejectsLateFragment() { + QueryId queryId = new QueryId("partial_attempt_dispatch"); + long deadline = System.currentTimeMillis() + TimeUnit.MINUTES.toMillis(1); + FragmentInstanceManager manager = FragmentInstanceManager.getInstance(); + org.apache.iotdb.db.queryengine.common.QueryAttemptId key = + new org.apache.iotdb.db.queryengine.common.QueryAttemptId(queryId, 0); + try { + for (int i = 0; i < 2; i++) { + FragmentInstance instance = mock(FragmentInstance.class); + PlanFragment fragment = mock(PlanFragment.class); + FragmentInstanceId id = new FragmentInstanceId(new PlanFragmentId(queryId, i), "0"); + when(instance.getId()).thenReturn(id); + when(instance.getFragment()).thenReturn(fragment); + when(instance.hasQueryAttempt()).thenReturn(true); + when(instance.getQueryAttemptId()).thenReturn(0); + when(instance.getQueryDeadline()).thenReturn(deadline); + when(instance.getAdmissionDeadline()).thenReturn(deadline); + when(instance.getDataNodeFINum()).thenReturn(3); + when(instance.getTimeOut()).thenReturn(TimeUnit.MINUTES.toMillis(1)); + when(fragment.getPlanNodeTree()) + .thenThrow(new MetadataLeaseFencedException(FAILURE_MESSAGE, RETRY_UNTIL_SUCCESS)); + FragmentInstanceInfo result = + manager.execDataQueryFragmentInstance(instance, mock(IDataRegionForQuery.class)); + assertTrue(result.getState().isFailed()); + } + assertTrue(manager.getQueryAttempts().hasSharedContext(key)); + manager.closeQueryAttempt(queryId, 0, deadline, true); + org.junit.Assert.assertFalse(manager.getQueryAttempts().hasSharedContext(key)); + FragmentInstance late = mock(FragmentInstance.class); + when(late.getId()).thenReturn(new FragmentInstanceId(new PlanFragmentId(queryId, 2), "0")); + when(late.hasQueryAttempt()).thenReturn(true); + when(late.getQueryAttemptId()).thenReturn(0); + when(late.getQueryDeadline()).thenReturn(deadline); + when(late.getAdmissionDeadline()).thenReturn(deadline); + org.junit.Assert.assertThrows( + org.apache.iotdb.commons.exception.IoTDBRuntimeException.class, + () -> manager.execDataQueryFragmentInstance(late, mock(IDataRegionForQuery.class))); + } finally { + manager.closeQueryAttempt(queryId, 0, deadline, true); + for (int i = 0; i < 3; i++) { + manager.cancelTask(new FragmentInstanceId(new PlanFragmentId(queryId, i), "0"), true); + } + } + } + private void assertFailureInfoAfterConcurrentContextRemoval(boolean dataQuery) throws Exception { String queryType = dataQuery ? "data" : "schema"; FragmentInstanceId instanceId = diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java new file mode 100644 index 0000000000000..ddf282306fd9b --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java @@ -0,0 +1,242 @@ +/* + * 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.execution.fragment; + +import org.apache.iotdb.commons.exception.IoTDBRuntimeException; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryAttemptId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.junit.Test; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class QueryAttemptRegistryTest { + private final AtomicLong clock = new AtomicLong(1000); + private final QueryAttemptRegistry registry = new QueryAttemptRegistry(clock::get); + private final QueryId queryId = new QueryId("attempt_test"); + + @Test + public void keepsSharedStateBetweenSequentiallyDispatchedFragments() { + QueryAttemptRegistry.Registration first = register(0, 0, 2); + DataNodeQueryContext context = first.getContext(); + first.close(); + assertTrue(registry.hasSharedContext(key(0))); + QueryAttemptRegistry.Registration second = register(0, 1, 2); + assertSame(context, second.getContext()); + second.close(); + assertFalse(registry.hasSharedContext(key(0))); + assertEquals(1, registry.getRetainedEntryCount()); + } + + @Test + public void closesPartiallyDispatchedAttemptAfterActualUsersRelease() { + FragmentInstance instance = instance(0, 0, 3); + FragmentInstanceStateMachine state = state(instance); + QueryAttemptRegistry.Registration first = registry.register(instance, state, true); + QueryAttemptRegistry.Registration failed = register(0, 1, 3); + failed.close(); + registry.close(key(0), 2000, true); + assertEquals(FragmentInstanceState.CANCELLED, state.getState()); + assertTrue(registry.hasSharedContext(key(0))); + first.close(); + first.close(); + registry.close(key(0), 2000, true); + assertFalse(registry.hasSharedContext(key(0))); + // The third FI never arrived; it cannot keep the shared LAST/filter state alive. + assertEquals(1, registry.getRetainedEntryCount()); + } + + @Test + public void lateOldCancellationAndReleaseDoNotAffectRetry() { + QueryAttemptRegistry.Registration old = register(0, 0, 3); + QueryAttemptRegistry.Registration retry = register(1, 3, 1); + assertNotSame(old.getContext(), retry.getContext()); + registry.close(key(0), 2000, true); + old.close(); + registry.close(key(0), 2000, true); + assertTrue(registry.hasSharedContext(key(1))); + assertFalse(registry.hasSharedContext(key(0))); + retry.close(); + } + + @Test + public void closedBeforeFirstDispatchRemainsFencedAfterRecordExpires() { + registry.close(key(0), 2000, true); + assertEquals( + TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), + assertThrows(IoTDBRuntimeException.class, () -> register(0, 0, 1)).getErrorCode()); + clock.set(2000); + registry.removeExpired(); + assertEquals(0, registry.getRetainedEntryCount()); + assertEquals( + TSStatusCode.QUERY_TIMEOUT.getStatusCode(), + assertThrows(IoTDBRuntimeException.class, () -> register(0, 0, 1)).getErrorCode()); + clock.set(1500); + assertThrows(IoTDBRuntimeException.class, () -> register(0, 0, 1)); + assertEquals(0, registry.getRetainedEntryCount()); + } + + @Test + public void deadlineReclaimsPartialAttemptWhenCoordinatorDoesNotCancel() { + FragmentInstance instance = instance(0, 0, 3); + FragmentInstanceStateMachine state = state(instance); + QueryAttemptRegistry.Registration first = registry.register(instance, state, true); + clock.set(2000); + registry.removeExpired(); + assertEquals(FragmentInstanceState.CANCELLED, state.getState()); + first.close(); + assertEquals(0, registry.getRetainedEntryCount()); + } + + @Test + public void repeatedDispatchDoesNotReplaceExistingRegistration() { + FragmentInstance instance = instance(0, 0, 2); + QueryAttemptRegistry.Registration first = registry.register(instance, state(instance), true); + assertEquals( + TSStatusCode.REPEATED_RPC_CALL.getStatusCode(), + assertThrows( + IoTDBRuntimeException.class, + () -> registry.register(instance, state(instance), true)) + .getErrorCode()); + first.close(); + assertTrue(registry.hasSharedContext(key(0))); + registry.close(key(0), 2000, true); + assertFalse(registry.hasSharedContext(key(0))); + } + + @Test + public void legacyCompletionDoesNotFenceRetryUsingTheSameQueryId() { + FragmentInstance instance = instance(0, 0, 1); + when(instance.hasQueryAttempt()).thenReturn(false); + QueryAttemptRegistry.Registration first = registry.register(instance, state(instance), true); + first.close(); + assertEquals(0, registry.getRetainedEntryCount()); + QueryAttemptRegistry.Registration second = registry.register(instance, state(instance), true); + assertNotSame(first.getContext(), second.getContext()); + second.close(); + } + + @Test + public void concurrentCloseCannotMissRegistrationInProgress() throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(2); + CountDownLatch registering = new CountDownLatch(1); + CountDownLatch allowRegistration = new CountDownLatch(1); + FragmentInstance instance = instance(0, 0, 3); + FragmentInstanceStateMachine state = state(instance); + when(instance.getDataNodeFINum()) + .thenAnswer( + ignored -> { + registering.countDown(); + assertTrue(allowRegistration.await(10, TimeUnit.SECONDS)); + return 3; + }); + try { + Future registration = + executor.submit(() -> registry.register(instance, state, true)); + assertTrue(registering.await(10, TimeUnit.SECONDS)); + Future cancellation = executor.submit(() -> registry.close(key(0), 2000, true)); + allowRegistration.countDown(); + QueryAttemptRegistry.Registration result = registration.get(10, TimeUnit.SECONDS); + cancellation.get(10, TimeUnit.SECONDS); + assertEquals(FragmentInstanceState.CANCELLED, state.getState()); + result.close(); + assertFalse(registry.hasSharedContext(key(0))); + assertThrows(IoTDBRuntimeException.class, () -> register(0, 1, 3)); + } finally { + allowRegistration.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + @Test + public void admissionExpiryDoesNotCancelAnAlreadyRunningLongQuery() { + FragmentInstance instance = instance(0, 0, 3); + when(instance.getQueryDeadline()).thenReturn(Long.MAX_VALUE); + when(instance.getAdmissionDeadline()).thenReturn(1200L); + FragmentInstanceStateMachine state = state(instance); + QueryAttemptRegistry.Registration running = registry.register(instance, state, true); + clock.set(1300); + registry.removeExpired(); + assertEquals(FragmentInstanceState.RUNNING, state.getState()); + assertTrue(registry.hasSharedContext(key(0))); + registry.close(key(0), 1200, true); + running.close(); + assertEquals(0, registry.getRetainedEntryCount()); + assertThrows( + IoTDBRuntimeException.class, () -> registry.register(instance, state(instance), true)); + } + + @Test + public void completedLongQueryRetainsOnlyABoundedAdmissionFence() { + FragmentInstance instance = instance(0, 0, 1); + when(instance.getQueryDeadline()).thenReturn(Long.MAX_VALUE); + when(instance.getAdmissionDeadline()).thenReturn(1200L); + registry.register(instance, state(instance), true).close(); + assertEquals(1, registry.getRetainedEntryCount()); + assertFalse(registry.hasSharedContext(key(0))); + clock.set(1200); + registry.removeExpired(); + assertEquals(0, registry.getRetainedEntryCount()); + } + + private QueryAttemptId key(int attempt) { + return new QueryAttemptId(queryId, attempt); + } + + private QueryAttemptRegistry.Registration register(int attempt, int fragment, int planned) { + FragmentInstance instance = instance(attempt, fragment, planned); + return registry.register(instance, state(instance), true); + } + + private FragmentInstanceStateMachine state(FragmentInstance instance) { + return new FragmentInstanceStateMachine(instance.getId(), Runnable::run); + } + + private FragmentInstance instance(int attempt, int fragment, int planned) { + FragmentInstance instance = mock(FragmentInstance.class); + when(instance.getId()) + .thenReturn(new FragmentInstanceId(new PlanFragmentId(queryId, fragment), "0")); + when(instance.hasQueryAttempt()).thenReturn(true); + when(instance.getQueryAttemptId()).thenReturn(attempt); + when(instance.getQueryDeadline()).thenReturn(2000L); + when(instance.getAdmissionDeadline()).thenReturn(2000L); + when(instance.getTimeOut()).thenReturn(1000L); + when(instance.getDataNodeFINum()).thenReturn(planned); + return instance; + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/NonAlignedTreeDeviceViewScanOperatorTreeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/NonAlignedTreeDeviceViewScanOperatorTreeTest.java index cc7f97984715f..3c3478d323a0c 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/NonAlignedTreeDeviceViewScanOperatorTreeTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/NonAlignedTreeDeviceViewScanOperatorTreeTest.java @@ -660,7 +660,7 @@ private Operator getOperator( LocalExecutionPlanContext localExecutionPlanContext = new LocalExecutionPlanContext( - typeProvider, fragmentInstanceContext, new DataNodeQueryContext(1)); + typeProvider, fragmentInstanceContext, new DataNodeQueryContext()); Operator operator = tableOperatorGenerator.visitTreeNonAlignedDeviceViewScan(node, localExecutionPlanContext); ((DataDriverContext) localExecutionPlanContext.getDriverContext()) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/sink/SinkOperatorDownStreamNodeIdTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/sink/SinkOperatorDownStreamNodeIdTest.java index 03c93dcf0ca6a..f7dd36b804b05 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/sink/SinkOperatorDownStreamNodeIdTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/sink/SinkOperatorDownStreamNodeIdTest.java @@ -85,7 +85,7 @@ private LocalExecutionPlanContext createLocalExecutionPlanContext( fragmentInstanceContext.setDataRegion(dataRegion); return new LocalExecutionPlanContext( - typeProvider, fragmentInstanceContext, new DataNodeQueryContext(1)); + typeProvider, fragmentInstanceContext, new DataNodeQueryContext()); } // for table model diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionLifecycleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionLifecycleTest.java new file mode 100644 index 0000000000000..9ab62a22d80c8 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionLifecycleTest.java @@ -0,0 +1,135 @@ +/* + * 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.plan.execution; + +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.QueryStateMachine; +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.planner.IPlanner; +import org.apache.iotdb.db.queryengine.plan.planner.plan.DistributedQueryPlan; +import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan; +import org.apache.iotdb.db.queryengine.plan.scheduler.IScheduler; + +import org.junit.Test; + +import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class QueryExecutionLifecycleTest { + @Test + public void stopBeforeFactoryReturnsIsDeliveredToCreatedScheduler() throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(2); + CountDownLatch creating = new CountDownLatch(1); + CountDownLatch allowCreation = new CountDownLatch(1); + IScheduler scheduler = mock(IScheduler.class); + IPlanner planner = planner(); + when(planner.createScheduler(any(), any(), any(), any())) + .thenAnswer( + ignored -> { + creating.countDown(); + assertTrue(allowCreation.await(10, TimeUnit.SECONDS)); + return scheduler; + }); + QueryExecution execution = new QueryExecution(planner, context(), executor); + Throwable cause = new RuntimeException("stop before scheduler publication"); + try { + Future start = executor.submit(execution::start); + assertTrue(creating.await(10, TimeUnit.SECONDS)); + execution.stop(cause); + allowCreation.countDown(); + start.get(10, TimeUnit.SECONDS); + execution.stop(cause); + verify(scheduler).stop(cause); + verify(scheduler, never()).start(); + } finally { + allowCreation.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + @Test + public void asynchronousFailureStopsSchedulerBeforeStartReturns() throws Exception { + ExecutorService executor = Executors.newSingleThreadExecutor(); + CountDownLatch stopped = new CountDownLatch(1); + Throwable cause = new RuntimeException("non-retryable dispatch failure"); + IScheduler scheduler = mock(IScheduler.class); + IPlanner planner = planner(); + doAnswer( + ignored -> { + stopped.countDown(); + return null; + }) + .when(scheduler) + .stop(cause); + when(planner.createScheduler(any(), any(), any(), any())) + .thenAnswer( + invocation -> { + QueryStateMachine state = invocation.getArgument(3); + doAnswer( + ignored -> { + state.transitionToFailed(cause); + assertTrue(stopped.await(10, TimeUnit.SECONDS)); + return null; + }) + .when(scheduler) + .start(); + return scheduler; + }); + QueryExecution execution = new QueryExecution(planner, context(), executor); + try { + execution.start(); + execution.stopAndCleanup(cause); + verify(scheduler).stop(cause); + } finally { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + private static IPlanner planner() { + IPlanner planner = mock(IPlanner.class); + when(planner.analyze(any())).thenReturn(mock(IAnalysis.class)); + when(planner.doLogicalPlan(any(), any())).thenReturn(mock(LogicalQueryPlan.class)); + DistributedQueryPlan plan = mock(DistributedQueryPlan.class); + when(plan.getInstances()).thenReturn(Collections.emptyList()); + when(planner.doDistributionPlan(any(), any(), any())).thenReturn(plan); + return planner; + } + + private static MPPQueryContext context() { + MPPQueryContext context = new MPPQueryContext(new QueryId("scheduler_publication")); + context.setQueryType(QueryType.OTHER); + return context; + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionRetryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionRetryTest.java index 8761346df5e6c..5a8bb325b21e1 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionRetryTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionRetryTest.java @@ -88,6 +88,7 @@ public void testAnalysisAttemptIsRolledBackOnlyWhenRetryStarts() { assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), result.status.code); assertEquals(2, planner.getScheduleAttempts()); + assertEquals(1, execution.getContext().getQueryAttemptId()); assertEquals( Arrays.asList("begin", "analyze", "rollback", "begin", "analyze", "commit"), planner.getAnalysisEvents()); @@ -223,19 +224,25 @@ public DistributedQueryPlan doDistributionPlan( } @Override - public IScheduler doSchedule( + public IScheduler createScheduler( IAnalysis analysis, DistributedQueryPlan distributedPlan, MPPQueryContext context, QueryStateMachine stateMachine) { - scheduleAttempts++; - stateMachine.transitionToDispatching(); - if (alwaysFailDispatch || scheduleAttempts == 1) { - TSStatus dispatchFailure = RpcUtils.getStatus(TSStatusCode.DISPATCH_ERROR); - stateMachine.transitionToPendingRetry(dispatchFailure); - } else { - stateMachine.transitionToRunning(); - } + org.mockito.Mockito.doAnswer( + invocation -> { + scheduleAttempts++; + stateMachine.transitionToDispatching(); + if (alwaysFailDispatch || scheduleAttempts == 1) { + TSStatus dispatchFailure = RpcUtils.getStatus(TSStatusCode.DISPATCH_ERROR); + stateMachine.transitionToPendingRetry(dispatchFailure); + } else { + stateMachine.transitionToRunning(); + } + return null; + }) + .when(scheduler) + .start(); return scheduler; } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LastQueryCacheTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LastQueryCacheTest.java index 865c882c98063..486bdb63ea016 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LastQueryCacheTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LastQueryCacheTest.java @@ -59,7 +59,6 @@ import java.util.Arrays; import java.util.Collection; import java.util.Collections; -import java.util.HashMap; import java.util.HashSet; import static org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext; @@ -181,9 +180,7 @@ private LocalExecutionPlanContext createContext(IDataRegionForQuery dataRegion) QueryId queryId = new QueryId("last_query_cache_test"); FragmentInstanceId instanceId = new FragmentInstanceId(new PlanFragmentId(queryId, 0), "instance"); - DataNodeQueryContext queryContext = new DataNodeQueryContext(1); - HashMap queryContexts = new HashMap<>(); - queryContexts.put(queryId, queryContext); + DataNodeQueryContext queryContext = new DataNodeQueryContext(); FragmentInstanceContext instanceContext = spy( createFragmentInstanceContext( @@ -193,7 +190,7 @@ private LocalExecutionPlanContext createContext(IDataRegionForQuery dataRegion) 1, new UserEntity(666, "test", "127.0.0.1"), ZoneId.systemDefault()), dataRegion, (TimePredicate) null, - queryContexts, + queryContext, Long.MAX_VALUE, false, false)); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java index ea9323fdfdbd3..55ff4365362ce 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java @@ -118,7 +118,7 @@ public void testTablePlanEnablesSharingBeforeMemoryEstimationAndRootFactory() th long freeBefore = planner.getFreeMemoryForOperators(); try { List factories = - planner.plan(plan, new TypeProvider(), instance, new DataNodeQueryContext(1)); + planner.plan(plan, new TypeProvider(), instance, new DataNodeQueryContext()); assertEquals(1, factories.size()); CollectOperator collect = (CollectOperator) factories.get(0).getOperation(); // Check the actual reservation, not only the formula after planning has completed. @@ -182,7 +182,7 @@ public void testNullSessionRetainsTreeGeneratorAndIndependentInputQuotas() throw long freeBefore = planner.getFreeMemoryForOperators(); try { List factories = - planner.plan(plan, new TypeProvider(), instance, new DataNodeQueryContext(1)); + planner.plan(plan, new TypeProvider(), instance, new DataNodeQueryContext()); assertEquals(1, factories.size()); CollectOperator collect = (CollectOperator) factories.get(0).getOperation(); assertEquals(150 + collect.ramBytesUsed(), freeBefore - planner.getFreeMemoryForOperators()); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/PipelineBuilderTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/PipelineBuilderTest.java index 585418a48d3e6..8b83f4e473286 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/PipelineBuilderTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/PipelineBuilderTest.java @@ -1592,7 +1592,7 @@ private LocalExecutionPlanContext createLocalExecutionPlanContext(TypeProvider t fragmentInstanceContext.setDataRegion(dataRegion); return new LocalExecutionPlanContext( - typeProvider, fragmentInstanceContext, new DataNodeQueryContext(1)); + typeProvider, fragmentInstanceContext, new DataNodeQueryContext()); } private LocalExecutionPlanContext createLocalExecutionPlanContextWithQueryId( @@ -1610,7 +1610,7 @@ private LocalExecutionPlanContext createLocalExecutionPlanContextWithQueryId( fragmentInstanceContext.setDataRegion(dataRegion); return new LocalExecutionPlanContext( - typeProvider, fragmentInstanceContext, new DataNodeQueryContext(1)); + typeProvider, fragmentInstanceContext, new DataNodeQueryContext()); } /** diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterSchedulerLifecycleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterSchedulerLifecycleTest.java new file mode 100644 index 0000000000000..4caf1af1d6899 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterSchedulerLifecycleTest.java @@ -0,0 +1,116 @@ +/* + * 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.plan.scheduler; + +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.QueryStateMachine; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.junit.Test; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; + +import static com.google.common.util.concurrent.Futures.immediateFuture; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class ClusterSchedulerLifecycleTest { + private final QueryStateMachine state = + new QueryStateMachine( + new QueryId("scheduler_test"), + com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService()); + private final IFragInstanceDispatcher dispatcher = mock(IFragInstanceDispatcher.class); + private final IFragInstanceStateTracker tracker = mock(IFragInstanceStateTracker.class); + private final IQueryTerminator terminator = mock(IQueryTerminator.class); + private final ClusterScheduler scheduler = + new ClusterScheduler(state, dispatcher, tracker, terminator); + + @Test + public void stopBeforeStartPreventsDispatchAndTerminatesOnce() { + Throwable cause = new RuntimeException("stopped before dispatch"); + scheduler.stop(cause); + scheduler.start(); + scheduler.stop(cause); + verify(dispatcher).abort(); + verify(dispatcher, never()).dispatch(any(), any()); + verify(terminator).terminateImmediately(cause); + verify(tracker, never()).start(); + } + + @Test + public void stopDuringDispatchDoesNotWaitForDispatchOrStartTracker() throws Exception { + ExecutorService executor = Executors.newSingleThreadExecutor(); + CountDownLatch dispatching = new CountDownLatch(1); + CountDownLatch finishDispatch = new CountDownLatch(1); + when(dispatcher.dispatch(any(), any())) + .thenAnswer( + ignored -> { + dispatching.countDown(); + assertTrue(finishDispatch.await(10, TimeUnit.SECONDS)); + return immediateFuture(new FragInstanceDispatchResult(true)); + }); + try { + Future start = executor.submit(scheduler::start); + assertTrue(dispatching.await(10, TimeUnit.SECONDS)); + Throwable cause = new RuntimeException("stopped during dispatch"); + scheduler.stop(cause); + verify(dispatcher).abort(); + verify(terminator).terminateImmediately(cause); + finishDispatch.countDown(); + start.get(10, TimeUnit.SECONDS); + verify(tracker, never()).start(); + } finally { + finishDispatch.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + @Test + public void nonRetryableDispatchFailureTerminatesImmediately() { + when(dispatcher.dispatch(any(), any())) + .thenReturn( + immediateFuture( + new FragInstanceDispatchResult( + new TSStatus(TSStatusCode.QUERY_EXECUTION_MEMORY_NOT_ENOUGH.getStatusCode())))); + scheduler.start(); + Throwable cause = new RuntimeException("dispatch failed"); + scheduler.stop(cause); + verify(terminator).terminateImmediately(cause); + } + + @Test + public void normalCompletionKeepsGracePeriod() { + when(dispatcher.dispatch(any(), any())) + .thenReturn(immediateFuture(new FragInstanceDispatchResult(true))); + scheduler.start(); + scheduler.stop(null); + verify(terminator).terminate(null); + verify(terminator, never()).terminateImmediately(any()); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java new file mode 100644 index 0000000000000..50c7e63da967a --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java @@ -0,0 +1,133 @@ +/* + * 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.plan.scheduler; + +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.client.IClientManager; +import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient; +import org.apache.iotdb.commons.partition.ExecutorType; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceReq; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceResp; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.thrift.TException; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class FragmentInstanceDispatcherLifecycleTest { + @Test + public void abortPreventsDispatchAndRetainsOriginalAttemptMetadata() throws Exception { + QueryId queryId = new QueryId("aborted_dispatch"); + MPPQueryContext context = new MPPQueryContext(queryId); + context.setStartTime(1000); + context.setTimeOut(2000); + FragmentInstanceDispatcherImpl dispatcher = + new FragmentInstanceDispatcherImpl(QueryType.READ, context, null, null); + context.prepareForRetry(); + FragmentInstance instance = mock(FragmentInstance.class); + when(instance.getId()).thenReturn(new FragmentInstanceId(new PlanFragmentId(queryId, 0), "0")); + PlanFragment fragment = mock(PlanFragment.class); + when(instance.getFragment()).thenReturn(fragment); + dispatcher.abort(); + FragInstanceDispatchResult result = + dispatcher.dispatch(null, Collections.singletonList(instance)).get(); + assertFalse(result.isSuccessful()); + assertEquals( + TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), result.getFailureStatus().getCode()); + verify(instance, never()).getHostDataNode(); + verify(instance).setQueryAttempt(0, 3000); + } + + @Test + public void transportRetryReusesAdmissionDeadlineAndStopFreezesCutoff() throws Exception { + QueryId queryId = new QueryId("retry_admission"); + MPPQueryContext context = new MPPQueryContext(queryId); + context.setStartTime(System.currentTimeMillis()); + context.setTimeOut(Long.MAX_VALUE); + PlanFragment fragment = mock(PlanFragment.class); + FragmentInstance instance = + spy( + new FragmentInstance( + fragment, + new FragmentInstanceId(new PlanFragmentId(queryId, 0), "0"), + null, + QueryType.READ, + Long.MAX_VALUE, + null, + false, + false)); + doReturn(ByteBuffer.wrap(new byte[] {1})).when(instance).serializeToByteBuffer(); + TEndPoint endpoint = new TEndPoint("remote-host", 10730); + doReturn(new TDataNodeLocation().setInternalEndPoint(endpoint)) + .when(instance) + .getHostDataNode(); + instance.setExecutorType(mock(ExecutorType.class)); + IClientManager clients = + mock(IClientManager.class); + SyncDataNodeInternalServiceClient client = mock(SyncDataNodeInternalServiceClient.class); + when(clients.borrowClient(endpoint)).thenReturn(client); + List requests = new ArrayList<>(); + when(client.sendFragmentInstance(any())) + .thenAnswer( + invocation -> { + TSendFragmentInstanceReq request = invocation.getArgument(0); + requests.add(request.deepCopy()); + if (requests.size() == 1) { + throw new TException("response lost"); + } + return new TSendFragmentInstanceResp(true); + }); + FragmentInstanceDispatcherImpl dispatcher = + new FragmentInstanceDispatcherImpl(QueryType.READ, context, clients, null); + assertTrue(dispatcher.dispatch(null, Collections.singletonList(instance)).get().isSuccessful()); + assertEquals(2, requests.size()); + assertEquals(Long.MAX_VALUE, requests.get(0).getQueryDeadline()); + assertTrue(requests.get(0).getAdmissionDeadline() < Long.MAX_VALUE); + assertEquals(requests.get(0).getAdmissionDeadline(), requests.get(1).getAdmissionDeadline()); + long cutoff = dispatcher.getLastAdmissionDeadline(); + dispatcher.abort(); + assertFalse( + dispatcher.dispatch(null, Collections.singletonList(instance)).get().isSuccessful()); + assertEquals(2, requests.size()); + assertEquals(cutoff, dispatcher.getLastAdmissionDeadline()); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryAttemptProtocolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryAttemptProtocolTest.java new file mode 100644 index 0000000000000..d5a12286286d6 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryAttemptProtocolTest.java @@ -0,0 +1,94 @@ +/* + * 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.plan.scheduler; + +import org.apache.iotdb.mpp.rpc.thrift.TCancelQueryReq; +import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstance; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceReq; + +import org.apache.thrift.TDeserializer; +import org.apache.thrift.TSerializer; +import org.apache.thrift.protocol.TBinaryProtocol; +import org.apache.thrift.protocol.TField; +import org.apache.thrift.protocol.TProtocolUtil; +import org.apache.thrift.protocol.TType; +import org.apache.thrift.transport.TMemoryInputTransport; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.util.Collections; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; + +public class QueryAttemptProtocolTest { + @Test + public void legacyRequestsRemainReadableWithNoAttemptMetadata() throws Exception { + TSendFragmentInstanceReq original = + new TSendFragmentInstanceReq(new TFragmentInstance(ByteBuffer.wrap(new byte[] {1, 2, 3}))); + TSendFragmentInstanceReq received = new TSendFragmentInstanceReq(); + new TDeserializer().deserialize(received, new TSerializer().serialize(original)); + assertEquals(original, received); + assertFalse(received.isSetQueryAttemptId()); + assertFalse(received.isSetQueryDeadline()); + TCancelQueryReq cancel = new TCancelQueryReq("query", Collections.emptyList(), true); + TCancelQueryReq receivedCancel = new TCancelQueryReq(); + new TDeserializer().deserialize(receivedCancel, new TSerializer().serialize(cancel)); + assertEquals(cancel, receivedCancel); + assertFalse(receivedCancel.isSetQueryAttemptId()); + assertFalse(receivedCancel.isSetAdmissionDeadline()); + } + + @Test + public void attemptMetadataRoundTripsAndLegacyReaderCanSkipIt() throws Exception { + TFragmentInstance body = new TFragmentInstance(ByteBuffer.wrap(new byte[] {1, 2, 3})); + TSendFragmentInstanceReq request = + new TSendFragmentInstanceReq(body) + .setQueryAttemptId(2) + .setQueryDeadline(123456L) + .setAdmissionDeadline(123450L); + byte[] serialized = new TSerializer().serialize(request); + TSendFragmentInstanceReq received = new TSendFragmentInstanceReq(); + new TDeserializer().deserialize(received, serialized); + assertEquals(request, received); + // Simulate a receiver that only knows the original binary-body field. + TBinaryProtocol protocol = new TBinaryProtocol(new TMemoryInputTransport(serialized)); + TFragmentInstance legacyBody = new TFragmentInstance(); + protocol.readStructBegin(); + for (TField field = protocol.readFieldBegin(); + field.type != TType.STOP; + field = protocol.readFieldBegin()) { + if (field.id == 1) { + legacyBody.read(protocol); + } else { + TProtocolUtil.skip(protocol, field.type); + } + protocol.readFieldEnd(); + } + protocol.readStructEnd(); + assertEquals(body, legacyBody); + TCancelQueryReq cancel = + new TCancelQueryReq("query", Collections.emptyList(), true) + .setQueryAttemptId(2) + .setAdmissionDeadline(123450L); + TCancelQueryReq receivedCancel = new TCancelQueryReq(); + new TDeserializer().deserialize(receivedCancel, new TSerializer().serialize(cancel)); + assertEquals(cancel, receivedCancel); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java index 911074f2eca35..972031f9f4236 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java @@ -19,17 +19,28 @@ package org.apache.iotdb.db.queryengine.plan.scheduler; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.client.IClientManager; +import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.mpp.rpc.thrift.TCancelQueryReq; import org.junit.Before; import org.junit.Test; +import org.mockito.ArgumentCaptor; import java.util.Collections; import java.util.concurrent.Callable; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; @@ -53,7 +64,8 @@ public void setUp() { queryContext, Collections.emptyList(), null, - mock(IFragInstanceStateTracker.class)); + mock(IFragInstanceStateTracker.class), + () -> 0L); } @Test @@ -69,4 +81,42 @@ public void terminateImmediatelySkipsTheGracePeriod() { verify(scheduledExecutor).schedule(any(Callable.class), eq(0L), eq(TimeUnit.MILLISECONDS)); } + + @Test + public void closesAttemptEvenWhenNoFragmentIsReportedUnfinished() throws Exception { + final MPPQueryContext queryContext = mock(MPPQueryContext.class); + final QueryId queryId = new QueryId("attempt_to_close"); + when(queryContext.getQueryId()).thenReturn(queryId); + when(queryContext.getQueryAttemptId()).thenReturn(2); + when(queryContext.getQueryDeadline()).thenReturn(123456L); + when(queryContext.getEndPointBlackList()).thenReturn(Collections.emptySet()); + final TEndPoint endpoint = new TEndPoint("remote-host", 10730); + final FragmentInstance instance = mock(FragmentInstance.class); + when(instance.getHostDataNode()) + .thenReturn(new TDataNodeLocation().setInternalEndPoint(endpoint)); + when(instance.getId()).thenReturn(new FragmentInstanceId(new PlanFragmentId(queryId, 0), "0")); + final IClientManager clients = + mock(IClientManager.class); + final SyncDataNodeInternalServiceClient client = mock(SyncDataNodeInternalServiceClient.class); + when(clients.borrowClient(endpoint)).thenReturn(client); + final IFragInstanceStateTracker tracker = mock(IFragInstanceStateTracker.class); + when(tracker.filterUnFinishedFIs(any())).thenReturn(Collections.emptyList()); + final SimpleQueryTerminator terminator = + new SimpleQueryTerminator( + scheduledExecutor, + queryContext, + Collections.singletonList(instance), + clients, + tracker, + () -> 123456L); + // A delayed termination must retain the old attempt, even after the context starts a retry. + when(queryContext.getQueryAttemptId()).thenReturn(3); + terminator.terminateImmediately(new RuntimeException("dispatch failure")); + terminator.syncTerminate(); + ArgumentCaptor request = ArgumentCaptor.forClass(TCancelQueryReq.class); + verify(client).cancelQuery(request.capture()); + assertEquals(2, request.getValue().getQueryAttemptId()); + assertEquals(123456L, request.getValue().getAdmissionDeadline()); + assertTrue(request.getValue().getFragmentInstanceIds().isEmpty()); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java index 10daf11c721e8..153c23df4ade4 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java @@ -72,7 +72,7 @@ public void tt() { new LoadTsFileScheduler( distributedQueryPlan, mock(MPPQueryContext.class), - mock(QueryStateMachine.class), + newStateMachine("load_start"), mock(IClientManager.class), mock(IPartitionFetcher.class), false)); @@ -114,6 +114,25 @@ public void testStopWithAFailureFailsTheLoad() { Assert.assertEquals("the staging file is gone", stateMachine.getFailureMessage()); } + @Test + public void testStopBeforeStartReleasesConstructorMemoryOnce() throws Exception { + QueryStateMachine state = newStateMachine("load_stop_before_start"); + LoadTsFileScheduler scheduler = newScheduler(state); + java.lang.reflect.Field field = LoadTsFileScheduler.class.getDeclaredField("memoryBlock"); + field.setAccessible(true); + LoadTsFileDataCacheMemoryBlock block = (LoadTsFileDataCacheMemoryBlock) field.get(scheduler); + Method referenceCount = + LoadTsFileDataCacheMemoryBlock.class.getDeclaredMethod("getReferenceCount"); + referenceCount.setAccessible(true); + int before = (int) referenceCount.invoke(block); + scheduler.stop(null); + Assert.assertEquals(before - 1, referenceCount.invoke(block)); + scheduler.start(); + scheduler.stop(null); + Assert.assertEquals(before - 1, referenceCount.invoke(block)); + Assert.assertEquals(QueryState.QUEUED, state.getState()); + } + private LoadTsFileScheduler newScheduler(final QueryStateMachine stateMachine) { return new LoadTsFileScheduler( distributedQueryPlan, diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/transformation/builder/EvaluationDAGBuilderTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/transformation/builder/EvaluationDAGBuilderTest.java index 64bc431a7b9f2..956c4902d7ca8 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/transformation/builder/EvaluationDAGBuilderTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/transformation/builder/EvaluationDAGBuilderTest.java @@ -113,7 +113,7 @@ private Operator generateOperatorTree(String sql) { instance.getFragment().getPlanNodeTree(), instance.getFragment().getTypeProvider(), mockFIContext(queryId), - new DataNodeQueryContext(1)); + new DataNodeQueryContext()); return driverFactories.get(0).getOperation(); } catch (Exception e) { e.printStackTrace(); diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index 57afa70e72846..eb3753f07ccdf 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -142,6 +142,11 @@ struct TPlanNode { struct TSendFragmentInstanceReq { 1: required TFragmentInstance fragmentInstance 2: optional common.TConsensusGroupId consensusGroupId + // Attempt metadata is absent for a legacy coordinator. Do not change the binary FI body. + 3: optional i32 queryAttemptId + 4: optional i64 queryDeadline + // Admission expires independently of a potentially long query execution timeout. + 5: optional i64 admissionDeadline } struct TSendFragmentInstanceResp { @@ -187,6 +192,9 @@ struct TCancelQueryReq { 1: required string queryId 2: required list fragmentInstanceIds 3: required bool hasThrowable + // Close this attempt even if some of its planned FIs were never dispatched. + 4: optional i32 queryAttemptId + 5: optional i64 admissionDeadline } struct TCancelPlanFragmentReq { From a77d7ee68e560ed7ec5719752940bdb2f3385a34 Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Thu, 12 Nov 2026 11:44:07 +0800 Subject: [PATCH 2/4] [Query] Fix terminal cleanup and receiver-local admission Wake query status waiters on terminal cleanup while retaining retry-only cleanup. Replace coordinator wall-clock deadlines with receiver-issued admission tokens and local monotonic execution timeouts. Preserve reconnect retries, timeout status 720, and legacy peer fallback. Cover the review regressions, including the existing reporter CI cases; 126 targeted tests and both full-reactor locale compiles passed. --- .../iotdb/db/i18n/DataNodeQueryMessages.java | 5 +- .../iotdb/db/i18n/DataNodeQueryMessages.java | 5 +- .../impl/DataNodeInternalRPCServiceImpl.java | 15 +- .../queryengine/common/MPPQueryContext.java | 5 - .../db/queryengine/common/QueryAttemptId.java | 1 + .../execution/QueryStateMachine.java | 3 +- .../fragment/FragmentInstanceManager.java | 9 +- .../fragment/QueryAttemptRegistry.java | 196 ++++++++++------- .../plan/execution/QueryExecution.java | 37 +++- .../plan/planner/plan/FragmentInstance.java | 22 +- .../plan/scheduler/ClusterScheduler.java | 8 +- .../FragmentInstanceDispatcherImpl.java | 154 +++++++++++--- .../plan/scheduler/IQueryTerminator.java | 6 +- .../plan/scheduler/SimpleQueryTerminator.java | 29 +-- .../fragment/FragmentInstanceManagerTest.java | 13 +- .../fragment/QueryAttemptAdmissionTest.java | 184 ++++++++++++++++ .../fragment/QueryAttemptRegistryTest.java | 124 +++++++++-- .../QueryExecutionLifecycleTest.java | 36 ++++ .../ClusterSchedulerLifecycleTest.java | 14 ++ ...agmentInstanceDispatcherLifecycleTest.java | 198 ++++++++++++++---- .../scheduler/QueryAttemptProtocolTest.java | 11 +- .../scheduler/QueryTerminalCleanupTest.java | 152 ++++++++++++++ .../scheduler/SimpleQueryTerminatorTest.java | 9 +- .../src/main/thrift/datanode.thrift | 16 +- 24 files changed, 1008 insertions(+), 244 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptAdmissionTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryTerminalCleanupTest.java diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 7e38eaa85fda6..27e1ae6de3320 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -20,7 +20,10 @@ 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 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"; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index c6aba75034517..e719b5e8a06fd 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -20,7 +20,10 @@ 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 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"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java index d4b2d24647389..4b7e0bd03dc02 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java @@ -323,6 +323,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 +511,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()); @@ -533,9 +539,8 @@ 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()); + if (req.isSetQueryAttemptId() && req.isSetAdmissionToken()) { + fragmentInstance.setQueryAttempt(req.getQueryAttemptId(), req.getAdmissionToken()); } } catch (final Exception e) { LOGGER.warn(DataNodeMiscMessages.DESERIALIZE_FRAGMENT_INSTANCE_FAILED, e); @@ -616,12 +621,12 @@ public TFragmentInstanceInfoResp fetchFragmentInstanceInfo(TFetchFragmentInstanc @Override public TCancelResp cancelQuery(TCancelQueryReq req) { try (SetThreadName threadName = new SetThreadName(req.getQueryId())) { - if (req.isSetQueryAttemptId() && req.isSetAdmissionDeadline()) { + if (req.isSetQueryAttemptId() && req.isSetAdmissionToken()) { FragmentInstanceManager.getInstance() .closeQueryAttempt( new QueryId(req.getQueryId()), req.getQueryAttemptId(), - req.getAdmissionDeadline(), + req.getAdmissionToken(), req.hasThrowable); } List taskIds = 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 9c1f34bc16b8c..ad9c932164cbe 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 @@ -55,7 +55,6 @@ 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; @@ -311,10 +310,6 @@ public int getQueryAttemptId() { return queryAttemptId; } - public long getQueryDeadline() { - return LongMath.saturatedAdd(startTime, timeOut); - } - public ExternalTsFileQueryResource createExternalTsFileQueryResource( String tableName, List tsFilePaths, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/QueryAttemptId.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/QueryAttemptId.java index a76cdd7a28153..802505d695828 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/QueryAttemptId.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/QueryAttemptId.java @@ -16,6 +16,7 @@ * 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. */ 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 153c17536d541..7cfaee867a02e 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 @@ -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; } 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 e794e3b4a0dde..2ea86ce9493de 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 @@ -49,6 +49,7 @@ import org.apache.iotdb.db.storageengine.dataregion.IDataRegionForQuery; import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.mpp.rpc.thrift.TFetchFragmentInstanceStatisticsResp; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.iotdb.rpc.TSStatusCode; import org.apache.iotdb.udf.api.exception.UDFException; import org.apache.iotdb.udf.api.exception.UDFTypeMismatchException; @@ -450,8 +451,12 @@ public FragmentInstanceInfo cancelTask(FragmentInstanceId instanceId, boolean ha } public void closeQueryAttempt( - QueryId queryId, int attemptId, long deadline, boolean hasThrowable) { - queryAttempts.close(new QueryAttemptId(queryId, attemptId), deadline, hasThrowable); + QueryId queryId, int attemptId, TQueryAdmissionToken token, boolean hasThrowable) { + queryAttempts.close(new QueryAttemptId(queryId, attemptId), token, hasThrowable); + } + + public TQueryAdmissionToken fetchQueryAdmissionToken(long validityMs) { + return queryAttempts.issueAdmissionToken(validityMs); } QueryAttemptRegistry getQueryAttempts() { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java index 1b36d4e05664e..0624cf7dec135 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java @@ -16,13 +16,16 @@ * specific language governing permissions and limitations * under the License. */ + package org.apache.iotdb.db.queryengine.execution.fragment; import org.apache.iotdb.commons.exception.IoTDBRuntimeException; +import org.apache.iotdb.commons.exception.QueryTimeoutException; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.queryengine.common.QueryAttemptId; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.iotdb.rpc.TSStatusCode; import org.slf4j.Logger; @@ -32,39 +35,50 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; import java.util.function.LongSupplier; -import static com.google.common.math.LongMath.saturatedAdd; - /** - * Owns shared DataNode state for one execution attempt. All registration/closing bookkeeping is - * serialized by compute on the attempt key; planning and state transitions run outside that lock. - * Closed entries retain no query data after their last actual user releases it. They only fence - * late dispatch until the request's immutable admission deadline, after which the request itself is - * rejected, even if the entry has been removed. + * Owns shared state for a query attempt. Registration and closure serialize on the attempt key; + * planning and state transitions run outside that lock. Admission tokens are issued and checked + * against this DataNode process's monotonic clock, without comparing clocks across nodes. Token + * issuance is stateless: a delayed issuance RPC cannot create or reopen an attempt. */ class QueryAttemptRegistry { private static final Logger LOGGER = LoggerFactory.getLogger(QueryAttemptRegistry.class); + private static final long MAX_ADMISSION_VALIDITY_MS = TimeUnit.DAYS.toMillis(1); private final Map attempts = new ConcurrentHashMap<>(); - private final LongSupplier clock; - // Expiring a fence must remain irreversible even if the wall clock subsequently goes backwards. - private final AtomicLong observedTime = new AtomicLong(); + private final LongSupplier ticker; + private final String issuer = UUID.randomUUID().toString(); QueryAttemptRegistry() { - this(System::currentTimeMillis); + this(System::nanoTime); + } + + QueryAttemptRegistry(LongSupplier ticker) { + this.ticker = ticker; } - QueryAttemptRegistry(LongSupplier clock) { - this.clock = clock; + TQueryAdmissionToken issueAdmissionToken(long validityMs) { + long boundedValidityMs = Math.max(1, Math.min(validityMs, MAX_ADMISSION_VALIDITY_MS)); + return new TQueryAdmissionToken( + issuer, + ticker.getAsLong() + TimeUnit.MILLISECONDS.toNanos(boundedValidityMs), + boundedValidityMs); } - private long now() { - return observedTime.accumulateAndGet(clock.getAsLong(), Math::max); + private boolean issuedHere(TQueryAdmissionToken token) { + return token != null && issuer.equals(token.getIssuer()); + } + + private static boolean expired(long deadline, long now) { + // Subtraction also handles nanoTime wraparound; token validity is bounded to one day. + return now - deadline >= 0; } Registration register( @@ -75,31 +89,50 @@ Registration register( instance.hasQueryAttempt() ? instance.getQueryAttemptId() : QueryAttemptId.LEGACY_ATTEMPT); - long executionDeadline = - instance.hasQueryAttempt() - ? instance.getQueryDeadline() - : saturatedAdd(now(), instance.getTimeOut()); - long admissionDeadline = - instance.hasQueryAttempt() - ? Math.min(executionDeadline, instance.getAdmissionDeadline()) - : executionDeadline; + TQueryAdmissionToken token = instance.hasQueryAttempt() ? instance.getAdmissionToken() : null; + if (!key.isLegacy() && !issuedHere(token)) { + throw new IoTDBRuntimeException( + String.format( + DataNodeQueryMessages + .EXCEPTION_INVALID_QUERY_ADMISSION_ISSUER_FOR_ATTEMPT_ARG_9D1C4E9C, + key), + TSStatusCode.DISPATCH_ERROR.getStatusCode()); + } AtomicReference result = new AtomicReference<>(); attempts.compute( key, (ignored, existing) -> { - long currentTime = now(); - if (admissionDeadline <= currentTime - || (existing != null - && (existing.closed || existing.executionDeadline <= currentTime))) { - throw rejected(key, admissionDeadline <= currentTime); + long now = ticker.getAsLong(); + if (token != null && expired(token.getDeadlineNanos(), now)) { + // Admission expiry does not mean the query exhausted its execution budget. Retry the + // query with a fresh attempt/token instead of renewing the old request's cutoff. + throw new IoTDBRuntimeException( + String.format( + DataNodeQueryMessages + .EXCEPTION_QUERY_ATTEMPT_ARG_ADMISSION_EXPIRED_DEADLINE_ARG_NS_NOW_ARG_NS_16A9846C, + key, + token.getDeadlineNanos(), + now), + TSStatusCode.DISPATCH_ERROR.getStatusCode()); + } + if (existing != null && existing.closed) { + throw new IoTDBRuntimeException( + String.format( + DataNodeQueryMessages.EXCEPTION_QUERY_ATTEMPT_ARG_IS_CLOSED_12143E89, key), + TSStatusCode.QUERY_WAS_KILLED.getStatusCode()); + } + long timeoutNanos = TimeUnit.MILLISECONDS.toNanos(Math.max(0, instance.getTimeOut())); + if (timeoutNanos == 0 || (existing != null && existing.executionExpired(now))) { + throw new QueryTimeoutException(); } Entry entry = existing == null - ? new Entry(executionDeadline, Math.max(1, instance.getDataNodeFINum())) + ? new Entry(now, timeoutNanos, Math.max(1, instance.getDataNodeFINum())) : existing; - // A transport retry must never extend the lifetime of an existing attempt. - entry.executionDeadline = Math.min(entry.executionDeadline, executionDeadline); - entry.fenceDeadline = Math.max(entry.fenceDeadline, admissionDeadline); + entry.shortenExecutionTimeout(now, timeoutNanos); + if (token != null) { + entry.extendFence(token.getDeadlineNanos()); + } if (entry.active.containsKey(instance.getId())) { throw new IoTDBRuntimeException( String.format( @@ -120,83 +153,81 @@ Registration register( return result.get(); } - void close(QueryAttemptId key, long deadline, boolean hasThrowable) { + void close(QueryAttemptId key, TQueryAdmissionToken token, boolean hasThrowable) { + // A token from before a restart must not close a new incarnation's state. + if (!issuedHere(token)) { + return; + } List toCancel = new ArrayList<>(); - long currentTime = now(); attempts.compute( key, (ignored, existing) -> { - if (existing == null && deadline <= currentTime) { + long now = ticker.getAsLong(); + if (existing == null && expired(token.getDeadlineNanos(), now)) { return null; } - Entry entry = existing == null ? new Entry(deadline, 0) : existing; - entry.closed = true; - entry.fenceDeadline = Math.max(entry.fenceDeadline, deadline); - if (entry.active != null) { + Entry entry = existing == null ? new Entry(now, 0, 0) : existing; + if (!entry.closed && entry.active != null) { toCancel.addAll(entry.active.values()); } - return retireIfUnused(key, entry, currentTime); + entry.closed = true; + entry.extendFence(token.getDeadlineNanos()); + return retireIfUnused(key, entry, now); }); - terminate(toCancel, hasThrowable); + terminate(toCancel, hasThrowable, false); } void removeExpired() { - long currentTime = now(); - List toCancel = new ArrayList<>(); + List timedOut = new ArrayList<>(); for (QueryAttemptId key : attempts.keySet()) { attempts.computeIfPresent( key, (ignored, entry) -> { - if ((entry.closed ? entry.fenceDeadline : entry.executionDeadline) > currentTime) { - return entry; - } - entry.closed = true; - if (entry.active != null) { - toCancel.addAll(entry.active.values()); + long now = ticker.getAsLong(); + if (!entry.closed && entry.executionExpired(now)) { + entry.closed = true; + timedOut.addAll(entry.active.values()); } - return retireIfUnused(key, entry, currentTime); + // Closed entries were already terminated by their first closer. Fence expiry must not + // add a timeout cause to an explicitly cancelled or successfully completed FI. + return retireIfUnused(key, entry, now); }); } - terminate(toCancel, true); + terminate(timedOut, true, true); } - private static void terminate(List registrations, boolean hasThrowable) { + private static void terminate( + List registrations, boolean hasThrowable, boolean timeout) { for (Registration registration : registrations) { try { - if (hasThrowable) { + if (timeout) { + if (!registration.stateMachine.getState().isDone()) { + registration.stateMachine.failed(new QueryTimeoutException()); + } + } else if (hasThrowable) { registration.stateMachine.cancel(); } else { registration.stateMachine.finished(); } } catch (RuntimeException e) { - // One failed notification must not prevent cancellation of the other FIs. LOGGER.warn(DataNodeQueryMessages.ERRORS_FINISHING_FI_PROCESS, e); } } } - private static Entry retireIfUnused(QueryAttemptId key, Entry entry, long currentTime) { + private static Entry retireIfUnused(QueryAttemptId key, Entry entry, long now) { if ((entry.active == null || entry.active.isEmpty()) && (entry.closed || entry.remaining == 0)) { entry.closed = true; entry.context = null; entry.active = null; - // Legacy requests have no admission deadline on the wire and may reuse the queryId on - // retry. Preserve their existing behavior instead of fencing a later legacy execution. - if (key.isLegacy() || entry.fenceDeadline <= currentTime) { + if (key.isLegacy() || expired(entry.fenceDeadline, now)) { return null; } } return entry; } - private static IoTDBRuntimeException rejected(QueryAttemptId key, boolean expired) { - return new IoTDBRuntimeException( - String.format( - DataNodeQueryMessages.EXCEPTION_QUERY_ATTEMPT_ARG_IS_CLOSED_OR_EXPIRED_F45E3DCF, key), - (expired ? TSStatusCode.QUERY_TIMEOUT : TSStatusCode.QUERY_WAS_KILLED).getStatusCode()); - } - int getRetainedEntryCount() { return attempts.size(); } @@ -213,17 +244,42 @@ boolean hasSharedContext(QueryAttemptId key) { } private static class Entry { - private long executionDeadline; + private long executionStart; + private long executionTimeout; private long fenceDeadline; + private boolean hasFence; private int remaining; private boolean closed; private DataNodeQueryContext context; private Map active = new HashMap<>(); - private Entry(long deadline, int remaining) { - this.executionDeadline = deadline; + private Entry(long start, long timeout, int remaining) { + this.executionStart = start; + this.executionTimeout = timeout; this.remaining = remaining; } + + private boolean executionExpired(long now) { + return executionTimeout != Long.MAX_VALUE && now - executionStart >= executionTimeout; + } + + private void shortenExecutionTimeout(long now, long timeout) { + long remainingNanos = + executionTimeout == Long.MAX_VALUE + ? Long.MAX_VALUE + : Math.max(0, executionTimeout - (now - executionStart)); + if (timeout < remainingNanos) { + executionStart = now; + executionTimeout = timeout; + } + } + + private void extendFence(long deadline) { + if (!hasFence || deadline - fenceDeadline > 0) { + fenceDeadline = deadline; + hasFence = true; + } + } } final class Registration implements AutoCloseable { @@ -263,7 +319,7 @@ public void close() { return current; } entry.remaining--; - return retireIfUnused(key, entry, now()); + return retireIfUnused(key, entry, ticker.getAsLong()); }); } } 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 1dbffd48ff2f8..fb7c393564ba2 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 @@ -70,6 +70,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Throwables.throwIfUnchecked; @@ -94,7 +95,8 @@ public class QueryExecution implements IQueryExecution { private static final long RETRY_INTERVAL_IN_MS = 2000; private int retryCount = 0; private final MPPQueryContext context; - private volatile SchedulingAttempt schedulingAttempt = new SchedulingAttempt(); + private final AtomicReference schedulingAttempt = + new AtomicReference<>(new SchedulingAttempt()); private final QueryStateMachine stateMachine; private final IPlanner planner; @@ -267,14 +269,14 @@ private ExecutionResult retry() { LOGGER.warn(DataNodeQueryMessages.REACHMAXRETRYCOUNT); // Keep the same ownership order as an ordinary retry: plans and scheduler stop using the // analysis working state before the journal restores parser-owned objects. - this.stopAndCleanup(stateMachine.getFailureException()); + this.stopAttemptAndRelease(stateMachine.getFailureException()); planner.rollbackAnalysisAttempt(); stateMachine.transitionToFailed(); return getStatus(); } LOGGER.warn(DataNodeQueryMessages.ERROR_WHEN_EXECUTING_QUERY, stateMachine.getFailureMessage()); // stop and clean up resources the QueryExecution used - this.stopAndCleanup(stateMachine.getFailureException()); + this.stopAttemptAndRelease(stateMachine.getFailureException()); // The generated plans no longer read the statement after cleanup. Restore the parser-owned // state before the next analysis attempt so retry observes the same SQL as the first attempt. planner.rollbackAnalysisAttempt(); @@ -293,7 +295,7 @@ private ExecutionResult retry() { // clear runtime variables in MPPQueryContext context.prepareForRetry(); // re-stop - this.schedulingAttempt = new SchedulingAttempt(); + this.schedulingAttempt.set(new SchedulingAttempt()); this.resultHandleCleanUp.compareAndSet(true, false); // re-analyze the query planner.beginAnalysisAttempt(); @@ -336,7 +338,7 @@ private IAnalysis analyze(MPPQueryContext context) { private void schedule() { final long startTime = System.nanoTime(); - SchedulingAttempt attempt = schedulingAttempt; + SchedulingAttempt attempt = schedulingAttempt.get(); IScheduler scheduler = planner.createScheduler(analysis, distributedPlan, context, stateMachine); if (stateMachine.getState().isDone()) { @@ -392,7 +394,7 @@ private String printFragmentInstances(List instances) { // Stop the workers for this query @Override public void stop(Throwable t) { - schedulingAttempt.stop(t); + schedulingAttempt.get().stop(t); } /** One-shot handoff between scheduler creation and concurrent stop requests. */ @@ -470,8 +472,27 @@ private void cleanUpResultHandle() { // Stop the query and clean up all the resources this query occupied @Override public void stopAndCleanup(Throwable t) { - stop(t); - releaseResource(t); + try { + // Terminal cleanup (including the stale-query sweep) must wake getStatus even if the + // dispatcher has not returned yet. Retry cleanup only stops the old execution attempt. + if (!stateMachine.getState().isDone()) { + if (t == null) { + stateMachine.transitionToAborted(); + } else { + stateMachine.transitionToFailed(t); + } + } + } finally { + stopAttemptAndRelease(t); + } + } + + private void stopAttemptAndRelease(Throwable t) { + try { + stop(t); + } finally { + releaseResource(t); + } } /** Release the resources that current QueryExecution hold with a specified exception */ diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java index 295ff71dd1cb9..c4b53ff8aba4a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/FragmentInstance.java @@ -33,6 +33,7 @@ import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeUtil; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.ReadWriteIOUtils; @@ -74,32 +75,23 @@ public class FragmentInstance implements IConsensusRequest { // Admission metadata travels in the optional RPC envelope, not the binary FI body, so older // DataNodes can still deserialize the plan. Local dispatch sets the same fields directly. private int queryAttemptId = -1; - private long queryDeadline; - private long admissionDeadline; + private TQueryAdmissionToken admissionToken; - public void setQueryAttempt(int queryAttemptId, long queryDeadline) { + public void setQueryAttempt(int queryAttemptId, TQueryAdmissionToken admissionToken) { this.queryAttemptId = queryAttemptId; - this.queryDeadline = queryDeadline; + this.admissionToken = admissionToken; } public boolean hasQueryAttempt() { - return queryAttemptId >= 0 && queryDeadline > 0; + return queryAttemptId >= 0 && admissionToken != null; } public int getQueryAttemptId() { return queryAttemptId; } - public long getQueryDeadline() { - return queryDeadline; - } - - public void setAdmissionDeadline(long admissionDeadline) { - this.admissionDeadline = admissionDeadline; - } - - public long getAdmissionDeadline() { - return admissionDeadline; + public TQueryAdmissionToken getAdmissionToken() { + return admissionToken; } private boolean isHighestPriority; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java index c15de059af9b9..589923d26810d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java @@ -105,7 +105,7 @@ public ClusterScheduler( instances, syncInternalServiceClientManager, stateTracker, - fragmentDispatcher::getLastAdmissionDeadline); + fragmentDispatcher::getIssuedAdmissionTokens); } } @@ -134,7 +134,7 @@ && isQuery() @Override public void start() { synchronized (this) { - if (started || stopped || stateMachine.getState().isDone()) { + if (started || stopped) { return; } started = true; @@ -206,8 +206,8 @@ public void stop(Throwable t) { } if (queryTerminator != null) { if (terminateImmediately) { - // The FragmentInstances dispatched before the failure can't finish on their own and hold - // operator memory until they are cancelled, so don't wait for the grace period. + // Dispatch failed, or stop arrived before dispatch completed. The FIs already sent may + // depend on unsent FIs and hold operator memory, so do not wait for the grace period. queryTerminator.terminateImmediately(t); } else { queryTerminator.terminate(t); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java index 76f2542f7f295..46d7119911b63 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java @@ -46,6 +46,7 @@ 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; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceManager; import org.apache.iotdb.db.queryengine.plan.Coordinator; import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; @@ -55,6 +56,7 @@ import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstance; import org.apache.iotdb.mpp.rpc.thrift.TPlanNode; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.iotdb.mpp.rpc.thrift.TSendBatchPlanNodeReq; import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceReq; import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceResp; @@ -63,7 +65,7 @@ import org.apache.iotdb.rpc.RpcUtils; import org.apache.iotdb.rpc.TSStatusCode; -import com.google.common.math.LongMath; +import org.apache.thrift.TApplicationException; import org.apache.thrift.TException; import org.apache.thrift.transport.TTransportException; import org.apache.tsfile.external.commons.lang3.exception.ExceptionUtils; @@ -74,13 +76,18 @@ import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.LongSupplier; import java.util.stream.Collectors; import static com.google.common.util.concurrent.Futures.immediateFuture; @@ -96,11 +103,14 @@ public class FragmentInstanceDispatcherImpl implements IFragInstanceDispatcher { private final QueryType type; private final MPPQueryContext queryContext; private final int queryAttemptId; - private final long queryDeadline; private final long admissionWindowMs; - private long lastAdmissionDeadline; + private final LongSupplier ticker; + private final Map admissionCache = new ConcurrentHashMap<>(); + // Protected by this, also used to serialize token publication against abort(). + private final Map issuedAdmissionTokens = new HashMap<>(); private volatile boolean aborted; - private volatile LinkedBlockingQueue> dispatchQueue; + private final AtomicReference>> dispatchQueue = + new AtomicReference<>(); private final String localhostIpAddr; private final int localhostInternalPort; private final IClientManager @@ -126,15 +136,32 @@ public FragmentInstanceDispatcherImpl( IClientManager syncInternalServiceClientManager, IClientManager asyncInternalServiceClientManager) { + this( + type, + queryContext, + syncInternalServiceClientManager, + asyncInternalServiceClientManager, + System::nanoTime); + } + + FragmentInstanceDispatcherImpl( + QueryType type, + MPPQueryContext queryContext, + IClientManager syncInternalServiceClientManager, + IClientManager + asyncInternalServiceClientManager, + LongSupplier ticker) { this.type = type; this.queryContext = queryContext; this.queryAttemptId = queryContext.getQueryAttemptId(); - this.queryDeadline = queryContext.getQueryDeadline(); - // A zero socket timeout means unlimited waiting, but admission fences must remain finite. + this.ticker = ticker; + // Budget both the initial RPC and the supported reconnect retry. A zero socket timeout + // still needs a finite token lifetime; actual query execution keeps its own timeout. this.admissionWindowMs = - COMMON_CONFIG.getDnConnectionTimeoutInMS() > 0 - ? COMMON_CONFIG.getDnConnectionTimeoutInMS() - : TimeUnit.MINUTES.toMillis(1); + 2L + * (COMMON_CONFIG.getDnConnectionTimeoutInMS() > 0 + ? COMMON_CONFIG.getDnConnectionTimeoutInMS() + : TimeUnit.MINUTES.toMillis(1)); this.syncInternalServiceClientManager = syncInternalServiceClientManager; this.asyncInternalServiceClientManager = asyncInternalServiceClientManager; this.localhostIpAddr = IoTDBDescriptor.getInstance().getConfig().getInternalAddress(); @@ -149,9 +176,6 @@ public FragmentInstanceDispatcherImpl( public Future dispatch( SubPlan root, List instances) { if (isQuery()) { - for (FragmentInstance instance : instances) { - instance.setQueryAttempt(queryAttemptId, queryDeadline); - } return instances.size() == 1 || root == null ? dispatchRead(instances) : topologicalParallelDispatchRead(root, instances); @@ -164,7 +188,7 @@ private Future topologicalParallelDispatchRead( SubPlan root, List instances) { long startTime = System.nanoTime(); LinkedBlockingQueue> queue = new LinkedBlockingQueue<>(); - dispatchQueue = queue; + dispatchQueue.set(queue); List> futures = new ArrayList<>(instances.size()); queue.add(new Pair<>(root, true)); try { @@ -207,7 +231,7 @@ private Future topologicalParallelDispatchRead( RpcUtils.getStatus( TSStatusCode.INTERNAL_SERVER_ERROR, UNEXPECTED_ERRORS + t.getMessage()))); } finally { - dispatchQueue = null; + dispatchQueue.set(null); long queryDispatchReadTime = System.nanoTime() - startTime; QUERY_EXECUTION_METRIC_SET.recordExecutionCost(DISPATCH_READ, queryDispatchReadTime); queryContext.recordDispatchCost(queryDispatchReadTime); @@ -482,11 +506,11 @@ private long getRetrySleepTime(int retryTimes) { } private void dispatchOneInstance(FragmentInstance instance) - throws FragmentInstanceDispatchException { + throws FragmentInstanceDispatchException, TException { checkNotAborted(); TEndPoint endPoint = instance.getHostDataNode().getInternalEndPoint(); if (isDispatchedToLocal(endPoint)) { - prepareAdmission(instance); + prepareAdmission(instance, endPoint, null); dispatchLocally(instance); } else { dispatchRemote(instance, endPoint); @@ -509,18 +533,26 @@ private void dispatchRemoteHelper(final FragmentInstance instance, final TEndPoi instance.getFragment().getPlanNodeTree(), instance.getSessionInfo() == null ? null : instance.getSessionInfo().getUserName()); boolean transferAttemptRecorded = false; + boolean legacyAdmissionFallback = false; try (final SyncDataNodeInternalServiceClient client = syncInternalServiceClientManager.borrowClient(endPoint)) { switch (instance.getType()) { case READ: final TSendFragmentInstanceReq sendFragmentInstanceReq = new TSendFragmentInstanceReq(new TFragmentInstance(instance.serializeToByteBuffer())); - prepareAdmission(instance); + try { + prepareAdmission(instance, endPoint, client); + } catch (TApplicationException e) { + if (e.getType() != TApplicationException.UNKNOWN_METHOD) { + throw e; + } + legacyAdmissionFallback = true; + throw e; + } if (instance.hasQueryAttempt()) { sendFragmentInstanceReq .setQueryAttemptId(instance.getQueryAttemptId()) - .setQueryDeadline(instance.getQueryDeadline()) - .setAdmissionDeadline(instance.getAdmissionDeadline()); + .setAdmissionToken(instance.getAdmissionToken()); } if (instance.getExecutorType().isStorageExecutor()) { sendFragmentInstanceReq.setConsensusGroupId( @@ -626,6 +658,13 @@ private void dispatchRemoteHelper(final FragmentInstance instance, final TEndPoi } throw e; } catch (TException e) { + if (legacyAdmissionFallback) { + // UNKNOWN_METHOD invalidates the pooled connection. The resource scope above must + // return it before borrowing again, including when the pool has only one slot. + // Capability negotiation does not consume the FI's normal reconnect retry. + dispatchRemoteHelper(instance, endPoint); + return; + } if (!transferAttemptRecorded) { recordWriteTransferFailureIfNecessary(endPoint, containsUserData, e); } @@ -799,27 +838,78 @@ private void dispatchLocally(final FragmentInstance instance) @Override public synchronized void abort() { aborted = true; - LinkedBlockingQueue> queue = dispatchQueue; + LinkedBlockingQueue> queue = dispatchQueue.get(); if (queue != null) { - queue.offer(new Pair<>(null, false)); + queue.add(new Pair<>(null, false)); } } - private synchronized void prepareAdmission(FragmentInstance instance) - throws FragmentInstanceDispatchException { + private void prepareAdmission( + FragmentInstance instance, TEndPoint endPoint, SyncDataNodeInternalServiceClient client) + throws FragmentInstanceDispatchException, TException { checkNotAborted(); - long deadline = instance.getAdmissionDeadline(); - if (deadline == 0) { - deadline = - Math.min( - queryDeadline, LongMath.saturatedAdd(System.currentTimeMillis(), admissionWindowMs)); - instance.setAdmissionDeadline(deadline); + if (instance.hasQueryAttempt()) { + // A resend must retain its old token, even after a socket timeout or a receiver restart. + // Expired/old-incarnation tokens are retryable dispatch failures, not execution timeouts. + return; + } + CachedAdmission cached = + admissionCache.computeIfAbsent(endPoint, ignored -> new CachedAdmission()); + // Serialize token acquisition per destination, but never hold the abort/publication lock + // across an RPC. Other FIs on this host reuse the token while it is conservatively valid. + synchronized (cached) { + checkNotAborted(); + if (cached.legacy) { + return; + } + if (cached.token == null + || ticker.getAsLong() - cached.requestStart + >= TimeUnit.MILLISECONDS.toNanos(cached.token.getValidityMs())) { + long requestStart = ticker.getAsLong(); + try { + cached.token = + client == null + ? FragmentInstanceManager.getInstance() + .fetchQueryAdmissionToken(admissionWindowMs) + : client.fetchQueryAdmissionToken(admissionWindowMs); + cached.requestStart = requestStart; + } catch (TException e) { + Throwable cause = ExceptionUtils.getRootCause(e); + if (!(cause instanceof TApplicationException) + || ((TApplicationException) cause).getType() + != TApplicationException.UNKNOWN_METHOD) { + throw e; + } + // The client interceptor wraps application exceptions in TException. Mark this peer + // as legacy, then let the caller replace the connection invalidated by that handler. + checkNotAborted(); + cached.legacy = true; + throw (TApplicationException) cause; + } + } + synchronized (this) { + checkNotAborted(); + instance.setQueryAttempt(queryAttemptId, cached.token); + TQueryAdmissionToken previous = issuedAdmissionTokens.get(endPoint); + if (previous == null + || !previous.getIssuer().equals(cached.token.getIssuer()) + || cached.token.getDeadlineNanos() - previous.getDeadlineNanos() > 0) { + issuedAdmissionTokens.put(endPoint, cached.token); + } + } } - lastAdmissionDeadline = Math.max(lastAdmissionDeadline, deadline); } - public synchronized long getLastAdmissionDeadline() { - return lastAdmissionDeadline; + public synchronized Map getIssuedAdmissionTokens() { + Map snapshot = new HashMap<>(); + issuedAdmissionTokens.forEach((endPoint, token) -> snapshot.put(endPoint, token.deepCopy())); + return snapshot; + } + + private static class CachedAdmission { + private TQueryAdmissionToken token; + private long requestStart; + private boolean legacy; } private void checkNotAborted() throws FragmentInstanceDispatchException { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java index 3be8958684732..bd36a7d0046aa 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java @@ -26,9 +26,9 @@ public interface IQueryTerminator { Future terminate(Throwable t); /** - * Cancels the unfinished FragmentInstances of the query without a grace period. Used when the - * dispatch failed: the dispatched FragmentInstances can't finish on their own and hold their - * resources until they are cancelled. + * Cancels the unfinished FragmentInstances without a grace period when dispatch failed or was + * stopped before completing. FIs already sent may depend on unsent FIs and hold resources until + * they are cancelled. */ Future terminateImmediately(Throwable t); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java index d0b6da402e266..f01a5c4974f53 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java @@ -31,11 +31,13 @@ import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceManager; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; import org.apache.iotdb.mpp.rpc.thrift.TCancelQueryReq; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.thrift.TException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -43,7 +45,7 @@ import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; -import java.util.function.LongSupplier; +import java.util.function.Supplier; import java.util.stream.Collectors; public class SimpleQueryTerminator implements IQueryTerminator { @@ -52,8 +54,8 @@ public class SimpleQueryTerminator implements IQueryTerminator { protected ScheduledExecutorService scheduledExecutor; private final QueryId queryId; private final int queryAttemptId; - private final LongSupplier admissionDeadlineSupplier; - private long admissionDeadline; + private final Supplier> admissionTokensSupplier; + private Map admissionTokens = Collections.emptyMap(); private final MPPQueryContext queryContext; private final IFragInstanceStateTracker stateTracker; @@ -72,11 +74,11 @@ public SimpleQueryTerminator( final IClientManager internalServiceClientManager, final IFragInstanceStateTracker stateTracker, - final LongSupplier admissionDeadlineSupplier) { + final Supplier> admissionTokensSupplier) { this.scheduledExecutor = scheduledExecutor; this.queryId = queryContext.getQueryId(); this.queryAttemptId = queryContext.getQueryAttemptId(); - this.admissionDeadlineSupplier = admissionDeadlineSupplier; + this.admissionTokensSupplier = admissionTokensSupplier; this.queryContext = queryContext; this.internalServiceClientManager = internalServiceClientManager; this.stateTracker = stateTracker; @@ -108,9 +110,9 @@ private Future terminate(final Throwable t, final long delayInMs) { .filter(endPoint -> !queryContext.getEndPointBlackList().contains(endPoint)) .collect(Collectors.toList()); this.hasThrowable = Objects.nonNull(t); - // ClusterScheduler aborts the dispatcher before termination: no later request can obtain an - // admission deadline beyond this cutoff. Transport retries reuse their original deadline. - this.admissionDeadline = admissionDeadlineSupplier.getAsLong(); + // Aborting dispatch first freezes the tokens used on each receiver. A token acquisition that + // finishes afterwards is stateless and cannot publish an FI or reopen the closed attempt. + this.admissionTokens = admissionTokensSupplier.get(); return scheduledExecutor.schedule(this::syncTerminate, delayInMs, TimeUnit.MILLISECONDS); } @@ -122,7 +124,8 @@ public Boolean syncTerminate() { // can still keep the shared DataNode context alive. Legacy receivers only use the FI list. final List unfinishedFIs = stateTracker.filterUnFinishedFIs(ownedFragmentInstance.get(endPoint)); - if (unfinishedFIs.isEmpty() && admissionDeadline <= 0) { + TQueryAdmissionToken token = admissionTokens.get(endPoint); + if (unfinishedFIs.isEmpty() && token == null) { continue; } @@ -130,9 +133,9 @@ public Boolean syncTerminate() { final int internalPort = IoTDBDescriptor.getInstance().getConfig().getInternalPort(); if (internalAddress.equalsIgnoreCase(endPoint.getIp()) && internalPort == endPoint.getPort()) { - if (admissionDeadline > 0) { + if (token != null) { FragmentInstanceManager.getInstance() - .closeQueryAttempt(queryId, queryAttemptId, admissionDeadline, hasThrowable); + .closeQueryAttempt(queryId, queryAttemptId, token, hasThrowable); } for (final FragmentInstanceId insId : unfinishedFIs) { FragmentInstanceManager.getInstance().cancelTask(insId, hasThrowable); @@ -149,8 +152,8 @@ public Boolean syncTerminate() { .map(FragmentInstanceId::toThrift) .collect(Collectors.toList()), hasThrowable); - if (admissionDeadline > 0) { - request.setQueryAttemptId(queryAttemptId).setAdmissionDeadline(admissionDeadline); + if (token != null) { + request.setQueryAttemptId(queryAttemptId).setAdmissionToken(token); } client.cancelQuery(request); } catch (final ClientManagerException e) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java index 6fe3c2232231b..7eeb56bd159f0 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManagerTest.java @@ -28,6 +28,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment; import org.apache.iotdb.db.schemaengine.schemaregion.ISchemaRegion; import org.apache.iotdb.db.storageengine.dataregion.IDataRegionForQuery; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.iotdb.rpc.TSStatusCode; import org.junit.BeforeClass; @@ -69,8 +70,8 @@ public void testSchemaQueryFailureInfoAfterConcurrentContextRemoval() throws Exc @Test public void testPartialDispatchReleasesAttemptContextAndRejectsLateFragment() { QueryId queryId = new QueryId("partial_attempt_dispatch"); - long deadline = System.currentTimeMillis() + TimeUnit.MINUTES.toMillis(1); FragmentInstanceManager manager = FragmentInstanceManager.getInstance(); + TQueryAdmissionToken token = manager.fetchQueryAdmissionToken(TimeUnit.MINUTES.toMillis(1)); org.apache.iotdb.db.queryengine.common.QueryAttemptId key = new org.apache.iotdb.db.queryengine.common.QueryAttemptId(queryId, 0); try { @@ -82,8 +83,7 @@ public void testPartialDispatchReleasesAttemptContextAndRejectsLateFragment() { when(instance.getFragment()).thenReturn(fragment); when(instance.hasQueryAttempt()).thenReturn(true); when(instance.getQueryAttemptId()).thenReturn(0); - when(instance.getQueryDeadline()).thenReturn(deadline); - when(instance.getAdmissionDeadline()).thenReturn(deadline); + when(instance.getAdmissionToken()).thenReturn(token); when(instance.getDataNodeFINum()).thenReturn(3); when(instance.getTimeOut()).thenReturn(TimeUnit.MINUTES.toMillis(1)); when(fragment.getPlanNodeTree()) @@ -93,19 +93,18 @@ public void testPartialDispatchReleasesAttemptContextAndRejectsLateFragment() { assertTrue(result.getState().isFailed()); } assertTrue(manager.getQueryAttempts().hasSharedContext(key)); - manager.closeQueryAttempt(queryId, 0, deadline, true); + manager.closeQueryAttempt(queryId, 0, token, true); org.junit.Assert.assertFalse(manager.getQueryAttempts().hasSharedContext(key)); FragmentInstance late = mock(FragmentInstance.class); when(late.getId()).thenReturn(new FragmentInstanceId(new PlanFragmentId(queryId, 2), "0")); when(late.hasQueryAttempt()).thenReturn(true); when(late.getQueryAttemptId()).thenReturn(0); - when(late.getQueryDeadline()).thenReturn(deadline); - when(late.getAdmissionDeadline()).thenReturn(deadline); + when(late.getAdmissionToken()).thenReturn(token); org.junit.Assert.assertThrows( org.apache.iotdb.commons.exception.IoTDBRuntimeException.class, () -> manager.execDataQueryFragmentInstance(late, mock(IDataRegionForQuery.class))); } finally { - manager.closeQueryAttempt(queryId, 0, deadline, true); + manager.closeQueryAttempt(queryId, 0, token, true); for (int i = 0; i < 3; i++) { manager.cancelTask(new FragmentInstanceId(new PlanFragmentId(queryId, i), "0"), true); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptAdmissionTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptAdmissionTest.java new file mode 100644 index 0000000000000..d3592a0164a70 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptAdmissionTest.java @@ -0,0 +1,184 @@ +/* + * 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.execution.fragment; + +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.client.IClientManager; +import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient; +import org.apache.iotdb.commons.conf.CommonConfig; +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.exception.IoTDBRuntimeException; +import org.apache.iotdb.commons.partition.ExecutorType; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment; +import org.apache.iotdb.db.queryengine.plan.scheduler.FragInstanceDispatchResult; +import org.apache.iotdb.db.queryengine.plan.scheduler.FragmentInstanceDispatcherImpl; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceReq; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceResp; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.thrift.transport.TTransportException; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +public class QueryAttemptAdmissionTest { + @Test + public void fullSocketTimeoutStillAllowsTheReconnectRetry() throws Exception { + assertTransportTimeout(20, true); + } + + @Test + public void expiredAdmissionIsRetryableInsteadOfReportingQueryTimeout() throws Exception { + assertTransportTimeout(41, false); + } + + private void assertTransportTimeout(long firstRpcElapsedMs, boolean expectedSuccess) + throws Exception { + CommonConfig config = CommonDescriptor.getInstance().getConfig(); + int previousTimeout = config.getDnConnectionTimeoutInMS(); + config.setDnConnectionTimeoutInMS(20); + try { + QueryId queryId = new QueryId("socket_timeout_admission"); + MPPQueryContext context = new MPPQueryContext(queryId); + context.setStartTime(System.currentTimeMillis()); + context.setTimeOut(TimeUnit.HOURS.toMillis(1)); + // An arbitrary receiver monotonic origin, unrelated to the coordinator's wall clock. + AtomicLong receiverTick = new AtomicLong(-TimeUnit.DAYS.toNanos(123)); + QueryAttemptRegistry receiver = new QueryAttemptRegistry(receiverTick::get); + FragmentInstance instance = + spy( + new FragmentInstance( + mock(PlanFragment.class), + new FragmentInstanceId(new PlanFragmentId(queryId, 0), "0"), + null, + QueryType.READ, + TimeUnit.HOURS.toMillis(1), + null, + false, + false)); + instance.setDataNodeFINum(1); + instance.setExecutorType(mock(ExecutorType.class)); + doReturn(ByteBuffer.wrap(new byte[] {1})).when(instance).serializeToByteBuffer(); + TEndPoint endpoint = new TEndPoint("remote-host", 10730); + doReturn(new TDataNodeLocation().setInternalEndPoint(endpoint)) + .when(instance) + .getHostDataNode(); + IClientManager clients = + mock(IClientManager.class); + SyncDataNodeInternalServiceClient client = mock(SyncDataNodeInternalServiceClient.class); + when(clients.borrowClient(endpoint)).thenReturn(client); + when(client.fetchQueryAdmissionToken(anyLong())) + .thenAnswer(invocation -> receiver.issueAdmissionToken(invocation.getArgument(0))); + List requests = new ArrayList<>(); + when(client.sendFragmentInstance(any())) + .thenAnswer( + invocation -> { + TSendFragmentInstanceReq request = invocation.getArgument(0); + requests.add(request.deepCopy()); + if (requests.size() == 1) { + receiverTick.addAndGet(TimeUnit.MILLISECONDS.toNanos(firstRpcElapsedMs)); + throw new TTransportException(TTransportException.TIMED_OUT, "socket timed out"); + } + FragmentInstanceId receivedId = instance.getId(); + FragmentInstance received = mock(FragmentInstance.class); + when(received.getId()).thenReturn(receivedId); + when(received.hasQueryAttempt()).thenReturn(true); + when(received.getQueryAttemptId()).thenReturn(request.getQueryAttemptId()); + when(received.getAdmissionToken()).thenReturn(request.getAdmissionToken()); + when(received.getDataNodeFINum()).thenReturn(1); + when(received.getTimeOut()).thenReturn(TimeUnit.HOURS.toMillis(1)); + try { + receiver + .register( + received, + new FragmentInstanceStateMachine(received.getId(), Runnable::run), + true) + .close(); + return new TSendFragmentInstanceResp(true); + } catch (IoTDBRuntimeException e) { + return new TSendFragmentInstanceResp(false) + .setNeedRetry(true) + .setStatus(new TSStatus(e.getErrorCode()).setMessage(e.getMessage())); + } + }); + FragmentInstanceDispatcherImpl dispatcher = + new FragmentInstanceDispatcherImpl(QueryType.READ, context, clients, null); + FragInstanceDispatchResult result = + dispatcher.dispatch(null, Collections.singletonList(instance)).get(); + assertEquals(expectedSuccess, result.isSuccessful()); + assertEquals(2, requests.size()); + assertEquals(requests.get(0).getAdmissionToken(), requests.get(1).getAdmissionToken()); + if (!expectedSuccess) { + assertEquals( + TSStatusCode.DISPATCH_ERROR.getStatusCode(), result.getFailureStatus().getCode()); + assertTrue( + result + .getFailureStatus() + .getMessage() + .contains(String.valueOf(requests.get(0).getAdmissionToken().getDeadlineNanos()))); + assertTrue( + result.getFailureStatus().getMessage().contains(String.valueOf(receiverTick.get()))); + } + } finally { + config.setDnConnectionTimeoutInMS(previousTimeout); + } + } + + @Test + public void tokenIssuanceAfterFenceExpiryCannotRecreateAnAttempt() { + AtomicLong tick = new AtomicLong(); + QueryAttemptRegistry receiver = new QueryAttemptRegistry(tick::get); + org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken token = receiver.issueAdmissionToken(1); + receiver.close( + new org.apache.iotdb.db.queryengine.common.QueryAttemptId(new QueryId("closed"), 0), + token, + true); + tick.set(TimeUnit.MILLISECONDS.toNanos(1)); + receiver.removeExpired(); + assertEquals(0, receiver.getRetainedEntryCount()); + // A delayed issuance RPC is deliberately query-agnostic and allocates no attempt entry. + receiver.issueAdmissionToken(Long.MAX_VALUE); + assertEquals(0, receiver.getRetainedEntryCount()); + assertTrue( + receiver.issueAdmissionToken(Long.MAX_VALUE).getValidityMs() <= TimeUnit.DAYS.toMillis(1)); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java index ddf282306fd9b..1d39fbf7cfdca 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java @@ -16,14 +16,17 @@ * specific language governing permissions and limitations * under the License. */ + package org.apache.iotdb.db.queryengine.execution.fragment; import org.apache.iotdb.commons.exception.IoTDBRuntimeException; +import org.apache.iotdb.commons.exception.QueryTimeoutException; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.queryengine.common.PlanFragmentId; import org.apache.iotdb.db.queryengine.common.QueryAttemptId; import org.apache.iotdb.db.queryengine.common.QueryId; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.iotdb.rpc.TSStatusCode; import org.junit.Test; @@ -45,8 +48,9 @@ import static org.mockito.Mockito.when; public class QueryAttemptRegistryTest { - private final AtomicLong clock = new AtomicLong(1000); + private final AtomicLong clock = new AtomicLong(TimeUnit.MILLISECONDS.toNanos(1000)); private final QueryAttemptRegistry registry = new QueryAttemptRegistry(clock::get); + private final TQueryAdmissionToken token = registry.issueAdmissionToken(1000); private final QueryId queryId = new QueryId("attempt_test"); @Test @@ -69,12 +73,12 @@ public void closesPartiallyDispatchedAttemptAfterActualUsersRelease() { QueryAttemptRegistry.Registration first = registry.register(instance, state, true); QueryAttemptRegistry.Registration failed = register(0, 1, 3); failed.close(); - registry.close(key(0), 2000, true); + registry.close(key(0), token, true); assertEquals(FragmentInstanceState.CANCELLED, state.getState()); assertTrue(registry.hasSharedContext(key(0))); first.close(); first.close(); - registry.close(key(0), 2000, true); + registry.close(key(0), token, true); assertFalse(registry.hasSharedContext(key(0))); // The third FI never arrived; it cannot keep the shared LAST/filter state alive. assertEquals(1, registry.getRetainedEntryCount()); @@ -85,9 +89,9 @@ public void lateOldCancellationAndReleaseDoNotAffectRetry() { QueryAttemptRegistry.Registration old = register(0, 0, 3); QueryAttemptRegistry.Registration retry = register(1, 3, 1); assertNotSame(old.getContext(), retry.getContext()); - registry.close(key(0), 2000, true); + registry.close(key(0), token, true); old.close(); - registry.close(key(0), 2000, true); + registry.close(key(0), token, true); assertTrue(registry.hasSharedContext(key(1))); assertFalse(registry.hasSharedContext(key(0))); retry.close(); @@ -95,18 +99,16 @@ public void lateOldCancellationAndReleaseDoNotAffectRetry() { @Test public void closedBeforeFirstDispatchRemainsFencedAfterRecordExpires() { - registry.close(key(0), 2000, true); + registry.close(key(0), token, true); assertEquals( TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), assertThrows(IoTDBRuntimeException.class, () -> register(0, 0, 1)).getErrorCode()); - clock.set(2000); + clock.set(TimeUnit.MILLISECONDS.toNanos(2000)); registry.removeExpired(); assertEquals(0, registry.getRetainedEntryCount()); assertEquals( - TSStatusCode.QUERY_TIMEOUT.getStatusCode(), + TSStatusCode.DISPATCH_ERROR.getStatusCode(), assertThrows(IoTDBRuntimeException.class, () -> register(0, 0, 1)).getErrorCode()); - clock.set(1500); - assertThrows(IoTDBRuntimeException.class, () -> register(0, 0, 1)); assertEquals(0, registry.getRetainedEntryCount()); } @@ -115,9 +117,15 @@ public void deadlineReclaimsPartialAttemptWhenCoordinatorDoesNotCancel() { FragmentInstance instance = instance(0, 0, 3); FragmentInstanceStateMachine state = state(instance); QueryAttemptRegistry.Registration first = registry.register(instance, state, true); - clock.set(2000); + clock.set(TimeUnit.MILLISECONDS.toNanos(2000)); registry.removeExpired(); - assertEquals(FragmentInstanceState.CANCELLED, state.getState()); + assertEquals(FragmentInstanceState.FAILED, state.getState()); + FragmentInstanceContext context = + FragmentInstanceContext.createFragmentInstanceContext( + instance.getId(), state, null, Long.MAX_VALUE, false, false); + assertEquals( + TSStatusCode.QUERY_TIMEOUT.getStatusCode(), + context.getInstanceInfo().getErrorCode().get().getCode()); first.close(); assertEquals(0, registry.getRetainedEntryCount()); } @@ -134,7 +142,7 @@ public void repeatedDispatchDoesNotReplaceExistingRegistration() { .getErrorCode()); first.close(); assertTrue(registry.hasSharedContext(key(0))); - registry.close(key(0), 2000, true); + registry.close(key(0), token, true); assertFalse(registry.hasSharedContext(key(0))); } @@ -168,7 +176,7 @@ public void concurrentCloseCannotMissRegistrationInProgress() throws Exception { Future registration = executor.submit(() -> registry.register(instance, state, true)); assertTrue(registering.await(10, TimeUnit.SECONDS)); - Future cancellation = executor.submit(() -> registry.close(key(0), 2000, true)); + Future cancellation = executor.submit(() -> registry.close(key(0), token, true)); allowRegistration.countDown(); QueryAttemptRegistry.Registration result = registration.get(10, TimeUnit.SECONDS); cancellation.get(10, TimeUnit.SECONDS); @@ -186,15 +194,16 @@ public void concurrentCloseCannotMissRegistrationInProgress() throws Exception { @Test public void admissionExpiryDoesNotCancelAnAlreadyRunningLongQuery() { FragmentInstance instance = instance(0, 0, 3); - when(instance.getQueryDeadline()).thenReturn(Long.MAX_VALUE); - when(instance.getAdmissionDeadline()).thenReturn(1200L); + when(instance.getTimeOut()).thenReturn(Long.MAX_VALUE); + TQueryAdmissionToken shortToken = registry.issueAdmissionToken(200); + when(instance.getAdmissionToken()).thenReturn(shortToken); FragmentInstanceStateMachine state = state(instance); QueryAttemptRegistry.Registration running = registry.register(instance, state, true); - clock.set(1300); + clock.set(TimeUnit.MILLISECONDS.toNanos(1300)); registry.removeExpired(); assertEquals(FragmentInstanceState.RUNNING, state.getState()); assertTrue(registry.hasSharedContext(key(0))); - registry.close(key(0), 1200, true); + registry.close(key(0), shortToken, true); running.close(); assertEquals(0, registry.getRetainedEntryCount()); assertThrows( @@ -204,14 +213,84 @@ public void admissionExpiryDoesNotCancelAnAlreadyRunningLongQuery() { @Test public void completedLongQueryRetainsOnlyABoundedAdmissionFence() { FragmentInstance instance = instance(0, 0, 1); - when(instance.getQueryDeadline()).thenReturn(Long.MAX_VALUE); - when(instance.getAdmissionDeadline()).thenReturn(1200L); + when(instance.getTimeOut()).thenReturn(Long.MAX_VALUE); + TQueryAdmissionToken shortToken = registry.issueAdmissionToken(200); + when(instance.getAdmissionToken()).thenReturn(shortToken); registry.register(instance, state(instance), true).close(); assertEquals(1, registry.getRetainedEntryCount()); assertFalse(registry.hasSharedContext(key(0))); - clock.set(1200); + clock.set(TimeUnit.MILLISECONDS.toNanos(1200)); + registry.removeExpired(); + assertEquals(0, registry.getRetainedEntryCount()); + } + + @Test + public void expiryOfClosedFenceDoesNotAddATimeoutCause() { + FragmentInstance instance = instance(0, 0, 3); + FragmentInstanceStateMachine state = state(instance); + QueryAttemptRegistry.Registration registration = registry.register(instance, state, true); + registry.close(key(0), token, true); + clock.addAndGet(TimeUnit.SECONDS.toNanos(2)); + registry.removeExpired(); + assertEquals(FragmentInstanceState.CANCELLED, state.getState()); + assertTrue(state.getFailureCauses().isEmpty()); + registration.close(); + assertEquals(0, registry.getRetainedEntryCount()); + } + + @Test + public void receiverClockOriginDoesNotShortenExecutionBudget() { + // The receiver can have an entirely different clock origin from the coordinator. Only the + // receiver's issued token and local relative execution timeout enter admission/expiry. + AtomicLong receiverTick = new AtomicLong(TimeUnit.DAYS.toNanos(123)); + QueryAttemptRegistry receiver = new QueryAttemptRegistry(receiverTick::get); + FragmentInstance instance = instance(0, 0, 1); + when(instance.getAdmissionToken()).thenReturn(receiver.issueAdmissionToken(20)); + FragmentInstanceStateMachine state = state(instance); + QueryAttemptRegistry.Registration registration = receiver.register(instance, state, true); + receiverTick.addAndGet(TimeUnit.MILLISECONDS.toNanos(999)); + receiver.removeExpired(); + assertEquals(FragmentInstanceState.RUNNING, state.getState()); + receiverTick.addAndGet(TimeUnit.MILLISECONDS.toNanos(1)); + receiver.removeExpired(); + assertEquals(FragmentInstanceState.FAILED, state.getState()); + assertTrue(state.getFailureCauses().peek() instanceof QueryTimeoutException); + registration.close(); + } + + @Test + public void oldProcessTokenCannotAdmitOrCloseANewProcessAttempt() { + QueryAttemptRegistry restarted = new QueryAttemptRegistry(clock::get); + FragmentInstance instance = instance(0, 0, 1); + assertEquals( + TSStatusCode.DISPATCH_ERROR.getStatusCode(), + assertThrows( + IoTDBRuntimeException.class, + () -> restarted.register(instance, state(instance), true)) + .getErrorCode()); + when(instance.getAdmissionToken()).thenReturn(restarted.issueAdmissionToken(1000)); + QueryAttemptRegistry.Registration registration = + restarted.register(instance, state(instance), true); + restarted.close(key(0), token, true); + assertTrue(restarted.hasSharedContext(key(0))); + registration.close(); + } + + @Test + public void freshAttemptWorksAfterOldFenceExpiresAndMonotonicCounterWraps() { + clock.set(Long.MAX_VALUE - TimeUnit.MILLISECONDS.toNanos(5)); + TQueryAdmissionToken issued = registry.issueAdmissionToken(10); + FragmentInstance instance = instance(0, 0, 1); + when(instance.getAdmissionToken()).thenReturn(issued); + registry.register(instance, state(instance), true).close(); + clock.addAndGet(TimeUnit.MILLISECONDS.toNanos(10)); registry.removeExpired(); assertEquals(0, registry.getRetainedEntryCount()); + assertThrows( + IoTDBRuntimeException.class, () -> registry.register(instance, state(instance), true)); + FragmentInstance fresh = instance(1, 1, 1); + when(fresh.getAdmissionToken()).thenReturn(registry.issueAdmissionToken(1000)); + registry.register(fresh, state(fresh), true).close(); } private QueryAttemptId key(int attempt) { @@ -233,8 +312,7 @@ private FragmentInstance instance(int attempt, int fragment, int planned) { .thenReturn(new FragmentInstanceId(new PlanFragmentId(queryId, fragment), "0")); when(instance.hasQueryAttempt()).thenReturn(true); when(instance.getQueryAttemptId()).thenReturn(attempt); - when(instance.getQueryDeadline()).thenReturn(2000L); - when(instance.getAdmissionDeadline()).thenReturn(2000L); + when(instance.getAdmissionToken()).thenReturn(token); when(instance.getTimeOut()).thenReturn(1000L); when(instance.getDataNodeFINum()).thenReturn(planned); return instance; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionLifecycleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionLifecycleTest.java index 9ab62a22d80c8..3dfbb37695158 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionLifecycleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionLifecycleTest.java @@ -16,6 +16,7 @@ * specific language governing permissions and limitations * under the License. */ + package org.apache.iotdb.db.queryengine.plan.execution; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; @@ -117,6 +118,41 @@ public void asynchronousFailureStopsSchedulerBeforeStartReturns() throws Excepti } } + @Test + public void terminalCleanupPreservesAnExistingStatusOnlyFailure() throws Exception { + ExecutorService notifications = Executors.newSingleThreadExecutor(); + IPlanner planner = planner(); + org.apache.iotdb.common.rpc.thrift.TSStatus original = + new org.apache.iotdb.common.rpc.thrift.TSStatus( + org.apache.iotdb.rpc.TSStatusCode.QUERY_EXECUTION_MEMORY_NOT_ENOUGH.getStatusCode()); + original.setMessage("original failure"); + when(planner.createScheduler(any(), any(), any(), any())) + .thenAnswer( + invocation -> { + QueryStateMachine state = invocation.getArgument(3); + IScheduler scheduler = mock(IScheduler.class); + doAnswer( + ignored -> { + state.transitionToFailed(original); + return null; + }) + .when(scheduler) + .start(); + return scheduler; + }); + QueryExecution execution = new QueryExecution(planner, context(), notifications); + try { + execution.start(); + execution.stopAndCleanup( + new org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException(0, 100, 1)); + org.junit.Assert.assertEquals(original, execution.getStatus().status); + } finally { + execution.stopAndCleanup(null); + notifications.shutdownNow(); + assertTrue(notifications.awaitTermination(10, TimeUnit.SECONDS)); + } + } + private static IPlanner planner() { IPlanner planner = mock(IPlanner.class); when(planner.analyze(any())).thenReturn(mock(IAnalysis.class)); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterSchedulerLifecycleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterSchedulerLifecycleTest.java index 4caf1af1d6899..6fc2541bb7264 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterSchedulerLifecycleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterSchedulerLifecycleTest.java @@ -16,6 +16,7 @@ * specific language governing permissions and limitations * under the License. */ + package org.apache.iotdb.db.queryengine.plan.scheduler; import org.apache.iotdb.common.rpc.thrift.TSStatus; @@ -104,6 +105,19 @@ public void nonRetryableDispatchFailureTerminatesImmediately() { verify(terminator).terminateImmediately(cause); } + @Test + public void admissionFailureCanStartAFreshQueryAttempt() { + state.transitionToPlanned(); + when(dispatcher.dispatch(any(), any())) + .thenReturn( + immediateFuture( + new FragInstanceDispatchResult( + new TSStatus(TSStatusCode.DISPATCH_ERROR.getStatusCode())))); + scheduler.start(); + org.junit.Assert.assertEquals( + org.apache.iotdb.db.queryengine.execution.QueryState.PENDING_RETRY, state.getState()); + } + @Test public void normalCompletionKeepsGracePeriod() { when(dispatcher.dispatch(any(), any())) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java index 50c7e63da967a..d291e5831449d 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java @@ -16,6 +16,7 @@ * specific language governing permissions and limitations * under the License. */ + package org.apache.iotdb.db.queryengine.plan.scheduler; import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; @@ -30,22 +31,32 @@ import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceReq; import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceResp; import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.thrift.TApplicationException; import org.apache.thrift.TException; import org.junit.Test; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -54,41 +65,29 @@ import static org.mockito.Mockito.when; public class FragmentInstanceDispatcherLifecycleTest { - @Test - public void abortPreventsDispatchAndRetainsOriginalAttemptMetadata() throws Exception { - QueryId queryId = new QueryId("aborted_dispatch"); - MPPQueryContext context = new MPPQueryContext(queryId); - context.setStartTime(1000); - context.setTimeOut(2000); - FragmentInstanceDispatcherImpl dispatcher = - new FragmentInstanceDispatcherImpl(QueryType.READ, context, null, null); - context.prepareForRetry(); - FragmentInstance instance = mock(FragmentInstance.class); - when(instance.getId()).thenReturn(new FragmentInstanceId(new PlanFragmentId(queryId, 0), "0")); - PlanFragment fragment = mock(PlanFragment.class); - when(instance.getFragment()).thenReturn(fragment); - dispatcher.abort(); - FragInstanceDispatchResult result = - dispatcher.dispatch(null, Collections.singletonList(instance)).get(); - assertFalse(result.isSuccessful()); - assertEquals( - TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), result.getFailureStatus().getCode()); - verify(instance, never()).getHostDataNode(); - verify(instance).setQueryAttempt(0, 3000); - } + private final QueryId queryId = new QueryId("dispatch_admission"); + private final MPPQueryContext context = new MPPQueryContext(queryId); + private final TEndPoint endpoint = new TEndPoint("remote-host", 10730); + private final IClientManager clients = + mock(IClientManager.class); + private final SyncDataNodeInternalServiceClient client = + mock(SyncDataNodeInternalServiceClient.class); + private final TQueryAdmissionToken token = new TQueryAdmissionToken("receiver", 123456L, 60000L); - @Test - public void transportRetryReusesAdmissionDeadlineAndStopFreezesCutoff() throws Exception { - QueryId queryId = new QueryId("retry_admission"); - MPPQueryContext context = new MPPQueryContext(queryId); + private FragmentInstanceDispatcherImpl dispatcher() throws Exception { context.setStartTime(System.currentTimeMillis()); context.setTimeOut(Long.MAX_VALUE); - PlanFragment fragment = mock(PlanFragment.class); + when(clients.borrowClient(endpoint)).thenReturn(client); + when(client.fetchQueryAdmissionToken(anyLong())).thenReturn(token); + return new FragmentInstanceDispatcherImpl(QueryType.READ, context, clients, null); + } + + private FragmentInstance instance(int fragmentId) { FragmentInstance instance = spy( new FragmentInstance( - fragment, - new FragmentInstanceId(new PlanFragmentId(queryId, 0), "0"), + mock(PlanFragment.class), + new FragmentInstanceId(new PlanFragmentId(queryId, fragmentId), "0"), null, QueryType.READ, Long.MAX_VALUE, @@ -96,15 +95,30 @@ public void transportRetryReusesAdmissionDeadlineAndStopFreezesCutoff() throws E false, false)); doReturn(ByteBuffer.wrap(new byte[] {1})).when(instance).serializeToByteBuffer(); - TEndPoint endpoint = new TEndPoint("remote-host", 10730); doReturn(new TDataNodeLocation().setInternalEndPoint(endpoint)) .when(instance) .getHostDataNode(); instance.setExecutorType(mock(ExecutorType.class)); - IClientManager clients = - mock(IClientManager.class); - SyncDataNodeInternalServiceClient client = mock(SyncDataNodeInternalServiceClient.class); - when(clients.borrowClient(endpoint)).thenReturn(client); + return instance; + } + + @Test + public void abortPreventsTokenAcquisitionAndDispatch() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + dispatcher.abort(); + FragInstanceDispatchResult result = + dispatcher.dispatch(null, Collections.singletonList(instance(0))).get(); + assertFalse(result.isSuccessful()); + assertEquals( + TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), result.getFailureStatus().getCode()); + verify(client, never()).fetchQueryAdmissionToken(anyLong()); + verify(client, never()).sendFragmentInstance(any()); + } + + @Test + public void transportRetryReusesTokenAndStopFreezesIssuedTokens() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + context.prepareForRetry(); List requests = new ArrayList<>(); when(client.sendFragmentInstance(any())) .thenAnswer( @@ -116,18 +130,120 @@ public void transportRetryReusesAdmissionDeadlineAndStopFreezesCutoff() throws E } return new TSendFragmentInstanceResp(true); }); - FragmentInstanceDispatcherImpl dispatcher = - new FragmentInstanceDispatcherImpl(QueryType.READ, context, clients, null); + FragmentInstance instance = instance(0); assertTrue(dispatcher.dispatch(null, Collections.singletonList(instance)).get().isSuccessful()); assertEquals(2, requests.size()); - assertEquals(Long.MAX_VALUE, requests.get(0).getQueryDeadline()); - assertTrue(requests.get(0).getAdmissionDeadline() < Long.MAX_VALUE); - assertEquals(requests.get(0).getAdmissionDeadline(), requests.get(1).getAdmissionDeadline()); - long cutoff = dispatcher.getLastAdmissionDeadline(); + assertEquals(0, requests.get(0).getQueryAttemptId()); + assertEquals(token, requests.get(0).getAdmissionToken()); + assertEquals(requests.get(0).getAdmissionToken(), requests.get(1).getAdmissionToken()); + verify(client).fetchQueryAdmissionToken(anyLong()); + Map issued = dispatcher.getIssuedAdmissionTokens(); dispatcher.abort(); assertFalse( dispatcher.dispatch(null, Collections.singletonList(instance)).get().isSuccessful()); assertEquals(2, requests.size()); - assertEquals(cutoff, dispatcher.getLastAdmissionDeadline()); + assertEquals(issued, dispatcher.getIssuedAdmissionTokens()); + } + + @Test + public void fragmentsOnTheSameNodeReuseOneTokenRpc() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + when(client.sendFragmentInstance(any())).thenReturn(new TSendFragmentInstanceResp(true)); + assertTrue( + dispatcher.dispatch(null, Arrays.asList(instance(0), instance(1))).get().isSuccessful()); + verify(client).fetchQueryAdmissionToken(anyLong()); + } + + @Test + public void legacyReceiverUsesTheOriginalProtocol() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + when(client.fetchQueryAdmissionToken(anyLong())) + .thenThrow( + new TApplicationException(TApplicationException.UNKNOWN_METHOD, "unknown method")); + when(client.sendFragmentInstance(any())) + .thenAnswer( + invocation -> { + TSendFragmentInstanceReq request = invocation.getArgument(0); + assertFalse(request.isSetQueryAttemptId()); + assertFalse(request.isSetAdmissionToken()); + return new TSendFragmentInstanceResp(true); + }); + assertTrue( + dispatcher.dispatch(null, Arrays.asList(instance(0), instance(1))).get().isSuccessful()); + verify(client).fetchQueryAdmissionToken(anyLong()); + assertTrue(dispatcher.getIssuedAdmissionTokens().isEmpty()); + } + + @Test + public void wrappedUnknownMethodUsesAFreshClientForLegacyDispatch() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + SyncDataNodeInternalServiceClient fresh = mock(SyncDataNodeInternalServiceClient.class); + java.util.concurrent.atomic.AtomicBoolean returned = + new java.util.concurrent.atomic.AtomicBoolean(); + java.util.concurrent.atomic.AtomicInteger borrows = + new java.util.concurrent.atomic.AtomicInteger(); + org.mockito.Mockito.doAnswer( + ignored -> { + returned.set(true); + return null; + }) + .when(client) + .close(); + when(clients.borrowClient(endpoint)) + .thenAnswer( + ignored -> { + if (borrows.getAndIncrement() == 0) { + return client; + } + assertTrue(returned.get()); + return fresh; + }); + when(client.fetchQueryAdmissionToken(anyLong())) + .thenThrow( + new TException( + new TApplicationException(TApplicationException.UNKNOWN_METHOD, "unknown method"))); + when(fresh.sendFragmentInstance(any())) + .thenAnswer( + invocation -> { + TSendFragmentInstanceReq request = invocation.getArgument(0); + assertFalse(request.isSetAdmissionToken()); + return new TSendFragmentInstanceResp(true); + }); + assertTrue( + dispatcher.dispatch(null, Collections.singletonList(instance(0))).get().isSuccessful()); + verify(client, never()).sendFragmentInstance(any()); + verify(fresh, never()).fetchQueryAdmissionToken(anyLong()); + verify(fresh).sendFragmentInstance(any()); + } + + @Test + public void abortDuringTokenRpcCannotPublishALateFragment() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + ExecutorService executor = Executors.newSingleThreadExecutor(); + CountDownLatch acquiring = new CountDownLatch(1); + CountDownLatch returnToken = new CountDownLatch(1); + when(client.fetchQueryAdmissionToken(anyLong())) + .thenAnswer( + invocation -> { + acquiring.countDown(); + assertTrue(returnToken.await(10, TimeUnit.SECONDS)); + return token; + }); + try { + Future result = + executor.submit( + () -> dispatcher.dispatch(null, Collections.singletonList(instance(0))).get()); + assertTrue(acquiring.await(10, TimeUnit.SECONDS)); + dispatcher.abort(); + assertTrue(dispatcher.getIssuedAdmissionTokens().isEmpty()); + returnToken.countDown(); + assertFalse(result.get(10, TimeUnit.SECONDS).isSuccessful()); + verify(client, never()).sendFragmentInstance(any()); + assertTrue(dispatcher.getIssuedAdmissionTokens().isEmpty()); + } finally { + returnToken.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryAttemptProtocolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryAttemptProtocolTest.java index d5a12286286d6..d77cdc8e1a752 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryAttemptProtocolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryAttemptProtocolTest.java @@ -16,10 +16,12 @@ * specific language governing permissions and limitations * under the License. */ + package org.apache.iotdb.db.queryengine.plan.scheduler; import org.apache.iotdb.mpp.rpc.thrift.TCancelQueryReq; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstance; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceReq; import org.apache.thrift.TDeserializer; @@ -46,13 +48,13 @@ public void legacyRequestsRemainReadableWithNoAttemptMetadata() throws Exception new TDeserializer().deserialize(received, new TSerializer().serialize(original)); assertEquals(original, received); assertFalse(received.isSetQueryAttemptId()); - assertFalse(received.isSetQueryDeadline()); + assertFalse(received.isSetAdmissionToken()); TCancelQueryReq cancel = new TCancelQueryReq("query", Collections.emptyList(), true); TCancelQueryReq receivedCancel = new TCancelQueryReq(); new TDeserializer().deserialize(receivedCancel, new TSerializer().serialize(cancel)); assertEquals(cancel, receivedCancel); assertFalse(receivedCancel.isSetQueryAttemptId()); - assertFalse(receivedCancel.isSetAdmissionDeadline()); + assertFalse(receivedCancel.isSetAdmissionToken()); } @Test @@ -61,8 +63,7 @@ public void attemptMetadataRoundTripsAndLegacyReaderCanSkipIt() throws Exception TSendFragmentInstanceReq request = new TSendFragmentInstanceReq(body) .setQueryAttemptId(2) - .setQueryDeadline(123456L) - .setAdmissionDeadline(123450L); + .setAdmissionToken(new TQueryAdmissionToken("receiver", 123450L, 1000L)); byte[] serialized = new TSerializer().serialize(request); TSendFragmentInstanceReq received = new TSendFragmentInstanceReq(); new TDeserializer().deserialize(received, serialized); @@ -86,7 +87,7 @@ public void attemptMetadataRoundTripsAndLegacyReaderCanSkipIt() throws Exception TCancelQueryReq cancel = new TCancelQueryReq("query", Collections.emptyList(), true) .setQueryAttemptId(2) - .setAdmissionDeadline(123450L); + .setAdmissionToken(new TQueryAdmissionToken("receiver", 123450L, 1000L)); TCancelQueryReq receivedCancel = new TCancelQueryReq(); new TDeserializer().deserialize(receivedCancel, new TSerializer().serialize(cancel)); assertEquals(cancel, receivedCancel); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryTerminalCleanupTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryTerminalCleanupTest.java new file mode 100644 index 0000000000000..9a462fae86dac --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryTerminalCleanupTest.java @@ -0,0 +1,152 @@ +/* + * 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.plan.scheduler; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.queryengine.common.SessionInfo; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.QueryState; +import org.apache.iotdb.db.queryengine.execution.QueryStateMachine; +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.ExecutionResult; +import org.apache.iotdb.db.queryengine.plan.execution.QueryExecution; +import org.apache.iotdb.db.queryengine.plan.planner.IPlanner; +import org.apache.iotdb.db.queryengine.plan.planner.plan.DistributedQueryPlan; +import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.junit.Test; + +import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import static com.google.common.util.concurrent.Futures.immediateFuture; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class QueryTerminalCleanupTest { + @Test + public void staleQueryCleanupDuringDispatchMustUnblockGetStatus() throws Exception { + assertTerminalCleanup( + new QueryTimeoutRuntimeException(0, 120_001, 60_000), + QueryState.FAILED, + TSStatusCode.QUERY_TIMEOUT.getStatusCode()); + } + + @Test + public void nullCauseCleanupDuringDispatchAbortsAndUnblocksGetStatus() throws Exception { + assertTerminalCleanup( + null, QueryState.ABORTED, TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode()); + } + + private void assertTerminalCleanup(Throwable cause, QueryState expectedState, int expectedStatus) + throws Exception { + ExecutorService notifications = Executors.newSingleThreadExecutor(); + ExecutorService rpc = Executors.newSingleThreadExecutor(); + ExecutorService statusWaiter = Executors.newSingleThreadExecutor(); + CountDownLatch dispatchStarted = new CountDownLatch(1); + CountDownLatch finishDispatch = new CountDownLatch(1); + AtomicReference state = new AtomicReference<>(); + IFragInstanceDispatcher dispatcher = mock(IFragInstanceDispatcher.class); + when(dispatcher.dispatch(any(), any())) + .thenAnswer( + ignored -> { + dispatchStarted.countDown(); + assertTrue(finishDispatch.await(10, TimeUnit.SECONDS)); + return immediateFuture(new FragInstanceDispatchResult(true)); + }); + IPlanner planner = mock(IPlanner.class); + when(planner.analyze(any())).thenReturn(mock(IAnalysis.class)); + when(planner.doLogicalPlan(any(), any())).thenReturn(mock(LogicalQueryPlan.class)); + DistributedQueryPlan plan = mock(DistributedQueryPlan.class); + when(plan.getInstances()).thenReturn(Collections.emptyList()); + when(planner.doDistributionPlan(any(), any(), any())).thenReturn(plan); + when(planner.createScheduler(any(), any(), any(), any())) + .thenAnswer( + invocation -> { + QueryStateMachine machine = invocation.getArgument(3); + state.set(machine); + return new ClusterScheduler( + machine, + dispatcher, + mock(IFragInstanceStateTracker.class), + mock(IQueryTerminator.class)); + }); + QueryId queryId = new QueryId("stale_cleanup_during_dispatch"); + TEndPoint localExchange = new TEndPoint("127.0.0.1", 10740); + MPPQueryContext context = + new MPPQueryContext( + "", queryId, mock(SessionInfo.class), localExchange, new TEndPoint("127.0.0.1", 10730)); + context.setQueryType(QueryType.READ); + context.setStartTime(System.currentTimeMillis()); + context.setTimeOut(60_000L); + context + .getResultNodeContext() + .setUpStream( + localExchange, + queryId.genPlanFragmentId().genFragmentInstanceId(), + new PlanNodeId("upstream")); + QueryExecution execution = new QueryExecution(planner, context, notifications); + Future result = null; + try { + result = + rpc.submit( + () -> { + execution.start(); + return execution.getStatus(); + }); + assertTrue(dispatchStarted.await(10, TimeUnit.SECONDS)); + assertEquals(QueryState.DISPATCHING, state.get().getState()); + // Coordinator.cleanUpStaleQueries calls stopAndCleanup directly, without cancel(). + execution.stopAndCleanup(cause); + // Status waiters must wake even before the in-flight dispatch returns. + assertEquals( + expectedStatus, + statusWaiter.submit(execution::getStatus).get(10, TimeUnit.SECONDS).status.getCode()); + finishDispatch.countDown(); + assertEquals(expectedStatus, result.get(10, TimeUnit.SECONDS).status.getCode()); + assertEquals(expectedState, state.get().getState()); + } finally { + finishDispatch.countDown(); + execution.stopAndCleanup(null); + if (result != null) { + result.cancel(true); + } + statusWaiter.shutdownNow(); + assertTrue(statusWaiter.awaitTermination(10, TimeUnit.SECONDS)); + rpc.shutdownNow(); + notifications.shutdownNow(); + assertTrue(rpc.awaitTermination(10, TimeUnit.SECONDS)); + assertTrue(notifications.awaitTermination(10, TimeUnit.SECONDS)); + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java index 972031f9f4236..3dfb1c4202b36 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java @@ -29,6 +29,7 @@ import org.apache.iotdb.db.queryengine.common.QueryId; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; import org.apache.iotdb.mpp.rpc.thrift.TCancelQueryReq; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.junit.Before; import org.junit.Test; @@ -65,7 +66,7 @@ public void setUp() { Collections.emptyList(), null, mock(IFragInstanceStateTracker.class), - () -> 0L); + Collections::emptyMap); } @Test @@ -88,9 +89,9 @@ public void closesAttemptEvenWhenNoFragmentIsReportedUnfinished() throws Excepti final QueryId queryId = new QueryId("attempt_to_close"); when(queryContext.getQueryId()).thenReturn(queryId); when(queryContext.getQueryAttemptId()).thenReturn(2); - when(queryContext.getQueryDeadline()).thenReturn(123456L); when(queryContext.getEndPointBlackList()).thenReturn(Collections.emptySet()); final TEndPoint endpoint = new TEndPoint("remote-host", 10730); + final TQueryAdmissionToken token = new TQueryAdmissionToken("receiver", 123456L, 1000L); final FragmentInstance instance = mock(FragmentInstance.class); when(instance.getHostDataNode()) .thenReturn(new TDataNodeLocation().setInternalEndPoint(endpoint)); @@ -108,7 +109,7 @@ public void closesAttemptEvenWhenNoFragmentIsReportedUnfinished() throws Excepti Collections.singletonList(instance), clients, tracker, - () -> 123456L); + () -> Collections.singletonMap(endpoint, token)); // A delayed termination must retain the old attempt, even after the context starts a retry. when(queryContext.getQueryAttemptId()).thenReturn(3); terminator.terminateImmediately(new RuntimeException("dispatch failure")); @@ -116,7 +117,7 @@ public void closesAttemptEvenWhenNoFragmentIsReportedUnfinished() throws Excepti ArgumentCaptor request = ArgumentCaptor.forClass(TCancelQueryReq.class); verify(client).cancelQuery(request.capture()); assertEquals(2, request.getValue().getQueryAttemptId()); - assertEquals(123456L, request.getValue().getAdmissionDeadline()); + assertEquals(token, request.getValue().getAdmissionToken()); assertTrue(request.getValue().getFragmentInstanceIds().isEmpty()); } } diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index eb3753f07ccdf..aaf76ff73418b 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -139,14 +139,19 @@ struct TPlanNode { 1: required binary body } +// Issued and checked by one DataNode process using its own monotonic clock. Opaque to peers. +struct TQueryAdmissionToken { + 1: required string issuer + 2: required i64 deadlineNanos + 3: required i64 validityMs +} + struct TSendFragmentInstanceReq { 1: required TFragmentInstance fragmentInstance 2: optional common.TConsensusGroupId consensusGroupId // Attempt metadata is absent for a legacy coordinator. Do not change the binary FI body. 3: optional i32 queryAttemptId - 4: optional i64 queryDeadline - // Admission expires independently of a potentially long query execution timeout. - 5: optional i64 admissionDeadline + 4: optional TQueryAdmissionToken admissionToken } struct TSendFragmentInstanceResp { @@ -194,7 +199,7 @@ struct TCancelQueryReq { 3: required bool hasThrowable // Close this attempt even if some of its planned FIs were never dispatched. 4: optional i32 queryAttemptId - 5: optional i64 admissionDeadline + 5: optional TQueryAdmissionToken admissionToken } struct TCancelPlanFragmentReq { @@ -946,6 +951,9 @@ service IDataNodeRPCService { /** * dispatch FragmentInstance to remote node for query request */ + // Stateless: a delayed token request cannot create or reopen a query attempt. + TQueryAdmissionToken fetchQueryAdmissionToken(1: i64 validityMs); + TSendFragmentInstanceResp sendFragmentInstance(TSendFragmentInstanceReq req); /** From f787a87252ed811bd04ee220b2a56e85579660aa Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Thu, 12 Nov 2026 20:37:44 +0800 Subject: [PATCH 3/4] [Query] Reduce closed-attempt retention and expiry overhead Separate compact closed markers from live attempt state and expire them by deadline on a dedicated executor. Cap admission validity by the remaining query budget, keep extended fences safe from stale expiry events, and avoid duplicate queue entries. Use an explicit legacy-admission signal for connection fallback and cover serialized RPC admission and cancellation. Validated 142 targeted tests and full-reactor compilation in both locales; measured retained heap and expiry costs with a 1.2-million-record synthetic workload. --- .../fragment/FragmentInstanceManager.java | 13 +- .../fragment/QueryAttemptRegistry.java | 77 +++++++++- .../FragmentInstanceDispatcherImpl.java | 41 +++-- .../fragment/QueryAttemptRegistryTest.java | 59 ++++++++ .../fragment/QueryAttemptRpcTest.java | 140 ++++++++++++++++++ ...agmentInstanceDispatcherLifecycleTest.java | 14 ++ .../iotdb/commons/concurrent/ThreadName.java | 2 + 7 files changed, 317 insertions(+), 29 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRpcTest.java 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 bf3ffee88585f..a1bb09ca6f57d 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 @@ -87,6 +87,7 @@ public class FragmentInstanceManager { private final IDriverScheduler scheduler = DriverScheduler.getInstance(); private final ScheduledExecutorService instanceManagementExecutor; + private final ScheduledExecutorService queryAttemptExpirationExecutor; public final ExecutorService instanceNotificationExecutor; private final Duration infoCacheTime; @@ -110,6 +111,9 @@ private FragmentInstanceManager() { this.instanceManagementExecutor = IoTDBThreadPoolFactory.newScheduledThreadPool( 1, ThreadName.FRAGMENT_INSTANCE_MANAGEMENT.getName()); + this.queryAttemptExpirationExecutor = + IoTDBThreadPoolFactory.newScheduledThreadPool( + 1, ThreadName.QUERY_ATTEMPT_EXPIRATION.getName()); this.instanceNotificationExecutor = IoTDBThreadPoolFactory.newFixedThreadPool( IoTDBDescriptor.getInstance().getConfig().getFragmentInstanceNotificationThreadCount(), @@ -117,6 +121,14 @@ private FragmentInstanceManager() { this.infoCacheTime = new Duration(5, TimeUnit.MINUTES); + // Large bursts of expired admission fences must not delay stale-query or flushing cleanup. + ScheduledExecutorUtil.safelyScheduleWithFixedDelay( + queryAttemptExpirationExecutor, + queryAttempts::removeExpired, + 2000, + 2000, + TimeUnit.MILLISECONDS); + ScheduledExecutorUtil.safelyScheduleWithFixedDelay( instanceManagementExecutor, this::removeOldInstances, 2000, 2000, TimeUnit.MILLISECONDS); ScheduledExecutorUtil.safelyScheduleWithFixedDelay( @@ -519,7 +531,6 @@ public TFetchFragmentInstanceStatisticsResp getFragmentInstanceStatistics( } private void removeOldInstances() { - queryAttempts.removeExpired(); long oldestAllowedInstance = System.currentTimeMillis() - infoCacheTime.toMillis(); instanceContext .entrySet() diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java index 0624cf7dec135..fcce3d453d1ac 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java @@ -37,6 +37,8 @@ import java.util.Map; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.DelayQueue; +import java.util.concurrent.Delayed; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; @@ -52,7 +54,11 @@ class QueryAttemptRegistry { private static final Logger LOGGER = LoggerFactory.getLogger(QueryAttemptRegistry.class); private static final long MAX_ADMISSION_VALIDITY_MS = TimeUnit.DAYS.toMillis(1); + // Only live attempts are scanned for execution timeout. Completed attempts retain a compact + // marker, indexed by deadline, rather than a full Entry on the management sweep's hot path. private final Map attempts = new ConcurrentHashMap<>(); + private final Map closedAttempts = new ConcurrentHashMap<>(); + private final DelayQueue closedExpirations = new DelayQueue<>(); private final LongSupplier ticker; private final String issuer = UUID.randomUUID().toString(); @@ -115,7 +121,7 @@ Registration register( now), TSStatusCode.DISPATCH_ERROR.getStatusCode()); } - if (existing != null && existing.closed) { + if ((existing != null && existing.closed) || closedAttempts.containsKey(key)) { throw new IoTDBRuntimeException( String.format( DataNodeQueryMessages.EXCEPTION_QUERY_ATTEMPT_ARG_IS_CLOSED_12143E89, key), @@ -163,10 +169,11 @@ void close(QueryAttemptId key, TQueryAdmissionToken token, boolean hasThrowable) key, (ignored, existing) -> { long now = ticker.getAsLong(); - if (existing == null && expired(token.getDeadlineNanos(), now)) { + if (existing == null) { + retainClosure(key, token.getDeadlineNanos(), now); return null; } - Entry entry = existing == null ? new Entry(now, 0, 0) : existing; + Entry entry = existing; if (!entry.closed && entry.active != null) { toCancel.addAll(entry.active.values()); } @@ -194,6 +201,13 @@ void removeExpired() { }); } terminate(timedOut, true, true); + ClosedAttempt fence; + while ((fence = closedExpirations.poll()) != null) { + ClosedAttempt expiredFence = fence; + // A later close may have extended the cutoff. An old queue event must not remove it. + closedAttempts.computeIfPresent( + fence.key, (key, current) -> current == expiredFence ? null : current); + } } private static void terminate( @@ -215,23 +229,49 @@ private static void terminate( } } - private static Entry retireIfUnused(QueryAttemptId key, Entry entry, long now) { + private Entry retireIfUnused(QueryAttemptId key, Entry entry, long now) { if ((entry.active == null || entry.active.isEmpty()) && (entry.closed || entry.remaining == 0)) { entry.closed = true; entry.context = null; entry.active = null; - if (key.isLegacy() || expired(entry.fenceDeadline, now)) { - return null; - } + retainClosure(key, entry.fenceDeadline, now); + return null; } return entry; } + private void retainClosure(QueryAttemptId key, long deadline, long now) { + if (key.isLegacy() || expired(deadline, now)) { + return; + } + // Called while holding the active-map key lock, before removing that Entry. Registration + // checks this map under the same key lock, so there is no gap in the admission fence. + closedAttempts.compute( + key, + (ignored, current) -> { + if (current != null && current.deadline - deadline >= 0) { + return current; + } + ClosedAttempt next = + new ClosedAttempt(current == null ? key : current.key, deadline, ticker); + closedExpirations.add(next); + return next; + }); + } + int getRetainedEntryCount() { + return attempts.size() + closedAttempts.size(); + } + + int getActiveEntryCount() { return attempts.size(); } + int getPendingExpirationCount() { + return closedExpirations.size(); + } + boolean hasSharedContext(QueryAttemptId key) { AtomicBoolean present = new AtomicBoolean(); attempts.computeIfPresent( @@ -243,6 +283,29 @@ boolean hasSharedContext(QueryAttemptId key) { return present.get(); } + private static class ClosedAttempt implements Delayed { + private final QueryAttemptId key; + private final long deadline; + private final LongSupplier ticker; + + private ClosedAttempt(QueryAttemptId key, long deadline, LongSupplier ticker) { + this.key = key; + this.deadline = deadline; + this.ticker = ticker; + } + + @Override + public long getDelay(TimeUnit unit) { + return unit.convert(deadline - ticker.getAsLong(), TimeUnit.NANOSECONDS); + } + + @Override + public int compareTo(Delayed other) { + // All queued cutoffs are within the bounded token window, including nanoTime wraparound. + return Long.compare(deadline - ((ClosedAttempt) other).deadline, 0); + } + } + private static class Entry { private long executionStart; private long executionTimeout; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java index 46d7119911b63..35b8c70b3287f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java @@ -533,22 +533,13 @@ private void dispatchRemoteHelper(final FragmentInstance instance, final TEndPoi instance.getFragment().getPlanNodeTree(), instance.getSessionInfo() == null ? null : instance.getSessionInfo().getUserName()); boolean transferAttemptRecorded = false; - boolean legacyAdmissionFallback = false; try (final SyncDataNodeInternalServiceClient client = syncInternalServiceClientManager.borrowClient(endPoint)) { switch (instance.getType()) { case READ: final TSendFragmentInstanceReq sendFragmentInstanceReq = new TSendFragmentInstanceReq(new TFragmentInstance(instance.serializeToByteBuffer())); - try { - prepareAdmission(instance, endPoint, client); - } catch (TApplicationException e) { - if (e.getType() != TApplicationException.UNKNOWN_METHOD) { - throw e; - } - legacyAdmissionFallback = true; - throw e; - } + prepareAdmission(instance, endPoint, client); if (instance.hasQueryAttempt()) { sendFragmentInstanceReq .setQueryAttemptId(instance.getQueryAttemptId()) @@ -657,14 +648,11 @@ private void dispatchRemoteHelper(final FragmentInstance instance, final TEndPoi recordWriteTransferFailureIfNecessary(endPoint, containsUserData, e); } throw e; + } catch (LegacyAdmissionException ignored) { + // The resource scope has returned the invalidated connection. Capability negotiation + // needs a fresh connection without consuming the FI's normal reconnect retry. + dispatchRemoteHelper(instance, endPoint); } catch (TException e) { - if (legacyAdmissionFallback) { - // UNKNOWN_METHOD invalidates the pooled connection. The resource scope above must - // return it before borrowing again, including when the pool has only one slot. - // Capability negotiation does not consume the FI's normal reconnect retry. - dispatchRemoteHelper(instance, endPoint); - return; - } if (!transferAttemptRecorded) { recordWriteTransferFailureIfNecessary(endPoint, containsUserData, e); } @@ -866,12 +854,14 @@ private void prepareAdmission( || ticker.getAsLong() - cached.requestStart >= TimeUnit.MILLISECONDS.toNanos(cached.token.getValidityMs())) { long requestStart = ticker.getAsLong(); + // A short query must not leave a fence for the full socket-timeout window. This is a + // relative duration, so capping it does not compare coordinator and receiver clocks. + long validityMs = Math.min(admissionWindowMs, Math.max(1, instance.getTimeOut())); try { cached.token = client == null - ? FragmentInstanceManager.getInstance() - .fetchQueryAdmissionToken(admissionWindowMs) - : client.fetchQueryAdmissionToken(admissionWindowMs); + ? FragmentInstanceManager.getInstance().fetchQueryAdmissionToken(validityMs) + : client.fetchQueryAdmissionToken(validityMs); cached.requestStart = requestStart; } catch (TException e) { Throwable cause = ExceptionUtils.getRootCause(e); @@ -884,7 +874,7 @@ private void prepareAdmission( // as legacy, then let the caller replace the connection invalidated by that handler. checkNotAborted(); cached.legacy = true; - throw (TApplicationException) cause; + throw new LegacyAdmissionException(cause); } } synchronized (this) { @@ -906,6 +896,15 @@ public synchronized Map getIssuedAdmissionToken return snapshot; } + /** Distinguishes capability negotiation from an error returned by the FI RPC itself. */ + private static class LegacyAdmissionException extends TException { + private static final long serialVersionUID = 1L; + + private LegacyAdmissionException(Throwable cause) { + super(cause); + } + } + private static class CachedAdmission { private TQueryAdmissionToken token; private long requestStart; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java index 1d39fbf7cfdca..5adda94bd00b0 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java @@ -293,6 +293,65 @@ public void freshAttemptWorksAfterOldFenceExpiresAndMonotonicCounterWraps() { registry.register(fresh, state(fresh), true).close(); } + @Test + public void completedAttemptsAreNotScannedBeforeTheirExpiry() { + java.util.concurrent.atomic.AtomicInteger reads = + new java.util.concurrent.atomic.AtomicInteger(); + QueryAttemptRegistry many = + new QueryAttemptRegistry( + () -> { + reads.incrementAndGet(); + return clock.get(); + }); + TQueryAdmissionToken issued = many.issueAdmissionToken(1000); + for (int i = 0; i < 10000; i++) { + many.close(new QueryAttemptId(new QueryId("closed_" + i), 0), issued, true); + } + assertEquals(0, many.getActiveEntryCount()); + assertEquals(10000, many.getRetainedEntryCount()); + reads.set(0); + many.removeExpired(); + // Only the queue head needs a clock check, independently of the retained marker count. + assertTrue(reads.get() <= 2); + assertEquals(10000, many.getRetainedEntryCount()); + clock.set(issued.getDeadlineNanos()); + many.removeExpired(); + assertEquals(0, many.getRetainedEntryCount()); + assertEquals(0, many.getPendingExpirationCount()); + } + + @Test + public void earlierExpiryCannotDeleteAnExtendedFence() { + TQueryAdmissionToken shortToken = registry.issueAdmissionToken(100); + registry.close(key(0), shortToken, true); + clock.addAndGet(TimeUnit.MILLISECONDS.toNanos(50)); + TQueryAdmissionToken extended = registry.issueAdmissionToken(1000); + registry.close(key(0), extended, true); + clock.set(shortToken.getDeadlineNanos()); + registry.removeExpired(); + assertEquals(1, registry.getRetainedEntryCount()); + FragmentInstance late = instance(0, 0, 1); + when(late.getAdmissionToken()).thenReturn(extended); + assertEquals( + TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), + assertThrows(IoTDBRuntimeException.class, () -> registry.register(late, state(late), true)) + .getErrorCode()); + clock.set(extended.getDeadlineNanos()); + registry.removeExpired(); + assertEquals(0, registry.getRetainedEntryCount()); + } + + @Test + public void repeatedCloseDoesNotAccumulateExpiryEvents() { + registry.close(key(0), token, true); + for (int i = 0; i < 1000; i++) { + registry.close(key(0), token, true); + } + assertEquals(1, registry.getPendingExpirationCount()); + assertEquals(1, registry.getRetainedEntryCount()); + assertEquals(0, registry.getActiveEntryCount()); + } + private QueryAttemptId key(int attempt) { return new QueryAttemptId(queryId, attempt); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRpcTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRpcTest.java new file mode 100644 index 0000000000000..f6c8d569dae41 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRpcTest.java @@ -0,0 +1,140 @@ +/* + * 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.execution.fragment; + +import org.apache.iotdb.commons.path.MeasurementPath; +import org.apache.iotdb.commons.queryengine.common.SessionInfo; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.protocol.thrift.impl.DataNodeInternalRPCServiceImpl; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.SeriesScanNode; +import org.apache.iotdb.db.service.DataNode.DataNodeContext; +import org.apache.iotdb.mpp.rpc.thrift.TCancelQueryReq; +import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstance; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceReq; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceResp; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.junit.BeforeClass; +import org.junit.Test; + +import java.time.ZoneId; +import java.util.Collections; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; + +public class QueryAttemptRpcTest { + private final DataNodeInternalRPCServiceImpl service = + new DataNodeInternalRPCServiceImpl(mock(DataNodeContext.class)); + + @BeforeClass + public static void configureNodeId() { + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0); + } + + @Test + public void serializedDispatchHonorsAnAttemptClosedByRpc() throws Exception { + QueryId queryId = new QueryId("rpc_closed_attempt"); + TQueryAdmissionToken token = service.fetchQueryAdmissionToken(30000); + assertTrue( + service + .cancelQuery( + new TCancelQueryReq(queryId.getId(), Collections.emptyList(), true) + .setQueryAttemptId(0) + .setAdmissionToken(token)) + .isCancelled()); + TSendFragmentInstanceResp response = + service.sendFragmentInstance(request(instance(queryId, 0), 0, token)); + assertFalse(response.isAccepted()); + assertEquals(TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), response.getStatus().getCode()); + } + + @Test + public void serializedExpiredAdmissionReturnsRetryableDispatchError() throws Exception { + TQueryAdmissionToken expired = + service + .fetchQueryAdmissionToken(30000) + .setDeadlineNanos(System.nanoTime() - TimeUnit.SECONDS.toNanos(1)); + TSendFragmentInstanceResp response = + service.sendFragmentInstance( + request(instance(new QueryId("rpc_expired_attempt"), 0), 0, expired)); + assertFalse(response.isAccepted()); + assertEquals(TSStatusCode.DISPATCH_ERROR.getStatusCode(), response.getStatus().getCode()); + } + + @Test + public void oldCancellationRpcDoesNotCloseRetry() throws Exception { + QueryId queryId = new QueryId("rpc_retry_attempt"); + TQueryAdmissionToken token = service.fetchQueryAdmissionToken(30000); + FragmentInstance retry = instance(queryId, 1); + retry.setQueryAttempt(1, token); + FragmentInstanceStateMachine state = + new FragmentInstanceStateMachine(retry.getId(), Runnable::run); + QueryAttemptRegistry registry = FragmentInstanceManager.getInstance().getQueryAttempts(); + try (QueryAttemptRegistry.Registration ignored = registry.register(retry, state, false)) { + service.cancelQuery( + new TCancelQueryReq(queryId.getId(), Collections.emptyList(), true) + .setQueryAttemptId(0) + .setAdmissionToken(token)); + assertEquals(FragmentInstanceState.RUNNING, state.getState()); + service.cancelQuery( + new TCancelQueryReq(queryId.getId(), Collections.emptyList(), true) + .setQueryAttemptId(1) + .setAdmissionToken(token)); + assertEquals(FragmentInstanceState.CANCELLED, state.getState()); + } + } + + private static FragmentInstance instance(QueryId queryId, int index) throws Exception { + PlanFragmentId fragmentId = new PlanFragmentId(queryId, index); + FragmentInstance instance = + new FragmentInstance( + new PlanFragment( + fragmentId, + new SeriesScanNode( + new PlanNodeId("scan"), new MeasurementPath("root.rpc.device.value"))), + fragmentId.genFragmentInstanceId(), + null, + QueryType.READ, + 30000, + new SessionInfo(0, "root", ZoneId.systemDefault()), + false, + false); + instance.setDataNodeFINum(3); + return instance; + } + + private static TSendFragmentInstanceReq request( + FragmentInstance instance, int attempt, TQueryAdmissionToken token) { + return new TSendFragmentInstanceReq(new TFragmentInstance(instance.serializeToByteBuffer())) + .setQueryAttemptId(attempt) + .setAdmissionToken(token); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java index d291e5831449d..5846254faad74 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java @@ -154,6 +154,20 @@ public void fragmentsOnTheSameNodeReuseOneTokenRpc() throws Exception { verify(client).fetchQueryAdmissionToken(anyLong()); } + @Test + public void admissionValidityDoesNotOutliveRemainingQueryBudget() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + FragmentInstance shortQuery = instance(0); + doReturn(5L).when(shortQuery).getTimeOut(); + when(client.fetchQueryAdmissionToken(anyLong())) + .thenAnswer( + invocation -> new TQueryAdmissionToken("receiver", 123456L, invocation.getArgument(0))); + when(client.sendFragmentInstance(any())).thenReturn(new TSendFragmentInstanceResp(true)); + assertTrue( + dispatcher.dispatch(null, Collections.singletonList(shortQuery)).get().isSuccessful()); + verify(client).fetchQueryAdmissionToken(5L); + } + @Test public void legacyReceiverUsesTheOriginalProtocol() throws Exception { FragmentInstanceDispatcherImpl dispatcher = dispatcher(); diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/concurrent/ThreadName.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/concurrent/ThreadName.java index f0d3b02dceead..906dfcdda77a5 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/concurrent/ThreadName.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/concurrent/ThreadName.java @@ -34,6 +34,7 @@ public enum ThreadName { QUERY_SENTINEL("Query-Sentinel-Thread"), TIMED_QUERY_SQL_COUNT("Timed-Query-SQL-Count"), FRAGMENT_INSTANCE_MANAGEMENT("Fragment-Instance-Management"), + QUERY_ATTEMPT_EXPIRATION("Query-Attempt-Expiration"), FRAGMENT_INSTANCE_NOTIFICATION("Fragment-Instance-Notification"), FRAGMENT_INSTANCE_DISPATCH("Fragment-Instance-Dispatch"), DRIVER_TASK_SCHEDULER_NOTIFICATION("Driver-Task-Scheduler-Notification"), @@ -225,6 +226,7 @@ public enum ThreadName { QUERY_SENTINEL, TIMED_QUERY_SQL_COUNT, FRAGMENT_INSTANCE_MANAGEMENT, + QUERY_ATTEMPT_EXPIRATION, FRAGMENT_INSTANCE_NOTIFICATION, DRIVER_TASK_SCHEDULER_NOTIFICATION)); private static final Set mppThreadNames = From 2608feb6b1f638bdb6b8362574982296ff529c55 Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Fri, 13 Nov 2026 15:02:54 +0800 Subject: [PATCH 4/4] [Query] Reuse fragment management executor for attempt expiry Schedule query attempt expiration on instanceManagementExecutor and remove the unused dedicated executor and thread name. Validated 142 targeted tests and full-reactor test compilation for both locales. --- .../execution/fragment/FragmentInstanceManager.java | 7 +------ .../org/apache/iotdb/commons/concurrent/ThreadName.java | 2 -- 2 files changed, 1 insertion(+), 8 deletions(-) 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 a1bb09ca6f57d..7926d2b501a76 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 @@ -87,7 +87,6 @@ public class FragmentInstanceManager { private final IDriverScheduler scheduler = DriverScheduler.getInstance(); private final ScheduledExecutorService instanceManagementExecutor; - private final ScheduledExecutorService queryAttemptExpirationExecutor; public final ExecutorService instanceNotificationExecutor; private final Duration infoCacheTime; @@ -111,9 +110,6 @@ private FragmentInstanceManager() { this.instanceManagementExecutor = IoTDBThreadPoolFactory.newScheduledThreadPool( 1, ThreadName.FRAGMENT_INSTANCE_MANAGEMENT.getName()); - this.queryAttemptExpirationExecutor = - IoTDBThreadPoolFactory.newScheduledThreadPool( - 1, ThreadName.QUERY_ATTEMPT_EXPIRATION.getName()); this.instanceNotificationExecutor = IoTDBThreadPoolFactory.newFixedThreadPool( IoTDBDescriptor.getInstance().getConfig().getFragmentInstanceNotificationThreadCount(), @@ -121,9 +117,8 @@ private FragmentInstanceManager() { this.infoCacheTime = new Duration(5, TimeUnit.MINUTES); - // Large bursts of expired admission fences must not delay stale-query or flushing cleanup. ScheduledExecutorUtil.safelyScheduleWithFixedDelay( - queryAttemptExpirationExecutor, + instanceManagementExecutor, queryAttempts::removeExpired, 2000, 2000, diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/concurrent/ThreadName.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/concurrent/ThreadName.java index 906dfcdda77a5..f0d3b02dceead 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/concurrent/ThreadName.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/concurrent/ThreadName.java @@ -34,7 +34,6 @@ public enum ThreadName { QUERY_SENTINEL("Query-Sentinel-Thread"), TIMED_QUERY_SQL_COUNT("Timed-Query-SQL-Count"), FRAGMENT_INSTANCE_MANAGEMENT("Fragment-Instance-Management"), - QUERY_ATTEMPT_EXPIRATION("Query-Attempt-Expiration"), FRAGMENT_INSTANCE_NOTIFICATION("Fragment-Instance-Notification"), FRAGMENT_INSTANCE_DISPATCH("Fragment-Instance-Dispatch"), DRIVER_TASK_SCHEDULER_NOTIFICATION("Driver-Task-Scheduler-Notification"), @@ -226,7 +225,6 @@ public enum ThreadName { QUERY_SENTINEL, TIMED_QUERY_SQL_COUNT, FRAGMENT_INSTANCE_MANAGEMENT, - QUERY_ATTEMPT_EXPIRATION, FRAGMENT_INSTANCE_NOTIFICATION, DRIVER_TASK_SCHEDULER_NOTIFICATION)); private static final Set mppThreadNames =