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 78e4928fdc3ed..e54062158a3f1 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.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 = 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 completion = new CompletableFuture<>(); + } + public ShuffleSinkHandle( TFragmentInstanceId localFragmentInstanceId, List 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(); } 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 01725611b6eca..15ef07241d404 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.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 closeResult = executor.submit(shuffleSinkHandle::close); Assert.assertTrue(closeStarted.await(5, TimeUnit.SECONDS)); - Assert.assertFalse(shuffleSinkHandle.abort()); + CountDownLatch abortCalled = new CountDownLatch(1); + Future 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 firstAbortResult = executor.submit(shuffleSinkHandle::abort); + Assert.assertTrue(abortStarted.await(5, TimeUnit.SECONDS)); + + CountDownLatch secondAbortCalled = new CountDownLatch(1); + Future 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 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 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";