diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasks.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasks.java index 3349c9e07fe8..2fb76bf6ac86 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasks.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasks.java @@ -25,6 +25,7 @@ import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.function.Consumer; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder; @@ -100,7 +101,9 @@ void shutdown() { executor.shutdownNow(); } - private void drainFinished(Consumer onDone) throws Exception { + /** Delivers every task that has finished, in queue order, without blocking. */ + @VisibleForTesting + void drainFinished(Consumer onDone) throws Exception { Iterator> iterator = active.iterator(); while (iterator.hasNext()) { Future future = iterator.next(); diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasksTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasksTest.java index 09471c1f414f..478a616d58a2 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasksTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BoundedAsyncTasksTest.java @@ -88,14 +88,17 @@ public void testOutOfOrderCompletionDrainsFinishedTasks() throws Exception { }, delivered::add); - // Wait until fast task has finished running in the background + // The latch fires inside the task body; the future is marked done a moment later, so a + // single drain can miss it. Drain until it shows up while "slow" is still blocked. fastTaskDone.await(); - - // Trigger a drain by submitting another task tasks.submit(() -> "noop", delivered::add); - - // "fast" should have been drained while "slow" is still blocked - assertTrue(delivered.contains("fast")); + int retries = 0; + while (!delivered.contains("fast") && retries < 1000) { + Thread.sleep(1); + tasks.drainFinished(delivered::add); + retries++; + } + assertTrue("fast task was never drained", delivered.contains("fast")); assertFalse(delivered.contains("slow")); slowTaskHold.countDown();