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 {