This is an automated email from the ASF dual-hosted git repository.
JackieTien97 pushed a commit to branch codex/fix-shutdown-interruption-logs
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to
refs/heads/codex/fix-shutdown-interruption-logs by this push:
new a37eade0251 Propagate remote write interruption provenance
a37eade0251 is described below
commit a37eade0251adcfaa3263eba15597d69a4ac1f3b
Author: JackieTien97 <[email protected]>
AuthorDate: Tue Oct 20 20:40:23 2026 +0800
Propagate remote write interruption provenance
---
.../iotdb/db/i18n/DataNodeQueryMessages.java | 1 +
.../iotdb/db/i18n/DataNodeQueryMessages.java | 1 +
.../queryengine/execution/QueryStateMachine.java | 49 ++++++-
.../plan/execution/ExecutionResult.java | 2 +-
.../queryengine/plan/execution/QueryExecution.java | 18 ++-
.../plan/scheduler/ClusterScheduler.java | 3 +-
.../plan/scheduler/FragInstanceDispatchResult.java | 12 ++
.../scheduler/FragmentInstanceDispatcherImpl.java | 6 +-
.../execution/QueryStateMachineTest.java | 35 +++++
.../metrics/IoTDBInternalLocalReporterTest.java | 157 ++++++++++++++++++++-
10 files changed, 265 insertions(+), 19 deletions(-)
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 123018f0eb0..35a2cf57f40 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
@@ -1243,6 +1243,7 @@ public final class DataNodeQueryMessages {
"Interrupted when dispatching read async";
public static final String INTERRUPTED_WHEN_DISPATCHING_WRITE_ASYNC =
"Interrupted when dispatching write async";
+ public static final String MESSAGE_INTERRUPTED_ERRORS_ARG_876F1B4B =
"Interrupted errors: %s";
public static final String DESERIALIZE_CONSENSUSGROUPID_FAILED =
"Deserialize ConsensusGroupId failed. ";
public static final String CAN_T_CONNECT_TO_NODE =
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 7c77e7218f3..a32008d51ad 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
@@ -1223,6 +1223,7 @@ public final class DataNodeQueryMessages {
"异步分发读取操作时被中断";
public static final String INTERRUPTED_WHEN_DISPATCHING_WRITE_ASYNC =
"异步分发写入操作时被中断";
+ public static final String MESSAGE_INTERRUPTED_ERRORS_ARG_876F1B4B =
"中断错误：%s";
public static final String DESERIALIZE_CONSENSUSGROUPID_FAILED =
"反序列化 ConsensusGroupId 失败。";
public static final String CAN_T_CONNECT_TO_NODE =
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 1bc3e2b0f32..2683fec32cf 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
@@ -54,7 +54,7 @@ public class QueryStateMachine {
private final StateMachine<QueryState> queryState;
private final AtomicReference<Throwable> failureException = new
AtomicReference<>();
- private final AtomicReference<TSStatus> failureStatus = new
AtomicReference<>();
+ private final AtomicReference<FailureStatus> failureStatus = new
AtomicReference<>();
public QueryStateMachine(QueryId queryId, ExecutorService executor) {
this.queryState =
@@ -90,7 +90,7 @@ public class QueryStateMachine {
}
public void transitionToPendingRetry(TSStatus failureStatus) {
- this.failureStatus.compareAndSet(null, failureStatus);
+ recordFailureStatus(failureStatus, false);
queryState.setIf(PENDING_RETRY, currentState -> currentState ==
DISPATCHING);
}
@@ -110,7 +110,7 @@ public class QueryStateMachine {
}
public boolean transitionToCanceled(Throwable throwable, TSStatus
failureStatus) {
- this.failureStatus.compareAndSet(null, failureStatus);
+ recordFailureStatus(failureStatus, false);
this.failureException.compareAndSet(null, throwable);
return transitionToDoneState(CANCELED);
}
@@ -129,10 +129,34 @@ public class QueryStateMachine {
}
public void transitionToFailed(TSStatus failureStatus) {
- this.failureStatus.compareAndSet(null, failureStatus);
+ transitionToFailed(failureStatus, false);
+ }
+
+ public void transitionToFailed(TSStatus failureStatus, boolean
failureCausedByInterruption) {
+ recordFailureStatus(failureStatus, failureCausedByInterruption);
transitionToDoneState(FAILED);
}
+ private void recordFailureStatus(TSStatus status, boolean
failureCausedByInterruption) {
+ if (status != null) {
+ // Publish the status and its origin together. A later interruption must
not reclassify
+ // the first recorded failure.
+ failureStatus.compareAndSet(null, new FailureStatus(status,
failureCausedByInterruption));
+ }
+ }
+
+ /** Only the exact recorded status carries the interruption provenance
stored with it. */
+ public boolean isFailureCausedByInterruption(TSStatus status) {
+ FailureStatus failure = failureStatus.get();
+ if (failure == null || failure.status != status ||
!failure.failureCausedByInterruption) {
+ return false;
+ }
+ // Statuses and exceptions can be recorded independently. Keep an
unrelated exception
+ // visible even if a subsequent dispatch interruption supplied the status.
+ Throwable throwable = failureException.get();
+ return throwable == null || getRootCause(throwable) instanceof
InterruptedException;
+ }
+
private boolean transitionToDoneState(QueryState doneState) {
requireNonNull(doneState,
DataNodeQueryMessages.EXCEPTION_DONESTATE_IS_NULL_D88F77E5);
checkArgument(
@@ -161,9 +185,9 @@ public class QueryStateMachine {
}
public TSStatus getFailureStatus() {
- TSStatus status = failureStatus.get();
+ FailureStatus status = failureStatus.get();
if (status != null) {
- return status;
+ return status.status;
} else {
Throwable throwable = failureException.get();
if (throwable != null) {
@@ -174,7 +198,18 @@ public class QueryStateMachine {
return RpcUtils.getStatus(((IoTDBException) t).getErrorCode(),
t.getMessage());
}
}
- return failureStatus.get();
+ status = failureStatus.get();
+ return status == null ? null : status.status;
+ }
+ }
+
+ private static class FailureStatus {
+ private final TSStatus status;
+ private final boolean failureCausedByInterruption;
+
+ private FailureStatus(TSStatus status, boolean
failureCausedByInterruption) {
+ this.status = status;
+ this.failureCausedByInterruption = failureCausedByInterruption;
}
}
}
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/ExecutionResult.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/ExecutionResult.java
index c06ee8b3954..2f46460635a 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/ExecutionResult.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/ExecutionResult.java
@@ -26,7 +26,7 @@ public class ExecutionResult {
public final QueryId queryId;
public final TSStatus status;
- /** Whether this status was generated by an interrupted wait without an
existing query failure. */
+ /** Whether this failure originated from an interrupted wait rather than
another query error. */
public final boolean failureCausedByInterruption;
public ExecutionResult(QueryId queryId, TSStatus status) {
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 49a7a375a33..b8f753d70a0 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
@@ -644,11 +644,14 @@ public class QueryExecution implements IQueryExecution {
} catch (InterruptedException e) {
// TODO: (xingtanzjr) use more accurate error handling
Thread.currentThread().interrupt();
- // An interrupted waiter can still observe a failure recorded before the
interruption.
- // Preserve its identity so callers do not mistake that failure for
shutdown cancellation.
+ // Dispatch may already have converted an interruption to a failure
status. Preserve its
+ // recorded origin without reclassifying an unrelated failure.
TSStatus failureStatus = stateMachine.getFailureStatus();
if (failureStatus != null) {
- return new ExecutionResult(context.getQueryId(), failureStatus);
+ return new ExecutionResult(
+ context.getQueryId(),
+ failureStatus,
+ stateMachine.isFailureCausedByInterruption(failureStatus));
}
return new ExecutionResult(
context.getQueryId(),
@@ -656,12 +659,14 @@ public class QueryExecution implements IQueryExecution {
// A terminal query may have failed with an exception that has no
dedicated status code.
!stateMachine.getState().isDone());
} catch (ExecutionException e) {
+ TSStatus failureStatus = stateMachine.getFailureStatus();
return new ExecutionResult(
context.getQueryId(),
- stateMachine.getFailureStatus() == null
+ failureStatus == null
? RpcUtils.getStatus(
TSStatusCode.INTERNAL_SERVER_ERROR,
stateMachine.getFailureMessage())
- : stateMachine.getFailureStatus());
+ : failureStatus,
+ stateMachine.isFailureCausedByInterruption(failureStatus));
}
}
@@ -732,7 +737,8 @@ public class QueryExecution implements IQueryExecution {
||
IoTDBConstant.ClientVersion.V_1_0.equals(context.getSession().getVersion())) {
planner.setRedirectInfo(analysis, CONFIG.getAddressAndPort(), tsstatus);
}
- return new ExecutionResult(context.getQueryId(), tsstatus);
+ return new ExecutionResult(
+ context.getQueryId(), tsstatus,
stateMachine.isFailureCausedByInterruption(tsstatus));
}
public DistributedQueryPlan getDistributedPlan() {
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 5e0aa02f5c0..4c828c3808c 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
@@ -125,7 +125,8 @@ public class ClusterScheduler implements IScheduler {
if (needRetry(result.getFailureStatus())) {
stateMachine.transitionToPendingRetry(result.getFailureStatus());
} else {
- stateMachine.transitionToFailed(result.getFailureStatus());
+ stateMachine.transitionToFailed(
+ result.getFailureStatus(),
result.isFailureCausedByInterruption());
}
return;
}
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragInstanceDispatchResult.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragInstanceDispatchResult.java
index 6a621139455..dd9d39694e8 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragInstanceDispatchResult.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragInstanceDispatchResult.java
@@ -26,13 +26,21 @@ public class FragInstanceDispatchResult {
private TSStatus failureStatus;
+ private final boolean failureCausedByInterruption;
+
public FragInstanceDispatchResult(boolean successful) {
this.successful = successful;
+ this.failureCausedByInterruption = false;
}
public FragInstanceDispatchResult(TSStatus failureStatus) {
+ this(failureStatus, false);
+ }
+
+ public FragInstanceDispatchResult(TSStatus failureStatus, boolean
failureCausedByInterruption) {
this.successful = false;
this.failureStatus = failureStatus;
+ this.failureCausedByInterruption = failureCausedByInterruption;
}
public boolean isSuccessful() {
@@ -42,4 +50,8 @@ public class FragInstanceDispatchResult {
public TSStatus getFailureStatus() {
return failureStatus;
}
+
+ public boolean isFailureCausedByInterruption() {
+ return failureCausedByInterruption;
+ }
}
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 24c42a0c16f..4894a8bc106 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
@@ -330,7 +330,11 @@ public class FragmentInstanceDispatcherImpl implements
IFragInstanceDispatcher {
return immediateFuture(
new FragInstanceDispatchResult(
RpcUtils.getStatus(
- TSStatusCode.INTERNAL_SERVER_ERROR, "Interrupted errors: " +
e.getMessage())));
+ TSStatusCode.INTERNAL_SERVER_ERROR,
+ String.format(
+
DataNodeQueryMessages.MESSAGE_INTERRUPTED_ERRORS_ARG_876F1B4B,
+ e.getMessage())),
+ true));
}
if (dispatchFailures.isEmpty()) {
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachineTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachineTest.java
index ad7ab094748..118ff7b5025 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachineTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/QueryStateMachineTest.java
@@ -31,13 +31,17 @@ import org.junit.Assert;
import org.junit.Test;
import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicInteger;
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
public class QueryStateMachineTest {
@@ -158,6 +162,37 @@ public class QueryStateMachineTest {
stateMachine.getFailureStatus().code);
}
+ @Test
+ public void testInterruptionProvenanceBelongsToTheRecordedStatus() {
+ QueryStateMachine stateMachine =
+ new QueryStateMachine(genQueryId(), mock(ExecutorService.class));
+ TSStatus interruption = new
TSStatus(TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode());
+
+ stateMachine.transitionToFailed(interruption, true);
+
+ assertSame(interruption, stateMachine.getFailureStatus());
+ assertTrue(stateMachine.isFailureCausedByInterruption(interruption));
+ assertFalse(
+ stateMachine.isFailureCausedByInterruption(
+ new TSStatus(TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode())));
+ }
+
+ @Test
+ public void testLaterInterruptionDoesNotReclassifyExistingFailure() {
+ QueryStateMachine stateMachine =
+ new QueryStateMachine(genQueryId(), mock(ExecutorService.class));
+ TSStatus original = new
TSStatus(TSStatusCode.WRITE_PROCESS_ERROR.getStatusCode());
+ TSStatus interruption = new
TSStatus(TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode());
+
+ stateMachine.transitionToFailed(original);
+ stateMachine.transitionToFailed(interruption, true);
+ stateMachine.transitionToFailed(original, true);
+
+ assertSame(original, stateMachine.getFailureStatus());
+ assertFalse(stateMachine.isFailureCausedByInterruption(original));
+ assertFalse(stateMachine.isFailureCausedByInterruption(interruption));
+ }
+
private QueryStateMachine genQueryStateMachine() {
return new QueryStateMachine(
genQueryId(),
IoTDBThreadPoolFactory.newSingleThreadExecutor("TestQueryStateMachine"));
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporterTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporterTest.java
index 46ab03ac685..a94cefc54c0 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporterTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporterTest.java
@@ -19,9 +19,17 @@
package org.apache.iotdb.db.service.metrics;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
+import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.client.IClientManager;
+import
org.apache.iotdb.commons.client.async.AsyncDataNodeInternalServiceClient;
import org.apache.iotdb.commons.conf.CommonDescriptor;
import org.apache.iotdb.commons.queryengine.common.SessionInfo;
+import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode;
import org.apache.iotdb.db.conf.IoTDBDescriptor;
import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
import org.apache.iotdb.db.queryengine.common.QueryId;
@@ -29,13 +37,19 @@ import
org.apache.iotdb.db.queryengine.execution.QueryStateMachine;
import org.apache.iotdb.db.queryengine.plan.Coordinator;
import org.apache.iotdb.db.queryengine.plan.analyze.IAnalysis;
import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher;
+import org.apache.iotdb.db.queryengine.plan.analyze.QueryType;
import org.apache.iotdb.db.queryengine.plan.analyze.schema.ISchemaFetcher;
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.FragmentInstance;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment;
+import org.apache.iotdb.db.queryengine.plan.scheduler.ClusterScheduler;
import org.apache.iotdb.db.queryengine.plan.statement.Statement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement;
import
org.apache.iotdb.db.queryengine.plan.statement.metadata.CreateMultiTimeSeriesStatement;
+import org.apache.iotdb.mpp.rpc.thrift.TSendBatchPlanNodeReq;
import org.apache.iotdb.rpc.TSStatusCode;
import ch.qos.logback.classic.Level;
@@ -52,6 +66,7 @@ import org.mockito.stubbing.OngoingStubbing;
import org.slf4j.LoggerFactory;
import java.lang.reflect.Field;
+import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
@@ -356,6 +371,138 @@ public class IoTDBInternalLocalReporterTest {
verifyReloadInterrupt(null, null, false);
}
+ @Test
+ public void testRemoteWriteInterruptedDuringReloadDoesNotWarn() throws
Exception {
+ verifyRemoteDispatchInterruption(false, false, null);
+ }
+
+ @Test
+ public void testRemoteAutoCreateInterruptedDuringReloadDoesNotWarn() throws
Exception {
+ verifyRemoteDispatchInterruption(true, false, null);
+ }
+
+ @Test
+ public void testRemoteInterruptionSurvivesClearedThreadFlag() throws
Exception {
+ verifyRemoteDispatchInterruption(false, true, null);
+ }
+
+ @Test
+ public void testRemoteInterruptionDoesNotReplaceExistingFailure() throws
Exception {
+ verifyRemoteDispatchInterruption(
+ false, false, new
TSStatus(TSStatusCode.WRITE_PROCESS_ERROR.getStatusCode()));
+ }
+
+ @Test
+ public void testRemoteInterruptionDoesNotHideExistingException() throws
Exception {
+ verifyRemoteDispatchInterruption(
+ false, false, null, new IllegalStateException("existing write
failure"));
+ }
+
+ private void verifyRemoteDispatchInterruption(
+ boolean autoCreate, boolean clearInterrupt, TSStatus recordedFailure)
throws Exception {
+ verifyRemoteDispatchInterruption(autoCreate, clearInterrupt,
recordedFailure, null);
+ }
+
+ private void verifyRemoteDispatchInterruption(
+ boolean autoCreate,
+ boolean clearInterrupt,
+ TSStatus recordedFailure,
+ Throwable recordedException)
+ throws Exception {
+ service = Executors.newSingleThreadScheduledExecutor();
+ reporter = newReporter(service);
+ MPPQueryContext context = mock(MPPQueryContext.class);
+ when(context.getQueryId()).thenReturn(new QueryId("remote_metric"));
+ IPlanner planner = mock(IPlanner.class);
+ when(planner.analyze(context)).thenReturn(mock(IAnalysis.class));
+ ExecutorService notifications = mock(ExecutorService.class);
+ QueryExecution queryExecution = new QueryExecution(planner, context,
notifications);
+ QueryStateMachine stateMachine = stateMachine(queryExecution);
+ if (recordedFailure != null) {
+ stateMachine.transitionToFailed(recordedFailure);
+ } else if (recordedException != null) {
+ stateMachine.transitionToFailed(recordedException);
+ } else {
+ stateMachine.transitionToPlanned();
+ }
+
+ TEndPoint endpoint = new TEndPoint("192.0.2.1", 10730);
+ TDataNodeLocation remote = new
TDataNodeLocation().setInternalEndPoint(endpoint);
+ TRegionReplicaSet region =
+ new TRegionReplicaSet(
+ new TConsensusGroupId(TConsensusGroupType.DataRegion, 1),
+ Collections.singletonList(remote));
+ PlanNode node = mock(PlanNode.class);
+ when(node.serializeToByteBuffer()).thenReturn(ByteBuffer.allocate(0));
+ when(node.getChildren()).thenReturn(Collections.emptyList());
+ PlanFragment fragment = mock(PlanFragment.class);
+ when(fragment.getPlanNodeTree()).thenReturn(node);
+ FragmentInstance instance = mock(FragmentInstance.class);
+ when(instance.getHostDataNode()).thenReturn(remote);
+ when(instance.getRegionReplicaSet()).thenReturn(region);
+ when(instance.getFragment()).thenReturn(fragment);
+ DistributedQueryPlan plan = mock(DistributedQueryPlan.class);
+ when(plan.getInstances()).thenReturn(Collections.singletonList(instance));
+
+ @SuppressWarnings("unchecked")
+ IClientManager<TEndPoint, AsyncDataNodeInternalServiceClient> clients =
+ mock(IClientManager.class);
+ AsyncDataNodeInternalServiceClient client =
mock(AsyncDataNodeInternalServiceClient.class);
+ when(clients.borrowClient(endpoint)).thenReturn(client);
+ CountDownLatch dispatched = new CountDownLatch(1);
+ doAnswer(
+ invocation -> {
+ // Leave the RPC response pending so the real
AsyncPlanNodeSender waits.
+ dispatched.countDown();
+ return null;
+ })
+ .when(client)
+ .sendBatchPlanNode(any(TSendBatchPlanNodeReq.class), any());
+ ClusterScheduler scheduler =
+ new ClusterScheduler(context, stateMachine, plan, QueryType.WRITE,
null, null, clients);
+ AtomicReference<ExecutionResult> returnedResult = new AtomicReference<>();
+ OngoingStubbing<ExecutionResult> execution = execution();
+ if (autoCreate) {
+ execution = execution.thenReturn(result(TSStatusCode.METADATA_ERROR));
+ }
+ execution.thenAnswer(
+ invocation -> {
+ scheduler.start();
+ if (clearInterrupt) {
+ Thread.interrupted();
+ }
+ // Deliver the getStatus() listener immediately, also exercising its
normal return
+ // path when the interrupt flag has been consumed after dispatch.
+ doAnswer(
+ notification -> {
+ notification.<Runnable>getArgument(0).run();
+ return null;
+ })
+ .when(notifications)
+ .execute(any(Runnable.class));
+ ExecutionResult result = queryExecution.getStatus();
+ returnedResult.set(result);
+ return result;
+ });
+
+ writeMetric();
+ assertTrue(dispatched.await(5, TimeUnit.SECONDS));
+ reporter.stop();
+ assertTrue(service.awaitTermination(5, TimeUnit.SECONDS));
+
+ ExecutionResult result = returnedResult.get();
+ if (recordedFailure != null) {
+ assertSame(recordedFailure, result.status);
+ } else {
+ assertEquals(TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode(),
result.status.getCode());
+ }
+ boolean hasRecordedFailure = recordedFailure != null || recordedException
!= null;
+ assertEquals(hasRecordedFailure, hasWarning());
+ assertEquals(!hasRecordedFailure, result.failureCausedByInterruption);
+ executedStatements(autoCreate ? 2 : 1);
+ verify(client).sendBatchPlanNode(any(TSendBatchPlanNodeReq.class), any());
+ }
+
private void verifyReloadInterrupt(
TSStatus recordedFailure, Throwable recordedException, boolean
autoCreate) throws Exception {
service = Executors.newSingleThreadScheduledExecutor();
@@ -408,9 +555,7 @@ public class IoTDBInternalLocalReporterTest {
when(planner.analyze(context)).thenReturn(mock(IAnalysis.class));
ExecutorService notifications = mock(ExecutorService.class);
QueryExecution execution = new QueryExecution(planner, context,
notifications);
- Field stateMachineField =
QueryExecution.class.getDeclaredField("stateMachine");
- stateMachineField.setAccessible(true);
- QueryStateMachine stateMachine = (QueryStateMachine)
stateMachineField.get(execution);
+ QueryStateMachine stateMachine = stateMachine(execution);
if (recordedFailure != null) {
stateMachine.transitionToFailed(recordedFailure);
} else if (recordedException != null) {
@@ -428,6 +573,12 @@ public class IoTDBInternalLocalReporterTest {
return execution;
}
+ private QueryStateMachine stateMachine(QueryExecution execution) throws
Exception {
+ Field stateMachineField =
QueryExecution.class.getDeclaredField("stateMachine");
+ stateMachineField.setAccessible(true);
+ return (QueryStateMachine) stateMachineField.get(execution);
+ }
+
private IoTDBInternalLocalReporter newReporter(ScheduledExecutorService
executor) {
return new IoTDBInternalLocalReporter(
executor, coordinator, sessionInfo, partitionFetcher, schemaFetcher);
