Stabilize query execution concurrency unit tests

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java

@@ -29,15 +29,17 @@
 
 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;
 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;
 
@@ -119,8 +121,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,120 +142,57 @@ 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<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;
-        }
-      }
-
-      if (blockedOnMemory != null) {
-        numOfTimesBlocked.updateAndGet(v -> v + 1);
-        blockedOnMemory.addListener(
-            new SendTask(queue, mockTsBlockSize, numOfTsBlocksToSend, numOfTimesBlocked, executor),
-            executor);
-      } else {
-        synchronized (queue) {
-          queue.setNoMoreTsBlocks(true);
-        }
+    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();
       }
-    }
-  }
-
-  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);
-      }
     }
   }
 

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());

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<RuntimeException> 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();
@@ -426,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));
+        }
+      }
     }
   }