[Java Harness] Handle swallowed exceptions across fused DoFns

sdks/java/core/src/main/java/org/apache/beam/sdk/options/ExperimentalOptions.java

@@ -35,6 +35,9 @@ public interface ExperimentalOptions extends PipelineOptions {
 
   String STATE_SAMPLING_PERIOD_MILLIS = "state_sampling_period_millis";
 
+  /** Experiment controlling exception handling at fused Java DoFn boundaries. */
+  String FUSED_DOFN_EXCEPTION_HANDLING = "fused_dofn_exception_handling";
+
   @Description(
       "Apache Beam provides a number of experimental features that can "
           + "be enabled with this flag. If executing against a managed service, please contact the "

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java

@@ -79,6 +79,7 @@
 import org.apache.beam.sdk.metrics.MetricsContainer;
 import org.apache.beam.sdk.metrics.MetricsEnvironment;
 import org.apache.beam.sdk.metrics.MetricsEnvironment.MetricsEnvironmentState;
+import org.apache.beam.sdk.options.ExperimentalOptions;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.SdkHarnessOptions;
 import org.apache.beam.sdk.transforms.DoFn.BundleFinalizer;
@@ -560,13 +561,24 @@ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest re
                 bundleProcessor.getInboundEndpointApiServiceDescriptors());
           }
 
-          // Need to reverse this since we want to call finish in topological order.
+          // Do not run @FinishBundle after a downstream failure was swallowed during element
+          // processing. Finish functions can flush state or write to external sinks.
+          bundleProcessor.getpCollectionConsumerRegistry().failBundleIfDownstreamFailed();
+
+          // Need to reverse this since we want to call finish in topological order. Check after
+          // each function as well, so a swallowed failure during @FinishBundle prevents later
+          // finish functions from running.
           for (ThrowingRunnable finishFunction :
               Lists.reverse(finishFunctionRegistry.getFunctions())) {
             LOG.debug("Finishing function {}", finishFunction);
             finishFunction.run();
+            bundleProcessor.getpCollectionConsumerRegistry().failBundleIfDownstreamFailed();
           }
 
+          // Surface a failure raised by the final @FinishBundle function before outbound data is
+          // finalized.
+          bundleProcessor.getpCollectionConsumerRegistry().failBundleIfDownstreamFailed();
+
           // If bundleProcessor has not flushed any elements, embed them in response.
           embedOutboundElementsIfApplicable(response, bundleProcessor);
 
@@ -598,6 +610,9 @@ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest re
         }
       }
 
+      // A warning is appropriate only once all bundle work has succeeded.
+      bundleProcessor.getpCollectionConsumerRegistry().logWarningIfDownstreamFailed();
+
       // Mark the bundle processor as re-usable.
       bundleProcessorCache.release(
           request.getProcessBundle().getProcessBundleDescriptorId(), bundleProcessor);
@@ -615,7 +630,40 @@ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest re
       // Ensure that if more data arrives for the instruction it is discarded.
       beamFnDataClient.poisonInstructionId(instructionId);
       throw e;
+    } catch (Error e) {
+      LOG.debug(
+          "Error processing bundle {} with bundleProcessor for {} after error",
+          instructionId,
+          request.getProcessBundle().getProcessBundleDescriptorId(),
+          e);
+      if (bundleProcessor != null) {
+        // Make sure we clean up from the active set of bundle processors.
+        bundleProcessorCache.discard(bundleProcessor);
+      }
+      // Ensure that if more data arrives for the instruction it is discarded.
+      beamFnDataClient.poisonInstructionId(instructionId);
+      throw e;
+    }
+  }
+
+  private static PCollectionConsumerRegistry.FusedDoFnExceptionHandling
+      getFusedDoFnExceptionHandling(PipelineOptions options) {
+    String value =
+        ExperimentalOptions.getExperimentValue(
+            options, ExperimentalOptions.FUSED_DOFN_EXCEPTION_HANDLING);
+    if (value == null || value.equalsIgnoreCase("legacy")) {
+      return PCollectionConsumerRegistry.FusedDoFnExceptionHandling.LEGACY;
+    }
+    if (value.equalsIgnoreCase("warn")) {
+      return PCollectionConsumerRegistry.FusedDoFnExceptionHandling.WARN;
+    }
+    if (value.equalsIgnoreCase("fail")) {
+      return PCollectionConsumerRegistry.FusedDoFnExceptionHandling.FAIL;
     }
+    throw new IllegalArgumentException(
+        String.format(
+            "Unsupported value '%s' for experiment %s. Expected legacy, warn, or fail.",
+            value, ExperimentalOptions.FUSED_DOFN_EXCEPTION_HANDLING));
   }
 
   private void embedOutboundElementsIfApplicable(
@@ -803,20 +851,24 @@ private BundleProcessor createBundleProcessor(
         new MetricsEnvironmentStateForBundle();
     ExecutionStateTracker stateTracker = executionStateSampler.create();
     bundleProgressReporterAndRegistrar.register(stateTracker);
+    PCollectionConsumerRegistry.FusedDoFnExceptionHandling fusedDoFnExceptionHandling =
+        getFusedDoFnExceptionHandling(options);
     PCollectionConsumerRegistry pCollectionConsumerRegistry =
         new PCollectionConsumerRegistry(
             stateTracker,
             shortIds,
             bundleProgressReporterAndRegistrar,
             bundleDescriptor,
-            dataSampler);
+            dataSampler,
+            fusedDoFnExceptionHandling);
     HashSet<String> processedPTransformIds = new HashSet<>();
 
     PTransformFunctionRegistry startFunctionRegistry =
         new PTransformFunctionRegistry(shortIds, stateTracker, Urns.START_BUNDLE_MSECS);
     PTransformFunctionRegistry finishFunctionRegistry =
         new PTransformFunctionRegistry(shortIds, stateTracker, Urns.FINISH_BUNDLE_MSECS);
     List<ThrowingRunnable> resetFunctions = new ArrayList<>();
+    resetFunctions.add(pCollectionConsumerRegistry::reset);
     List<ThrowingRunnable> tearDownFunctions = new ArrayList<>();
 
     // Build a multimap of PCollection ids to PTransform ids which consume said PCollections