This is an automated email from the ASF dual-hosted git repository.
JackieTien97 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 9b2784c3ad4 Fix data exchange failure propagation (#18468)
9b2784c3ad4 is described below
commit 9b2784c3ad47f215d140aa49e25493c89f16a74d
Author: Jackie Tien <[email protected]>
AuthorDate: Fri Aug 14 12:14:44 2026 +0800
Fix data exchange failure propagation (#18468)
---
.../execution/exchange/MPPDataExchangeManager.java | 54 ++-
.../execution/exchange/sink/ShuffleSinkHandle.java | 121 ++++---
.../execution/exchange/sink/SinkChannel.java | 24 +-
.../execution/exchange/source/SourceHandle.java | 39 ++-
.../exchange/MPPDataExchangeManagerTest.java | 80 +++++
.../execution/exchange/ShuffleSinkHandleTest.java | 96 ++++++
.../SinkChannelFailurePropagationTest.java | 364 +++++++++++++++++++++
7 files changed, 699 insertions(+), 79 deletions(-)
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java
index 2d3447dee69..5a9371570a9 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java
@@ -463,17 +463,21 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
private final AtomicInteger cnt;
+ private final AtomicBoolean hasChannelFailedOrAborted;
+
private final AtomicBoolean hasDecremented = new AtomicBoolean(false);
public ISinkChannelListenerImpl(
TFragmentInstanceId localFragmentInstanceId,
FragmentInstanceContext context,
IMPPDataExchangeManagerCallback<Throwable> onFailureCallback,
- AtomicInteger cnt) {
+ AtomicInteger cnt,
+ AtomicBoolean hasChannelFailedOrAborted) {
this.shuffleSinkHandleId = localFragmentInstanceId;
this.context = context;
this.onFailureCallback = onFailureCallback;
this.cnt = cnt;
+ this.hasChannelFailedOrAborted = hasChannelFailedOrAborted;
}
@Override
@@ -496,6 +500,7 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
if (LOGGER.isDebugEnabled()) {
LOGGER.debug(DataNodeQueryMessages.SKH_LISTENER_ON_ABORT);
}
+ hasChannelFailedOrAborted.set(true);
decrementCnt();
return context.getFailureCause();
}
@@ -503,21 +508,25 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
@Override
public void onFailure(ISink sink, Throwable t) {
LOGGER.warn(DataNodeQueryMessages.ISINKCHANNEL_FAILED_DUE_TO, t);
- decrementCnt();
- if (onFailureCallback != null) {
- onFailureCallback.call(t);
+ hasChannelFailedOrAborted.set(true);
+ try {
+ if (onFailureCallback != null) {
+ onFailureCallback.call(t);
+ }
+ } finally {
+ decrementCnt();
}
}
private void decrementCnt() {
if (hasDecremented.compareAndSet(false, true) && (cnt.decrementAndGet()
== 0)) {
- closeShuffleSinkHandle();
+ releaseShuffleSinkHandle();
}
}
- private void closeShuffleSinkHandle() {
+ private void releaseShuffleSinkHandle() {
ISinkHandle sinkHandle = shuffleSinkHandles.remove(shuffleSinkHandleId);
- if (sinkHandle != null) {
+ if (sinkHandle != null && !hasChannelFailedOrAborted.get()) {
if (LOGGER.isDebugEnabled()) {
LOGGER.debug(DataNodeQueryMessages.CLOSE_SHUFFLE_SINK_HANDLE,
shuffleSinkHandleId);
}
@@ -651,7 +660,8 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
// TODO: replace with callbacks to decouple MPPDataExchangeManager from
// FragmentInstanceContext
FragmentInstanceContext instanceContext,
- AtomicInteger cnt) {
+ AtomicInteger cnt,
+ AtomicBoolean hasChannelFailedOrAborted) {
if (LOGGER.isDebugEnabled()) {
LOGGER.debug(
@@ -687,7 +697,11 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
localFragmentInstanceId,
queue,
new ISinkChannelListenerImpl(
- localFragmentInstanceId, instanceContext, instanceContext::failed,
cnt));
+ localFragmentInstanceId,
+ instanceContext,
+ instanceContext::failed,
+ cnt,
+ hasChannelFailedOrAborted));
}
/**
@@ -723,7 +737,8 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
// TODO: replace with callbacks to decouple MPPDataExchangeManager from
// FragmentInstanceContext
FragmentInstanceContext instanceContext,
- AtomicInteger cnt) {
+ AtomicInteger cnt,
+ AtomicBoolean hasChannelFailedOrAborted) {
if (LOGGER.isDebugEnabled()) {
LOGGER.debug(
@@ -743,7 +758,11 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
executorService,
tsBlockSerdeFactory.get(),
new ISinkChannelListenerImpl(
- localFragmentInstanceId, instanceContext, instanceContext::failed,
cnt),
+ localFragmentInstanceId,
+ instanceContext,
+ instanceContext::failed,
+ cnt,
+ hasChannelFailedOrAborted),
instanceContext.isHighestPriority(),
mppDataExchangeServiceClientManager);
}
@@ -767,6 +786,7 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
int channelNum = downStreamChannelLocationList.size();
AtomicInteger cnt = new AtomicInteger(channelNum);
+ AtomicBoolean hasChannelFailedOrAborted = new AtomicBoolean(false);
List<ISinkChannel> downStreamChannelList =
downStreamChannelLocationList.stream()
.map(
@@ -776,7 +796,8 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
localPlanNodeId,
downStreamChannelLocation,
instanceContext,
- cnt))
+ cnt,
+ hasChannelFailedOrAborted))
.collect(Collectors.toList());
ShuffleSinkHandle shuffleSinkHandle =
@@ -795,7 +816,8 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
String localPlanNodeId,
DownStreamChannelLocation downStreamChannelLocation,
FragmentInstanceContext instanceContext,
- AtomicInteger cnt) {
+ AtomicInteger cnt,
+ AtomicBoolean hasChannelFailedOrAborted) {
if (isSameNode(downStreamChannelLocation.getRemoteEndpoint())) {
return createLocalSinkChannel(
localFragmentInstanceId,
@@ -803,7 +825,8 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
downStreamChannelLocation.getRemotePlanNodeId(),
localPlanNodeId,
instanceContext,
- cnt);
+ cnt,
+ hasChannelFailedOrAborted);
} else {
return createSinkChannel(
localFragmentInstanceId,
@@ -812,7 +835,8 @@ public class MPPDataExchangeManager implements
IMPPDataExchangeManager {
downStreamChannelLocation.getRemotePlanNodeId(),
localPlanNodeId,
instanceContext,
- cnt);
+ cnt,
+ hasChannelFailedOrAborted);
}
}
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java
index 11c9d4d36bf..78e4928fdc3 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java
@@ -33,6 +33,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
@@ -63,6 +64,10 @@ public class ShuffleSinkHandle implements ISinkHandle {
private volatile boolean closed = false;
+ // close() and abort() invoke channel callbacks, so they cannot be protected
by this handle's
+ // lock.
+ private final AtomicBoolean terminationClaimed = new AtomicBoolean(false);
+
private static final DataExchangeCostMetricSet DATA_EXCHANGE_COST_METRIC_SET
=
DataExchangeCostMetricSet.getInstance();
private final Lock lock = new ReentrantLock();
@@ -201,37 +206,43 @@ public class ShuffleSinkHandle implements ISinkHandle {
@Override
public boolean abort() {
- if (aborted || closed) {
+ if (aborted || closed || !terminationClaimed.compareAndSet(false, true)) {
return false;
}
- if (LOGGER.isDebugEnabled()) {
- LOGGER.debug(DataNodeQueryMessages.START_ABORT_SHUFFLE_SINK_HANDLE);
- }
- boolean meetError = false;
- Exception firstException = null;
- boolean selfAborted = true;
- for (ISink channel : downStreamChannelList) {
- try {
- selfAborted = channel.abort();
- } catch (Exception e) {
- if (!meetError) {
- firstException = e;
- meetError = true;
+ try {
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug(DataNodeQueryMessages.START_ABORT_SHUFFLE_SINK_HANDLE);
+ }
+ boolean meetError = false;
+ Exception firstException = null;
+ boolean selfAborted = true;
+ for (ISink channel : downStreamChannelList) {
+ try {
+ selfAborted = channel.abort();
+ } catch (Exception e) {
+ if (!meetError) {
+ firstException = e;
+ meetError = true;
+ }
}
}
- }
- if (meetError) {
-
LOGGER.warn(DataNodeQueryMessages.ERROR_OCCURRED_WHEN_TRY_TO_ABORT_CHANNEL,
firstException);
- }
- if (selfAborted) {
- sinkListener.onAborted(this);
- aborted = true;
- if (LOGGER.isDebugEnabled()) {
- LOGGER.debug(DataNodeQueryMessages.END_ABORT_SHUFFLE_SINK_HANDLE);
+ if (meetError) {
+
LOGGER.warn(DataNodeQueryMessages.ERROR_OCCURRED_WHEN_TRY_TO_ABORT_CHANNEL,
firstException);
+ }
+ if (selfAborted) {
+ sinkListener.onAborted(this);
+ aborted = true;
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug(DataNodeQueryMessages.END_ABORT_SHUFFLE_SINK_HANDLE);
+ }
+ return true;
+ } else {
+ return false;
+ }
+ } finally {
+ if (!aborted) {
+ terminationClaimed.set(false);
}
- return true;
- } else {
- return false;
}
}
@@ -241,37 +252,43 @@ public class ShuffleSinkHandle implements ISinkHandle {
// Lock ShuffleSinkHandle and wait to lock LocalSinkChannel
@Override
public boolean close() {
- if (closed || aborted) {
+ if (closed || aborted || !terminationClaimed.compareAndSet(false, true)) {
return false;
}
- if (LOGGER.isDebugEnabled()) {
- LOGGER.debug(DataNodeQueryMessages.START_CLOSE_SHUFFLE_SINK_HANDLE);
- }
- boolean meetError = false;
- Exception firstException = null;
- boolean selfClosed = true;
- for (ISink channel : downStreamChannelList) {
- try {
- selfClosed = channel.close();
- } catch (Exception e) {
- if (!meetError) {
- firstException = e;
- meetError = true;
+ try {
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug(DataNodeQueryMessages.START_CLOSE_SHUFFLE_SINK_HANDLE);
+ }
+ boolean meetError = false;
+ Exception firstException = null;
+ boolean selfClosed = true;
+ for (ISink channel : downStreamChannelList) {
+ try {
+ selfClosed = channel.close();
+ } catch (Exception e) {
+ if (!meetError) {
+ firstException = e;
+ meetError = true;
+ }
}
}
- }
- if (meetError) {
-
LOGGER.warn(DataNodeQueryMessages.ERROR_OCCURRED_WHEN_TRY_TO_CLOSE_CHANNEL,
firstException);
- }
- if (selfClosed) {
- sinkListener.onFinish(this);
- closed = true;
- if (LOGGER.isDebugEnabled()) {
- LOGGER.debug(DataNodeQueryMessages.END_CLOSE_SHUFFLE_SINK_HANDLE);
+ if (meetError) {
+
LOGGER.warn(DataNodeQueryMessages.ERROR_OCCURRED_WHEN_TRY_TO_CLOSE_CHANNEL,
firstException);
+ }
+ if (selfClosed) {
+ sinkListener.onFinish(this);
+ closed = true;
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug(DataNodeQueryMessages.END_CLOSE_SHUFFLE_SINK_HANDLE);
+ }
+ return true;
+ } else {
+ return false;
+ }
+ } finally {
+ if (!closed) {
+ terminationClaimed.set(false);
}
- return true;
- } else {
- return false;
}
}
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java
index ec6c9c680a1..daaff9fdbcd 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java
@@ -39,6 +39,7 @@ import org.apache.iotdb.mpp.rpc.thrift.TNewDataBlockEvent;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.thrift.TApplicationException;
import org.apache.tsfile.common.conf.TSFileDescriptor;
import org.apache.tsfile.external.commons.lang3.Validate;
import org.apache.tsfile.read.common.block.TsBlock;
@@ -523,6 +524,26 @@ public class SinkChannel implements ISinkChannel {
// endregion
+ /**
+ * A {@link TApplicationException} means the downstream received the RPC but
failed while handling
+ * it. Retrying the same data-block notification cannot recover that
failure. The sync client may
+ * wrap it in multiple {@code TException}s, so inspect the complete cause
chain.
+ */
+ private static boolean containsTApplicationException(Throwable throwable) {
+ Throwable current = throwable;
+ while (current != null) {
+ if (current instanceof TApplicationException) {
+ return true;
+ }
+ Throwable cause = current.getCause();
+ if (cause == current) {
+ return false;
+ }
+ current = cause;
+ }
+ return false;
+ }
+
// region ============ TestOnly ============
@TestOnly
public void setRetryIntervalInMs(long retryIntervalInMs) {
@@ -581,8 +602,9 @@ public class SinkChannel implements ISinkChannel {
} catch (Exception e) {
LOGGER.warn(
DataNodeQueryMessages.FAILED_TO_SEND_NEW_DATA_BLOCK_EVENT_ATTEMPT, attempt, e);
- if (attempt == MAX_ATTEMPT_TIMES) {
+ if (containsTApplicationException(e) || attempt ==
MAX_ATTEMPT_TIMES) {
sinkListener.onFailure(SinkChannel.this, e);
+ return;
}
try {
Thread.sleep(retryIntervalInMs);
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java
index 4f90db0caed..2cd33b2d769 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java
@@ -384,18 +384,35 @@ public class SourceHandle implements ISourceHandle {
public synchronized void updatePendingDataBlockInfo(
int startSequenceId, List<Long> dataBlockSizes) {
- if (LOGGER.isDebugEnabled()) {
- LOGGER.debug(
-
DataNodeQueryMessages.RECEIVENEWTSBLOCKNOTIFICATION_ARG_ARG_EACH_SIZE_IS_ARG,
- startSequenceId,
- startSequenceId + dataBlockSizes.size(),
- dataBlockSizes);
- }
- for (int i = 0; i < dataBlockSizes.size(); i++) {
- sequenceIdToDataBlockSize.put(i + startSequenceId,
dataBlockSizes.get(i));
+ try {
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug(
+
DataNodeQueryMessages.RECEIVENEWTSBLOCKNOTIFICATION_ARG_ARG_EACH_SIZE_IS_ARG,
+ startSequenceId,
+ startSequenceId + dataBlockSizes.size(),
+ dataBlockSizes);
+ }
+ for (int i = 0; i < dataBlockSizes.size(); i++) {
+ sequenceIdToDataBlockSize.put(i + startSequenceId,
dataBlockSizes.get(i));
+ }
+ if (canGetTsBlockFromRemote) {
+ trySubmitGetDataBlocksTask();
+ }
+ } catch (RuntimeException | Error t) {
+ // This method runs in the inbound RPC thread. Mark the local FI failed
before the same
+ // exception is returned to the upstream FI through the existing Thrift
error channel.
+ notifyFailure(t);
+ throw t;
}
- if (canGetTsBlockFromRemote) {
- trySubmitGetDataBlocksTask();
+ }
+
+ private void notifyFailure(Throwable t) {
+ try {
+ sourceHandleListener.onFailure(this, t);
+ } catch (Throwable callbackFailure) {
+ if (callbackFailure != t) {
+ t.addSuppressed(callbackFailure);
+ }
}
}
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java
index 448a094c046..ce491042481 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java
@@ -25,6 +25,9 @@ import org.apache.iotdb.commons.client.IClientManager;
import
org.apache.iotdb.commons.client.sync.SyncDataNodeMPPDataExchangeServiceClient;
import org.apache.iotdb.commons.memory.MemoryManager;
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.QueryId;
import
org.apache.iotdb.db.queryengine.execution.exchange.sink.DownStreamChannelIndex;
import
org.apache.iotdb.db.queryengine.execution.exchange.sink.DownStreamChannelLocation;
import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISinkHandle;
@@ -33,18 +36,95 @@ import
org.apache.iotdb.db.queryengine.execution.exchange.sink.ShuffleSinkHandle
import org.apache.iotdb.db.queryengine.execution.exchange.source.ISourceHandle;
import
org.apache.iotdb.db.queryengine.execution.exchange.source.LocalSourceHandle;
import
org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
+import
org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceState;
+import
org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine;
import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager;
import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool;
import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId;
+import org.apache.iotdb.mpp.rpc.thrift.TNewDataBlockEvent;
import org.junit.Assert;
import org.junit.Test;
import org.mockito.Mockito;
import java.util.Collections;
+import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
public class MPPDataExchangeManagerTest {
+ @Test
+ public void testNewDataBlockMemoryFailureMarksTargetFragmentFailed() throws
Exception {
+ final String queryId = "q_memory_reservation_failure";
+ final String targetPlanNodeId = "root-exchange";
+ final long dataBlockSize = 1024L;
+ final TFragmentInstanceId targetFragmentInstanceId = new
TFragmentInstanceId(queryId, 0, "0");
+ final TFragmentInstanceId sourceFragmentInstanceId = new
TFragmentInstanceId(queryId, 1, "0");
+ final TEndPoint sourceEndpoint = new TEndPoint("remote-exchange", 10740);
+ ExecutorService exchangeExecutor = Executors.newSingleThreadExecutor();
+ ExecutorService notificationExecutor = Executors.newSingleThreadExecutor();
+
+ try {
+ LocalMemoryManager localMemoryManager =
Mockito.mock(LocalMemoryManager.class);
+ MemoryPool memoryPool = Mockito.mock(MemoryPool.class);
+ Mockito.when(localMemoryManager.getQueryPool()).thenReturn(memoryPool);
+ IllegalArgumentException expectedFailure =
+ new IllegalArgumentException("injected memory reservation failure");
+ Mockito.when(
+ memoryPool.reserveWithPriority(
+ Mockito.eq(queryId),
+ Mockito.anyString(),
+ Mockito.eq(targetPlanNodeId),
+ Mockito.eq(dataBlockSize),
+ Mockito.anyLong(),
+ Mockito.eq(false)))
+ .thenThrow(expectedFailure);
+
+ FragmentInstanceId targetId =
+ new FragmentInstanceId(new PlanFragmentId(new QueryId(queryId), 0),
"0");
+ FragmentInstanceStateMachine stateMachine =
+ new FragmentInstanceStateMachine(targetId, notificationExecutor);
+ FragmentInstanceContext context =
+ FragmentInstanceContext.createFragmentInstanceContext(targetId,
stateMachine);
+ MPPDataExchangeManager exchangeManager =
+ new MPPDataExchangeManager(
+ localMemoryManager,
+ new TsBlockSerdeFactory(),
+ exchangeExecutor,
+ Mockito.mock(IClientManager.class));
+ ISourceHandle sourceHandle =
+ exchangeManager.createSourceHandle(
+ targetFragmentInstanceId,
+ targetPlanNodeId,
+ 0,
+ sourceEndpoint,
+ sourceFragmentInstanceId,
+ context::failed);
+ sourceHandle.isBlocked();
+
+ TNewDataBlockEvent event =
+ new TNewDataBlockEvent(
+ targetFragmentInstanceId,
+ targetPlanNodeId,
+ sourceFragmentInstanceId,
+ 0,
+ Collections.singletonList(dataBlockSize));
+ IllegalArgumentException actualFailure =
+ Assert.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ exchangeManager
+ .getOrCreateMPPDataExchangeServiceImpl()
+ .onNewDataBlockEvent(event));
+
+ Assert.assertSame(expectedFailure, actualFailure);
+ Assert.assertEquals(FragmentInstanceState.FAILED,
stateMachine.getState());
+ Assert.assertSame(expectedFailure,
stateMachine.getFailureCauses().peek());
+ } finally {
+ exchangeExecutor.shutdownNow();
+ notificationExecutor.shutdownNow();
+ }
+ }
+
@Test
public void testCreateLocalSinkHandle() {
final TFragmentInstanceId localFragmentInstanceId = new
TFragmentInstanceId("q0", 1, "0");
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/ShuffleSinkHandleTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/ShuffleSinkHandleTest.java
index 75d204acade..01725611b6e 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/ShuffleSinkHandleTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/ShuffleSinkHandleTest.java
@@ -22,6 +22,7 @@ package org.apache.iotdb.db.queryengine.execution.exchange;
import org.apache.iotdb.commons.memory.MemoryManager;
import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
import
org.apache.iotdb.db.queryengine.execution.exchange.sink.DownStreamChannelIndex;
+import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISinkChannel;
import
org.apache.iotdb.db.queryengine.execution.exchange.sink.LocalSinkChannel;
import
org.apache.iotdb.db.queryengine.execution.exchange.sink.ShuffleSinkHandle;
import
org.apache.iotdb.db.queryengine.execution.exchange.source.LocalSourceHandle;
@@ -35,10 +36,105 @@ import org.junit.Test;
import org.mockito.Mockito;
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
com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
public class ShuffleSinkHandleTest {
+ @Test
+ public void testClosePreventsConcurrentAbort() throws Exception {
+ TFragmentInstanceId fragmentInstanceId = new TFragmentInstanceId("q0", 0,
"0");
+ ISinkChannel channel = Mockito.mock(ISinkChannel.class);
+ MPPDataExchangeManager.SinkListener sinkListener =
+ Mockito.mock(MPPDataExchangeManager.SinkListener.class);
+ CountDownLatch closeStarted = new CountDownLatch(1);
+ CountDownLatch allowCloseToFinish = new CountDownLatch(1);
+ Mockito.when(channel.close())
+ .thenAnswer(
+ invocation -> {
+ closeStarted.countDown();
+ Assert.assertTrue(allowCloseToFinish.await(5, TimeUnit.SECONDS));
+ return true;
+ });
+
+ ShuffleSinkHandle shuffleSinkHandle =
+ new ShuffleSinkHandle(
+ fragmentInstanceId,
+ Collections.singletonList(channel),
+ new DownStreamChannelIndex(0),
+ ShuffleSinkHandle.ShuffleStrategyEnum.PLAIN,
+ sinkListener);
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+
+ try {
+ Future<Boolean> closeResult = executor.submit(shuffleSinkHandle::close);
+ Assert.assertTrue(closeStarted.await(5, TimeUnit.SECONDS));
+
+ Assert.assertFalse(shuffleSinkHandle.abort());
+ allowCloseToFinish.countDown();
+
+ Assert.assertTrue(closeResult.get(5, TimeUnit.SECONDS));
+ Assert.assertTrue(shuffleSinkHandle.isClosed());
+ Assert.assertFalse(shuffleSinkHandle.isAborted());
+ Mockito.verify(channel).close();
+ Mockito.verify(channel, Mockito.never()).abort();
+ Mockito.verify(sinkListener).onFinish(shuffleSinkHandle);
+ Mockito.verify(sinkListener,
Mockito.never()).onAborted(shuffleSinkHandle);
+ } finally {
+ allowCloseToFinish.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testAbortPreventsConcurrentClose() throws Exception {
+ TFragmentInstanceId fragmentInstanceId = new TFragmentInstanceId("q0", 0,
"0");
+ ISinkChannel channel = Mockito.mock(ISinkChannel.class);
+ MPPDataExchangeManager.SinkListener sinkListener =
+ Mockito.mock(MPPDataExchangeManager.SinkListener.class);
+ CountDownLatch abortStarted = new CountDownLatch(1);
+ CountDownLatch allowAbortToFinish = new CountDownLatch(1);
+ Mockito.when(channel.abort())
+ .thenAnswer(
+ invocation -> {
+ abortStarted.countDown();
+ Assert.assertTrue(allowAbortToFinish.await(5, TimeUnit.SECONDS));
+ return true;
+ });
+
+ ShuffleSinkHandle shuffleSinkHandle =
+ new ShuffleSinkHandle(
+ fragmentInstanceId,
+ Collections.singletonList(channel),
+ new DownStreamChannelIndex(0),
+ ShuffleSinkHandle.ShuffleStrategyEnum.PLAIN,
+ sinkListener);
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+
+ try {
+ Future<Boolean> abortResult = executor.submit(shuffleSinkHandle::abort);
+ Assert.assertTrue(abortStarted.await(5, TimeUnit.SECONDS));
+
+ Assert.assertFalse(shuffleSinkHandle.close());
+ allowAbortToFinish.countDown();
+
+ Assert.assertTrue(abortResult.get(5, TimeUnit.SECONDS));
+ Assert.assertTrue(shuffleSinkHandle.isAborted());
+ Assert.assertFalse(shuffleSinkHandle.isClosed());
+ Mockito.verify(channel).abort();
+ Mockito.verify(channel, Mockito.never()).close();
+ Mockito.verify(sinkListener).onAborted(shuffleSinkHandle);
+ Mockito.verify(sinkListener,
Mockito.never()).onFinish(shuffleSinkHandle);
+ } finally {
+ allowAbortToFinish.countDown();
+ executor.shutdownNow();
+ }
+ }
+
@Test
public void testAbort() {
final String queryId = "q0";
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SinkChannelFailurePropagationTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SinkChannelFailurePropagationTest.java
new file mode 100644
index 00000000000..9564b104766
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SinkChannelFailurePropagationTest.java
@@ -0,0 +1,364 @@
+/*
+ * 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.exchange;
+
+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.client.sync.SyncDataNodeMPPDataExchangeServiceClient;
+import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
+import org.apache.iotdb.db.queryengine.common.PlanFragmentId;
+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.execution.exchange.sink.DownStreamChannelIndex;
+import
org.apache.iotdb.db.queryengine.execution.exchange.sink.DownStreamChannelLocation;
+import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink;
+import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISinkHandle;
+import
org.apache.iotdb.db.queryengine.execution.exchange.sink.ShuffleSinkHandle;
+import org.apache.iotdb.db.queryengine.execution.exchange.sink.SinkChannel;
+import
org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
+import
org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceFailureInfo;
+import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceInfo;
+import
org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceState;
+import
org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine;
+import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager;
+import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance;
+import
org.apache.iotdb.db.queryengine.plan.scheduler.FixedRateFragInsStateTracker;
+import org.apache.iotdb.mpp.rpc.thrift.TFetchFragmentInstanceInfoReq;
+import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId;
+import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceInfoResp;
+import org.apache.iotdb.mpp.rpc.thrift.TNewDataBlockEvent;
+
+import org.apache.thrift.TApplicationException;
+import org.apache.thrift.TException;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class SinkChannelFailurePropagationTest {
+
+ @Test
+ public void testFailureCallbackRunsBeforeChannelAccounting() {
+ ExecutorService exchangeExecutor = Executors.newSingleThreadExecutor();
+ try {
+ MPPDataExchangeManager exchangeManager =
+ new MPPDataExchangeManager(
+ Mockito.mock(LocalMemoryManager.class),
+ new TsBlockSerdeFactory(),
+ exchangeExecutor,
+ Mockito.mock(IClientManager.class));
+ AtomicInteger remainingChannels = new AtomicInteger(1);
+ AtomicBoolean hasChannelFailedOrAborted = new AtomicBoolean(false);
+ AtomicBoolean callbackInvoked = new AtomicBoolean(false);
+ IllegalStateException callbackFailure =
+ new IllegalStateException("injected callback failure");
+ MPPDataExchangeManager.ISinkChannelListenerImpl listener =
+ exchangeManager
+ .new ISinkChannelListenerImpl(
+ new TFragmentInstanceId("q_failure_callback_order", 0, "0"),
+ Mockito.mock(FragmentInstanceContext.class),
+ failure -> {
+ Assert.assertEquals(1, remainingChannels.get());
+ callbackInvoked.set(true);
+ throw callbackFailure;
+ },
+ remainingChannels,
+ hasChannelFailedOrAborted);
+
+ IllegalStateException actualFailure =
+ Assert.assertThrows(
+ IllegalStateException.class,
+ () ->
+ listener.onFailure(
+ Mockito.mock(ISink.class), new TException("injected
channel failure")));
+
+ Assert.assertSame(callbackFailure, actualFailure);
+ Assert.assertTrue(callbackInvoked.get());
+ Assert.assertTrue(hasChannelFailedOrAborted.get());
+ Assert.assertEquals(0, remainingChannels.get());
+ } finally {
+ exchangeExecutor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void
testApplicationExceptionFailsWithoutRetryAndPreventsNormalClose() throws
Exception {
+ final String queryId = "q_failure_with_two_channels";
+ final TFragmentInstanceId upstreamThriftId = new
TFragmentInstanceId(queryId, 1, "0");
+ final FragmentInstanceId upstreamId =
+ new FragmentInstanceId(new PlanFragmentId(new QueryId(queryId), 1),
"0");
+ final TEndPoint firstEndpoint = new TEndPoint("remote-exchange-0", 10740);
+ final TEndPoint secondEndpoint = new TEndPoint("remote-exchange-1", 10740);
+
+ ExecutorService exchangeExecutor = Executors.newSingleThreadExecutor();
+ ExecutorService fragmentNotificationExecutor =
Executors.newSingleThreadExecutor();
+
+ try {
+ LocalMemoryManager localMemoryManager =
Mockito.mock(LocalMemoryManager.class);
+ MemoryPool memoryPool = Utils.createMockNonBlockedMemoryPool();
+ Mockito.when(localMemoryManager.getQueryPool()).thenReturn(memoryPool);
+
+ IClientManager<TEndPoint, SyncDataNodeMPPDataExchangeServiceClient>
exchangeClientManager =
+ Mockito.mock(IClientManager.class);
+ SyncDataNodeMPPDataExchangeServiceClient exchangeClient =
+ Mockito.mock(SyncDataNodeMPPDataExchangeServiceClient.class);
+ TException expectedFailure =
+ new TException(
+ "wrapped application failure",
+ new TException(
+ "wrapped receive failure",
+ new TApplicationException(
+ TApplicationException.INTERNAL_ERROR,
+ "Internal error processing onNewDataBlockEvent")));
+
Mockito.when(exchangeClientManager.borrowClient(secondEndpoint)).thenReturn(exchangeClient);
+ Mockito.doThrow(expectedFailure)
+ .when(exchangeClient)
+ .onNewDataBlockEvent(Mockito.any(TNewDataBlockEvent.class));
+
+ FragmentInstanceStateMachine fragmentStateMachine =
+ new FragmentInstanceStateMachine(upstreamId,
fragmentNotificationExecutor);
+ FragmentInstanceContext fragmentContext =
+ FragmentInstanceContext.createFragmentInstanceContext(upstreamId,
fragmentStateMachine);
+ MPPDataExchangeManager exchangeManager =
+ new MPPDataExchangeManager(
+ localMemoryManager,
+ new TsBlockSerdeFactory(),
+ exchangeExecutor,
+ exchangeClientManager);
+ ISinkHandle sinkHandle =
+ exchangeManager.createShuffleSinkHandle(
+ Arrays.asList(
+ new DownStreamChannelLocation(
+ firstEndpoint, new TFragmentInstanceId(queryId, 0, "0"),
"root-exchange-0"),
+ new DownStreamChannelLocation(
+ secondEndpoint, new TFragmentInstanceId(queryId, 0,
"1"), "root-exchange-1")),
+ new DownStreamChannelIndex(1),
+ ShuffleSinkHandle.ShuffleStrategyEnum.PLAIN,
+ upstreamThriftId,
+ "upstream-sink",
+ fragmentContext);
+ SinkChannel normallyFinishedChannel = (SinkChannel)
sinkHandle.getChannel(0);
+ SinkChannel failedChannel = (SinkChannel) sinkHandle.getChannel(1);
+ failedChannel.setRetryIntervalInMs(0);
+
+ Assert.assertTrue(sinkHandle.isFull().isDone());
+ sinkHandle.send(Utils.createMockTsBlocks(1, 1024).get(0));
+ exchangeExecutor.submit(() -> {}).get(5, TimeUnit.SECONDS);
+ Mockito.verify(exchangeClient, Mockito.times(1))
+ .onNewDataBlockEvent(Mockito.any(TNewDataBlockEvent.class));
+ long waitStartNanos = System.nanoTime();
+ while (!fragmentStateMachine.getState().isDone()
+ && TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() -
waitStartNanos) < 5) {
+ Thread.sleep(10);
+ }
+
+ Assert.assertEquals(FragmentInstanceState.FAILED,
fragmentStateMachine.getState());
+ Assert.assertEquals(expectedFailure,
fragmentStateMachine.getFailureCauses().peek());
+ Assert.assertTrue(normallyFinishedChannel.close());
+
+ Assert.assertEquals(0, exchangeManager.getShuffleSinkHandleSize());
+ Assert.assertFalse(sinkHandle.isClosed());
+ Assert.assertFalse(sinkHandle.isAborted());
+
+ Assert.assertTrue(sinkHandle.abort());
+ Assert.assertTrue(sinkHandle.isAborted());
+ Assert.assertFalse(sinkHandle.isClosed());
+ Assert.assertTrue(failedChannel.isAborted());
+ } finally {
+ exchangeExecutor.shutdownNow();
+ fragmentNotificationExecutor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testSingleChannelFailurePropagation() throws Exception {
+ final String queryId = "q_failure_propagation";
+ final TFragmentInstanceId upstreamThriftId = new
TFragmentInstanceId(queryId, 1, "0");
+ final TFragmentInstanceId rootThriftId = new TFragmentInstanceId(queryId,
0, "0");
+ final FragmentInstanceId upstreamId =
+ new FragmentInstanceId(new PlanFragmentId(new QueryId(queryId), 1),
"0");
+ final FragmentInstanceId rootId =
+ new FragmentInstanceId(new PlanFragmentId(new QueryId(queryId), 0),
"0");
+ final TEndPoint exchangeEndpoint = new TEndPoint("remote-exchange", 10740);
+ final TEndPoint stateEndpoint = new TEndPoint("remote-state", 10730);
+
+ ExecutorService exchangeExecutor = Executors.newSingleThreadExecutor();
+ ExecutorService fragmentNotificationExecutor =
Executors.newSingleThreadExecutor();
+ ExecutorService queryNotificationExecutor =
Executors.newSingleThreadExecutor();
+ ScheduledExecutorService stateTrackerExecutor =
Executors.newSingleThreadScheduledExecutor();
+ FixedRateFragInsStateTracker stateTracker = null;
+
+ try {
+ LocalMemoryManager localMemoryManager =
Mockito.mock(LocalMemoryManager.class);
+ MemoryPool memoryPool = Utils.createMockNonBlockedMemoryPool();
+ Mockito.when(localMemoryManager.getQueryPool()).thenReturn(memoryPool);
+
+ IClientManager<TEndPoint, SyncDataNodeMPPDataExchangeServiceClient>
exchangeClientManager =
+ Mockito.mock(IClientManager.class);
+ SyncDataNodeMPPDataExchangeServiceClient exchangeClient =
+ Mockito.mock(SyncDataNodeMPPDataExchangeServiceClient.class);
+ TException expectedFailure = new TException("injected
onNewDataBlockEvent failure");
+
Mockito.when(exchangeClientManager.borrowClient(exchangeEndpoint)).thenReturn(exchangeClient);
+ Mockito.doThrow(expectedFailure)
+ .when(exchangeClient)
+ .onNewDataBlockEvent(Mockito.any(TNewDataBlockEvent.class));
+
+ FragmentInstanceStateMachine fragmentStateMachine =
+ new FragmentInstanceStateMachine(upstreamId,
fragmentNotificationExecutor);
+ FragmentInstanceContext fragmentContext =
+ FragmentInstanceContext.createFragmentInstanceContext(upstreamId,
fragmentStateMachine);
+
+ MPPDataExchangeManager exchangeManager =
+ new MPPDataExchangeManager(
+ localMemoryManager,
+ new TsBlockSerdeFactory(),
+ exchangeExecutor,
+ exchangeClientManager);
+ ISinkHandle sinkHandle =
+ exchangeManager.createShuffleSinkHandle(
+ Collections.singletonList(
+ new DownStreamChannelLocation(exchangeEndpoint,
rootThriftId, "root-exchange")),
+ new DownStreamChannelIndex(0),
+ ShuffleSinkHandle.ShuffleStrategyEnum.PLAIN,
+ upstreamThriftId,
+ "upstream-sink",
+ fragmentContext);
+ SinkChannel sinkChannel = (SinkChannel) sinkHandle.getChannel(0);
+ sinkChannel.setRetryIntervalInMs(0);
+ fragmentStateMachine.addStateChangeListener(
+ newState -> {
+ if (newState.isFailed()) {
+ sinkHandle.abort();
+ } else if (newState.isDone()) {
+ sinkHandle.close();
+ }
+ });
+
+ QueryStateMachine queryStateMachine =
+ new QueryStateMachine(new QueryId(queryId),
queryNotificationExecutor);
+ queryStateMachine.transitionToRunning();
+
+ FragmentInstance upstreamInstance = Mockito.mock(FragmentInstance.class);
+ FragmentInstance rootInstance = Mockito.mock(FragmentInstance.class);
+ TDataNodeLocation stateLocation =
+ new
TDataNodeLocation().setDataNodeId(1).setInternalEndPoint(stateEndpoint);
+ Mockito.when(upstreamInstance.getId()).thenReturn(upstreamId);
+
Mockito.when(upstreamInstance.getHostDataNode()).thenReturn(stateLocation);
+ Mockito.when(upstreamInstance.isRoot()).thenReturn(false);
+ Mockito.when(rootInstance.getId()).thenReturn(rootId);
+ Mockito.when(rootInstance.getHostDataNode()).thenReturn(stateLocation);
+ Mockito.when(rootInstance.isRoot()).thenReturn(true);
+
+ IClientManager<TEndPoint, SyncDataNodeInternalServiceClient>
stateClientManager =
+ Mockito.mock(IClientManager.class);
+ SyncDataNodeInternalServiceClient stateClient =
+ Mockito.mock(SyncDataNodeInternalServiceClient.class);
+
Mockito.when(stateClientManager.borrowClient(stateEndpoint)).thenReturn(stateClient);
+ Mockito.when(
+ stateClient.fetchFragmentInstanceInfo(
+ Mockito.any(TFetchFragmentInstanceInfoReq.class)))
+ .thenAnswer(
+ invocation -> {
+ TFetchFragmentInstanceInfoReq request =
invocation.getArgument(0);
+ if (upstreamThriftId.equals(request.getFragmentInstanceId())) {
+ return toThriftResponse(fragmentContext.getInstanceInfo());
+ }
+ return new
TFragmentInstanceInfoResp(FragmentInstanceState.RUNNING.toString());
+ });
+
+ stateTracker =
+ new FixedRateFragInsStateTracker(
+ queryStateMachine,
+ stateTrackerExecutor,
+ Arrays.asList(upstreamInstance, rootInstance),
+ stateClientManager);
+ stateTracker.start();
+
+ Assert.assertTrue(sinkHandle.isFull().isDone());
+ long failureStartNanos = System.nanoTime();
+ sinkHandle.send(Utils.createMockTsBlocks(1, 1024).get(0));
+
+ Mockito.verify(exchangeClient,
Mockito.timeout(5_000).times(SinkChannel.MAX_ATTEMPT_TIMES))
+ .onNewDataBlockEvent(Mockito.any(TNewDataBlockEvent.class));
+ while (!fragmentStateMachine.getState().isDone()
+ && TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() -
failureStartNanos) < 5) {
+ Thread.sleep(10);
+ }
+
+ Assert.assertEquals(FragmentInstanceState.FAILED,
fragmentStateMachine.getState());
+ Assert.assertEquals(expectedFailure,
fragmentStateMachine.getFailureCauses().peek());
+ while (!sinkChannel.isAborted()
+ && TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() -
failureStartNanos) < 5) {
+ Thread.sleep(10);
+ }
+ Assert.assertTrue(sinkHandle.isAborted());
+ Assert.assertFalse(sinkHandle.isClosed());
+ Assert.assertTrue(sinkChannel.isAborted());
+ Assert.assertFalse(sinkChannel.isClosed());
+ Assert.assertEquals(0, exchangeManager.getShuffleSinkHandleSize());
+
+ while (!queryStateMachine.getState().isDone()
+ && TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() -
failureStartNanos) < 5) {
+ Thread.sleep(10);
+ }
+ Mockito.verify(stateClient, Mockito.timeout(3_000).atLeastOnce())
+
.fetchFragmentInstanceInfo(Mockito.any(TFetchFragmentInstanceInfoReq.class));
+ Assert.assertEquals(QueryState.FAILED, queryStateMachine.getState());
+ } finally {
+ if (stateTracker != null) {
+ stateTracker.abort();
+ }
+ exchangeExecutor.shutdownNow();
+ fragmentNotificationExecutor.shutdownNow();
+ queryNotificationExecutor.shutdownNow();
+ stateTrackerExecutor.shutdownNow();
+ }
+ }
+
+ private static TFragmentInstanceInfoResp
toThriftResponse(FragmentInstanceInfo info)
+ throws IOException {
+ TFragmentInstanceInfoResp response = new
TFragmentInstanceInfoResp(info.getState().toString());
+ response.setEndTime(info.getEndTime());
+ response.setFailedMessages(Collections.singletonList(info.getMessage()));
+ List<ByteBuffer> failureInfoList = new ArrayList<>();
+ for (FragmentInstanceFailureInfo failureInfo : info.getFailureInfoList()) {
+ failureInfoList.add(failureInfo.serialize());
+ }
+ response.setFailureInfoList(failureInfoList);
+ info.getErrorCode().ifPresent(response::setErrorCode);
+ return response;
+ }
+}
