From 13fb856686ed935683e15a223e79f397d9cb8d2e Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 2 Nov 2026 14:48:24 +0800 Subject: [PATCH 1/3] Stabilize query execution concurrency unit tests --- .../exchange/SharedTsBlockQueueTest.java | 161 ++++++------------ .../FragmentInstanceExecutionTest.java | 5 +- .../SeriesScanUtilCancellationTest.java | 37 ++-- .../schedule/queue/L1PriorityQueueTest.java | 42 +++-- 4 files changed, 97 insertions(+), 148 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 391393cf5fa0e..29fe431b32781 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 @@ -29,7 +29,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; @@ -37,7 +36,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; @@ -119,8 +119,8 @@ 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; @@ -140,123 +140,58 @@ public void concurrencyTest() { queue.getCanAddTsBlock().set(null); queue.setMaxBytesCanReserve(Long.MAX_VALUE); + // Reuse an immutable block so mock generation is outside the concurrency deadline. + TsBlock block = Utils.createMockTsBlock(mockTsBlockSize); + int blockCount = 1000; 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 { + Future sender = + executor.submit( + () -> { + for (int i = 0; i < blockCount; i++) { + ListenableFuture blocked; + synchronized (queue) { + blocked = queue.add(block); + } + blocked.get(10, TimeUnit.SECONDS); + } + synchronized (queue) { + queue.setNoMoreTsBlocks(true); + } + return null; + }); + Future receiver = + executor.submit( + () -> { + for (int i = 0; i < blockCount; i++) { + ListenableFuture blocked; + synchronized (queue) { + blocked = queue.isBlocked(); + } + blocked.get(10, TimeUnit.SECONDS); + synchronized (queue) { + Assert.assertSame(block, queue.remove()); + } + } + return null; + }); + sender.get(10, TimeUnit.SECONDS); + receiver.get(10, TimeUnit.SECONDS); + Assert.assertTrue(queue.hasNoMoreTsBlocks()); + Assert.assertTrue(queue.isEmpty()); + Assert.assertEquals(0, spyMemoryPool.getReservedBytes()); + } finally { + executor.shutdownNow(); 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 { + Assert.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } finally { synchronized (queue) { - queue.setNoMoreTsBlocks(true); + queue.close(); } } } } - 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); - } - } - } - @Test public void testDependentQueuesKeepIndependentBudgetsByDefault() { MemoryPool pool = new MemoryPool("pipeline-chain", new MemoryManager(1024), 512); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java index 0caf1692f635b..237aacd72e45b 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java @@ -134,7 +134,7 @@ public void testFragmentInstanceExecution() throws InterruptedException { } @Test - public void testTVListOwnerTransfer() throws InterruptedException { + public void testTVListOwnerTransfer() throws Exception { ExecutorService instanceNotificationExecutor = IoTDBThreadPoolFactory.newFixedThreadPool(1, "test-instance-notification"); @@ -173,8 +173,9 @@ public void testTVListOwnerTransfer() throws InterruptedException { fragmentInstanceContext2.decrementNumOfUnClosedDriver(); fragmentInstanceContext1.getStateMachine().finished(); - Thread.sleep(100); + instanceNotificationExecutor.submit(() -> {}).get(10, TimeUnit.SECONDS); fragmentInstanceContext2.getStateMachine().finished(); + instanceNotificationExecutor.submit(() -> {}).get(10, TimeUnit.SECONDS); assertTrue(execution1.getInstanceState().isDone()); assertTrue(execution2.getInstanceState().isDone()); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtilCancellationTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtilCancellationTest.java index e347699cb6d93..62e5a26fe3a3e 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtilCancellationTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtilCancellationTest.java @@ -244,12 +244,12 @@ public void testCompactionContextWithoutStateMachine() throws IOException { assertEquals(2, loadedFiles); } - @Test(timeout = 15000) + @Test public void testTimeoutUnblocksDriverResourceCleanup() throws Exception { assertFailureUnblocksDriverResourceCleanup(new QueryTimeoutException()); } - @Test(timeout = 15000) + @Test public void testSemanticFailureUnblocksDriverResourceCleanup() throws Exception { assertFailureUnblocksDriverResourceCleanup( new SemanticException("Scalar sub-query has returned multiple rows.")); @@ -260,6 +260,7 @@ private void assertFailureUnblocksDriverResourceCleanup(RuntimeException failure ExecutorService notifications = Executors.newSingleThreadExecutor(); FragmentInstanceContext context = newContext(notifications); context.initializeNumOfDrivers(1); + ExecutorService worker = Executors.newSingleThreadExecutor(); try { CountDownLatch closeRequested = new CountDownLatch(1); SeriesScanUtil scanner = @@ -314,10 +315,14 @@ public void close() { false, exchangeManager); - assertSame( - failure, - assertThrows( - RuntimeException.class, () -> driver.processFor(new Duration(1, TimeUnit.SECONDS)))); + // Bound the driver operation rather than cold context and mock initialization in a CI fork. + Future execution = + worker.submit( + () -> + assertThrows( + RuntimeException.class, + () -> driver.processFor(new Duration(1, TimeUnit.SECONDS)))); + assertSame(failure, execution.get(5, TimeUnit.SECONDS)); // This can complete only after the preceding cleanup callback gets past allDriversClosed. notifications.submit(() -> {}).get(5, TimeUnit.SECONDS); assertEquals(2, loadedFiles); @@ -329,12 +334,17 @@ public void close() { } finally { // Ensure a failed assertion cannot leave the notification thread waiting on the latch. context.decrementNumOfUnClosedDriver(); - notifications.shutdownNow(); - assertTrue(notifications.awaitTermination(5, TimeUnit.SECONDS)); + worker.shutdownNow(); + notifications.shutdown(); + try { + assertTrue(worker.awaitTermination(5, TimeUnit.SECONDS)); + } finally { + assertTrue(notifications.awaitTermination(5, TimeUnit.SECONDS)); + } } } - @Test(timeout = 60000) + @Test public void testFinishedScanUnblocksDriverResourceCleanup() throws Exception { for (boolean sequence : new boolean[] {true, false}) { for (boolean waitForClose : new boolean[] {false, true}) { @@ -352,7 +362,6 @@ private void assertFinishedScanUnblocksDriverResourceCleanup( CountDownLatch resumeScan = new CountDownLatch(1); CountDownLatch closeRequested = new CountDownLatch(1); // Cover both orderings: the FI is FINISHED before driver.close(), and close is already pending. - notifications.submit(() -> await(allowNotifications)); FragmentInstanceContext context = newContext(notifications); context.initializeNumOfDrivers(1); Operator operator = mock(Operator.class); @@ -401,6 +410,13 @@ public void close() { false, exchangeManager); try { + // Release this gate in finally so a timeout cannot change the ordering under test. + Future notificationsAllowed = + notifications.submit( + () -> { + allowNotifications.await(); + return null; + }); Future execution = worker.submit(() -> driver.processFor(new Duration(1, TimeUnit.SECONDS))); await(scanEntered); @@ -415,6 +431,7 @@ public void close() { assertEquals(FragmentInstanceState.FINISHED, context.getStateMachine().getState()); assertTrue(context.getStateMachine().getFailureCauses().isEmpty()); allowNotifications.countDown(); + notificationsAllowed.get(5, TimeUnit.SECONDS); notifications.submit(() -> {}).get(5, TimeUnit.SECONDS); verify(operator).close(); verify(context.getMemoryReservationContext()).releaseAllReservedMemory(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L1PriorityQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L1PriorityQueueTest.java index 4038bd51c2809..2b1cd8ce0096f 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L1PriorityQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L1PriorityQueueTest.java @@ -23,13 +23,15 @@ import org.junit.Assert; import org.junit.Test; -import java.util.ArrayList; -import java.util.List; +import java.util.concurrent.FutureTask; +import java.util.concurrent.TimeUnit; + +import static org.awaitility.Awaitility.await; public class L1PriorityQueueTest { @Test - public void testPollBlocked() throws InterruptedException { + public void testPollBlocked() throws Exception { IndexedBlockingQueue queue = new L1PriorityQueue<>( 10, @@ -40,27 +42,21 @@ public void testPollBlocked() throws InterruptedException { return Integer.compare(o1.getValue(), o2.getValue()); }, new QueueElement(new QueueElement.QueueElementID(0), 0)); - List res = new ArrayList<>(); - Thread t1 = - new Thread( - () -> { - try { - QueueElement e = queue.poll(); - res.add(e); - } catch (InterruptedException e) { - e.printStackTrace(); - Assert.fail(); - } - }); + FutureTask result = new FutureTask<>(queue::poll); + Thread t1 = new Thread(result); t1.start(); - Thread.sleep(100); - Assert.assertEquals(Thread.State.WAITING, t1.getState()); - QueueElement e2 = new QueueElement(new QueueElement.QueueElementID(1), 1); - queue.push(e2); - Thread.sleep(100); - Assert.assertEquals(Thread.State.TERMINATED, t1.getState()); - Assert.assertEquals(1, res.size()); - Assert.assertEquals(e2.getDriverTaskId().toString(), res.get(0).getDriverTaskId().toString()); + try { + await() + .atMost(10, TimeUnit.SECONDS) + .untilAsserted(() -> Assert.assertEquals(Thread.State.WAITING, t1.getState())); + QueueElement e2 = new QueueElement(new QueueElement.QueueElementID(1), 1); + queue.push(e2); + Assert.assertSame(e2, result.get(10, TimeUnit.SECONDS)); + } finally { + t1.interrupt(); + t1.join(TimeUnit.SECONDS.toMillis(10)); + Assert.assertFalse(t1.isAlive()); + } } @Test From a717a68912d5e04d9ab94fbe4ee46ab63dd49fb0 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 2 Nov 2026 15:29:15 +0800 Subject: [PATCH 2/3] Stabilize blocked priority queue unit tests --- .../schedule/queue/L2PriorityQueueTest.java | 41 ++++++++----------- .../queue/MultilevelPriorityQueueTest.java | 38 ++++++----------- 2 files changed, 28 insertions(+), 51 deletions(-) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L2PriorityQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L2PriorityQueueTest.java index 537358d9a49e8..7809c75c288e1 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L2PriorityQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L2PriorityQueueTest.java @@ -23,8 +23,7 @@ import org.junit.Assert; import org.junit.Test; -import java.util.ArrayList; -import java.util.List; +import java.util.concurrent.FutureTask; import java.util.concurrent.TimeUnit; import static org.awaitility.Awaitility.await; @@ -32,7 +31,7 @@ public class L2PriorityQueueTest { @Test - public void testPollBlocked() throws InterruptedException { + public void testPollBlocked() throws Exception { IndexedBlockingQueue queue = new L2PriorityQueue<>( 10, @@ -43,29 +42,21 @@ public void testPollBlocked() throws InterruptedException { return Integer.compare(o1.getValue(), o2.getValue()); }, new QueueElement(new QueueElement.QueueElementID(0), 0)); - List res = new ArrayList<>(); - Thread t1 = - new Thread( - () -> { - try { - QueueElement e = queue.poll(); - res.add(e); - } catch (InterruptedException e) { - e.printStackTrace(); - fail(); - } - }); + FutureTask result = new FutureTask<>(queue::poll); + Thread t1 = new Thread(result); t1.start(); - await() - .atMost(1, TimeUnit.MINUTES) - .untilAsserted(() -> Assert.assertEquals(Thread.State.WAITING, t1.getState())); - QueueElement e2 = new QueueElement(new QueueElement.QueueElementID(1), 1); - queue.push(e2); - await() - .atMost(1, TimeUnit.MINUTES) - .untilAsserted(() -> Assert.assertEquals(Thread.State.TERMINATED, t1.getState())); - Assert.assertEquals(1, res.size()); - Assert.assertEquals(e2.getDriverTaskId().toString(), res.get(0).getDriverTaskId().toString()); + try { + await() + .atMost(10, TimeUnit.SECONDS) + .untilAsserted(() -> Assert.assertEquals(Thread.State.WAITING, t1.getState())); + QueueElement e2 = new QueueElement(new QueueElement.QueueElementID(1), 1); + queue.push(e2); + Assert.assertSame(e2, result.get(10, TimeUnit.SECONDS)); + } finally { + t1.interrupt(); + t1.join(TimeUnit.SECONDS.toMillis(10)); + Assert.assertFalse(t1.isAlive()); + } } @Test diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/MultilevelPriorityQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/MultilevelPriorityQueueTest.java index 42bb75a2dc59d..87f9a8db73e47 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/MultilevelPriorityQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/MultilevelPriorityQueueTest.java @@ -35,44 +35,30 @@ import org.junit.Test; import org.mockito.Mockito; -import java.util.ArrayList; -import java.util.List; import java.util.OptionalInt; +import java.util.concurrent.FutureTask; import java.util.concurrent.TimeUnit; import static org.awaitility.Awaitility.await; public class MultilevelPriorityQueueTest { @Test - public void testPollBlocked() { + public void testPollBlocked() throws Exception { + IndexedBlockingQueue queue = new MultilevelPriorityQueue(2, 1000, new DriverTask()); + FutureTask result = new FutureTask<>(queue::poll); + Thread t1 = new Thread(result); + t1.start(); try { - IndexedBlockingQueue queue = - new MultilevelPriorityQueue(2, 1000, new DriverTask()); - List res = new ArrayList<>(); - Thread t1 = - new Thread( - () -> { - try { - DriverTask e = queue.poll(); - res.add(e); - } catch (InterruptedException e) { - e.printStackTrace(); - Assert.fail(); - } - }); - t1.start(); await() - .atMost(1, TimeUnit.MINUTES) + .atMost(10, TimeUnit.SECONDS) .untilAsserted(() -> Assert.assertEquals(Thread.State.WAITING, t1.getState())); DriverTask e2 = mockDriverTask(mockDriverTaskId(), false); queue.push(e2); - await() - .atMost(1, TimeUnit.MINUTES) - .untilAsserted(() -> Assert.assertEquals(Thread.State.TERMINATED, t1.getState())); - Assert.assertEquals(1, res.size()); - Assert.assertEquals(e2.getDriverTaskId().toString(), res.get(0).getDriverTaskId().toString()); - } catch (Exception e) { - Assert.fail(); + Assert.assertSame(e2, result.get(10, TimeUnit.SECONDS)); + } finally { + t1.interrupt(); + t1.join(TimeUnit.SECONDS.toMillis(10)); + Assert.assertFalse(t1.isAlive()); } } From d3d726cd39bd00ad63f25b48d1b6fead83fcb96b Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 2 Nov 2026 16:04:03 +0800 Subject: [PATCH 3/3] Propagate worker failures and guarantee unit test cleanup --- .../exchange/SharedTsBlockQueueTest.java | 68 ++++++++++--------- .../SeriesScanUtilCancellationTest.java | 19 ++++-- .../schedule/queue/L1PriorityQueueTest.java | 5 +- .../schedule/queue/L2PriorityQueueTest.java | 5 +- .../queue/MultilevelPriorityQueueTest.java | 5 +- 5 files changed, 61 insertions(+), 41 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 29fe431b32781..18a3094507eaf 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 @@ -34,6 +34,8 @@ import org.junit.Test; import org.mockito.Mockito; +import java.util.concurrent.CompletionService; +import java.util.concurrent.ExecutorCompletionService; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -144,39 +146,41 @@ public void concurrencyTest() throws Exception { TsBlock block = Utils.createMockTsBlock(mockTsBlockSize); int blockCount = 1000; ExecutorService executor = Executors.newFixedThreadPool(2); + CompletionService completedTasks = new ExecutorCompletionService<>(executor); try { - Future sender = - executor.submit( - () -> { - for (int i = 0; i < blockCount; i++) { - ListenableFuture blocked; - synchronized (queue) { - blocked = queue.add(block); - } - blocked.get(10, TimeUnit.SECONDS); - } - synchronized (queue) { - queue.setNoMoreTsBlocks(true); - } - return null; - }); - Future receiver = - executor.submit( - () -> { - for (int i = 0; i < blockCount; i++) { - ListenableFuture blocked; - synchronized (queue) { - blocked = queue.isBlocked(); - } - blocked.get(10, TimeUnit.SECONDS); - synchronized (queue) { - Assert.assertSame(block, queue.remove()); - } - } - return null; - }); - sender.get(10, TimeUnit.SECONDS); - receiver.get(10, TimeUnit.SECONDS); + completedTasks.submit( + () -> { + for (int i = 0; i < blockCount; i++) { + ListenableFuture blocked; + synchronized (queue) { + blocked = queue.add(block); + } + blocked.get(10, TimeUnit.SECONDS); + } + synchronized (queue) { + queue.setNoMoreTsBlocks(true); + } + return null; + }); + completedTasks.submit( + () -> { + for (int i = 0; i < blockCount; i++) { + ListenableFuture blocked; + synchronized (queue) { + blocked = queue.isBlocked(); + } + blocked.get(10, TimeUnit.SECONDS); + synchronized (queue) { + Assert.assertSame(block, queue.remove()); + } + } + return null; + }); + for (int i = 0; i < 2; i++) { + Future completed = completedTasks.poll(10, TimeUnit.SECONDS); + Assert.assertNotNull(completed); + completed.get(); + } Assert.assertTrue(queue.hasNoMoreTsBlocks()); Assert.assertTrue(queue.isEmpty()); Assert.assertEquals(0, spyMemoryPool.getReservedBytes()); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtilCancellationTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtilCancellationTest.java index 62e5a26fe3a3e..d0bda579acade 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtilCancellationTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtilCancellationTest.java @@ -443,12 +443,19 @@ public void close() { resumeScan.countDown(); allowNotifications.countDown(); worker.shutdownNow(); - assertTrue(worker.awaitTermination(5, TimeUnit.SECONDS)); - driver.close(); - context.decrementNumOfUnClosedDriver(); - // Let FI cleanup finish after opening its latch, without interrupting its driver-close wait. - notifications.shutdown(); - assertTrue(notifications.awaitTermination(5, TimeUnit.SECONDS)); + try { + assertTrue(worker.awaitTermination(5, TimeUnit.SECONDS)); + } finally { + try { + driver.close(); + } finally { + context.decrementNumOfUnClosedDriver(); + // Let FI cleanup finish after opening its latch, without interrupting its driver-close + // wait. + notifications.shutdown(); + assertTrue(notifications.awaitTermination(5, TimeUnit.SECONDS)); + } + } } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L1PriorityQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L1PriorityQueueTest.java index 2b1cd8ce0096f..cb5abc56abbe9 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L1PriorityQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L1PriorityQueueTest.java @@ -48,7 +48,10 @@ public void testPollBlocked() throws Exception { try { await() .atMost(10, TimeUnit.SECONDS) - .untilAsserted(() -> Assert.assertEquals(Thread.State.WAITING, t1.getState())); + .until(() -> result.isDone() || t1.getState() == Thread.State.WAITING); + if (result.isDone()) { + Assert.fail("poll completed before an element was pushed: " + result.get()); + } QueueElement e2 = new QueueElement(new QueueElement.QueueElementID(1), 1); queue.push(e2); Assert.assertSame(e2, result.get(10, TimeUnit.SECONDS)); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L2PriorityQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L2PriorityQueueTest.java index 7809c75c288e1..3a9e85fcc2fbc 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L2PriorityQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L2PriorityQueueTest.java @@ -48,7 +48,10 @@ public void testPollBlocked() throws Exception { try { await() .atMost(10, TimeUnit.SECONDS) - .untilAsserted(() -> Assert.assertEquals(Thread.State.WAITING, t1.getState())); + .until(() -> result.isDone() || t1.getState() == Thread.State.WAITING); + if (result.isDone()) { + Assert.fail("poll completed before an element was pushed: " + result.get()); + } QueueElement e2 = new QueueElement(new QueueElement.QueueElementID(1), 1); queue.push(e2); Assert.assertSame(e2, result.get(10, TimeUnit.SECONDS)); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/MultilevelPriorityQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/MultilevelPriorityQueueTest.java index 87f9a8db73e47..7d9de811af5cd 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/MultilevelPriorityQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/MultilevelPriorityQueueTest.java @@ -51,7 +51,10 @@ public void testPollBlocked() throws Exception { try { await() .atMost(10, TimeUnit.SECONDS) - .untilAsserted(() -> Assert.assertEquals(Thread.State.WAITING, t1.getState())); + .until(() -> result.isDone() || t1.getState() == Thread.State.WAITING); + if (result.isDone()) { + Assert.fail("poll completed before an element was pushed: " + result.get()); + } DriverTask e2 = mockDriverTask(mockDriverTaskId(), false); queue.push(e2); Assert.assertSame(e2, result.get(10, TimeUnit.SECONDS));