This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 402b41c726d Fix ThreadSanitizer data race and failed RPC calls (#39076)
402b41c726d is described below
commit 402b41c726d14d6bd2064629bf82353712563bdd
Author: Tejas Iyer <[email protected]>
AuthorDate: Mon Jun 29 07:28:15 2026 -0700
Fix ThreadSanitizer data race and failed RPC calls (#39076)
* Fix ThreadSanitizer data race in AsyncWrapperTest
* Fix failed future memory leak to support transient RPC retries
* Fix AsyncWrapper to not mark cancelled futures as finished
---
.../apache/beam/sdk/transforms/AsyncWrapper.java | 10 ++--
.../beam/sdk/transforms/AsyncWrapperTest.java | 66 +++++++++++++++++++++-
sdks/python/apache_beam/transforms/async_dofn.py | 4 +-
.../apache_beam/transforms/async_dofn_test.py | 47 +++++++++++++++
4 files changed, 121 insertions(+), 6 deletions(-)
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncWrapper.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncWrapper.java
index f6272ba4fe6..23372be2a2d 100644
---
a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncWrapper.java
+++
b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncWrapper.java
@@ -644,15 +644,17 @@ public class AsyncWrapper<K, InputT, OutputT> extends
DoFn<KV<K, InputT>, Output
if (activeElements.containsKey(elementId)) {
InFlightElement<OutputT> inFlight = activeElements.get(elementId);
+ // Future is either completed, cancelled, or throws an exception
if (inFlight.future.isDone()) {
+ // Remove from local active map before checking the result
+ activeElements.remove(elementId);
try {
if (!inFlight.future.isCancelled()) {
toReturn.add(inFlight.future.get());
+ // Only mark as finished if future was not cancelled
+ finishedElementIds.add(elementId);
+ itemsFinished++;
}
-
- finishedElementIds.add(elementId);
- activeElements.remove(elementId);
- itemsFinished++;
} catch (Exception e) {
LOG.error("Error executing async task for element {}", element,
e);
throw new RuntimeException("Error executing async task for
element " + element, e);
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 183b1851459..d098104a8cf 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
@@ -19,6 +19,7 @@ package org.apache.beam.sdk.transforms;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertThrows;
import java.io.Serializable;
import java.util.ArrayList;
@@ -179,7 +180,7 @@ public class AsyncWrapperTest implements Serializable {
// 4. Used for testing Timer mock implementations.
private static class FakeTimer implements Timer {
- private Instant time = Instant.EPOCH;
+ private volatile Instant time = Instant.EPOCH;
@Override
public void set(Instant absoluteTime) {
@@ -874,4 +875,67 @@ public class AsyncWrapperTest implements Serializable {
// Verify calling resetState() while background tasks are running finishes
cleanly
AsyncWrapper.resetState();
}
+
+ // Test 15: testTransientRpcFailureRetry
+ // Verify DoFn exceptions wipe local active state so bundle retries
reschedule work.
+ @Test
+ public void testTransientRpcFailureRetry() {
+ FlakyDoFn dofn = new FlakyDoFn();
+ AsyncWrapper<String, String, String> asyncWrapper =
+ new AsyncWrapper<>(
+ dofn, 1, Duration.standardSeconds(5), null, null, null, null,
useThreadPool);
+ asyncWrapper.setup(null);
+
+ FakeBagState<KV<String, String>> fakeBagState = new FakeBagState<>();
+ FakeTimer fakeTimer = new FakeTimer();
+ KV<String, String> msg = KV.of("key1", "1");
+
+ asyncWrapper.processDirect(msg, GlobalWindow.INSTANCE, Instant.now(),
fakeBagState, fakeTimer);
+ waitForEmpty(asyncWrapper);
+
+ // Attempt 1 should throw RuntimeException wrapping ExecutionException
+ assertThrows(
+ RuntimeException.class,
+ () ->
+ asyncWrapper.commitFinishedItemsDirect(
+ fakeTimer.getCurrentRelativeTime(), fakeBagState, fakeTimer));
+
+ // Verify the failed Future was removed from local active elements
+ assertEquals(0, asyncWrapper.getItemsInBufferCount());
+
+ // Simulate runner bundle retry: commitFinishedItemsDirect runs again with
msg still in state.
+ // Because the dead future was removed, it will reschedule msg and succeed
on Attempt 2.
+ waitForEmpty(asyncWrapper);
+ List<String> result =
+ asyncWrapper.commitFinishedItemsDirect(
+ fakeTimer.getCurrentRelativeTime(), fakeBagState, fakeTimer);
+ if (result.isEmpty()) {
+ waitForEmpty(asyncWrapper);
+ result =
+ asyncWrapper.commitFinishedItemsDirect(
+ fakeTimer.getCurrentRelativeTime(), fakeBagState, fakeTimer);
+ }
+
+ checkOutput(result, Collections.singletonList("1"));
+ assertEquals(0, fakeBagState.items.size());
+ }
+
+ private static class FlakyDoFn extends DoFn<String, String> {
+ private int attempts = 0;
+ private final ReentrantLock lock = new ReentrantLock();
+
+ @ProcessElement
+ public void processElement(@Element String element, OutputReceiver<String>
receiver) {
+ lock.lock();
+ try {
+ attempts++;
+ if (attempts == 1) {
+ throw new RuntimeException("Transient RPC Error");
+ }
+ } finally {
+ lock.unlock();
+ }
+ receiver.output(element);
+ }
+ }
}
diff --git a/sdks/python/apache_beam/transforms/async_dofn.py
b/sdks/python/apache_beam/transforms/async_dofn.py
index ad3d5bc6646..f104770d8c0 100644
--- a/sdks/python/apache_beam/transforms/async_dofn.py
+++ b/sdks/python/apache_beam/transforms/async_dofn.py
@@ -489,9 +489,11 @@ class AsyncWrapper(beam.DoFn):
if x_id in processing_elements:
_, future = processing_elements[x_id]
if future.done():
+ # Pop from local active map before checking the result
+ processing_elements.pop(x_id)
to_return.append(future.result())
+ # Only mark as finished if result() succeeded without exception
finished_items.append(x)
- processing_elements.pop(x_id)
items_finished += 1
else:
items_not_yet_finished += 1
diff --git a/sdks/python/apache_beam/transforms/async_dofn_test.py
b/sdks/python/apache_beam/transforms/async_dofn_test.py
index 39901d791fb..f9e07d57be9 100644
--- a/sdks/python/apache_beam/transforms/async_dofn_test.py
+++ b/sdks/python/apache_beam/transforms/async_dofn_test.py
@@ -522,6 +522,53 @@ class AsyncTest(unittest.TestCase):
else:
self.assertEqual(p.exitcode, 0)
+ def test_transient_rpc_failure_retry(self):
+ # Verify DoFn exceptions wipe local active state so retries reschedule
work.
+ class FlakyDoFn(beam.DoFn):
+ def __init__(self):
+ self.attempts = 0
+ self.lock = Lock()
+
+ def process(self, element):
+ with self.lock:
+ self.attempts += 1
+ current_attempt = self.attempts
+ if current_attempt == 1:
+ raise RuntimeError("Transient RPC Error")
+ yield element
+
+ dofn = FlakyDoFn()
+ async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
+ async_dofn.setup()
+ fake_bag_state = FakeBagState([])
+ fake_timer = FakeTimer(0)
+ msg = ('key1', 1)
+
+ async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer)
+ self.wait_for_empty(async_dofn)
+
+ # Attempt 1 should raise the RuntimeError stored in the future
+ with self.assertRaises(RuntimeError):
+ async_dofn.commit_finished_items(fake_bag_state, fake_timer)
+
+ # Verify the failed future was popped from local processing_elements
+ with async_lib.AsyncWrapper._lock:
+ self.assertNotIn(
+ async_dofn._id_fn(msg[1]),
+ async_lib.AsyncWrapper._processing_elements[async_dofn._uuid],
+ )
+
+ # Simulate runner bundle retry: commit_finished_items runs again with msg
still in state.
+ # Because the dead future was popped, it will reschedule msg and succeed
on Attempt 2.
+ self.wait_for_empty(async_dofn)
+ result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
+ if not result:
+ self.wait_for_empty(async_dofn)
+ result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
+
+ self.check_output(result, [msg])
+ self.assertEqual(fake_bag_state.items, [])
+
if __name__ == '__main__':
unittest.main()
