Stabilize query execution tests on dev/1.3

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

@@ -26,15 +26,17 @@
 
 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 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;
 
@@ -136,120 +138,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;
-        }
+    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);
-      }
     }
   }
 }

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceExecutionTest.java

@@ -68,6 +68,7 @@
 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 @@
   }
 
   @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 @@
         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

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L1PriorityQueueTest.java

@@ -21,13 +21,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<QueueElement> queue =
         new L1PriorityQueue<>(
             10,
@@ -38,27 +40,24 @@ public void testPollBlocked() throws InterruptedException {
               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

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/schedule/queue/L2PriorityQueueTest.java

@@ -21,16 +21,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;
 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 void testPollBlocked() throws InterruptedException {
               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