From 0030ac2c76a3b0d33f0306dea26534f28ac6feba Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Wed, 11 Nov 2026 17:02:31 +0800 Subject: [PATCH] [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 {