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 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 reader = + new RecordReader() { + @Nullable + @Override + public RecordIterator readBatch() { + return null; + } + + @Override + public void close() { + readerCloseStarted.countDown(); + allowReaderClose.acquireUninterruptibly(); + readerClosed.countDown(); + } + }; + + ParallelExecution 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 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 create(Queue> queue) { return new RecordReader() { @Nullable