From de60bcaf3164a7d2d97b9d6cb4f025af6febe30f Mon Sep 17 00:00:00 2001
From: QuakeWang <wangfuzheng0814@foxmail.com>
Date: Mon, 27 Jul 2026 11:08:02 +0800
Subject: [PATCH] [common] Ensure ParallelExecution waits for reader cleanup

ParallelExecution marks tasks complete before reader cleanup, while close only interrupts the executor. It can therefore return before input streams are released.

Wait for executor termination in close and cover the race with a controlled concurrency test.

Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com>
---
 .../paimon/utils/ParallelExecution.java       |  8 +++
 .../crosspartition/IndexBootstrapTest.java    |  5 --
 .../paimon/utils/ParallelExecutionTest.java   | 62 +++++++++++++++++++
 3 files changed, 70 insertions(+), 5 deletions(-)

diff --git a/paimon-common/src/main/java/org/apache/paimon/utils/ParallelExecution.java b/paimon-common/src/main/java/org/apache/paimon/utils/ParallelExecution.java
index cb8b1c1f650e..6341aec8b2ae 100644
--- a/paimon-common/src/main/java/org/apache/paimon/utils/ParallelExecution.java
+++ b/paimon-common/src/main/java/org/apache/paimon/utils/ParallelExecution.java
@@ -174,6 +174,14 @@ private void sendToResults(SimpleCollectingOutputView outputView, int count, E e
     @Override
     public void close() throws IOException {
         this.executorService.shutdownNow();
+        try {
+            if (!this.executorService.awaitTermination(1, TimeUnit.MINUTES)) {
+                throw new IOException("Timed out while closing parallel execution.");
+            }
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new IOException("Interrupted while closing parallel execution.", e);
+        }
     }
 
     private ParallelBatch<T, E> iterator(MemorySegment page, int numRecords, E extraMessage) {
diff --git a/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java b/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java
index 0921ddf565ea..0e352b63ec7d 100644
--- a/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java
+++ b/paimon-core/src/test/java/org/apache/paimon/crosspartition/IndexBootstrapTest.java
@@ -105,11 +105,6 @@ public void testBoostrap() throws Exception {
                 .containsExactlyInAnyOrder(
                         GenericRow.of(2, 1, 3), GenericRow.of(4, 2, 5), GenericRow.of(6, 3, 7));
         result.clear();
-
-        // In ParallelExecution, latch.countDown first, then close the reader, it may not be closed
-        // here, (this is good, beneficial for query speed) but TableTestBase.after will check leak
-        // streams. So sleep here to avoid unstable.
-        Thread.sleep(1000);
     }
 
     private Table createTable() throws Exception {
diff --git a/paimon-core/src/test/java/org/apache/paimon/utils/ParallelExecutionTest.java b/paimon-core/src/test/java/org/apache/paimon/utils/ParallelExecutionTest.java
index 4e3feab57891..0f66d31bb964 100644
--- a/paimon-core/src/test/java/org/apache/paimon/utils/ParallelExecutionTest.java
+++ b/paimon-core/src/test/java/org/apache/paimon/utils/ParallelExecutionTest.java
@@ -32,6 +32,13 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Queue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Supplier;
 
 import static java.util.Collections.singletonList;
@@ -142,6 +149,61 @@ public void testTooBigRecord() {
                                 + " Please increase the 'page-size' table option.");
     }
 
+    @Test
+    public void testCloseWaitsForReaderClose() throws Exception {
+        CountDownLatch readerCloseStarted = new CountDownLatch(1);
+        CountDownLatch readerClosed = new CountDownLatch(1);
+        Semaphore allowReaderClose = new Semaphore(0);
+        RecordReader<Integer> reader =
+                new RecordReader<Integer>() {
+                    @Nullable
+                    @Override
+                    public RecordIterator<Integer> readBatch() {
+                        return null;
+                    }
+
+                    @Override
+                    public void close() {
+                        readerCloseStarted.countDown();
+                        allowReaderClose.acquireUninterruptibly();
+                        readerClosed.countDown();
+                    }
+                };
+
+        ParallelExecution<Integer, Integer> execution =
+                new ParallelExecution<>(
+                        new IntSerializer(), 1024, 1, singletonList(() -> Pair.of(reader, 1)));
+        ExecutorService closeExecutor = Executors.newSingleThreadExecutor();
+
+        try {
+            assertThat(readerCloseStarted.await(10, TimeUnit.SECONDS)).isTrue();
+
+            CountDownLatch executionCloseStarted = new CountDownLatch(1);
+            Future<Void> closeFuture =
+                    closeExecutor.submit(
+                            () -> {
+                                executionCloseStarted.countDown();
+                                execution.close();
+                                return null;
+                            });
+            assertThat(executionCloseStarted.await(10, TimeUnit.SECONDS)).isTrue();
+            assertThatThrownBy(() -> closeFuture.get(1, TimeUnit.SECONDS))
+                    .isInstanceOf(TimeoutException.class);
+
+            allowReaderClose.release();
+            closeFuture.get(10, TimeUnit.SECONDS);
+            assertThat(readerClosed.await(10, TimeUnit.SECONDS)).isTrue();
+        } finally {
+            allowReaderClose.release();
+            try {
+                execution.close();
+                readerClosed.await(10, TimeUnit.SECONDS);
+            } finally {
+                closeExecutor.shutdownNow();
+            }
+        }
+    }
+
     private RecordReader<Integer> create(Queue<List<Integer>> queue) {
         return new RecordReader<Integer>() {
             @Nullable