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 e54062158a3f1..f118439fb34e8 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 @@ -277,7 +277,10 @@ public boolean close() { boolean selfClosed = true; for (ISink channel : downStreamChannelList) { try { - selfClosed = channel.close(); + // close() returns false for a channel that is already closed, for example when its + // consumer closed it and the channel's listener re-entered this method. Such a channel + // ended normally, so count it as closed; an aborted channel still does not count. + selfClosed = channel.close() || (channel.isClosed() && !channel.isAborted()); } catch (Exception e) { if (!meetError) { firstException = e; 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 a5573159a94e5..f12a195c5e7e4 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 @@ -41,18 +41,24 @@ import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; +import org.apache.iotdb.mpp.rpc.thrift.TCloseSinkChannelEvent; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; import org.apache.iotdb.mpp.rpc.thrift.TNewDataBlockEvent; +import com.google.common.util.concurrent.MoreExecutors; import org.apache.tsfile.read.common.block.TsBlock; import org.junit.Assert; import org.junit.Test; import org.mockito.Mockito; +import java.util.Arrays; import java.util.Collections; +import java.util.List; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; + public class MPPDataExchangeManagerTest { @Test public void testNewDataBlockMemoryFailureMarksTargetFragmentFailed() throws Exception { @@ -255,6 +261,228 @@ public void testCreateLocalSourceHandle() { localPlanNodeId); } + @Test + public void testConsumerCloseAfterNoMoreTsBlocksFinishesShuffleSinkHandle() { + final TFragmentInstanceId producerId = new TFragmentInstanceId("q0", 1, "0"); + final TFragmentInstanceId consumerId = new TFragmentInstanceId("q0", 0, "0"); + final FragmentInstanceContext producerContext = Mockito.mock(FragmentInstanceContext.class); + final MPPDataExchangeManager mppDataExchangeManager = createMPPDataExchangeManager(); + final ISinkHandle shuffleSinkHandle = + createShuffleSinkHandle( + mppDataExchangeManager, + producerId, + Collections.singletonList(localChannelLocation(consumerId)), + producerContext); + final ISourceHandle localSourceHandle = + createLocalSourceHandle(mppDataExchangeManager, consumerId, producerId, 0); + + // The consumer takes every block before the producer signals the end of data. + localSourceHandle.isBlocked(); + shuffleSinkHandle.send(Utils.createMockTsBlock(256)); + Assert.assertTrue(localSourceHandle.isBlocked().isDone()); + Assert.assertNotNull(localSourceHandle.receive()); + + // The consumer closes as soon as the queue reports that it is finished, before the producer's + // own LocalSinkChannel#checkAndInvokeOnFinished() runs. This is the order seen when a finished + // query's result handle is closed. + localSourceHandle.isBlocked().addListener(localSourceHandle::close, directExecutor()); + shuffleSinkHandle.setNoMoreTsBlocks(); + + Assert.assertTrue(shuffleSinkHandle.isChannelClosed(0)); + Assert.assertTrue( + "ShuffleSinkHandle did not finish after the consumer closed its only channel", + shuffleSinkHandle.isClosed()); + Mockito.verify(producerContext).finished(); + } + + @Test + public void testConsumerCloseBeforeNoMoreTsBlocksFinishesShuffleSinkHandle() { + final TFragmentInstanceId producerId = new TFragmentInstanceId("q0", 1, "0"); + final TFragmentInstanceId consumerId = new TFragmentInstanceId("q0", 0, "0"); + final FragmentInstanceContext producerContext = Mockito.mock(FragmentInstanceContext.class); + final MPPDataExchangeManager mppDataExchangeManager = createMPPDataExchangeManager(); + final ISinkHandle shuffleSinkHandle = + createShuffleSinkHandle( + mppDataExchangeManager, + producerId, + Collections.singletonList(localChannelLocation(consumerId)), + producerContext); + final ISourceHandle localSourceHandle = + createLocalSourceHandle(mppDataExchangeManager, consumerId, producerId, 0); + + localSourceHandle.isBlocked(); + shuffleSinkHandle.send(Utils.createMockTsBlock(256)); + // The consumer stops reading early, for example because a LIMIT is satisfied. + localSourceHandle.close(); + + Assert.assertTrue(shuffleSinkHandle.isChannelClosed(0)); + Assert.assertTrue( + "ShuffleSinkHandle did not finish after the consumer closed its only channel", + shuffleSinkHandle.isClosed()); + Mockito.verify(producerContext).finished(); + + // The producer's driver may still be running; its later calls must be harmless. + shuffleSinkHandle.send(Utils.createMockTsBlock(256)); + shuffleSinkHandle.setNoMoreTsBlocks(); + Mockito.verify(producerContext, Mockito.times(1)).finished(); + } + + @Test + public void testConsumerCloseOfRemoteChannelFinishesShuffleSinkHandle() throws Exception { + final TFragmentInstanceId producerId = new TFragmentInstanceId("q0", 1, "0"); + final TFragmentInstanceId consumerId = new TFragmentInstanceId("q0", 0, "0"); + final FragmentInstanceContext producerContext = Mockito.mock(FragmentInstanceContext.class); + final MPPDataExchangeManager mppDataExchangeManager = createMPPDataExchangeManager(); + final ISinkHandle shuffleSinkHandle = + createShuffleSinkHandle( + mppDataExchangeManager, + producerId, + Collections.singletonList( + new DownStreamChannelLocation( + new TEndPoint("remote-exchange", 10740), consumerId, "exchange_0")), + producerContext); + + // The remote consumer closes its source handle before the producer signals the end of data. + mppDataExchangeManager + .getOrCreateMPPDataExchangeServiceImpl() + .onCloseSinkChannelEvent(new TCloseSinkChannelEvent(producerId, 0)); + + Assert.assertTrue(shuffleSinkHandle.isChannelClosed(0)); + Assert.assertTrue(shuffleSinkHandle.isClosed()); + Mockito.verify(producerContext).finished(); + } + + @Test + public void testConsumerCloseOfAllLocalChannelsFinishesShuffleSinkHandle() { + final TFragmentInstanceId producerId = new TFragmentInstanceId("q0", 1, "0"); + final TFragmentInstanceId firstConsumerId = new TFragmentInstanceId("q0", 0, "0"); + final TFragmentInstanceId secondConsumerId = new TFragmentInstanceId("q0", 2, "0"); + final FragmentInstanceContext producerContext = Mockito.mock(FragmentInstanceContext.class); + final MPPDataExchangeManager mppDataExchangeManager = createMPPDataExchangeManager(); + final ISinkHandle shuffleSinkHandle = + createShuffleSinkHandle( + mppDataExchangeManager, + producerId, + Arrays.asList( + localChannelLocation(firstConsumerId), localChannelLocation(secondConsumerId)), + producerContext); + final ISourceHandle firstSourceHandle = + createLocalSourceHandle(mppDataExchangeManager, firstConsumerId, producerId, 0); + final ISourceHandle secondSourceHandle = + createLocalSourceHandle(mppDataExchangeManager, secondConsumerId, producerId, 1); + + firstSourceHandle.close(); + Assert.assertFalse(shuffleSinkHandle.isClosed()); + Mockito.verify(producerContext, Mockito.never()).finished(); + + secondSourceHandle.close(); + Assert.assertTrue(shuffleSinkHandle.isChannelClosed(0)); + Assert.assertTrue(shuffleSinkHandle.isChannelClosed(1)); + Assert.assertTrue( + "ShuffleSinkHandle did not finish after the consumers closed all of its channels", + shuffleSinkHandle.isClosed()); + Mockito.verify(producerContext).finished(); + } + + @Test + public void testAbortedLocalChannelKeepsShuffleSinkHandleFromFinishing() { + final TFragmentInstanceId producerId = new TFragmentInstanceId("q0", 1, "0"); + final TFragmentInstanceId consumerId = new TFragmentInstanceId("q0", 0, "0"); + final FragmentInstanceContext producerContext = Mockito.mock(FragmentInstanceContext.class); + final MPPDataExchangeManager mppDataExchangeManager = createMPPDataExchangeManager(); + final ShuffleSinkHandle shuffleSinkHandle = + (ShuffleSinkHandle) + createShuffleSinkHandle( + mppDataExchangeManager, + producerId, + Collections.singletonList(localChannelLocation(consumerId)), + producerContext); + createLocalSourceHandle(mppDataExchangeManager, consumerId, producerId, 0); + + Assert.assertTrue(shuffleSinkHandle.getChannel(0).abort()); + + Assert.assertFalse( + "ShuffleSinkHandle must not finish when its channel was aborted", + shuffleSinkHandle.close()); + Assert.assertFalse(shuffleSinkHandle.isClosed()); + Mockito.verify(producerContext, Mockito.never()).finished(); + } + + @Test + public void testAbortedLocalChannelKeepsShuffleSinkHandleFromFinishingOnConsumerClose() { + final TFragmentInstanceId producerId = new TFragmentInstanceId("q0", 1, "0"); + final TFragmentInstanceId firstConsumerId = new TFragmentInstanceId("q0", 0, "0"); + final TFragmentInstanceId secondConsumerId = new TFragmentInstanceId("q0", 2, "0"); + final FragmentInstanceContext producerContext = Mockito.mock(FragmentInstanceContext.class); + final MPPDataExchangeManager mppDataExchangeManager = createMPPDataExchangeManager(); + final ShuffleSinkHandle shuffleSinkHandle = + (ShuffleSinkHandle) + createShuffleSinkHandle( + mppDataExchangeManager, + producerId, + Arrays.asList( + localChannelLocation(firstConsumerId), localChannelLocation(secondConsumerId)), + producerContext); + createLocalSourceHandle(mppDataExchangeManager, firstConsumerId, producerId, 0); + final ISourceHandle secondSourceHandle = + createLocalSourceHandle(mppDataExchangeManager, secondConsumerId, producerId, 1); + + Assert.assertTrue(shuffleSinkHandle.getChannel(0).abort()); + secondSourceHandle.close(); + + Assert.assertTrue(shuffleSinkHandle.isChannelClosed(1)); + Assert.assertFalse( + "ShuffleSinkHandle must not finish when one of its channels was aborted", + shuffleSinkHandle.isClosed()); + Mockito.verify(producerContext, Mockito.never()).finished(); + } + + private static MPPDataExchangeManager createMPPDataExchangeManager() { + LocalMemoryManager mockLocalMemoryManager = Mockito.mock(LocalMemoryManager.class); + MemoryManager memoryManager = Mockito.spy(new MemoryManager(10240L)); + MemoryPool spyMemoryPool = Mockito.spy(new MemoryPool("test", memoryManager, 5120L)); + Mockito.when(mockLocalMemoryManager.getQueryPool()).thenReturn(spyMemoryPool); + return new MPPDataExchangeManager( + mockLocalMemoryManager, + new TsBlockSerdeFactory(), + MoreExecutors.newDirectExecutorService(), + new IClientManager.Factory() + .createClientManager( + new ClientPoolFactory.SyncDataNodeMPPDataExchangeServiceClientPoolFactory())); + } + + private static DownStreamChannelLocation localChannelLocation(TFragmentInstanceId consumerId) { + return new DownStreamChannelLocation( + new TEndPoint( + IoTDBDescriptor.getInstance().getConfig().getInternalAddress(), + IoTDBDescriptor.getInstance().getConfig().getMppDataExchangePort()), + consumerId, + "exchange_0"); + } + + private static ISinkHandle createShuffleSinkHandle( + MPPDataExchangeManager mppDataExchangeManager, + TFragmentInstanceId producerId, + List channelLocations, + FragmentInstanceContext producerContext) { + return mppDataExchangeManager.createShuffleSinkHandle( + channelLocations, + new DownStreamChannelIndex(0), + ShuffleSinkHandle.ShuffleStrategyEnum.PLAIN, + producerId, + "shuffleSink_0", + producerContext); + } + + private static ISourceHandle createLocalSourceHandle( + MPPDataExchangeManager mppDataExchangeManager, + TFragmentInstanceId consumerId, + TFragmentInstanceId producerId, + int channelIndex) { + return mppDataExchangeManager.createLocalSourceHandleForFragment( + consumerId, "exchange_0", "shuffleSink_0", producerId, channelIndex, t -> {}); + } + private static void assertProducerOwnsQueue( MemoryPool pool, ISinkHandle sink,