[common] Ensure ParallelExecution waits for reader cleanup

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) {

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 {

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