Adds bundleFinalizer support to Dataflow non-portable worker. (#37723)

* Adds bundleFinalizer support to non-portable worker.
* Removes check preventing stateful DoFn's with bundle finalizers from running on Dataflow streaming non-portable worker when using Streaming Engine

Author: acrites

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json

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

.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":  1,
+  "modification":  6,
 }

runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java

@@ -311,9 +311,14 @@ public DoFn<InputT, OutputT>.StartBundleContext startBundleContext(DoFn<InputT,
     public String getErrorContext() {
       return "SimpleDoFnRunner/StartBundle";
     }
+
+    @Override
+    public BundleFinalizer bundleFinalizer() {
+      return stepContext.bundleFinalizer();
+    }
   }
 
-  /** An {@link DoFnInvoker.ArgumentProvider} for {@link DoFn.StartBundle @StartBundle}. */
+  /** An {@link DoFnInvoker.ArgumentProvider} for {@link DoFn.FinishBundle @FinishBundle}. */
   private class DoFnFinishBundleArgumentProvider
       extends DoFnInvoker.BaseArgumentProvider<InputT, OutputT> {
     /** A concrete implementation of {@link DoFn.FinishBundleContext}. */
@@ -356,6 +361,11 @@ public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
     public String getErrorContext() {
       return "SimpleDoFnRunner/FinishBundle";
     }
+
+    @Override
+    public BundleFinalizer bundleFinalizer() {
+      return stepContext.bundleFinalizer();
+    }
   }
 
   /**
@@ -1030,7 +1040,7 @@ public <T> void outputWindowedValue(
     @Override
     public BundleFinalizer bundleFinalizer() {
       throw new UnsupportedOperationException(
-          "Bundle finalization is not supported in non-portable pipelines.");
+          "Bundle finalization is not supported in OnTimer calls.");
     }
   }
 
@@ -1289,7 +1299,7 @@ public <T> void outputWindowedValue(
     @Override
     public BundleFinalizer bundleFinalizer() {
       throw new UnsupportedOperationException(
-          "Bundle finalization is not supported in non-portable pipelines.");
+          "Bundle finalization is not supported in OnWindowExpiration calls.");
     }
   }
 

runners/google-cloud-dataflow-java/build.gradle

@@ -205,7 +205,6 @@ def commonLegacyExcludeCategories = [
   'org.apache.beam.sdk.testing.UsesGaugeMetrics',
   'org.apache.beam.sdk.testing.UsesTestStream',
   'org.apache.beam.sdk.testing.UsesMetricsPusher',
-  'org.apache.beam.sdk.testing.UsesBundleFinalizer',
   'org.apache.beam.sdk.testing.UsesBoundedTrieMetrics', // Dataflow QM as of now does not support returning back BoundedTrie in metric result.
 ]
 
@@ -456,7 +455,9 @@ task validatesRunner {
       'org.apache.beam.sdk.transforms.ParDoLifecycleTest.testTeardownCalledAfterExceptionInSetupStateful',
       'org.apache.beam.sdk.transforms.ParDoLifecycleTest.testTeardownCalledAfterExceptionInStartBundle',
       'org.apache.beam.sdk.transforms.ParDoLifecycleTest.testTeardownCalledAfterExceptionInStartBundleStateful',
-      ]
+    ],
+    // Batch legacy worker does not support bundle finalization.
+    excludedCategories: [ 'org.apache.beam.sdk.testing.UsesBundleFinalizer', ],
   ))
 }
 
@@ -490,6 +491,8 @@ task validatesRunnerStreaming {
   description "Validates Dataflow runner forcing streaming mode"
   dependsOn(createLegacyWorkerValidatesRunnerTest(validatesRunnerStreamingConfig + [
     name: 'validatesRunnerLegacyWorkerTestStreaming',
+    // Streaming appliance currently fails bundle finalizer tests.
+    excludedCategories: validatesRunnerStreamingConfig.excludedCategories + [ 'org.apache.beam.sdk.testing.UsesBundleFinalizer', ],
   ]))
 }
 

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java

@@ -2737,11 +2737,11 @@ static void verifyDoFnSupported(
               DataflowRunner.class.getSimpleName()));
     }
     boolean isUnifiedWorker = useUnifiedWorker(options);
-    if (DoFnSignatures.usesBundleFinalizer(fn) && !isUnifiedWorker) {
+    if (DoFnSignatures.usesBundleFinalizer(fn) && !isUnifiedWorker && !streaming) {
       throw new UnsupportedOperationException(
           String.format(
-              "%s does not currently support %s when not using unified worker because it uses "
-                  + "BundleFinalizers in its implementation. Set the `--experiments=use_runner_v2` "
+              "%s does not currently support %s in batch mode when not using unified worker because it "
+                  + "uses BundleFinalizers in its implementation. Set the `--experiments=use_runner_v2` "
                   + "option to use this DoFn.",
               DataflowRunner.class.getSimpleName(), fn.getClass().getSimpleName()));
     }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SplittableProcessFnFactory.java

@@ -156,10 +156,7 @@ public DoFnRunner<KeyedWorkItem<byte[], KV<InputT, RestrictionT>>, OutputT> crea
               // in the event of a crash.
               10000,
               Duration.standardSeconds(10),
-              () -> {
-                throw new UnsupportedOperationException(
-                    "BundleFinalizer unsupported by non-portable Dataflow.");
-              }));
+              stepContext::bundleFinalizer));
       DoFnRunner<KeyedWorkItem<byte[], KV<InputT, RestrictionT>>, OutputT> simpleRunner =
           new SimpleDoFnRunner<>(
               options,

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java

@@ -35,6 +35,7 @@
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicLong;
 import javax.annotation.concurrent.NotThreadSafe;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair;
 import org.apache.beam.runners.core.SideInputReader;
 import org.apache.beam.runners.core.StateInternals;
 import org.apache.beam.runners.core.StateNamespace;
@@ -73,6 +74,7 @@
 import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader;
 import org.apache.beam.sdk.metrics.MetricsContainer;
 import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.transforms.DoFn.BundleFinalizer;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.util.ByteStringOutputStream;
 import org.apache.beam.sdk.values.PCollectionView;
@@ -82,8 +84,10 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.FluentIterable;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.HashBasedTable;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.PeekingIterator;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
@@ -446,11 +450,27 @@ public void invalidateCache() {
     }
   }
 
-  public Map<Long, Runnable> flushState() {
-    Map<Long, Runnable> callbacks = new HashMap<>();
+  public Map<Long, Pair<Instant, Runnable>> flushState() {
+    Map<Long, Pair<Instant, Runnable>> callbacks = new HashMap<>();
 
     for (StepContext stepContext : getAllStepContexts()) {
       stepContext.flushState();
+      for (Pair<Instant, BundleFinalizer.Callback> bundleFinalizer :
+          stepContext.flushBundleFinalizerCallbacks()) {
+        long id = ThreadLocalRandom.current().nextLong();
+        callbacks.put(
+            id,
+            Pair.of(
+                bundleFinalizer.getLeft(),
+                () -> {
+                  try {
+                    bundleFinalizer.getRight().onBundleSuccess();
+                  } catch (Exception e) {
+                    throw new RuntimeException("Exception while running bundle finalizer", e);
+                  }
+                }));
+        outputBuilder.addFinalizeIds(id);
+      }
     }
 
     if (activeReader != null) {
@@ -462,13 +482,15 @@ public Map<Long, Runnable> flushState() {
       sourceStateBuilder.addFinalizeIds(id);
       callbacks.put(
           id,
-          () -> {
-            try {
-              checkpointMark.finalizeCheckpoint();
-            } catch (IOException e) {
-              throw new RuntimeException("Exception while finalizing checkpoint", e);
-            }
-          });
+          Pair.of(
+              Instant.now().plus(Duration.standardMinutes(5)),
+              () -> {
+                try {
+                  checkpointMark.finalizeCheckpoint();
+                } catch (IOException e) {
+                  throw new RuntimeException("Exception while finalizing checkpoint", e);
+                }
+              }));
 
       @SuppressWarnings("unchecked")
       Coder<UnboundedSource.CheckpointMark> checkpointCoder =
@@ -699,6 +721,11 @@ public <W extends BoundedWindow> void setStateCleanupTimer(
     public DataflowStepContext namespacedToUser() {
       return this;
     }
+
+    @Override
+    public BundleFinalizer bundleFinalizer() {
+      return wrapped.bundleFinalizer();
+    }
   }
 
   /** A {@link SideInputReader} that fetches side inputs from the streaming worker's cache. */
@@ -771,6 +798,7 @@ class StepContext extends DataflowExecutionContext.DataflowStepContext
     // A list of timer keys that were modified by user processing earlier in this bundle. This
     // serves a tombstone, so that we know not to fire any bundle timers that were modified.
     private Table<String, StateNamespace, TimerData> modifiedUserTimerKeys = null;
+    private final WindmillBundleFinalizer bundleFinalizer = new WindmillBundleFinalizer();
 
     public StepContext(DataflowOperationContext operationContext) {
       super(operationContext.nameContext());
@@ -1043,9 +1071,37 @@ public TimerInternals timerInternals() {
       return checkNotNull(systemTimerInternals);
     }
 
+    @Override
+    public BundleFinalizer bundleFinalizer() {
+      return bundleFinalizer;
+    }
+
     public TimerInternals userTimerInternals() {
       ensureStateful("Tried to access user timers");
       return checkNotNull(userTimerInternals);
     }
+
+    public ImmutableList<Pair<Instant, BundleFinalizer.Callback>> flushBundleFinalizerCallbacks() {
+      return bundleFinalizer.flushCallbacks();
+    }
+  }
+
+  private static class WindmillBundleFinalizer implements BundleFinalizer {
+    private ImmutableList.Builder<Pair<Instant, Callback>> callbacks = ImmutableList.builder();
+
+    private WindmillBundleFinalizer() {}
+
+    private ImmutableList<Pair<Instant, Callback>> flushCallbacks() {
+      ImmutableList<Pair<Instant, Callback>> flushedCallbacks = callbacks.build();
+      if (!Iterables.isEmpty(flushedCallbacks)) {
+        callbacks = ImmutableList.builder();
+      }
+      return flushedCallbacks;
+    }
+
+    @Override
+    public void afterBundleCommit(Instant callbackExpiry, Callback callback) {
+      callbacks.add(Pair.of(callbackExpiry, callback));
+    }
   }
 }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkResponseChunkAssembler.java

@@ -54,24 +54,27 @@ final class GetWorkResponseChunkAssembler {
   private final WorkItem.Builder workItemBuilder; // Reused to reduce GC overhead.
   private ByteString data;
   private long bufferedSize;
+  private final List<Long> appliedFinalizeIds;
 
   GetWorkResponseChunkAssembler() {
     workTimingInfosTracker = new GetWorkTimingInfosTracker(System::currentTimeMillis);
     data = ByteString.EMPTY;
     bufferedSize = 0;
     metadata = null;
     workItemBuilder = WorkItem.newBuilder();
+    appliedFinalizeIds = new ArrayList<>();
   }
 
   /**
-   * Appends the response chunk bytes to the {@link #data }byte buffer. Return the assembled
+   * Appends the response chunk bytes to the {@link #data} byte buffer. Return the assembled
    * WorkItem if all response chunks for a WorkItem have been received.
    */
   List<AssembledWorkItem> append(Windmill.StreamingGetWorkResponseChunk chunk) {
     if (chunk.hasComputationMetadata()) {
       metadata = ComputationMetadata.fromProto(chunk.getComputationMetadata());
     }
     workTimingInfosTracker.addTimingInfo(chunk.getPerWorkItemTimingInfosList());
+    appliedFinalizeIds.addAll(chunk.getAppliedFinalizeIdsList());
 
     List<AssembledWorkItem> response = new ArrayList<>();
     for (int i = 0; i < chunk.getSerializedWorkItemList().size(); i++) {
@@ -90,13 +93,14 @@ List<AssembledWorkItem> append(Windmill.StreamingGetWorkResponseChunk chunk) {
   }
 
   /**
-   * Attempt to flush the {@link #data} bytes into a {@link WorkItem} w/ it's metadata. Resets the
+   * Attempt to flush the {@link #data} bytes into a {@link WorkItem} w/ its metadata. Resets the
    * data byte string and tracking metadata afterwards, whether the {@link WorkItem} deserialization
    * was successful or not.
    */
   private Optional<AssembledWorkItem> flushToWorkItem() {
     try {
       workItemBuilder.mergeFrom(data);
+      workItemBuilder.addAllAppliedFinalizeIds(appliedFinalizeIds);
       return Optional.of(
           AssembledWorkItem.create(
               workItemBuilder.build(),
@@ -110,6 +114,7 @@ private Optional<AssembledWorkItem> flushToWorkItem() {
       workTimingInfosTracker.reset();
       data = ByteString.EMPTY;
       bufferedSize = 0;
+      appliedFinalizeIds.clear();
     }
 
     return Optional.empty();