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 df18eaa2e4d2a..574feb482cc56 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,11 @@ package org.apache.iotdb.db.i18n; public final class DataNodeQueryMessages { + public static final String EXCEPTION_QUERY_ATTEMPT_ARG_IS_CLOSED_12143E89 = "Query attempt %s is closed."; + public static final String EXCEPTION_QUERY_ATTEMPT_ARG_ADMISSION_EXPIRED_DEADLINE_ARG_NS_NOW_ARG_NS_16A9846C = "Query attempt %s admission expired: deadline=%d ns, now=%d ns."; + public static final String EXCEPTION_INVALID_QUERY_ADMISSION_ISSUER_FOR_ATTEMPT_ARG_9D1C4E9C = "Invalid query admission issuer for attempt %s."; + + public static final String MESSAGE_FAIL_TO_INSERT_MEASUREMENTS_ARG_CAUSED_BY_ARG_B9858E6B = "Fail to insert measurements %s caused by %s"; 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 b2514babdfcd2..3895194221b60 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,11 @@ package org.apache.iotdb.db.i18n; public final class DataNodeQueryMessages { + public static final String EXCEPTION_QUERY_ATTEMPT_ARG_IS_CLOSED_12143E89 = "查询执行轮次 %s 已关闭。"; + public static final String EXCEPTION_QUERY_ATTEMPT_ARG_ADMISSION_EXPIRED_DEADLINE_ARG_NS_NOW_ARG_NS_16A9846C = "查询执行轮次 %s 的接收凭证已过期：截止时间=%d ns，当前时间=%d ns。"; + public static final String EXCEPTION_INVALID_QUERY_ADMISSION_ISSUER_FOR_ATTEMPT_ARG_9D1C4E9C = "查询执行轮次 %s 的接收凭证签发节点进程不匹配。"; + + public static final String MESSAGE_FAIL_TO_INSERT_MEASUREMENTS_ARG_CAUSED_BY_ARG_B9858E6B = "测点 %s 插入失败，原因：%s"; 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 3ce117f767a80..79e0065fd4f0a 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 @@ -133,6 +133,7 @@ import org.apache.iotdb.db.protocol.session.SessionManager; import org.apache.iotdb.db.protocol.thrift.OperationType; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.QueryId; import org.apache.iotdb.db.queryengine.execution.executor.RegionExecutionResult; import org.apache.iotdb.db.queryengine.execution.executor.RegionReadExecutor; import org.apache.iotdb.db.queryengine.execution.executor.RegionWriteExecutor; @@ -323,6 +324,7 @@ import org.apache.iotdb.mpp.rpc.thrift.TPushTopicMetaResp; import org.apache.iotdb.mpp.rpc.thrift.TPushTopicMetaRespExceptionMessage; import org.apache.iotdb.mpp.rpc.thrift.TPushTopicOwnerLeaseReq; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.iotdb.mpp.rpc.thrift.TRegionLeaderChangeReq; import org.apache.iotdb.mpp.rpc.thrift.TRegionLeaderChangeResp; import org.apache.iotdb.mpp.rpc.thrift.TRegionMigrateResult; @@ -510,6 +512,11 @@ private TSStatus waitForConsensusStarted() { } } + @Override + public TQueryAdmissionToken fetchQueryAdmissionToken(long validityMs) { + return FragmentInstanceManager.getInstance().fetchQueryAdmissionToken(validityMs); + } + @Override public TSendFragmentInstanceResp sendFragmentInstance(final TSendFragmentInstanceReq req) { LOGGER.debug(DataNodeMiscMessages.RECEIVE_FRAGMENT_INSTANCE, req.getConsensusGroupId()); @@ -556,6 +563,9 @@ public TSendFragmentInstanceResp sendFragmentInstance(final TSendFragmentInstanc FragmentInstance.deserializeFrom( fragmentInstanceId, req.fragmentInstance.body)); fragmentInstance.attachUDFQueryScope(queryScope); + if (req.isSetQueryAttemptId() && req.isSetAdmissionToken()) { + fragmentInstance.setQueryAttempt(req.getQueryAttemptId(), req.getAdmissionToken()); + } } catch (final Exception e) { if (classLoaderManager != null && queryScope != null) { classLoaderManager.finalizeUDFQuery(queryScope); @@ -642,6 +652,14 @@ public TFragmentInstanceInfoResp fetchFragmentInstanceInfo(TFetchFragmentInstanc @Override public TCancelResp cancelQuery(TCancelQueryReq req) { try (SetThreadName threadName = new SetThreadName(req.getQueryId())) { + if (req.isSetQueryAttemptId() && req.isSetAdmissionToken()) { + FragmentInstanceManager.getInstance() + .closeQueryAttempt( + new QueryId(req.getQueryId()), + req.getQueryAttemptId(), + req.getAdmissionToken(), + req.hasThrowable); + } List 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..ad9c932164cbe 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java @@ -109,6 +109,8 @@ public enum ExplainType { // time unit is ms private long startTime; + private int queryAttemptId; + private TEndPoint localDataBlockEndpoint; private TEndPoint localInternalEndpoint; private ResultNodeContext resultNodeContext; @@ -280,6 +282,7 @@ public void releaseMemoryForSchemaTree() { } public void prepareForRetry() { + queryAttemptId++; if (!isInnerTriggeredQuery()) { cleanUpCte(); } @@ -303,6 +306,10 @@ public QueryId getQueryId() { return queryId; } + public int getQueryAttemptId() { + return queryAttemptId; + } + public ExternalTsFileQueryResource createExternalTsFileQueryResource( String tableName, List 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..802505d695828 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/QueryAttemptId.java @@ -0,0 +1,59 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.queryengine.common; + +/** Identifies one execution of a query, independently of transport-level RPC retries. */ +public final class QueryAttemptId { + public static final int LEGACY_ATTEMPT = -1; + + private final QueryId queryId; + private final int attemptId; + + public QueryAttemptId(QueryId queryId, int attemptId) { + this.queryId = queryId; + this.attemptId = attemptId; + } + + public boolean isLegacy() { + return attemptId == LEGACY_ATTEMPT; + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof QueryAttemptId)) { + return false; + } + QueryAttemptId that = (QueryAttemptId) other; + return attemptId == that.attemptId && queryId.equals(that.queryId); + } + + @Override + public int hashCode() { + return 31 * queryId.hashCode() + attemptId; + } + + @Override + public String toString() { + return queryId + "/" + attemptId; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachine.java index 153c17536d541..7cfaee867a02e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachine.java @@ -178,7 +178,8 @@ public String getFailureMessage() { public Throwable getFailureException() { Throwable throwable = failureException.get(); if (throwable == null) { - return new IoTDBException(getFailureStatus()); + TSStatus status = getFailureStatus(); + return status == null ? null : new IoTDBException(status); } else { return throwable; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/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 676618ed420fc..6d33ea28a6066 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 @@ -41,7 +41,6 @@ import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.DeviceContext; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; -import org.apache.iotdb.db.queryengine.common.QueryId; import org.apache.iotdb.db.queryengine.metric.DriverSchedulerMetricSet; import org.apache.iotdb.db.queryengine.metric.QueryRelatedResourceMetricSet; import org.apache.iotdb.db.queryengine.metric.QueryResourceMetricSet; @@ -183,7 +182,8 @@ 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, @@ -244,7 +244,7 @@ public static FragmentInstanceContext createFragmentInstanceContext( SessionInfo sessionInfo, IDataRegionForQuery dataRegion, TimePredicate globalTimePredicate, - Map dataNodeQueryContextMap, + DataNodeQueryContext dataNodeQueryContext, long outerQueryDeadlineMs, boolean debug, boolean isVerbose) { @@ -255,7 +255,7 @@ public static FragmentInstanceContext createFragmentInstanceContext( sessionInfo, dataRegion, globalTimePredicate, - dataNodeQueryContextMap, + dataNodeQueryContext, outerQueryDeadlineMs, debug, isVerbose); @@ -359,7 +359,7 @@ private FragmentInstanceContext( SessionInfo sessionInfo, IDataRegionForQuery dataRegion, TimePredicate globalTimePredicate, - Map dataNodeQueryContextMap, + DataNodeQueryContext dataNodeQueryContext, long outerQueryDeadlineMs, boolean debug, boolean verbose) { @@ -375,8 +375,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()); } @@ -392,7 +391,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()); @@ -410,7 +408,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; } @@ -430,7 +427,6 @@ private FragmentInstanceContext( this.sessionInfo = sessionInfo; this.dataRegion = dataRegion; this.globalTimeFilter = globalTimeFilter; - this.dataNodeQueryContextMap = null; this.memoryReservationManager = new ThreadSafeMemoryReservationManager(id.getQueryId(), this.getClass().getName()); } @@ -452,7 +448,6 @@ protected FragmentInstanceContext( this.queryId = queryId; this.id = null; this.stateMachine = null; - this.dataNodeQueryContextMap = null; this.dataNodeQueryContext = null; this.dataRegion = dataRegion; this.globalTimeFilter = timeFilter; @@ -547,12 +542,6 @@ 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; } } @@ -1352,7 +1341,15 @@ private void releaseTVListOwnedByQuery() { * be decreased. */ public synchronized void releaseResource() { - finalizeUDFQuery(); + try { + finalizeUDFQuery(); + releaseResourceInternal(); + } finally { + releaseDataNodeQueryContext(); + } + } + + private void releaseResourceInternal() { try { getQueryChunkReaders().close(); } catch (IOException e) { @@ -1486,16 +1483,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; } synchronized void adoptUDFQueryScope(UDFClassLoaderManager.QueryScope transferredScope) { 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 71635656aff4b..7926d2b501a76 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceManager.java @@ -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; @@ -48,6 +49,7 @@ import org.apache.iotdb.db.storageengine.dataregion.IDataRegionForQuery; import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.mpp.rpc.thrift.TFetchFragmentInstanceStatisticsResp; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.iotdb.rpc.TSStatusCode; import org.apache.iotdb.udf.api.exception.UDFException; import org.apache.iotdb.udf.api.exception.UDFTypeMismatchException; @@ -80,7 +82,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 +107,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()); @@ -116,6 +117,13 @@ private FragmentInstanceManager() { this.infoCacheTime = new Duration(5, TimeUnit.MINUTES); + ScheduledExecutorUtil.safelyScheduleWithFixedDelay( + instanceManagementExecutor, + queryAttempts::removeExpired, + 2000, + 2000, + TimeUnit.MILLISECONDS); + ScheduledExecutorUtil.safelyScheduleWithFixedDelay( instanceManagementExecutor, this::removeOldInstances, 2000, 2000, TimeUnit.MILLISECONDS); ScheduledExecutorUtil.safelyScheduleWithFixedDelay( @@ -152,42 +160,54 @@ 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()); - FragmentInstanceContext createdContext = - createFragmentInstanceContext( - fragmentInstanceId, - stateMachine, - instance.getSessionInfo(), - dataRegion, - instance.getGlobalTimePredicate(), - dataNodeQueryContextMap, - DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()), - instance.isDebug(), - instance.isVerbose()); - createdContext.adoptUDFQueryScope(instance.claimUDFQueryScope()); - return createdContext; + registrations[0] = queryAttempts.register(instance, stateMachine, true); + try { + FragmentInstanceContext createdContext = + createFragmentInstanceContext( + fragmentInstanceId, + stateMachine, + instance.getSessionInfo(), + dataRegion, + instance.getGlobalTimePredicate(), + registrations[0].getContext(), + DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()), + instance.isDebug(), + instance.isVerbose()); + createdContext.adoptUDFQueryScope(instance.claimUDFQueryScope()); + return createdContext; + } 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()); + try { + context.finalizeUDFQuery(); + } finally { + 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())); // For ShowQueries related instances, isHighestPriority == true @@ -237,8 +257,16 @@ 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()); - context.finalizeUDFQuery(); - clearFIRelatedResources(instanceId); + try { + context.finalizeUDFQuery(); + 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; } }); @@ -298,11 +326,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) { @@ -316,26 +339,44 @@ 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; - FragmentInstanceContext createdContext = - createFragmentInstanceContext( - fragmentInstanceId, - stateMachine, - instance.getSessionInfo(), - DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()), - instance.isDebug(), - instance.isVerbose()); - createdContext.adoptUDFQueryScope(instance.claimUDFQueryScope()); - return createdContext; + registrations[0] = queryAttempts.register(instance, stateMachine, false); + try { + FragmentInstanceContext createdContext = + createFragmentInstanceContext( + fragmentInstanceId, + stateMachine, + instance.getSessionInfo(), + DriverTask.computeDeadlineTimeInMs(instance.getTimeOut()), + instance.isDebug(), + instance.isVerbose()); + createdContext.adoptUDFQueryScope(instance.claimUDFQueryScope()); + return createdContext; + } 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()); + try { + context.finalizeUDFQuery(); + } finally { + context.releaseDataNodeQueryContext(); + } + return null; + } List driverFactories = planner.plan( instance.getFragment().getPlanNodeTree(), @@ -376,8 +417,14 @@ public FragmentInstanceInfo execSchemaQueryFragmentInstance( // See execDataQueryFragmentInstance for why this result must not be fetched from // instanceContext after computeIfAbsent returns. failedInstanceInfo.set(context.getInstanceInfo()); - context.finalizeUDFQuery(); - clearFIRelatedResources(instanceId); + try { + context.finalizeUDFQuery(); + clearFIRelatedResources(instanceId); + } finally { + if (!context.isExecutionInitialized()) { + context.releaseDataNodeQueryContext(); + } + } return null; } }); @@ -425,6 +472,19 @@ public FragmentInstanceInfo cancelTask(FragmentInstanceId instanceId, boolean ha return null; } + public void closeQueryAttempt( + QueryId queryId, int attemptId, TQueryAdmissionToken token, boolean hasThrowable) { + queryAttempts.close(new QueryAttemptId(queryId, attemptId), token, hasThrowable); + } + + public TQueryAdmissionToken fetchQueryAdmissionToken(long validityMs) { + return queryAttempts.issueAdmissionToken(validityMs); + } + + QueryAttemptRegistry getQueryAttempts() { + return queryAttempts; + } + /** * Gets the info for the specified fragment instance. * 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..fcce3d453d1ac --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistry.java @@ -0,0 +1,389 @@ +/* + * 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.commons.exception.QueryTimeoutException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.QueryAttemptId; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.DelayQueue; +import java.util.concurrent.Delayed; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.LongSupplier; + +/** + * Owns shared state for a query attempt. Registration and closure serialize on the attempt key; + * planning and state transitions run outside that lock. Admission tokens are issued and checked + * against this DataNode process's monotonic clock, without comparing clocks across nodes. Token + * issuance is stateless: a delayed issuance RPC cannot create or reopen an attempt. + */ +class QueryAttemptRegistry { + private static final Logger LOGGER = LoggerFactory.getLogger(QueryAttemptRegistry.class); + private static final long MAX_ADMISSION_VALIDITY_MS = TimeUnit.DAYS.toMillis(1); + + // Only live attempts are scanned for execution timeout. Completed attempts retain a compact + // marker, indexed by deadline, rather than a full Entry on the management sweep's hot path. + private final Map attempts = new ConcurrentHashMap<>(); + private final Map closedAttempts = new ConcurrentHashMap<>(); + private final DelayQueue closedExpirations = new DelayQueue<>(); + private final LongSupplier ticker; + private final String issuer = UUID.randomUUID().toString(); + + QueryAttemptRegistry() { + this(System::nanoTime); + } + + QueryAttemptRegistry(LongSupplier ticker) { + this.ticker = ticker; + } + + TQueryAdmissionToken issueAdmissionToken(long validityMs) { + long boundedValidityMs = Math.max(1, Math.min(validityMs, MAX_ADMISSION_VALIDITY_MS)); + return new TQueryAdmissionToken( + issuer, + ticker.getAsLong() + TimeUnit.MILLISECONDS.toNanos(boundedValidityMs), + boundedValidityMs); + } + + private boolean issuedHere(TQueryAdmissionToken token) { + return token != null && issuer.equals(token.getIssuer()); + } + + private static boolean expired(long deadline, long now) { + // Subtraction also handles nanoTime wraparound; token validity is bounded to one day. + return now - deadline >= 0; + } + + Registration register( + FragmentInstance instance, FragmentInstanceStateMachine stateMachine, boolean dataQuery) { + QueryAttemptId key = + new QueryAttemptId( + instance.getId().getQueryId(), + instance.hasQueryAttempt() + ? instance.getQueryAttemptId() + : QueryAttemptId.LEGACY_ATTEMPT); + TQueryAdmissionToken token = instance.hasQueryAttempt() ? instance.getAdmissionToken() : null; + if (!key.isLegacy() && !issuedHere(token)) { + throw new IoTDBRuntimeException( + String.format( + DataNodeQueryMessages + .EXCEPTION_INVALID_QUERY_ADMISSION_ISSUER_FOR_ATTEMPT_ARG_9D1C4E9C, + key), + TSStatusCode.DISPATCH_ERROR.getStatusCode()); + } + AtomicReference result = new AtomicReference<>(); + attempts.compute( + key, + (ignored, existing) -> { + long now = ticker.getAsLong(); + if (token != null && expired(token.getDeadlineNanos(), now)) { + // Admission expiry does not mean the query exhausted its execution budget. Retry the + // query with a fresh attempt/token instead of renewing the old request's cutoff. + throw new IoTDBRuntimeException( + String.format( + DataNodeQueryMessages + .EXCEPTION_QUERY_ATTEMPT_ARG_ADMISSION_EXPIRED_DEADLINE_ARG_NS_NOW_ARG_NS_16A9846C, + key, + token.getDeadlineNanos(), + now), + TSStatusCode.DISPATCH_ERROR.getStatusCode()); + } + if ((existing != null && existing.closed) || closedAttempts.containsKey(key)) { + throw new IoTDBRuntimeException( + String.format( + DataNodeQueryMessages.EXCEPTION_QUERY_ATTEMPT_ARG_IS_CLOSED_12143E89, key), + TSStatusCode.QUERY_WAS_KILLED.getStatusCode()); + } + long timeoutNanos = TimeUnit.MILLISECONDS.toNanos(Math.max(0, instance.getTimeOut())); + if (timeoutNanos == 0 || (existing != null && existing.executionExpired(now))) { + throw new QueryTimeoutException(); + } + Entry entry = + existing == null + ? new Entry(now, timeoutNanos, Math.max(1, instance.getDataNodeFINum())) + : existing; + entry.shortenExecutionTimeout(now, timeoutNanos); + if (token != null) { + entry.extendFence(token.getDeadlineNanos()); + } + 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, TQueryAdmissionToken token, boolean hasThrowable) { + // A token from before a restart must not close a new incarnation's state. + if (!issuedHere(token)) { + return; + } + List toCancel = new ArrayList<>(); + attempts.compute( + key, + (ignored, existing) -> { + long now = ticker.getAsLong(); + if (existing == null) { + retainClosure(key, token.getDeadlineNanos(), now); + return null; + } + Entry entry = existing; + if (!entry.closed && entry.active != null) { + toCancel.addAll(entry.active.values()); + } + entry.closed = true; + entry.extendFence(token.getDeadlineNanos()); + return retireIfUnused(key, entry, now); + }); + terminate(toCancel, hasThrowable, false); + } + + void removeExpired() { + List timedOut = new ArrayList<>(); + for (QueryAttemptId key : attempts.keySet()) { + attempts.computeIfPresent( + key, + (ignored, entry) -> { + long now = ticker.getAsLong(); + if (!entry.closed && entry.executionExpired(now)) { + entry.closed = true; + timedOut.addAll(entry.active.values()); + } + // Closed entries were already terminated by their first closer. Fence expiry must not + // add a timeout cause to an explicitly cancelled or successfully completed FI. + return retireIfUnused(key, entry, now); + }); + } + terminate(timedOut, true, true); + ClosedAttempt fence; + while ((fence = closedExpirations.poll()) != null) { + ClosedAttempt expiredFence = fence; + // A later close may have extended the cutoff. An old queue event must not remove it. + closedAttempts.computeIfPresent( + fence.key, (key, current) -> current == expiredFence ? null : current); + } + } + + private static void terminate( + List registrations, boolean hasThrowable, boolean timeout) { + for (Registration registration : registrations) { + try { + if (timeout) { + if (!registration.stateMachine.getState().isDone()) { + registration.stateMachine.failed(new QueryTimeoutException()); + } + } else if (hasThrowable) { + registration.stateMachine.cancel(); + } else { + registration.stateMachine.finished(); + } + } catch (RuntimeException e) { + LOGGER.warn(DataNodeQueryMessages.ERRORS_FINISHING_FI_PROCESS, e); + } + } + } + + private Entry retireIfUnused(QueryAttemptId key, Entry entry, long now) { + if ((entry.active == null || entry.active.isEmpty()) + && (entry.closed || entry.remaining == 0)) { + entry.closed = true; + entry.context = null; + entry.active = null; + retainClosure(key, entry.fenceDeadline, now); + return null; + } + return entry; + } + + private void retainClosure(QueryAttemptId key, long deadline, long now) { + if (key.isLegacy() || expired(deadline, now)) { + return; + } + // Called while holding the active-map key lock, before removing that Entry. Registration + // checks this map under the same key lock, so there is no gap in the admission fence. + closedAttempts.compute( + key, + (ignored, current) -> { + if (current != null && current.deadline - deadline >= 0) { + return current; + } + ClosedAttempt next = + new ClosedAttempt(current == null ? key : current.key, deadline, ticker); + closedExpirations.add(next); + return next; + }); + } + + int getRetainedEntryCount() { + return attempts.size() + closedAttempts.size(); + } + + int getActiveEntryCount() { + return attempts.size(); + } + + int getPendingExpirationCount() { + return closedExpirations.size(); + } + + boolean hasSharedContext(QueryAttemptId key) { + AtomicBoolean present = new AtomicBoolean(); + attempts.computeIfPresent( + key, + (ignored, entry) -> { + present.set(entry.context != null); + return entry; + }); + return present.get(); + } + + private static class ClosedAttempt implements Delayed { + private final QueryAttemptId key; + private final long deadline; + private final LongSupplier ticker; + + private ClosedAttempt(QueryAttemptId key, long deadline, LongSupplier ticker) { + this.key = key; + this.deadline = deadline; + this.ticker = ticker; + } + + @Override + public long getDelay(TimeUnit unit) { + return unit.convert(deadline - ticker.getAsLong(), TimeUnit.NANOSECONDS); + } + + @Override + public int compareTo(Delayed other) { + // All queued cutoffs are within the bounded token window, including nanoTime wraparound. + return Long.compare(deadline - ((ClosedAttempt) other).deadline, 0); + } + } + + private static class Entry { + private long executionStart; + private long executionTimeout; + private long fenceDeadline; + private boolean hasFence; + private int remaining; + private boolean closed; + private DataNodeQueryContext context; + private Map active = new HashMap<>(); + + private Entry(long start, long timeout, int remaining) { + this.executionStart = start; + this.executionTimeout = timeout; + this.remaining = remaining; + } + + private boolean executionExpired(long now) { + return executionTimeout != Long.MAX_VALUE && now - executionStart >= executionTimeout; + } + + private void shortenExecutionTimeout(long now, long timeout) { + long remainingNanos = + executionTimeout == Long.MAX_VALUE + ? Long.MAX_VALUE + : Math.max(0, executionTimeout - (now - executionStart)); + if (timeout < remainingNanos) { + executionStart = now; + executionTimeout = timeout; + } + } + + private void extendFence(long deadline) { + if (!hasFence || deadline - fenceDeadline > 0) { + fenceDeadline = deadline; + hasFence = true; + } + } + } + + final class Registration implements AutoCloseable { + 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, ticker.getAsLong()); + }); + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index b55df9547b9c1..5ab3b122e1c6c 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 @@ -71,6 +71,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Throwables.throwIfUnchecked; @@ -95,7 +96,8 @@ 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 final AtomicReference schedulingAttempt = + new AtomicReference<>(new SchedulingAttempt()); private final QueryStateMachine stateMachine; private final IPlanner planner; @@ -111,7 +113,6 @@ public class QueryExecution implements IQueryExecution { // used for cleaning resultHandle up exactly once private final AtomicBoolean resultHandleCleanUp; - private final AtomicBoolean stopped; private final UDFClassLoaderManager.QueryScope udfQueryScope; // cost time in ns of finished rpc @@ -178,7 +179,6 @@ public QueryExecution(IPlanner planner, MPPQueryContext context, ExecutorService } } }); - this.stopped = new AtomicBoolean(false); this.resultHandleCleanUp = new AtomicBoolean(false); } @@ -323,7 +323,7 @@ private ExecutionResult retry() { // clear runtime variables in MPPQueryContext context.prepareForRetry(); // re-stop - this.stopped.compareAndSet(true, false); + this.schedulingAttempt.set(new SchedulingAttempt()); this.resultHandleCleanUp.compareAndSet(true, false); if (stateMachine.getState().isDone()) { return getStatus(); @@ -381,7 +381,13 @@ private IAnalysis analyzeWithCleanupOnFailure(MPPQueryContext context) { private void schedule() { final long startTime = System.nanoTime(); - this.scheduler = planner.doSchedule(analysis, distributedPlan, context, stateMachine); + SchedulingAttempt attempt = schedulingAttempt.get(); + 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); } @@ -431,9 +437,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.get().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); + } } } @@ -474,13 +515,31 @@ private void cleanUpResultHandle() { // Stop the query and clean up all the resources this query occupied @Override public void stopAndCleanup(Throwable t) { - stop(t); - releaseResource(t, false); + try { + // Terminal cleanup must wake status waiters, while retry cleanup keeps the query and its + // UDF scope alive for the next execution attempt. + if (!stateMachine.getState().isDone()) { + if (t == null) { + stateMachine.transitionToAborted(); + } else { + stateMachine.transitionToFailed(t); + } + } + } finally { + try { + stop(t); + } finally { + releaseResource(t, false); + } + } } private void cleanupForRetry(Throwable t) { - stop(t); - releaseResource(t, true); + try { + stop(t); + } finally { + releaseResource(t, true); + } } /** 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 10bba4daedc13..477a48ad56c9f 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 @@ -34,6 +34,7 @@ import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeUtil; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.ReadWriteIOUtils; @@ -72,6 +73,28 @@ 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 TQueryAdmissionToken admissionToken; + + public void setQueryAttempt(int queryAttemptId, TQueryAdmissionToken admissionToken) { + this.queryAttemptId = queryAttemptId; + this.admissionToken = admissionToken; + } + + public boolean hasQueryAttempt() { + return queryAttemptId >= 0 && admissionToken != null; + } + + public int getQueryAttemptId() { + return queryAttemptId; + } + + public TQueryAdmissionToken getAdmissionToken() { + return admissionToken; + } + private boolean isHighestPriority; private transient UDFClassLoaderManager.QueryScope udfQueryScope; 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..589923d26810d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java @@ -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::getIssuedAdmissionTokens); } } + 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) { + 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,14 +192,22 @@ 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) { - // 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. + if (terminateImmediately) { + // Dispatch failed, or stop arrived before dispatch completed. The FIs already sent may + // depend on unsent FIs and hold operator memory, so do not wait for the grace period. queryTerminator.terminateImmediately(t); } else { queryTerminator.terminate(t); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java index 9af76e1b40121..35b8c70b3287f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java @@ -39,12 +39,14 @@ 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; import org.apache.iotdb.db.queryengine.execution.executor.RegionExecutionResult; import org.apache.iotdb.db.queryengine.execution.executor.RegionReadExecutor; import org.apache.iotdb.db.queryengine.execution.executor.RegionWriteExecutor; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceManager; import org.apache.iotdb.db.queryengine.plan.Coordinator; import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; @@ -54,6 +56,7 @@ import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstance; import org.apache.iotdb.mpp.rpc.thrift.TPlanNode; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.iotdb.mpp.rpc.thrift.TSendBatchPlanNodeReq; import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceReq; import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceResp; @@ -62,6 +65,7 @@ import org.apache.iotdb.rpc.RpcUtils; import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.thrift.TApplicationException; import org.apache.thrift.TException; import org.apache.thrift.transport.TTransportException; import org.apache.tsfile.external.commons.lang3.exception.ExceptionUtils; @@ -72,13 +76,18 @@ import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.LongSupplier; import java.util.stream.Collectors; import static com.google.common.util.concurrent.Futures.immediateFuture; @@ -93,6 +102,15 @@ 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 admissionWindowMs; + private final LongSupplier ticker; + private final Map admissionCache = new ConcurrentHashMap<>(); + // Protected by this, also used to serialize token publication against abort(). + private final Map issuedAdmissionTokens = new HashMap<>(); + private volatile boolean aborted; + private final AtomicReference>> dispatchQueue = + new AtomicReference<>(); private final String localhostIpAddr; private final int localhostInternalPort; private final IClientManager @@ -118,8 +136,32 @@ public FragmentInstanceDispatcherImpl( IClientManager syncInternalServiceClientManager, IClientManager asyncInternalServiceClientManager) { + this( + type, + queryContext, + syncInternalServiceClientManager, + asyncInternalServiceClientManager, + System::nanoTime); + } + + FragmentInstanceDispatcherImpl( + QueryType type, + MPPQueryContext queryContext, + IClientManager syncInternalServiceClientManager, + IClientManager + asyncInternalServiceClientManager, + LongSupplier ticker) { this.type = type; this.queryContext = queryContext; + this.queryAttemptId = queryContext.getQueryAttemptId(); + this.ticker = ticker; + // Budget both the initial RPC and the supported reconnect retry. A zero socket timeout + // still needs a finite token lifetime; actual query execution keeps its own timeout. + this.admissionWindowMs = + 2L + * (COMMON_CONFIG.getDnConnectionTimeoutInMS() > 0 + ? COMMON_CONFIG.getDnConnectionTimeoutInMS() + : TimeUnit.MINUTES.toMillis(1)); this.syncInternalServiceClientManager = syncInternalServiceClientManager; this.asyncInternalServiceClientManager = asyncInternalServiceClientManager; this.localhostIpAddr = IoTDBDescriptor.getInstance().getConfig().getInternalAddress(); @@ -145,11 +187,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.set(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 +231,7 @@ private Future topologicalParallelDispatchRead( RpcUtils.getStatus( TSStatusCode.INTERNAL_SERVER_ERROR, UNEXPECTED_ERRORS + t.getMessage()))); } finally { + dispatchQueue.set(null); long queryDispatchReadTime = System.nanoTime() - startTime; QUERY_EXECUTION_METRIC_SET.recordExecutionCost(DISPATCH_READ, queryDispatchReadTime); queryContext.recordDispatchCost(queryDispatchReadTime); @@ -218,7 +262,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)); } @@ -462,9 +506,11 @@ private long getRetrySleepTime(int retryTimes) { } private void dispatchOneInstance(FragmentInstance instance) - throws FragmentInstanceDispatchException { + throws FragmentInstanceDispatchException, TException { + checkNotAborted(); TEndPoint endPoint = instance.getHostDataNode().getInternalEndPoint(); if (isDispatchedToLocal(endPoint)) { + prepareAdmission(instance, endPoint, null); dispatchLocally(instance); } else { dispatchRemote(instance, endPoint); @@ -493,6 +539,12 @@ private void dispatchRemoteHelper(final FragmentInstance instance, final TEndPoi case READ: final TSendFragmentInstanceReq sendFragmentInstanceReq = new TSendFragmentInstanceReq(new TFragmentInstance(instance.serializeToByteBuffer())); + prepareAdmission(instance, endPoint, client); + if (instance.hasQueryAttempt()) { + sendFragmentInstanceReq + .setQueryAttemptId(instance.getQueryAttemptId()) + .setAdmissionToken(instance.getAdmissionToken()); + } if (instance.getExecutorType().isStorageExecutor()) { sendFragmentInstanceReq.setConsensusGroupId( instance.getRegionReplicaSet().getRegionId()); @@ -596,6 +648,10 @@ private void dispatchRemoteHelper(final FragmentInstance instance, final TEndPoi recordWriteTransferFailureIfNecessary(endPoint, containsUserData, e); } throw e; + } catch (LegacyAdmissionException ignored) { + // The resource scope has returned the invalidated connection. Capability negotiation + // needs a fresh connection without consuming the FI's normal reconnect retry. + dispatchRemoteHelper(instance, endPoint); } catch (TException e) { if (!transferAttemptRecorded) { recordWriteTransferFailureIfNecessary(endPoint, containsUserData, e); @@ -662,6 +718,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 +824,99 @@ private void dispatchLocally(final FragmentInstance instance) } @Override - public void abort() {} + public synchronized void abort() { + aborted = true; + LinkedBlockingQueue> queue = dispatchQueue.get(); + if (queue != null) { + queue.add(new Pair<>(null, false)); + } + } + + private void prepareAdmission( + FragmentInstance instance, TEndPoint endPoint, SyncDataNodeInternalServiceClient client) + throws FragmentInstanceDispatchException, TException { + checkNotAborted(); + if (instance.hasQueryAttempt()) { + // A resend must retain its old token, even after a socket timeout or a receiver restart. + // Expired/old-incarnation tokens are retryable dispatch failures, not execution timeouts. + return; + } + CachedAdmission cached = + admissionCache.computeIfAbsent(endPoint, ignored -> new CachedAdmission()); + // Serialize token acquisition per destination, but never hold the abort/publication lock + // across an RPC. Other FIs on this host reuse the token while it is conservatively valid. + synchronized (cached) { + checkNotAborted(); + if (cached.legacy) { + return; + } + if (cached.token == null + || ticker.getAsLong() - cached.requestStart + >= TimeUnit.MILLISECONDS.toNanos(cached.token.getValidityMs())) { + long requestStart = ticker.getAsLong(); + // A short query must not leave a fence for the full socket-timeout window. This is a + // relative duration, so capping it does not compare coordinator and receiver clocks. + long validityMs = Math.min(admissionWindowMs, Math.max(1, instance.getTimeOut())); + try { + cached.token = + client == null + ? FragmentInstanceManager.getInstance().fetchQueryAdmissionToken(validityMs) + : client.fetchQueryAdmissionToken(validityMs); + cached.requestStart = requestStart; + } catch (TException e) { + Throwable cause = ExceptionUtils.getRootCause(e); + if (!(cause instanceof TApplicationException) + || ((TApplicationException) cause).getType() + != TApplicationException.UNKNOWN_METHOD) { + throw e; + } + // The client interceptor wraps application exceptions in TException. Mark this peer + // as legacy, then let the caller replace the connection invalidated by that handler. + checkNotAborted(); + cached.legacy = true; + throw new LegacyAdmissionException(cause); + } + } + synchronized (this) { + checkNotAborted(); + instance.setQueryAttempt(queryAttemptId, cached.token); + TQueryAdmissionToken previous = issuedAdmissionTokens.get(endPoint); + if (previous == null + || !previous.getIssuer().equals(cached.token.getIssuer()) + || cached.token.getDeadlineNanos() - previous.getDeadlineNanos() > 0) { + issuedAdmissionTokens.put(endPoint, cached.token); + } + } + } + } + + public synchronized Map getIssuedAdmissionTokens() { + Map snapshot = new HashMap<>(); + issuedAdmissionTokens.forEach((endPoint, token) -> snapshot.put(endPoint, token.deepCopy())); + return snapshot; + } + + /** Distinguishes capability negotiation from an error returned by the FI RPC itself. */ + private static class LegacyAdmissionException extends TException { + private static final long serialVersionUID = 1L; + + private LegacyAdmissionException(Throwable cause) { + super(cause); + } + } + + private static class CachedAdmission { + private TQueryAdmissionToken token; + private long requestStart; + private boolean legacy; + } + + 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/IQueryTerminator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java index 3be8958684732..bd36a7d0046aa 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/IQueryTerminator.java @@ -26,9 +26,9 @@ public interface IQueryTerminator { Future terminate(Throwable t); /** - * Cancels the unfinished FragmentInstances of the query without a grace period. Used when the - * dispatch failed: the dispatched FragmentInstances can't finish on their own and hold their - * resources until they are cancelled. + * Cancels the unfinished FragmentInstances without a grace period when dispatch failed or was + * stopped before completing. FIs already sent may depend on unsent FIs and hold resources until + * they are cancelled. */ Future terminateImmediately(Throwable t); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/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..f01a5c4974f53 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminator.java @@ -31,11 +31,13 @@ import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceManager; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; import org.apache.iotdb.mpp.rpc.thrift.TCancelQueryReq; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.thrift.TException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -43,6 +45,7 @@ import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; import java.util.stream.Collectors; public class SimpleQueryTerminator implements IQueryTerminator { @@ -50,6 +53,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 Supplier> admissionTokensSupplier; + private Map admissionTokens = Collections.emptyMap(); private final MPPQueryContext queryContext; private final IFragInstanceStateTracker stateTracker; @@ -67,9 +73,12 @@ public SimpleQueryTerminator( final List fragmentInstances, final IClientManager internalServiceClientManager, - final IFragInstanceStateTracker stateTracker) { + final IFragInstanceStateTracker stateTracker, + final Supplier> admissionTokensSupplier) { this.scheduledExecutor = scheduledExecutor; this.queryId = queryContext.getQueryId(); + this.queryAttemptId = queryContext.getQueryAttemptId(); + this.admissionTokensSupplier = admissionTokensSupplier; this.queryContext = queryContext; this.internalServiceClientManager = internalServiceClientManager; this.stateTracker = stateTracker; @@ -101,6 +110,9 @@ private Future terminate(final Throwable t, final long delayInMs) { .filter(endPoint -> !queryContext.getEndPointBlackList().contains(endPoint)) .collect(Collectors.toList()); this.hasThrowable = Objects.nonNull(t); + // Aborting dispatch first freezes the tokens used on each receiver. A token acquisition that + // finishes afterwards is stateless and cannot publish an FI or reopen the closed attempt. + this.admissionTokens = admissionTokensSupplier.get(); return scheduledExecutor.schedule(this::syncTerminate, delayInMs, TimeUnit.MILLISECONDS); } @@ -108,10 +120,12 @@ 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()) { + TQueryAdmissionToken token = admissionTokens.get(endPoint); + if (unfinishedFIs.isEmpty() && token == null) { continue; } @@ -119,6 +133,10 @@ public Boolean syncTerminate() { final int internalPort = IoTDBDescriptor.getInstance().getConfig().getInternalPort(); if (internalAddress.equalsIgnoreCase(endPoint.getIp()) && internalPort == endPoint.getPort()) { + if (token != null) { + FragmentInstanceManager.getInstance() + .closeQueryAttempt(queryId, queryAttemptId, token, hasThrowable); + } for (final FragmentInstanceId insId : unfinishedFIs) { FragmentInstanceManager.getInstance().cancelTask(insId, hasThrowable); } @@ -127,13 +145,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 (token != null) { + request.setQueryAttemptId(queryAttemptId).setAdmissionToken(token); + } + 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..a25e6358f9ea5 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 @@ -20,14 +20,17 @@ package org.apache.iotdb.db.queryengine.execution.fragment; import org.apache.iotdb.commons.exception.MetadataLeaseFencedException; +import org.apache.iotdb.commons.udf.service.UDFClassLoaderManager; import org.apache.iotdb.db.conf.IoTDBDescriptor; 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.db.queryengine.plan.planner.plan.PlanFragment; import org.apache.iotdb.db.schemaengine.schemaregion.ISchemaRegion; import org.apache.iotdb.db.storageengine.dataregion.IDataRegionForQuery; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; import org.apache.iotdb.rpc.TSStatusCode; import org.junit.BeforeClass; @@ -66,6 +69,103 @@ public void testSchemaQueryFailureInfoAfterConcurrentContextRemoval() throws Exc assertFailureInfoAfterConcurrentContextRemoval(false); } + @Test + public void testPartialDispatchReleasesAttemptContextAndRejectsLateFragment() { + QueryId queryId = new QueryId("partial_attempt_dispatch"); + FragmentInstanceManager manager = FragmentInstanceManager.getInstance(); + TQueryAdmissionToken token = manager.fetchQueryAdmissionToken(TimeUnit.MINUTES.toMillis(1)); + org.apache.iotdb.db.queryengine.common.QueryAttemptId key = + new org.apache.iotdb.db.queryengine.common.QueryAttemptId(queryId, 0); + try { + 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.getAdmissionToken()).thenReturn(token); + 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, token, true); + org.junit.Assert.assertFalse(manager.getQueryAttempts().hasSharedContext(key)); + FragmentInstance late = mock(FragmentInstance.class); + when(late.getId()).thenReturn(new FragmentInstanceId(new PlanFragmentId(queryId, 2), "0")); + when(late.hasQueryAttempt()).thenReturn(true); + when(late.getQueryAttemptId()).thenReturn(0); + when(late.getAdmissionToken()).thenReturn(token); + org.junit.Assert.assertThrows( + org.apache.iotdb.commons.exception.IoTDBRuntimeException.class, + () -> manager.execDataQueryFragmentInstance(late, mock(IDataRegionForQuery.class))); + } finally { + manager.closeQueryAttempt(queryId, 0, token, true); + for (int i = 0; i < 3; i++) { + manager.cancelTask(new FragmentInstanceId(new PlanFragmentId(queryId, i), "0"), true); + } + } + } + + @Test + public void testCancelledDataAdmissionReleasesTransferredUdfScope() throws Exception { + assertCancelledAdmissionReleasesUdfScope(true); + } + + @Test + public void testCancelledSchemaAdmissionReleasesTransferredUdfScope() throws Exception { + assertCancelledAdmissionReleasesUdfScope(false); + } + + private void assertCancelledAdmissionReleasesUdfScope(boolean dataQuery) throws Exception { + UDFClassLoaderManager classLoaders = UDFClassLoaderManager.getInstance(); + if (classLoaders == null) { + classLoaders = UDFClassLoaderManager.setupAndGetInstance(); + } + UDFClassLoaderManager loader = classLoaders; + QueryId queryId = new QueryId("cancelled_admission_" + dataQuery); + FragmentInstanceId id = new FragmentInstanceId(new PlanFragmentId(queryId, 0), "0"); + UDFClassLoaderManager.QueryScope scope = loader.openQueryScope(id.getFullId()); + FragmentInstanceManager manager = FragmentInstanceManager.getInstance(); + TQueryAdmissionToken token = manager.fetchQueryAdmissionToken(30000); + FragmentInstance instance = mock(FragmentInstance.class); + when(instance.getId()).thenReturn(id); + when(instance.hasQueryAttempt()).thenReturn(true); + when(instance.getQueryAttemptId()).thenReturn(0); + when(instance.getAdmissionToken()).thenReturn(token); + when(instance.getTimeOut()).thenReturn(30000L); + when(instance.getDataNodeFINum()).thenReturn(3); + when(instance.claimUDFQueryScope()).thenReturn(scope); + // The attempt closes after registration, before the newly created FI context adopts the + // deserializer's scope. No execution listener will exist to perform its cleanup. + when(instance.getSessionInfo()) + .thenAnswer( + ignored -> { + manager.closeQueryAttempt(queryId, 0, token, true); + return null; + }); + try { + FragmentInstanceInfo info = + dataQuery + ? manager.execDataQueryFragmentInstance(instance, mock(IDataRegionForQuery.class)) + : manager.execSchemaQueryFragmentInstance(instance, mock(ISchemaRegion.class)); + assertEquals(FragmentInstanceState.CANCELLED, info.getState()); + org.junit.Assert.assertThrows( + IllegalStateException.class, () -> loader.withQueryScope(scope, () -> null)); + org.junit.Assert.assertFalse( + manager.getQueryAttempts().hasSharedContext(new QueryAttemptId(queryId, 0))); + } finally { + manager.cancelTask(id, true); + loader.finalizeUDFQuery(scope); + } + } + 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/QueryAttemptAdmissionTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptAdmissionTest.java new file mode 100644 index 0000000000000..d3592a0164a70 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptAdmissionTest.java @@ -0,0 +1,184 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.queryengine.execution.fragment; + +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.client.IClientManager; +import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient; +import org.apache.iotdb.commons.conf.CommonConfig; +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.exception.IoTDBRuntimeException; +import org.apache.iotdb.commons.partition.ExecutorType; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment; +import org.apache.iotdb.db.queryengine.plan.scheduler.FragInstanceDispatchResult; +import org.apache.iotdb.db.queryengine.plan.scheduler.FragmentInstanceDispatcherImpl; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceReq; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceResp; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.thrift.transport.TTransportException; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +public class QueryAttemptAdmissionTest { + @Test + public void fullSocketTimeoutStillAllowsTheReconnectRetry() throws Exception { + assertTransportTimeout(20, true); + } + + @Test + public void expiredAdmissionIsRetryableInsteadOfReportingQueryTimeout() throws Exception { + assertTransportTimeout(41, false); + } + + private void assertTransportTimeout(long firstRpcElapsedMs, boolean expectedSuccess) + throws Exception { + CommonConfig config = CommonDescriptor.getInstance().getConfig(); + int previousTimeout = config.getDnConnectionTimeoutInMS(); + config.setDnConnectionTimeoutInMS(20); + try { + QueryId queryId = new QueryId("socket_timeout_admission"); + MPPQueryContext context = new MPPQueryContext(queryId); + context.setStartTime(System.currentTimeMillis()); + context.setTimeOut(TimeUnit.HOURS.toMillis(1)); + // An arbitrary receiver monotonic origin, unrelated to the coordinator's wall clock. + AtomicLong receiverTick = new AtomicLong(-TimeUnit.DAYS.toNanos(123)); + QueryAttemptRegistry receiver = new QueryAttemptRegistry(receiverTick::get); + FragmentInstance instance = + spy( + new FragmentInstance( + mock(PlanFragment.class), + new FragmentInstanceId(new PlanFragmentId(queryId, 0), "0"), + null, + QueryType.READ, + TimeUnit.HOURS.toMillis(1), + null, + false, + false)); + instance.setDataNodeFINum(1); + instance.setExecutorType(mock(ExecutorType.class)); + doReturn(ByteBuffer.wrap(new byte[] {1})).when(instance).serializeToByteBuffer(); + TEndPoint endpoint = new TEndPoint("remote-host", 10730); + doReturn(new TDataNodeLocation().setInternalEndPoint(endpoint)) + .when(instance) + .getHostDataNode(); + IClientManager clients = + mock(IClientManager.class); + SyncDataNodeInternalServiceClient client = mock(SyncDataNodeInternalServiceClient.class); + when(clients.borrowClient(endpoint)).thenReturn(client); + when(client.fetchQueryAdmissionToken(anyLong())) + .thenAnswer(invocation -> receiver.issueAdmissionToken(invocation.getArgument(0))); + List requests = new ArrayList<>(); + when(client.sendFragmentInstance(any())) + .thenAnswer( + invocation -> { + TSendFragmentInstanceReq request = invocation.getArgument(0); + requests.add(request.deepCopy()); + if (requests.size() == 1) { + receiverTick.addAndGet(TimeUnit.MILLISECONDS.toNanos(firstRpcElapsedMs)); + throw new TTransportException(TTransportException.TIMED_OUT, "socket timed out"); + } + FragmentInstanceId receivedId = instance.getId(); + FragmentInstance received = mock(FragmentInstance.class); + when(received.getId()).thenReturn(receivedId); + when(received.hasQueryAttempt()).thenReturn(true); + when(received.getQueryAttemptId()).thenReturn(request.getQueryAttemptId()); + when(received.getAdmissionToken()).thenReturn(request.getAdmissionToken()); + when(received.getDataNodeFINum()).thenReturn(1); + when(received.getTimeOut()).thenReturn(TimeUnit.HOURS.toMillis(1)); + try { + receiver + .register( + received, + new FragmentInstanceStateMachine(received.getId(), Runnable::run), + true) + .close(); + return new TSendFragmentInstanceResp(true); + } catch (IoTDBRuntimeException e) { + return new TSendFragmentInstanceResp(false) + .setNeedRetry(true) + .setStatus(new TSStatus(e.getErrorCode()).setMessage(e.getMessage())); + } + }); + FragmentInstanceDispatcherImpl dispatcher = + new FragmentInstanceDispatcherImpl(QueryType.READ, context, clients, null); + FragInstanceDispatchResult result = + dispatcher.dispatch(null, Collections.singletonList(instance)).get(); + assertEquals(expectedSuccess, result.isSuccessful()); + assertEquals(2, requests.size()); + assertEquals(requests.get(0).getAdmissionToken(), requests.get(1).getAdmissionToken()); + if (!expectedSuccess) { + assertEquals( + TSStatusCode.DISPATCH_ERROR.getStatusCode(), result.getFailureStatus().getCode()); + assertTrue( + result + .getFailureStatus() + .getMessage() + .contains(String.valueOf(requests.get(0).getAdmissionToken().getDeadlineNanos()))); + assertTrue( + result.getFailureStatus().getMessage().contains(String.valueOf(receiverTick.get()))); + } + } finally { + config.setDnConnectionTimeoutInMS(previousTimeout); + } + } + + @Test + public void tokenIssuanceAfterFenceExpiryCannotRecreateAnAttempt() { + AtomicLong tick = new AtomicLong(); + QueryAttemptRegistry receiver = new QueryAttemptRegistry(tick::get); + org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken token = receiver.issueAdmissionToken(1); + receiver.close( + new org.apache.iotdb.db.queryengine.common.QueryAttemptId(new QueryId("closed"), 0), + token, + true); + tick.set(TimeUnit.MILLISECONDS.toNanos(1)); + receiver.removeExpired(); + assertEquals(0, receiver.getRetainedEntryCount()); + // A delayed issuance RPC is deliberately query-agnostic and allocates no attempt entry. + receiver.issueAdmissionToken(Long.MAX_VALUE); + assertEquals(0, receiver.getRetainedEntryCount()); + assertTrue( + receiver.issueAdmissionToken(Long.MAX_VALUE).getValidityMs() <= TimeUnit.DAYS.toMillis(1)); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java new file mode 100644 index 0000000000000..5adda94bd00b0 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRegistryTest.java @@ -0,0 +1,379 @@ +/* + * 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.commons.exception.QueryTimeoutException; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryAttemptId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.junit.Test; + +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(TimeUnit.MILLISECONDS.toNanos(1000)); + private final QueryAttemptRegistry registry = new QueryAttemptRegistry(clock::get); + private final TQueryAdmissionToken token = registry.issueAdmissionToken(1000); + private final QueryId queryId = new QueryId("attempt_test"); + + @Test + 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), token, true); + assertEquals(FragmentInstanceState.CANCELLED, state.getState()); + assertTrue(registry.hasSharedContext(key(0))); + first.close(); + first.close(); + registry.close(key(0), token, true); + assertFalse(registry.hasSharedContext(key(0))); + // The third FI never arrived; it cannot keep the shared LAST/filter state alive. + assertEquals(1, registry.getRetainedEntryCount()); + } + + @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), token, true); + old.close(); + registry.close(key(0), token, true); + assertTrue(registry.hasSharedContext(key(1))); + assertFalse(registry.hasSharedContext(key(0))); + retry.close(); + } + + @Test + public void closedBeforeFirstDispatchRemainsFencedAfterRecordExpires() { + registry.close(key(0), token, true); + assertEquals( + TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), + assertThrows(IoTDBRuntimeException.class, () -> register(0, 0, 1)).getErrorCode()); + clock.set(TimeUnit.MILLISECONDS.toNanos(2000)); + registry.removeExpired(); + assertEquals(0, registry.getRetainedEntryCount()); + assertEquals( + TSStatusCode.DISPATCH_ERROR.getStatusCode(), + assertThrows(IoTDBRuntimeException.class, () -> register(0, 0, 1)).getErrorCode()); + 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(TimeUnit.MILLISECONDS.toNanos(2000)); + registry.removeExpired(); + assertEquals(FragmentInstanceState.FAILED, state.getState()); + FragmentInstanceContext context = + FragmentInstanceContext.createFragmentInstanceContext( + instance.getId(), state, null, Long.MAX_VALUE, false, false); + assertEquals( + TSStatusCode.QUERY_TIMEOUT.getStatusCode(), + context.getInstanceInfo().getErrorCode().get().getCode()); + first.close(); + assertEquals(0, registry.getRetainedEntryCount()); + } + + @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), token, 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), token, 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.getTimeOut()).thenReturn(Long.MAX_VALUE); + TQueryAdmissionToken shortToken = registry.issueAdmissionToken(200); + when(instance.getAdmissionToken()).thenReturn(shortToken); + FragmentInstanceStateMachine state = state(instance); + QueryAttemptRegistry.Registration running = registry.register(instance, state, true); + clock.set(TimeUnit.MILLISECONDS.toNanos(1300)); + registry.removeExpired(); + assertEquals(FragmentInstanceState.RUNNING, state.getState()); + assertTrue(registry.hasSharedContext(key(0))); + registry.close(key(0), shortToken, 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.getTimeOut()).thenReturn(Long.MAX_VALUE); + TQueryAdmissionToken shortToken = registry.issueAdmissionToken(200); + when(instance.getAdmissionToken()).thenReturn(shortToken); + registry.register(instance, state(instance), true).close(); + assertEquals(1, registry.getRetainedEntryCount()); + assertFalse(registry.hasSharedContext(key(0))); + clock.set(TimeUnit.MILLISECONDS.toNanos(1200)); + registry.removeExpired(); + assertEquals(0, registry.getRetainedEntryCount()); + } + + @Test + public void expiryOfClosedFenceDoesNotAddATimeoutCause() { + FragmentInstance instance = instance(0, 0, 3); + FragmentInstanceStateMachine state = state(instance); + QueryAttemptRegistry.Registration registration = registry.register(instance, state, true); + registry.close(key(0), token, true); + clock.addAndGet(TimeUnit.SECONDS.toNanos(2)); + registry.removeExpired(); + assertEquals(FragmentInstanceState.CANCELLED, state.getState()); + assertTrue(state.getFailureCauses().isEmpty()); + registration.close(); + assertEquals(0, registry.getRetainedEntryCount()); + } + + @Test + public void receiverClockOriginDoesNotShortenExecutionBudget() { + // The receiver can have an entirely different clock origin from the coordinator. Only the + // receiver's issued token and local relative execution timeout enter admission/expiry. + AtomicLong receiverTick = new AtomicLong(TimeUnit.DAYS.toNanos(123)); + QueryAttemptRegistry receiver = new QueryAttemptRegistry(receiverTick::get); + FragmentInstance instance = instance(0, 0, 1); + when(instance.getAdmissionToken()).thenReturn(receiver.issueAdmissionToken(20)); + FragmentInstanceStateMachine state = state(instance); + QueryAttemptRegistry.Registration registration = receiver.register(instance, state, true); + receiverTick.addAndGet(TimeUnit.MILLISECONDS.toNanos(999)); + receiver.removeExpired(); + assertEquals(FragmentInstanceState.RUNNING, state.getState()); + receiverTick.addAndGet(TimeUnit.MILLISECONDS.toNanos(1)); + receiver.removeExpired(); + assertEquals(FragmentInstanceState.FAILED, state.getState()); + assertTrue(state.getFailureCauses().peek() instanceof QueryTimeoutException); + registration.close(); + } + + @Test + public void oldProcessTokenCannotAdmitOrCloseANewProcessAttempt() { + QueryAttemptRegistry restarted = new QueryAttemptRegistry(clock::get); + FragmentInstance instance = instance(0, 0, 1); + assertEquals( + TSStatusCode.DISPATCH_ERROR.getStatusCode(), + assertThrows( + IoTDBRuntimeException.class, + () -> restarted.register(instance, state(instance), true)) + .getErrorCode()); + when(instance.getAdmissionToken()).thenReturn(restarted.issueAdmissionToken(1000)); + QueryAttemptRegistry.Registration registration = + restarted.register(instance, state(instance), true); + restarted.close(key(0), token, true); + assertTrue(restarted.hasSharedContext(key(0))); + registration.close(); + } + + @Test + public void freshAttemptWorksAfterOldFenceExpiresAndMonotonicCounterWraps() { + clock.set(Long.MAX_VALUE - TimeUnit.MILLISECONDS.toNanos(5)); + TQueryAdmissionToken issued = registry.issueAdmissionToken(10); + FragmentInstance instance = instance(0, 0, 1); + when(instance.getAdmissionToken()).thenReturn(issued); + registry.register(instance, state(instance), true).close(); + clock.addAndGet(TimeUnit.MILLISECONDS.toNanos(10)); + registry.removeExpired(); + assertEquals(0, registry.getRetainedEntryCount()); + assertThrows( + IoTDBRuntimeException.class, () -> registry.register(instance, state(instance), true)); + FragmentInstance fresh = instance(1, 1, 1); + when(fresh.getAdmissionToken()).thenReturn(registry.issueAdmissionToken(1000)); + registry.register(fresh, state(fresh), true).close(); + } + + @Test + public void completedAttemptsAreNotScannedBeforeTheirExpiry() { + java.util.concurrent.atomic.AtomicInteger reads = + new java.util.concurrent.atomic.AtomicInteger(); + QueryAttemptRegistry many = + new QueryAttemptRegistry( + () -> { + reads.incrementAndGet(); + return clock.get(); + }); + TQueryAdmissionToken issued = many.issueAdmissionToken(1000); + for (int i = 0; i < 10000; i++) { + many.close(new QueryAttemptId(new QueryId("closed_" + i), 0), issued, true); + } + assertEquals(0, many.getActiveEntryCount()); + assertEquals(10000, many.getRetainedEntryCount()); + reads.set(0); + many.removeExpired(); + // Only the queue head needs a clock check, independently of the retained marker count. + assertTrue(reads.get() <= 2); + assertEquals(10000, many.getRetainedEntryCount()); + clock.set(issued.getDeadlineNanos()); + many.removeExpired(); + assertEquals(0, many.getRetainedEntryCount()); + assertEquals(0, many.getPendingExpirationCount()); + } + + @Test + public void earlierExpiryCannotDeleteAnExtendedFence() { + TQueryAdmissionToken shortToken = registry.issueAdmissionToken(100); + registry.close(key(0), shortToken, true); + clock.addAndGet(TimeUnit.MILLISECONDS.toNanos(50)); + TQueryAdmissionToken extended = registry.issueAdmissionToken(1000); + registry.close(key(0), extended, true); + clock.set(shortToken.getDeadlineNanos()); + registry.removeExpired(); + assertEquals(1, registry.getRetainedEntryCount()); + FragmentInstance late = instance(0, 0, 1); + when(late.getAdmissionToken()).thenReturn(extended); + assertEquals( + TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), + assertThrows(IoTDBRuntimeException.class, () -> registry.register(late, state(late), true)) + .getErrorCode()); + clock.set(extended.getDeadlineNanos()); + registry.removeExpired(); + assertEquals(0, registry.getRetainedEntryCount()); + } + + @Test + public void repeatedCloseDoesNotAccumulateExpiryEvents() { + registry.close(key(0), token, true); + for (int i = 0; i < 1000; i++) { + registry.close(key(0), token, true); + } + assertEquals(1, registry.getPendingExpirationCount()); + assertEquals(1, registry.getRetainedEntryCount()); + assertEquals(0, registry.getActiveEntryCount()); + } + + private QueryAttemptId key(int attempt) { + return new QueryAttemptId(queryId, attempt); + } + + 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.getAdmissionToken()).thenReturn(token); + when(instance.getTimeOut()).thenReturn(1000L); + when(instance.getDataNodeFINum()).thenReturn(planned); + return instance; + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRpcTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRpcTest.java new file mode 100644 index 0000000000000..f6c8d569dae41 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryAttemptRpcTest.java @@ -0,0 +1,140 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.queryengine.execution.fragment; + +import org.apache.iotdb.commons.path.MeasurementPath; +import org.apache.iotdb.commons.queryengine.common.SessionInfo; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.protocol.thrift.impl.DataNodeInternalRPCServiceImpl; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.SeriesScanNode; +import org.apache.iotdb.db.service.DataNode.DataNodeContext; +import org.apache.iotdb.mpp.rpc.thrift.TCancelQueryReq; +import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstance; +import org.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceReq; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceResp; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.junit.BeforeClass; +import org.junit.Test; + +import java.time.ZoneId; +import java.util.Collections; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; + +public class QueryAttemptRpcTest { + private final DataNodeInternalRPCServiceImpl service = + new DataNodeInternalRPCServiceImpl(mock(DataNodeContext.class)); + + @BeforeClass + public static void configureNodeId() { + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0); + } + + @Test + public void serializedDispatchHonorsAnAttemptClosedByRpc() throws Exception { + QueryId queryId = new QueryId("rpc_closed_attempt"); + TQueryAdmissionToken token = service.fetchQueryAdmissionToken(30000); + assertTrue( + service + .cancelQuery( + new TCancelQueryReq(queryId.getId(), Collections.emptyList(), true) + .setQueryAttemptId(0) + .setAdmissionToken(token)) + .isCancelled()); + TSendFragmentInstanceResp response = + service.sendFragmentInstance(request(instance(queryId, 0), 0, token)); + assertFalse(response.isAccepted()); + assertEquals(TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), response.getStatus().getCode()); + } + + @Test + public void serializedExpiredAdmissionReturnsRetryableDispatchError() throws Exception { + TQueryAdmissionToken expired = + service + .fetchQueryAdmissionToken(30000) + .setDeadlineNanos(System.nanoTime() - TimeUnit.SECONDS.toNanos(1)); + TSendFragmentInstanceResp response = + service.sendFragmentInstance( + request(instance(new QueryId("rpc_expired_attempt"), 0), 0, expired)); + assertFalse(response.isAccepted()); + assertEquals(TSStatusCode.DISPATCH_ERROR.getStatusCode(), response.getStatus().getCode()); + } + + @Test + public void oldCancellationRpcDoesNotCloseRetry() throws Exception { + QueryId queryId = new QueryId("rpc_retry_attempt"); + TQueryAdmissionToken token = service.fetchQueryAdmissionToken(30000); + FragmentInstance retry = instance(queryId, 1); + retry.setQueryAttempt(1, token); + FragmentInstanceStateMachine state = + new FragmentInstanceStateMachine(retry.getId(), Runnable::run); + QueryAttemptRegistry registry = FragmentInstanceManager.getInstance().getQueryAttempts(); + try (QueryAttemptRegistry.Registration ignored = registry.register(retry, state, false)) { + service.cancelQuery( + new TCancelQueryReq(queryId.getId(), Collections.emptyList(), true) + .setQueryAttemptId(0) + .setAdmissionToken(token)); + assertEquals(FragmentInstanceState.RUNNING, state.getState()); + service.cancelQuery( + new TCancelQueryReq(queryId.getId(), Collections.emptyList(), true) + .setQueryAttemptId(1) + .setAdmissionToken(token)); + assertEquals(FragmentInstanceState.CANCELLED, state.getState()); + } + } + + private static FragmentInstance instance(QueryId queryId, int index) throws Exception { + PlanFragmentId fragmentId = new PlanFragmentId(queryId, index); + FragmentInstance instance = + new FragmentInstance( + new PlanFragment( + fragmentId, + new SeriesScanNode( + new PlanNodeId("scan"), new MeasurementPath("root.rpc.device.value"))), + fragmentId.genFragmentInstanceId(), + null, + QueryType.READ, + 30000, + new SessionInfo(0, "root", ZoneId.systemDefault()), + false, + false); + instance.setDataNodeFINum(3); + return instance; + } + + private static TSendFragmentInstanceReq request( + FragmentInstance instance, int attempt, TQueryAdmissionToken token) { + return new TSendFragmentInstanceReq(new TFragmentInstance(instance.serializeToByteBuffer())) + .setQueryAttemptId(attempt) + .setAdmissionToken(token); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/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..3dfbb37695158 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecutionLifecycleTest.java @@ -0,0 +1,171 @@ +/* + * 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)); + } + } + + @Test + public void terminalCleanupPreservesAnExistingStatusOnlyFailure() throws Exception { + ExecutorService notifications = Executors.newSingleThreadExecutor(); + IPlanner planner = planner(); + org.apache.iotdb.common.rpc.thrift.TSStatus original = + new org.apache.iotdb.common.rpc.thrift.TSStatus( + org.apache.iotdb.rpc.TSStatusCode.QUERY_EXECUTION_MEMORY_NOT_ENOUGH.getStatusCode()); + original.setMessage("original failure"); + when(planner.createScheduler(any(), any(), any(), any())) + .thenAnswer( + invocation -> { + QueryStateMachine state = invocation.getArgument(3); + IScheduler scheduler = mock(IScheduler.class); + doAnswer( + ignored -> { + state.transitionToFailed(original); + return null; + }) + .when(scheduler) + .start(); + return scheduler; + }); + QueryExecution execution = new QueryExecution(planner, context(), notifications); + try { + execution.start(); + execution.stopAndCleanup( + new org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException(0, 100, 1)); + org.junit.Assert.assertEquals(original, execution.getStatus().status); + } finally { + execution.stopAndCleanup(null); + notifications.shutdownNow(); + assertTrue(notifications.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + private static IPlanner planner() { + IPlanner planner = mock(IPlanner.class); + when(planner.analyze(any())).thenReturn(mock(IAnalysis.class)); + 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 2698f81d047d0..ec2e570bfa5d0 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 @@ -91,6 +91,7 @@ public void testAnalysisAttemptIsRolledBackOnlyWhenRetryStarts() throws Exceptio 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()); @@ -281,19 +282,25 @@ public DistributedQueryPlan doDistributionPlan( } @Override - public IScheduler doSchedule( + public IScheduler createScheduler( IAnalysis analysis, DistributedQueryPlan distributedPlan, MPPQueryContext context, QueryStateMachine stateMachine) { - scheduleAttempts++; - stateMachine.transitionToDispatching(); - if (alwaysFailDispatch || (retryOnFirstDispatch && 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 || (retryOnFirstDispatch && 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..6fc2541bb7264 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterSchedulerLifecycleTest.java @@ -0,0 +1,130 @@ +/* + * 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 admissionFailureCanStartAFreshQueryAttempt() { + state.transitionToPlanned(); + when(dispatcher.dispatch(any(), any())) + .thenReturn( + immediateFuture( + new FragInstanceDispatchResult( + new TSStatus(TSStatusCode.DISPATCH_ERROR.getStatusCode())))); + scheduler.start(); + org.junit.Assert.assertEquals( + org.apache.iotdb.db.queryengine.execution.QueryState.PENDING_RETRY, state.getState()); + } + + @Test + public void normalCompletionKeepsGracePeriod() { + when(dispatcher.dispatch(any(), any())) + .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..5846254faad74 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherLifecycleTest.java @@ -0,0 +1,263 @@ +/* + * 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.TQueryAdmissionToken; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceReq; +import org.apache.iotdb.mpp.rpc.thrift.TSendFragmentInstanceResp; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.thrift.TApplicationException; +import org.apache.thrift.TException; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class FragmentInstanceDispatcherLifecycleTest { + private final QueryId queryId = new QueryId("dispatch_admission"); + private final MPPQueryContext context = new MPPQueryContext(queryId); + private final TEndPoint endpoint = new TEndPoint("remote-host", 10730); + private final IClientManager clients = + mock(IClientManager.class); + private final SyncDataNodeInternalServiceClient client = + mock(SyncDataNodeInternalServiceClient.class); + private final TQueryAdmissionToken token = new TQueryAdmissionToken("receiver", 123456L, 60000L); + + private FragmentInstanceDispatcherImpl dispatcher() throws Exception { + context.setStartTime(System.currentTimeMillis()); + context.setTimeOut(Long.MAX_VALUE); + when(clients.borrowClient(endpoint)).thenReturn(client); + when(client.fetchQueryAdmissionToken(anyLong())).thenReturn(token); + return new FragmentInstanceDispatcherImpl(QueryType.READ, context, clients, null); + } + + private FragmentInstance instance(int fragmentId) { + FragmentInstance instance = + spy( + new FragmentInstance( + mock(PlanFragment.class), + new FragmentInstanceId(new PlanFragmentId(queryId, fragmentId), "0"), + null, + QueryType.READ, + Long.MAX_VALUE, + null, + false, + false)); + doReturn(ByteBuffer.wrap(new byte[] {1})).when(instance).serializeToByteBuffer(); + doReturn(new TDataNodeLocation().setInternalEndPoint(endpoint)) + .when(instance) + .getHostDataNode(); + instance.setExecutorType(mock(ExecutorType.class)); + return instance; + } + + @Test + public void abortPreventsTokenAcquisitionAndDispatch() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + dispatcher.abort(); + FragInstanceDispatchResult result = + dispatcher.dispatch(null, Collections.singletonList(instance(0))).get(); + assertFalse(result.isSuccessful()); + assertEquals( + TSStatusCode.QUERY_WAS_KILLED.getStatusCode(), result.getFailureStatus().getCode()); + verify(client, never()).fetchQueryAdmissionToken(anyLong()); + verify(client, never()).sendFragmentInstance(any()); + } + + @Test + public void transportRetryReusesTokenAndStopFreezesIssuedTokens() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + context.prepareForRetry(); + List requests = new ArrayList<>(); + when(client.sendFragmentInstance(any())) + .thenAnswer( + invocation -> { + TSendFragmentInstanceReq request = invocation.getArgument(0); + requests.add(request.deepCopy()); + if (requests.size() == 1) { + throw new TException("response lost"); + } + return new TSendFragmentInstanceResp(true); + }); + FragmentInstance instance = instance(0); + assertTrue(dispatcher.dispatch(null, Collections.singletonList(instance)).get().isSuccessful()); + assertEquals(2, requests.size()); + assertEquals(0, requests.get(0).getQueryAttemptId()); + assertEquals(token, requests.get(0).getAdmissionToken()); + assertEquals(requests.get(0).getAdmissionToken(), requests.get(1).getAdmissionToken()); + verify(client).fetchQueryAdmissionToken(anyLong()); + Map issued = dispatcher.getIssuedAdmissionTokens(); + dispatcher.abort(); + assertFalse( + dispatcher.dispatch(null, Collections.singletonList(instance)).get().isSuccessful()); + assertEquals(2, requests.size()); + assertEquals(issued, dispatcher.getIssuedAdmissionTokens()); + } + + @Test + public void fragmentsOnTheSameNodeReuseOneTokenRpc() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + when(client.sendFragmentInstance(any())).thenReturn(new TSendFragmentInstanceResp(true)); + assertTrue( + dispatcher.dispatch(null, Arrays.asList(instance(0), instance(1))).get().isSuccessful()); + verify(client).fetchQueryAdmissionToken(anyLong()); + } + + @Test + public void admissionValidityDoesNotOutliveRemainingQueryBudget() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + FragmentInstance shortQuery = instance(0); + doReturn(5L).when(shortQuery).getTimeOut(); + when(client.fetchQueryAdmissionToken(anyLong())) + .thenAnswer( + invocation -> new TQueryAdmissionToken("receiver", 123456L, invocation.getArgument(0))); + when(client.sendFragmentInstance(any())).thenReturn(new TSendFragmentInstanceResp(true)); + assertTrue( + dispatcher.dispatch(null, Collections.singletonList(shortQuery)).get().isSuccessful()); + verify(client).fetchQueryAdmissionToken(5L); + } + + @Test + public void legacyReceiverUsesTheOriginalProtocol() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + when(client.fetchQueryAdmissionToken(anyLong())) + .thenThrow( + new TApplicationException(TApplicationException.UNKNOWN_METHOD, "unknown method")); + when(client.sendFragmentInstance(any())) + .thenAnswer( + invocation -> { + TSendFragmentInstanceReq request = invocation.getArgument(0); + assertFalse(request.isSetQueryAttemptId()); + assertFalse(request.isSetAdmissionToken()); + return new TSendFragmentInstanceResp(true); + }); + assertTrue( + dispatcher.dispatch(null, Arrays.asList(instance(0), instance(1))).get().isSuccessful()); + verify(client).fetchQueryAdmissionToken(anyLong()); + assertTrue(dispatcher.getIssuedAdmissionTokens().isEmpty()); + } + + @Test + public void wrappedUnknownMethodUsesAFreshClientForLegacyDispatch() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + SyncDataNodeInternalServiceClient fresh = mock(SyncDataNodeInternalServiceClient.class); + java.util.concurrent.atomic.AtomicBoolean returned = + new java.util.concurrent.atomic.AtomicBoolean(); + java.util.concurrent.atomic.AtomicInteger borrows = + new java.util.concurrent.atomic.AtomicInteger(); + org.mockito.Mockito.doAnswer( + ignored -> { + returned.set(true); + return null; + }) + .when(client) + .close(); + when(clients.borrowClient(endpoint)) + .thenAnswer( + ignored -> { + if (borrows.getAndIncrement() == 0) { + return client; + } + assertTrue(returned.get()); + return fresh; + }); + when(client.fetchQueryAdmissionToken(anyLong())) + .thenThrow( + new TException( + new TApplicationException(TApplicationException.UNKNOWN_METHOD, "unknown method"))); + when(fresh.sendFragmentInstance(any())) + .thenAnswer( + invocation -> { + TSendFragmentInstanceReq request = invocation.getArgument(0); + assertFalse(request.isSetAdmissionToken()); + return new TSendFragmentInstanceResp(true); + }); + assertTrue( + dispatcher.dispatch(null, Collections.singletonList(instance(0))).get().isSuccessful()); + verify(client, never()).sendFragmentInstance(any()); + verify(fresh, never()).fetchQueryAdmissionToken(anyLong()); + verify(fresh).sendFragmentInstance(any()); + } + + @Test + public void abortDuringTokenRpcCannotPublishALateFragment() throws Exception { + FragmentInstanceDispatcherImpl dispatcher = dispatcher(); + ExecutorService executor = Executors.newSingleThreadExecutor(); + CountDownLatch acquiring = new CountDownLatch(1); + CountDownLatch returnToken = new CountDownLatch(1); + when(client.fetchQueryAdmissionToken(anyLong())) + .thenAnswer( + invocation -> { + acquiring.countDown(); + assertTrue(returnToken.await(10, TimeUnit.SECONDS)); + return token; + }); + try { + Future result = + executor.submit( + () -> dispatcher.dispatch(null, Collections.singletonList(instance(0))).get()); + assertTrue(acquiring.await(10, TimeUnit.SECONDS)); + dispatcher.abort(); + assertTrue(dispatcher.getIssuedAdmissionTokens().isEmpty()); + returnToken.countDown(); + assertFalse(result.get(10, TimeUnit.SECONDS).isSuccessful()); + verify(client, never()).sendFragmentInstance(any()); + assertTrue(dispatcher.getIssuedAdmissionTokens().isEmpty()); + } finally { + returnToken.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryAttemptProtocolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryAttemptProtocolTest.java new file mode 100644 index 0000000000000..d77cdc8e1a752 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryAttemptProtocolTest.java @@ -0,0 +1,95 @@ +/* + * 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.TQueryAdmissionToken; +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.isSetAdmissionToken()); + TCancelQueryReq cancel = new TCancelQueryReq("query", Collections.emptyList(), true); + TCancelQueryReq receivedCancel = new TCancelQueryReq(); + new TDeserializer().deserialize(receivedCancel, new TSerializer().serialize(cancel)); + assertEquals(cancel, receivedCancel); + assertFalse(receivedCancel.isSetQueryAttemptId()); + assertFalse(receivedCancel.isSetAdmissionToken()); + } + + @Test + public void attemptMetadataRoundTripsAndLegacyReaderCanSkipIt() throws Exception { + TFragmentInstance body = new TFragmentInstance(ByteBuffer.wrap(new byte[] {1, 2, 3})); + TSendFragmentInstanceReq request = + new TSendFragmentInstanceReq(body) + .setQueryAttemptId(2) + .setAdmissionToken(new TQueryAdmissionToken("receiver", 123450L, 1000L)); + 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) + .setAdmissionToken(new TQueryAdmissionToken("receiver", 123450L, 1000L)); + TCancelQueryReq receivedCancel = new TCancelQueryReq(); + new TDeserializer().deserialize(receivedCancel, new TSerializer().serialize(cancel)); + assertEquals(cancel, receivedCancel); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryTerminalCleanupTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryTerminalCleanupTest.java new file mode 100644 index 0000000000000..9a462fae86dac --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/QueryTerminalCleanupTest.java @@ -0,0 +1,152 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.queryengine.plan.scheduler; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.queryengine.common.SessionInfo; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.QueryState; +import org.apache.iotdb.db.queryengine.execution.QueryStateMachine; +import org.apache.iotdb.db.queryengine.plan.analyze.IAnalysis; +import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; +import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult; +import org.apache.iotdb.db.queryengine.plan.execution.QueryExecution; +import org.apache.iotdb.db.queryengine.plan.planner.IPlanner; +import org.apache.iotdb.db.queryengine.plan.planner.plan.DistributedQueryPlan; +import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.junit.Test; + +import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import static com.google.common.util.concurrent.Futures.immediateFuture; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class QueryTerminalCleanupTest { + @Test + public void staleQueryCleanupDuringDispatchMustUnblockGetStatus() throws Exception { + assertTerminalCleanup( + new QueryTimeoutRuntimeException(0, 120_001, 60_000), + QueryState.FAILED, + TSStatusCode.QUERY_TIMEOUT.getStatusCode()); + } + + @Test + public void nullCauseCleanupDuringDispatchAbortsAndUnblocksGetStatus() throws Exception { + assertTerminalCleanup( + null, QueryState.ABORTED, TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode()); + } + + private void assertTerminalCleanup(Throwable cause, QueryState expectedState, int expectedStatus) + throws Exception { + ExecutorService notifications = Executors.newSingleThreadExecutor(); + ExecutorService rpc = Executors.newSingleThreadExecutor(); + ExecutorService statusWaiter = Executors.newSingleThreadExecutor(); + CountDownLatch dispatchStarted = new CountDownLatch(1); + CountDownLatch finishDispatch = new CountDownLatch(1); + AtomicReference state = new AtomicReference<>(); + IFragInstanceDispatcher dispatcher = mock(IFragInstanceDispatcher.class); + when(dispatcher.dispatch(any(), any())) + .thenAnswer( + ignored -> { + dispatchStarted.countDown(); + assertTrue(finishDispatch.await(10, TimeUnit.SECONDS)); + return immediateFuture(new FragInstanceDispatchResult(true)); + }); + IPlanner planner = mock(IPlanner.class); + when(planner.analyze(any())).thenReturn(mock(IAnalysis.class)); + when(planner.doLogicalPlan(any(), any())).thenReturn(mock(LogicalQueryPlan.class)); + DistributedQueryPlan plan = mock(DistributedQueryPlan.class); + when(plan.getInstances()).thenReturn(Collections.emptyList()); + when(planner.doDistributionPlan(any(), any(), any())).thenReturn(plan); + when(planner.createScheduler(any(), any(), any(), any())) + .thenAnswer( + invocation -> { + QueryStateMachine machine = invocation.getArgument(3); + state.set(machine); + return new ClusterScheduler( + machine, + dispatcher, + mock(IFragInstanceStateTracker.class), + mock(IQueryTerminator.class)); + }); + QueryId queryId = new QueryId("stale_cleanup_during_dispatch"); + TEndPoint localExchange = new TEndPoint("127.0.0.1", 10740); + MPPQueryContext context = + new MPPQueryContext( + "", queryId, mock(SessionInfo.class), localExchange, new TEndPoint("127.0.0.1", 10730)); + context.setQueryType(QueryType.READ); + context.setStartTime(System.currentTimeMillis()); + context.setTimeOut(60_000L); + context + .getResultNodeContext() + .setUpStream( + localExchange, + queryId.genPlanFragmentId().genFragmentInstanceId(), + new PlanNodeId("upstream")); + QueryExecution execution = new QueryExecution(planner, context, notifications); + Future result = null; + try { + result = + rpc.submit( + () -> { + execution.start(); + return execution.getStatus(); + }); + assertTrue(dispatchStarted.await(10, TimeUnit.SECONDS)); + assertEquals(QueryState.DISPATCHING, state.get().getState()); + // Coordinator.cleanUpStaleQueries calls stopAndCleanup directly, without cancel(). + execution.stopAndCleanup(cause); + // Status waiters must wake even before the in-flight dispatch returns. + assertEquals( + expectedStatus, + statusWaiter.submit(execution::getStatus).get(10, TimeUnit.SECONDS).status.getCode()); + finishDispatch.countDown(); + assertEquals(expectedStatus, result.get(10, TimeUnit.SECONDS).status.getCode()); + assertEquals(expectedState, state.get().getState()); + } finally { + finishDispatch.countDown(); + execution.stopAndCleanup(null); + if (result != null) { + result.cancel(true); + } + statusWaiter.shutdownNow(); + assertTrue(statusWaiter.awaitTermination(10, TimeUnit.SECONDS)); + rpc.shutdownNow(); + notifications.shutdownNow(); + assertTrue(rpc.awaitTermination(10, TimeUnit.SECONDS)); + assertTrue(notifications.awaitTermination(10, TimeUnit.SECONDS)); + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java index 911074f2eca35..3dfb1c4202b36 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/SimpleQueryTerminatorTest.java @@ -19,17 +19,29 @@ 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.apache.iotdb.mpp.rpc.thrift.TQueryAdmissionToken; 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 +65,8 @@ public void setUp() { queryContext, Collections.emptyList(), null, - mock(IFragInstanceStateTracker.class)); + mock(IFragInstanceStateTracker.class), + Collections::emptyMap); } @Test @@ -69,4 +82,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.getEndPointBlackList()).thenReturn(Collections.emptySet()); + final TEndPoint endpoint = new TEndPoint("remote-host", 10730); + final TQueryAdmissionToken token = new TQueryAdmissionToken("receiver", 123456L, 1000L); + final FragmentInstance instance = mock(FragmentInstance.class); + when(instance.getHostDataNode()) + .thenReturn(new TDataNodeLocation().setInternalEndPoint(endpoint)); + 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, + () -> Collections.singletonMap(endpoint, token)); + // A delayed termination must retain the old attempt, even after the context starts a retry. + when(queryContext.getQueryAttemptId()).thenReturn(3); + terminator.terminateImmediately(new RuntimeException("dispatch failure")); + terminator.syncTerminate(); + ArgumentCaptor request = ArgumentCaptor.forClass(TCancelQueryReq.class); + verify(client).cancelQuery(request.capture()); + assertEquals(2, request.getValue().getQueryAttemptId()); + assertEquals(token, request.getValue().getAdmissionToken()); + 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..aaf76ff73418b 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -139,9 +139,19 @@ struct TPlanNode { 1: required binary body } +// Issued and checked by one DataNode process using its own monotonic clock. Opaque to peers. +struct TQueryAdmissionToken { + 1: required string issuer + 2: required i64 deadlineNanos + 3: required i64 validityMs +} + struct TSendFragmentInstanceReq { 1: required TFragmentInstance fragmentInstance 2: optional common.TConsensusGroupId consensusGroupId + // Attempt metadata is absent for a legacy coordinator. Do not change the binary FI body. + 3: optional i32 queryAttemptId + 4: optional TQueryAdmissionToken admissionToken } struct TSendFragmentInstanceResp { @@ -187,6 +197,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 TQueryAdmissionToken admissionToken } struct TCancelPlanFragmentReq { @@ -938,6 +951,9 @@ service IDataNodeRPCService { /** * dispatch FragmentInstance to remote node for query request */ + // Stateless: a delayed token request cannot create or reopen a query attempt. + TQueryAdmissionToken fetchQueryAdmissionToken(1: i64 validityMs); + TSendFragmentInstanceResp sendFragmentInstance(TSendFragmentInstanceReq req); /**