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 c30554f743c Fix fragment failure propagation race (#18543)
c30554f743c is described below
commit c30554f743cbe801a0c2ed6cf59740cc3c63e4b5
Author: Jackie Tien <[email protected]>
AuthorDate: Sun Aug 30 14:07:32 2026 +0800
Fix fragment failure propagation race (#18543)
---
.../execution/exchange/sink/ShuffleSinkHandle.java | 63 ++++++++--
.../execution/exchange/ShuffleSinkHandleTest.java | 134 ++++++++++++++++++++-
2 files changed, 182 insertions(+), 15 deletions(-)
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 78e4928fdc3..e54062158a3 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,7 +33,8 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.List;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
@@ -65,8 +66,9 @@ 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);
+ // lock. An abort caller that loses the claim waits for the owner to finish;
otherwise it may
+ // release fragment memory while the owner is still releasing channel
buffers.
+ private final AtomicReference<TerminationClaim> terminationClaim = new
AtomicReference<>();
private static final DataExchangeCostMetricSet DATA_EXCHANGE_COST_METRIC_SET
=
DataExchangeCostMetricSet.getInstance();
@@ -81,6 +83,11 @@ public class ShuffleSinkHandle implements ISinkHandle {
+ RamUsageEstimator.shallowSizeOfInstance(TFragmentInstanceId.class)
+
RamUsageEstimator.shallowSizeOfInstance(DownStreamChannelIndex.class);
+ private static final class TerminationClaim {
+ private final Thread owner = Thread.currentThread();
+ private final CompletableFuture<Void> completion = new
CompletableFuture<>();
+ }
+
public ShuffleSinkHandle(
TFragmentInstanceId localFragmentInstanceId,
List<ISinkChannel> downStreamChannelList,
@@ -149,11 +156,14 @@ public class ShuffleSinkHandle implements ISinkHandle {
@Override
public void setNoMoreTsBlocks() {
- if (closed || aborted) {
+ if (closed || aborted || terminationClaim.get() != null) {
return;
}
try {
lock.lock();
+ if (closed || aborted || terminationClaim.get() != null) {
+ return;
+ }
for (int i = 0; i < downStreamChannelList.size(); i++) {
if (!hasSetNoMoreTsBlocks[i]) {
downStreamChannelList.get(i).setNoMoreTsBlocks();
@@ -168,13 +178,16 @@ public class ShuffleSinkHandle implements ISinkHandle {
@Override
public void setNoMoreTsBlocksOfOneChannel(int channelIndex) {
- if (closed || aborted) {
+ if (closed || aborted || terminationClaim.get() != null) {
// if this ShuffleSinkHandle has been closed, Driver.close() will
attempt to setNoMoreTsBlocks
// for all the channels
return;
}
try {
lock.lock();
+ if (closed || aborted || terminationClaim.get() != null) {
+ return;
+ }
if (!hasSetNoMoreTsBlocks[channelIndex]) {
downStreamChannelList.get(channelIndex).setNoMoreTsBlocks();
hasSetNoMoreTsBlocks[channelIndex] = true;
@@ -206,7 +219,8 @@ public class ShuffleSinkHandle implements ISinkHandle {
@Override
public boolean abort() {
- if (aborted || closed || !terminationClaimed.compareAndSet(false, true)) {
+ TerminationClaim claim = claimTermination(true);
+ if (claim == null) {
return false;
}
try {
@@ -240,9 +254,7 @@ public class ShuffleSinkHandle implements ISinkHandle {
return false;
}
} finally {
- if (!aborted) {
- terminationClaimed.set(false);
- }
+ releaseTerminationClaim(claim);
}
}
@@ -252,7 +264,8 @@ public class ShuffleSinkHandle implements ISinkHandle {
// Lock ShuffleSinkHandle and wait to lock LocalSinkChannel
@Override
public boolean close() {
- if (closed || aborted || !terminationClaimed.compareAndSet(false, true)) {
+ TerminationClaim claim = claimTermination(false);
+ if (claim == null) {
return false;
}
try {
@@ -286,9 +299,7 @@ public class ShuffleSinkHandle implements ISinkHandle {
return false;
}
} finally {
- if (!closed) {
- terminationClaimed.set(false);
- }
+ releaseTerminationClaim(claim);
}
}
@@ -316,6 +327,32 @@ public class ShuffleSinkHandle implements ISinkHandle {
}
}
+ private TerminationClaim claimTermination(boolean waitForCurrentClaim) {
+ TerminationClaim claim = new TerminationClaim();
+ while (!aborted && !closed) {
+ TerminationClaim currentClaim = terminationClaim.get();
+ if (currentClaim == null) {
+ if (terminationClaim.compareAndSet(null, claim)) {
+ return claim;
+ }
+ } else {
+ // A channel callback can re-enter close() while holding the channel
lock. Therefore close()
+ // must not wait for another termination operation. Abort callers are
not invoked under a
+ // channel lock and need the completion barrier before fragment memory
is deregistered.
+ if (!waitForCurrentClaim || currentClaim.owner ==
Thread.currentThread()) {
+ return null;
+ }
+ currentClaim.completion.join();
+ }
+ }
+ return null;
+ }
+
+ private void releaseTerminationClaim(TerminationClaim claim) {
+ terminationClaim.compareAndSet(claim, null);
+ claim.completion.complete(null);
+ }
+
private void switchChannelIfNecessary() {
shuffleStrategy.shuffle();
}
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 01725611b6e..15ef07241d4 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
@@ -41,6 +41,7 @@ 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.MoreExecutors.newDirectExecutorService;
@@ -68,16 +69,25 @@ public class ShuffleSinkHandleTest {
new DownStreamChannelIndex(0),
ShuffleSinkHandle.ShuffleStrategyEnum.PLAIN,
sinkListener);
- ExecutorService executor = Executors.newSingleThreadExecutor();
+ ExecutorService executor = Executors.newFixedThreadPool(2);
try {
Future<Boolean> closeResult = executor.submit(shuffleSinkHandle::close);
Assert.assertTrue(closeStarted.await(5, TimeUnit.SECONDS));
- Assert.assertFalse(shuffleSinkHandle.abort());
+ CountDownLatch abortCalled = new CountDownLatch(1);
+ Future<Boolean> abortResult =
+ executor.submit(
+ () -> {
+ abortCalled.countDown();
+ return shuffleSinkHandle.abort();
+ });
+ Assert.assertTrue(abortCalled.await(5, TimeUnit.SECONDS));
+ Assert.assertFalse(abortResult.isDone());
allowCloseToFinish.countDown();
Assert.assertTrue(closeResult.get(5, TimeUnit.SECONDS));
+ Assert.assertFalse(abortResult.get(5, TimeUnit.SECONDS));
Assert.assertTrue(shuffleSinkHandle.isClosed());
Assert.assertFalse(shuffleSinkHandle.isAborted());
Mockito.verify(channel).close();
@@ -135,6 +145,126 @@ public class ShuffleSinkHandleTest {
}
}
+ @Test
+ public void testConcurrentAbortWaitsForCompletion() 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.newFixedThreadPool(2);
+
+ try {
+ Future<Boolean> firstAbortResult =
executor.submit(shuffleSinkHandle::abort);
+ Assert.assertTrue(abortStarted.await(5, TimeUnit.SECONDS));
+
+ CountDownLatch secondAbortCalled = new CountDownLatch(1);
+ Future<Boolean> secondAbortResult =
+ executor.submit(
+ () -> {
+ secondAbortCalled.countDown();
+ return shuffleSinkHandle.abort();
+ });
+ Assert.assertTrue(secondAbortCalled.await(5, TimeUnit.SECONDS));
+ Assert.assertFalse(secondAbortResult.isDone());
+ allowAbortToFinish.countDown();
+
+ Assert.assertTrue(firstAbortResult.get(5, TimeUnit.SECONDS));
+ Assert.assertFalse(secondAbortResult.get(5, TimeUnit.SECONDS));
+ Mockito.verify(channel).abort();
+ Mockito.verify(sinkListener).onAborted(shuffleSinkHandle);
+ } finally {
+ allowAbortToFinish.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testAbortPreventsConcurrentSetNoMoreTsBlocks() 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));
+
+ shuffleSinkHandle.setNoMoreTsBlocks();
+
+ Mockito.verify(channel, Mockito.never()).setNoMoreTsBlocks();
+ Mockito.verify(sinkListener,
Mockito.never()).onEndOfBlocks(shuffleSinkHandle);
+ allowAbortToFinish.countDown();
+ Assert.assertTrue(abortResult.get(5, TimeUnit.SECONDS));
+ } finally {
+ allowAbortToFinish.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testCloseDoesNotWaitOnReentrantCall() {
+ TFragmentInstanceId fragmentInstanceId = new TFragmentInstanceId("q0", 0,
"0");
+ ISinkChannel channel = Mockito.mock(ISinkChannel.class);
+ MPPDataExchangeManager.SinkListener sinkListener =
+ Mockito.mock(MPPDataExchangeManager.SinkListener.class);
+ AtomicReference<ShuffleSinkHandle> shuffleSinkHandleReference = new
AtomicReference<>();
+ Mockito.when(channel.close())
+ .thenAnswer(
+ invocation -> {
+ Assert.assertFalse(shuffleSinkHandleReference.get().close());
+ return true;
+ });
+
+ ShuffleSinkHandle shuffleSinkHandle =
+ new ShuffleSinkHandle(
+ fragmentInstanceId,
+ Collections.singletonList(channel),
+ new DownStreamChannelIndex(0),
+ ShuffleSinkHandle.ShuffleStrategyEnum.PLAIN,
+ sinkListener);
+ shuffleSinkHandleReference.set(shuffleSinkHandle);
+
+ Assert.assertTrue(shuffleSinkHandle.close());
+ Assert.assertTrue(shuffleSinkHandle.isClosed());
+ Mockito.verify(channel).close();
+ Mockito.verify(sinkListener).onFinish(shuffleSinkHandle);
+ }
+
@Test
public void testAbort() {
final String queryId = "q0";
