From 30734910c6dbf835ed4b6bb1a68b2bb78f7fed92 Mon Sep 17 00:00:00 2001 From: Tejas Iyer Date: Mon, 15 Jun 2026 17:50:19 +0000 Subject: [PATCH 1/2] Fix AsyncWrapperTest timeout and flakiness issues. --- .../beam/sdk/transforms/AsyncWrapperTest.java | 38 ++++++++++--------- 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java index e95f586e8849..cfd992af3fdb 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java @@ -231,20 +231,22 @@ private void waitForEmpty(AsyncWrapper asyncWrapper) { private void waitForEmpty(AsyncWrapper asyncWrapper, int timeoutSeconds) { int count = 0; + // Poll every 5 milliseconds instead of 1000 milliseconds for instant response + int maxIterations = timeoutSeconds * 200; while (!asyncWrapper.isEmpty()) { try { - Thread.sleep(1000); + Thread.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new RuntimeException(e); } count += 1; - if (count > timeoutSeconds) { + if (count > maxIterations) { throw new RuntimeException("Timed out waiting for async dofn to be empty"); } } try { - Thread.sleep(1000); + Thread.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } @@ -419,7 +421,7 @@ public void testMultiKey() { // execution task has not finished processing yet. @Test public void testLongItem() { - BasicDofn dofn = new BasicDofn(1000); + BasicDofn dofn = new BasicDofn(10); AsyncWrapper asyncWrapper = new AsyncWrapper<>( dofn, 1, Duration.standardSeconds(5), null, null, null, null, useThreadPool); @@ -438,7 +440,7 @@ public void testLongItem() { assertEquals(0, dofn.getProcessed()); assertEquals(1, fakeBagState.items.size()); - waitForEmpty(asyncWrapper, 20); + waitForEmpty(asyncWrapper, 2); result = asyncWrapper.commitFinishedItemsDirect( @@ -538,7 +540,7 @@ public void testMultiElementDofn() { // Identical elements should not spawn multiple concurrent background executions. @Test public void testDuplicates() { - BasicDofn dofn = new BasicDofn(1000); + BasicDofn dofn = new BasicDofn(10); AsyncWrapper asyncWrapper = new AsyncWrapper<>( dofn, 1, Duration.standardSeconds(5), null, null, null, null, useThreadPool); @@ -568,7 +570,7 @@ public void testDuplicates() { // has cleared are correctly tracked and processed. @Test public void testSlowDuplicates() { - BasicDofn dofn = new BasicDofn(5000); + BasicDofn dofn = new BasicDofn(20); AsyncWrapper asyncWrapper = new AsyncWrapper<>( dofn, 1, Duration.standardSeconds(5), null, null, null, null, useThreadPool); @@ -581,7 +583,7 @@ public void testSlowDuplicates() { asyncWrapper.processDirect(msg, GlobalWindow.INSTANCE, Instant.now(), fakeBagState, fakeTimer); try { - Thread.sleep(10000); + Thread.sleep(100); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } @@ -610,7 +612,7 @@ public void testSlowDuplicates() { // and decrement immediately upon execution completion. @Test public void testBufferCount() { - BasicDofn dofn = new BasicDofn(1000); + BasicDofn dofn = new BasicDofn(10); AsyncWrapper asyncWrapper = new AsyncWrapper<>( dofn, 1, Duration.standardSeconds(5), null, null, null, null, useThreadPool); @@ -637,7 +639,7 @@ public void testBufferCount() { // the scheduler must block and delay submissions appropriately. @Test public void testBufferStopsAcceptingItems() { - BasicDofn dofn = new BasicDofn(1000); + BasicDofn dofn = new BasicDofn(10); AsyncWrapper asyncWrapper = new AsyncWrapper<>( dofn, @@ -670,7 +672,7 @@ public void testBufferStopsAcceptingItems() { } try { - Thread.sleep(200); + Thread.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } @@ -707,7 +709,7 @@ public void testBufferStopsAcceptingItems() { // Verifies actively cancelled elements are cleanly dropped from the buffer during throttling. @Test public void testBufferWithCancellation() { - BasicDofn dofn = new BasicDofn(1000); + BasicDofn dofn = new BasicDofn(10); AsyncWrapper asyncWrapper = new AsyncWrapper<>( dofn, 1, Duration.standardSeconds(5), null, null, null, null, useThreadPool); @@ -746,7 +748,7 @@ public void testBufferWithCancellation() { // across multiple keys correctly under heavy multi-threaded load. @Test public void testLoadCorrectness() { - BasicDofn dofn = new BasicDofn(1000); + BasicDofn dofn = new BasicDofn(10); AsyncWrapper asyncWrapper = new AsyncWrapper<>( dofn, @@ -791,14 +793,14 @@ public void testLoadCorrectness() { timers.get(key)); })); try { - Thread.sleep(random.nextInt(200)); + Thread.sleep(random.nextInt(2)); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } } try { - Thread.sleep(3000 + random.nextInt(2000)); + Thread.sleep(50 + random.nextInt(20)); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } @@ -834,7 +836,7 @@ public void testLoadCorrectness() { } } try { - Thread.sleep(1000 + random.nextInt(2000)); + Thread.sleep(10 + random.nextInt(20)); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } @@ -854,7 +856,7 @@ public void testLoadCorrectness() { // must complete cleanly without thread or lock deadlocks. @Test public void testResetStateConcurrentTeardown() { - BasicDofn dofn = new BasicDofn(500); + BasicDofn dofn = new BasicDofn(10); AsyncWrapper asyncWrapper = new AsyncWrapper<>( dofn, 1, Duration.standardSeconds(5), null, null, null, null, useThreadPool); @@ -867,7 +869,7 @@ public void testResetStateConcurrentTeardown() { KV.of("key1", "1"), GlobalWindow.INSTANCE, Instant.now(), fakeBagState, fakeTimer); try { - Thread.sleep(50); + Thread.sleep(2); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } From fa018c7ff71c1b59da3e46834b192ced416e874a Mon Sep 17 00:00:00 2001 From: Tejas Iyer Date: Mon, 15 Jun 2026 18:59:26 +0000 Subject: [PATCH 2/2] Fix AsyncWrapperTest to prevent new race conditions and premature timeouts. --- .../beam/sdk/transforms/AsyncWrapperTest.java | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java index cfd992af3fdb..183b1851459c 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java @@ -230,20 +230,17 @@ private void waitForEmpty(AsyncWrapper asyncWrapper) { } private void waitForEmpty(AsyncWrapper asyncWrapper, int timeoutSeconds) { - int count = 0; - // Poll every 5 milliseconds instead of 1000 milliseconds for instant response - int maxIterations = timeoutSeconds * 200; + long limit = System.currentTimeMillis() + timeoutSeconds * 1000L; while (!asyncWrapper.isEmpty()) { + if (System.currentTimeMillis() > limit) { + throw new RuntimeException("Timed out waiting for async dofn to be empty"); + } try { Thread.sleep(5); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new RuntimeException(e); } - count += 1; - if (count > maxIterations) { - throw new RuntimeException("Timed out waiting for async dofn to be empty"); - } } try { Thread.sleep(5); @@ -421,7 +418,7 @@ public void testMultiKey() { // execution task has not finished processing yet. @Test public void testLongItem() { - BasicDofn dofn = new BasicDofn(10); + BasicDofn dofn = new BasicDofn(500); AsyncWrapper asyncWrapper = new AsyncWrapper<>( dofn, 1, Duration.standardSeconds(5), null, null, null, null, useThreadPool); @@ -639,7 +636,7 @@ public void testBufferCount() { // the scheduler must block and delay submissions appropriately. @Test public void testBufferStopsAcceptingItems() { - BasicDofn dofn = new BasicDofn(10); + BasicDofn dofn = new BasicDofn(500); AsyncWrapper asyncWrapper = new AsyncWrapper<>( dofn, @@ -672,7 +669,7 @@ public void testBufferStopsAcceptingItems() { } try { - Thread.sleep(5); + Thread.sleep(100); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } @@ -800,7 +797,7 @@ public void testLoadCorrectness() { } try { - Thread.sleep(50 + random.nextInt(20)); + Thread.sleep(1000 + random.nextInt(1000)); } catch (InterruptedException e) { Thread.currentThread().interrupt(); }