Caideyipi commented on code in PR #18244:
URL: https://github.com/apache/iotdb/pull/18244#discussion_r3680529653


##########
iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java:
##########
@@ -139,119 +139,47 @@ 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 {
-        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;
-          }
-        }
-      }
-
-      if (blocked != null) {
-        numOfTimesBlocked.updateAndGet(v -> v + 1);
-        blocked.addListener(
-            new ReceiveTask(queue, numOfTsBlocksToReceive, numOfTimesBlocked, 
executor), executor);
-      }
+    try {
+      Future<Void> sender =
+          executor.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 =
+          executor.submit(
+              () -> {
+                for (int i = 0; i < numOfTsBlocks; i++) {
+                  ListenableFuture<Void> blocked;
+                  synchronized (queue) {
+                    blocked = queue.isBlocked();
+                  }
+                  blocked.get();
+                  synchronized (queue) {
+                    queue.remove();
+                  }
+                }
+                return null;
+              });
+
+      sender.get(30, TimeUnit.SECONDS);

Review Comment:
   Applied in af46d3df2f1. The test now consumes worker results through 
ExecutorCompletionService with one shared 30-second deadline, so the first 
completed worker is observed immediately. If either worker fails or the shared 
deadline expires, both futures are cancelled before rethrowing, preserving the 
root cause. SharedTsBlockQueueTest passes (2 tests).



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to