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

Author: tejasiyer-dev

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncWrapper.java

@@ -644,15 +644,17 @@ void commitFinishedItems(
 
         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);

sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java

@@ -19,6 +19,7 @@
 
 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 BagState<T> readLater() {
 
   // 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 void testResetStateConcurrentTeardown() {
     // 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);
+    }
+  }
 }

sdks/python/apache_beam/transforms/async_dofn.py

@@ -489,9 +489,11 @@ def commit_finished_items(
         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

sdks/python/apache_beam/transforms/async_dofn_test.py

@@ -522,6 +522,53 @@ def test_reset_state_concurrent_teardown(self):
     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()