Fix SDF bundle finalization timeout in streaming test (#38287)

* Fix SDF bundle finalization timeout in streaming test

* trigger workflow

* increased FINALIZATION_CALLBACK_TIMEOUT_SECS

* trigger workflow

* trigger workflow

* Increased timeout

* Fix SDF bundle-finalizer timeout constant in SplittableDoFnTest

* Stabilize bundle finalization SDF validates runner tests

Author: aIbrahiim

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run!",
-  "modification":  7,
+  "modification":  1,
 }

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

@@ -31,10 +31,10 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
@@ -958,16 +958,22 @@ public OffsetRange getInitialRestriction() {
   }
 
   /**
-   * While the finalization callback hasn't been invoked, this DoFn will keep requesting
-   * finalization, wait one second and then checkpoint upto MAX_ATTEMPTS amount of times. Once the
-   * callback has been invoked, the DoFn will output the element and stop.
+   * While the finalization callback hasn't been invoked, this DoFn repeatedly registers a
+   * finalization request, sleeps ~100ms, and checkpoints (via {@link ProcessContinuation#resume()})
+   * up to {@link #MAX_ATTEMPTS} times. Once the callback runs and flips {@code WAS_FINALIZED}, the
+   * DoFn outputs the element and stops.
+   *
+   * <p>Shared state must be thread-safe: {@link BundleFinalizer} callbacks can run on a different
+   * thread than {@code @ProcessElement}; use {@link ConcurrentHashMap} for {@code WAS_FINALIZED} so
+   * {@code computeIfAbsent} / updates are not racy (a plain {@link HashMap} can hang or corrupt
+   * under concurrent structural access).
    */
   public static class BundleFinalizingSplittableDoFn extends DoFn<String, String> {
+    /** Upper bound on {@link ProcessContinuation#resume()} iterations via restriction width. */
     private static final long MAX_ATTEMPTS = 3000;
-    // We use the UUID to uniquely identify this DoFn in case this test is run with
-    // other tests in the same JVM.
-    private static final Map<UUID, AtomicBoolean> WAS_FINALIZED = new HashMap();
-    private final UUID uuid = UUID.randomUUID();
+
+    private static final long FINALIZATION_CALLBACK_TIMEOUT_SECS = 300;
+    private static final Map<String, AtomicBoolean> WAS_FINALIZED = new ConcurrentHashMap<>();
 
     @NewTracker
     public RestrictionTracker<OffsetRange, Long> newTracker(@Restriction OffsetRange restriction) {
@@ -987,23 +993,29 @@ public ProcessContinuation process(
         RestrictionTracker<OffsetRange, Long> tracker,
         BundleFinalizer bundleFinalizer)
         throws InterruptedException {
-      if (WAS_FINALIZED.computeIfAbsent(uuid, (unused) -> new AtomicBoolean()).get()) {
+      AtomicBoolean wasFinalized =
+          WAS_FINALIZED.computeIfAbsent(element, (unused) -> new AtomicBoolean());
+      if (wasFinalized.get()) {
         tracker.tryClaim(tracker.currentRestriction().getFrom() + 1);
         receiver.output(element);
+        WAS_FINALIZED.remove(element);
         // Claim beyond the end now that we know we have been finalized.
         tracker.tryClaim(Long.MAX_VALUE);
         return stop();
       }
       if (tracker.tryClaim(tracker.currentRestriction().getFrom() + 1)) {
         bundleFinalizer.afterBundleCommit(
-            Instant.now().plus(Duration.standardSeconds(MAX_ATTEMPTS)),
-            () -> WAS_FINALIZED.computeIfAbsent(uuid, (unused) -> new AtomicBoolean()).set(true));
+            Instant.now().plus(Duration.standardSeconds(FINALIZATION_CALLBACK_TIMEOUT_SECS)),
+            () -> wasFinalized.set(true));
         // We sleep here instead of setting a resume time since the resume time doesn't need to
         // be honored.
         sleep(100L);
         return resume();
       }
-      return stop();
+      WAS_FINALIZED.remove(element);
+      throw new RuntimeException(
+          String.format(
+              "Bundle finalization callback was not observed after %d checkpoints.", MAX_ATTEMPTS));
     }
 
     @GetInitialRestriction
@@ -1017,9 +1029,10 @@ public OffsetRange getInitialRestriction() {
   public void testBundleFinalizationOccursOnBoundedSplittableDoFn() throws Exception {
     @BoundedPerElement
     class BoundedBundleFinalizingSplittableDoFn extends BundleFinalizingSplittableDoFn {}
-    PCollection<String> foo = p.apply(Create.of("foo"));
+    String element = "foo-" + UUID.randomUUID();
+    PCollection<String> foo = p.apply(Create.of(element));
     PCollection<String> res = foo.apply(ParDo.of(new BoundedBundleFinalizingSplittableDoFn()));
-    PAssert.that(res).containsInAnyOrder("foo");
+    PAssert.that(res).containsInAnyOrder(element);
     p.run();
   }
 
@@ -1028,9 +1041,10 @@ class BoundedBundleFinalizingSplittableDoFn extends BundleFinalizingSplittableDo
   public void testBundleFinalizationOccursOnUnboundedSplittableDoFn() throws Exception {
     @UnboundedPerElement
     class UnboundedBundleFinalizingSplittableDoFn extends BundleFinalizingSplittableDoFn {}
-    PCollection<String> foo = p.apply(Create.of("foo"));
+    String element = "foo-" + UUID.randomUUID();
+    PCollection<String> foo = p.apply(Create.of(element));
     PCollection<String> res = foo.apply(ParDo.of(new UnboundedBundleFinalizingSplittableDoFn()));
-    PAssert.that(res).containsInAnyOrder("foo");
+    PAssert.that(res).containsInAnyOrder(element);
     p.run();
   }