This is an automated email from the ASF dual-hosted git repository. Caideyipi pushed a commit to branch codex/cp-stabilize-query-concurrency-tests in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 22271c21cb8fa3eb2f0b9888366763ef389cada5 Author: Caideyipi <[email protected]> AuthorDate: Mon Nov 2 18:44:12 2026 +0800 Stabilize query execution concurrency unit tests (#19079) * Stabilize query execution concurrency unit tests * Stabilize blocked priority queue unit tests * Propagate worker failures and guarantee unit test cleanup (cherry picked from commit 24fd30e1ac6b43a32aaeca2066034b1ae81fbc44) --- .../execution/exchange/SharedTsBlockQueueTest.java | 165 +++++++-------------- .../fragment/FragmentInstanceExecutionTest.java | 8 +- .../schedule/queue/L1PriorityQueueTest.java | 45 +++--- .../schedule/queue/L2PriorityQueueTest.java | 44 +++--- .../queue/MultilevelPriorityQueueTest.java | 43 ++---- 5 files changed, 114 insertions(+), 191 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 a95fd35ba76..1e9080d3518 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 @@ -26,15 +26,17 @@ import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; -import org.apache.commons.lang3.Validate; import org.apache.tsfile.read.common.block.TsBlock; import org.junit.Assert; 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.atomic.AtomicReference; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService; @@ -116,8 +118,8 @@ public class SharedTsBlockQueueTest { 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; @@ -136,120 +138,57 @@ public class SharedTsBlockQueueTest { 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<Integer> numOfTimesSenderBlocked = new AtomicReference<>(0); - AtomicReference<Integer> numOfTimesReceiverBlocked = new AtomicReference<>(0); - AtomicReference<Integer> numOfTsBlocksToSend = new AtomicReference<>(1000); - AtomicReference<Integer> 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<Integer> numOfTsBlocksToSend; - private final AtomicReference<Integer> numOfTimesBlocked; - private final ExecutorService executor; - - public SendTask( - SharedTsBlockQueue queue, - long mockTsBlockSize, - AtomicReference<Integer> numOfTsBlocksToSend, - AtomicReference<Integer> 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<Void> blockedOnMemory = null; - while (numOfTsBlocksToSend.get() > 0) { - synchronized (queue) { - blockedOnMemory = queue.add(Utils.createMockTsBlock(mockTsBlockSize)); - } - numOfTsBlocksToSend.updateAndGet(v -> v - 1); - if (!blockedOnMemory.isDone()) { - break; - } + CompletionService<Void> completedTasks = new ExecutorCompletionService<>(executor); + try { + completedTasks.submit( + () -> { + for (int i = 0; i < blockCount; i++) { + ListenableFuture<Void> 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<Void> 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<Void> completed = completedTasks.poll(10, TimeUnit.SECONDS); + Assert.assertNotNull(completed); + completed.get(); } - - 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<Integer> numOfTsBlocksToReceive; - private final AtomicReference<Integer> numOfTimesBlocked; - private final ExecutorService executor; - - public ReceiveTask( - SharedTsBlockQueue queue, - AtomicReference<Integer> numOfTsBlocksToReceive, - AtomicReference<Integer> 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<Void> blocked = null; - while (numOfTsBlocksToReceive.get() > 0) { + Assert.assertTrue(queue.hasNoMoreTsBlocks()); + Assert.assertTrue(queue.isEmpty()); + Assert.assertEquals(0, spyMemoryPool.getReservedBytes()); + } finally { + executor.shutdownNow(); + try { + Assert.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } finally { synchronized (queue) { - blocked = queue.isBlocked(); - if (blocked.isDone()) { - queue.remove(); - numOfTsBlocksToReceive.updateAndGet(v -> v - 1); - } else { - break; - } + queue.close(); } } - - if (blocked != null) { - numOfTimesBlocked.updateAndGet(v -> v + 1); - blocked.addListener( - new ReceiveTask(queue, numOfTsBlocksToReceive, numOfTimesBlocked, executor), executor); - } } } } 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 8bde6b302d5..8f91e9abf92 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 @@ -68,6 +68,7 @@ import java.util.HashSet; import java.util.List; import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; import static org.apache.iotdb.db.queryengine.common.QueryId.MOCK_QUERY_ID; import static org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext.createFragmentInstanceContext; @@ -133,7 +134,7 @@ public class FragmentInstanceExecutionTest { } @Test - public void testTVListOwnerTransfer() throws InterruptedException { + public void testTVListOwnerTransfer() throws Exception { // Capture System.err to check for warning messages PrintStream systemOut = System.out; ByteArrayOutputStream logPrint = new ByteArrayOutputStream(); @@ -171,16 +172,17 @@ public class FragmentInstanceExecutionTest { 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()); - Thread.sleep(100); } catch (CpuNotEnoughException | MemoryNotEnoughException | IllegalArgumentException e) { fail(e.getMessage()); } finally { instanceNotificationExecutor.shutdown(); + assertTrue(instanceNotificationExecutor.awaitTermination(10, TimeUnit.SECONDS)); } } finally { // Restore original System.out 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 016328c42a5..2161fa5e345 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 @@ -21,13 +21,15 @@ package org.apache.iotdb.db.queryengine.execution.schedule.queue; 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<QueueElement> queue = new L1PriorityQueue<>( 10, @@ -38,27 +40,24 @@ public class L1PriorityQueueTest { return Integer.compare(o1.getValue(), o2.getValue()); }, new QueueElement(new QueueElement.QueueElementID(0), 0)); - List<QueueElement> res = new ArrayList<>(); - Thread t1 = - new Thread( - () -> { - try { - QueueElement e = queue.poll(); - res.add(e); - } catch (InterruptedException e) { - e.printStackTrace(); - Assert.fail(); - } - }); + FutureTask<QueueElement> 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) + .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)); + } 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/L2PriorityQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L2PriorityQueueTest.java index e7162b741ea..ca42d10249c 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 @@ -21,8 +21,7 @@ package org.apache.iotdb.db.queryengine.execution.schedule.queue; 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; @@ -30,7 +29,7 @@ import static org.junit.Assert.fail; public class L2PriorityQueueTest { @Test - public void testPollBlocked() throws InterruptedException { + public void testPollBlocked() throws Exception { IndexedBlockingQueue<QueueElement> queue = new L2PriorityQueue<>( 10, @@ -41,29 +40,24 @@ public class L2PriorityQueueTest { return Integer.compare(o1.getValue(), o2.getValue()); }, new QueueElement(new QueueElement.QueueElementID(0), 0)); - List<QueueElement> res = new ArrayList<>(); - Thread t1 = - new Thread( - () -> { - try { - QueueElement e = queue.poll(); - res.add(e); - } catch (InterruptedException e) { - e.printStackTrace(); - fail(); - } - }); + FutureTask<QueueElement> 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) + .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)); + } 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 0c7e40569d1..e247816bf4f 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 @@ -34,44 +34,33 @@ import org.junit.Assert; 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<DriverTask> queue = new MultilevelPriorityQueue(2, 1000, new DriverTask()); + FutureTask<DriverTask> result = new FutureTask<>(queue::poll); + Thread t1 = new Thread(result); + t1.start(); try { - IndexedBlockingQueue<DriverTask> queue = - new MultilevelPriorityQueue(2, 1000, new DriverTask()); - List<DriverTask> 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) - .untilAsserted(() -> Assert.assertEquals(Thread.State.WAITING, t1.getState())); + .atMost(10, TimeUnit.SECONDS) + .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); - 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()); } }
