Fix fragment failure propagation race

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java

@@ -33,7 +33,8 @@
 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 synchronized void send(TsBlock tsBlock) {
 
   @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 void setNoMoreTsBlocks() {
 
   @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 synchronized boolean isFinished() {
 
   @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 boolean abort() {
         return false;
       }
     } finally {
-      if (!aborted) {
-        terminationClaimed.set(false);
-      }
+      releaseTerminationClaim(claim);
     }
   }
 
@@ -252,7 +264,8 @@ public boolean abort() {
   // 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 boolean close() {
         return false;
       }
     } finally {
-      if (!closed) {
-        terminationClaimed.set(false);
-      }
+      releaseTerminationClaim(claim);
     }
   }
 
@@ -316,6 +327,32 @@ private void checkState() {
     }
   }
 
+  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();
   }

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/ShuffleSinkHandleTest.java

@@ -41,6 +41,7 @@
 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 void testClosePreventsConcurrentAbort() throws Exception {
             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 void testAbortPreventsConcurrentClose() throws Exception {
     }
   }
 
+  @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";