From d2a97d923ba78ce5385c6ea578ab3273a4b319fd Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 20 Jul 2026 10:08:31 +0800 Subject: [PATCH] Fix flaky SharedTsBlockQueue concurrency test --- .../exchange/SharedTsBlockQueueTest.java | 170 +++++------------- 1 file changed, 49 insertions(+), 121 deletions(-) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java index a92eaa5d0e283..f5faa1b5ce114 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java @@ -27,7 +27,6 @@ import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; -import org.apache.tsfile.external.commons.lang3.Validate; import org.apache.tsfile.read.common.block.TsBlock; import org.junit.Assert; import org.junit.Test; @@ -35,7 +34,8 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; @@ -117,18 +117,18 @@ public void testAsyncListenerAfterAbortDoesNotAddTsBlock() { Assert.assertTrue(addFuture.isDone()); } - @Test(timeout = 15000L) - public void concurrencyTest() { + @Test + public void concurrencyTest() throws Exception { final String queryId = "q0"; final long mockTsBlockSize = 1024L * 1024L; + final int numOfTsBlocks = 1000; // Construct a mock LocalMemoryManager with capacity 5 * mockTsBlockSize per // query. LocalMemoryManager mockLocalMemoryManager = Mockito.mock(LocalMemoryManager.class); - MemoryManager memoryManager = Mockito.spy(new MemoryManager(10 * mockTsBlockSize)); - MemoryPool spyMemoryPool = - Mockito.spy(new MemoryPool("test", memoryManager, 5 * mockTsBlockSize)); - Mockito.when(mockLocalMemoryManager.getQueryPool()).thenReturn(spyMemoryPool); + MemoryManager memoryManager = new MemoryManager(10 * mockTsBlockSize); + MemoryPool memoryPool = new MemoryPool("test", memoryManager, 5 * mockTsBlockSize); + Mockito.when(mockLocalMemoryManager.getQueryPool()).thenReturn(memoryPool); SharedTsBlockQueue queue = new SharedTsBlockQueue( new TFragmentInstanceId(queryId, 0, "0"), @@ -139,119 +139,47 @@ public void concurrencyTest() { queue.setMaxBytesCanReserve(Long.MAX_VALUE); ExecutorService executor = Executors.newFixedThreadPool(2); - AtomicReference numOfTimesSenderBlocked = new AtomicReference<>(0); - AtomicReference numOfTimesReceiverBlocked = new AtomicReference<>(0); - AtomicReference numOfTsBlocksToSend = new AtomicReference<>(1000); - AtomicReference numOfTsBlocksToReceive = new AtomicReference<>(1000); - executor.submit( - new SendTask( - queue, mockTsBlockSize, numOfTsBlocksToSend, numOfTimesSenderBlocked, executor)); - executor.submit( - new ReceiveTask(queue, numOfTsBlocksToReceive, numOfTimesReceiverBlocked, executor)); - - while (numOfTsBlocksToSend.get() != 0 && numOfTsBlocksToReceive.get() != 0) { - String message = - String.format( - "Sender %d: %d, Receiver %d: %d", - numOfTimesSenderBlocked.get(), - numOfTsBlocksToSend.get(), - numOfTimesReceiverBlocked.get(), - numOfTsBlocksToReceive.get()); - System.out.println(message); - try { - Thread.sleep(10L); - } catch (InterruptedException e) { - Assert.fail(e.getMessage()); - } - } - } - - private static class SendTask implements Runnable { - - private final SharedTsBlockQueue queue; - private final long mockTsBlockSize; - private final AtomicReference numOfTsBlocksToSend; - private final AtomicReference numOfTimesBlocked; - private final ExecutorService executor; - - public SendTask( - SharedTsBlockQueue queue, - long mockTsBlockSize, - AtomicReference numOfTsBlocksToSend, - AtomicReference numOfTimesBlocked, - ExecutorService executor) { - this.queue = Validate.notNull(queue); - Validate.isTrue(mockTsBlockSize > 0L); - this.mockTsBlockSize = mockTsBlockSize; - this.numOfTsBlocksToSend = Validate.notNull(numOfTsBlocksToSend); - this.numOfTimesBlocked = Validate.notNull(numOfTimesBlocked); - this.executor = Validate.notNull(executor); - } - - @Override - public void run() { - ListenableFuture blockedOnMemory = null; - while (numOfTsBlocksToSend.get() > 0) { - synchronized (queue) { - blockedOnMemory = queue.add(Utils.createMockTsBlock(mockTsBlockSize)); - } - numOfTsBlocksToSend.updateAndGet(v -> v - 1); - if (!blockedOnMemory.isDone()) { - break; - } - } - - if (blockedOnMemory != null) { - numOfTimesBlocked.updateAndGet(v -> v + 1); - blockedOnMemory.addListener( - new SendTask(queue, mockTsBlockSize, numOfTsBlocksToSend, numOfTimesBlocked, executor), - executor); - } else { - synchronized (queue) { - queue.setNoMoreTsBlocks(true); - } - } - } - } - - private static class ReceiveTask implements Runnable { - - private final SharedTsBlockQueue queue; - private final AtomicReference numOfTsBlocksToReceive; - private final AtomicReference numOfTimesBlocked; - private final ExecutorService executor; - - public ReceiveTask( - SharedTsBlockQueue queue, - AtomicReference numOfTsBlocksToReceive, - AtomicReference numOfTimesBlocked, - ExecutorService executor) { - this.queue = Validate.notNull(queue); - this.numOfTsBlocksToReceive = Validate.notNull(numOfTsBlocksToReceive); - this.numOfTimesBlocked = Validate.notNull(numOfTimesBlocked); - this.executor = Validate.notNull(executor); - } - - @Override - public void run() { - ListenableFuture blocked = null; - while (numOfTsBlocksToReceive.get() > 0) { - synchronized (queue) { - blocked = queue.isBlocked(); - if (blocked.isDone()) { - queue.remove(); - numOfTsBlocksToReceive.updateAndGet(v -> v - 1); - } else { - break; - } - } - } - - if (blocked != null) { - numOfTimesBlocked.updateAndGet(v -> v + 1); - blocked.addListener( - new ReceiveTask(queue, numOfTsBlocksToReceive, numOfTimesBlocked, executor), executor); - } + try { + Future sender = + executor.submit( + () -> { + for (int i = 0; i < numOfTsBlocks; i++) { + ListenableFuture blockedOnMemory; + synchronized (queue) { + blockedOnMemory = queue.add(Utils.createMockTsBlock(mockTsBlockSize)); + } + blockedOnMemory.get(); + } + synchronized (queue) { + queue.setNoMoreTsBlocks(true); + } + return null; + }); + Future receiver = + executor.submit( + () -> { + for (int i = 0; i < numOfTsBlocks; i++) { + ListenableFuture blocked; + synchronized (queue) { + blocked = queue.isBlocked(); + } + blocked.get(); + synchronized (queue) { + queue.remove(); + } + } + return null; + }); + + sender.get(30, TimeUnit.SECONDS); + receiver.get(30, TimeUnit.SECONDS); + + Assert.assertTrue(queue.hasNoMoreTsBlocks()); + Assert.assertTrue(queue.isEmpty()); + Assert.assertEquals(0L, memoryPool.getReservedBytes()); + } finally { + executor.shutdownNow(); + executor.awaitTermination(10, TimeUnit.SECONDS); } } }