Fix flaky SharedTsBlockQueue concurrency test

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

@@ -27,15 +27,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.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 java.util.concurrent.TimeoutException;
 
 import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
 
@@ -117,18 +119,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 +141,62 @@ public void concurrencyTest() {
     queue.setMaxBytesCanReserve(Long.MAX_VALUE);
 
     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 {
+      ExecutorCompletionService<Void> completionService = new ExecutorCompletionService<>(executor);
+      Future<Void> sender =
+          completionService.submit(
+              () -> {
+                for (int i = 0; i < numOfTsBlocks; i++) {
+                  ListenableFuture<Void> blockedOnMemory;
+                  synchronized (queue) {
+                    blockedOnMemory = queue.add(Utils.createMockTsBlock(mockTsBlockSize));
+                  }
+                  blockedOnMemory.get();
+                }
+                synchronized (queue) {
+                  queue.setNoMoreTsBlocks(true);
+                }
+                return null;
+              });
+      Future<Void> receiver =
+          completionService.submit(
+              () -> {
+                for (int i = 0; i < numOfTsBlocks; i++) {
+                  ListenableFuture<Void> blocked;
+                  synchronized (queue) {
+                    blocked = queue.isBlocked();
+                  }
+                  blocked.get();
+                  synchronized (queue) {
+                    queue.remove();
+                  }
+                }
+                return null;
+              });
+
+      final long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(30);
       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);
-        }
-      }
-    }
-  }
-
-  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) {
-        synchronized (queue) {
-          blocked = queue.isBlocked();
-          if (blocked.isDone()) {
-            queue.remove();
-            numOfTsBlocksToReceive.updateAndGet(v -> v - 1);
-          } else {
-            break;
+        for (int completedWorkerCount = 0; completedWorkerCount < 2; completedWorkerCount++) {
+          final long remainingNanos = deadline - System.nanoTime();
+          final Future<Void> completedWorker =
+              completionService.poll(Math.max(remainingNanos, 0), TimeUnit.NANOSECONDS);
+          if (completedWorker == null) {
+            throw new TimeoutException();
           }
+          completedWorker.get();
         }
+      } catch (Exception e) {
+        sender.cancel(true);
+        receiver.cancel(true);
+        throw e;
       }
 
-      if (blocked != null) {
-        numOfTimesBlocked.updateAndGet(v -> v + 1);
-        blocked.addListener(
-            new ReceiveTask(queue, numOfTsBlocksToReceive, numOfTimesBlocked, executor), executor);
-      }
+      Assert.assertTrue(queue.hasNoMoreTsBlocks());
+      Assert.assertTrue(queue.isEmpty());
+      Assert.assertEquals(0L, memoryPool.getReservedBytes());
+    } finally {
+      executor.shutdownNow();
+      executor.awaitTermination(10, TimeUnit.SECONDS);
     }
   }
 }