From c26947ddc92e62bd2b452151dd6aea916d5e6e7e Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:02:25 +0530 Subject: [PATCH 01/37] Track exceptions thrown by fused downstream consumers --- .../data/PCollectionConsumerRegistry.java | 84 ++++++++++++++++++- 1 file changed, 83 insertions(+), 1 deletion(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java index 58c06bac310f..a947754968bd 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java @@ -24,6 +24,7 @@ import java.util.List; import java.util.Map; import java.util.Random; +import java.util.concurrent.atomic.AtomicReference; import javax.annotation.Nullable; import org.apache.beam.fn.harness.HandlesSplits; import org.apache.beam.fn.harness.control.BundleProgressReporter; @@ -70,6 +71,16 @@ }) public class PCollectionConsumerRegistry { + /** Controls how exceptions from fused downstream consumers are handled at bundle completion. */ + public enum FusedDoFnExceptionHandling { + LEGACY, + WARN, + FAIL + } + + private final FusedDoFnExceptionHandling fusedDoFnExceptionHandling; + private final AtomicReference firstDownstreamException = new AtomicReference<>(); + /** Stores metadata about each consumer so that the appropriate metrics tracking can occur. */ @AutoValue @AutoValue.CopyAnnotations @@ -121,7 +132,13 @@ public PCollectionConsumerRegistry( ShortIdMap shortIdMap, BundleProgressReporter.Registrar bundleProgressReporterRegistrar, ProcessBundleDescriptor processBundleDescriptor) { - this(stateTracker, shortIdMap, bundleProgressReporterRegistrar, processBundleDescriptor, null); + this( + stateTracker, + shortIdMap, + bundleProgressReporterRegistrar, + processBundleDescriptor, + null, + FusedDoFnExceptionHandling.LEGACY); } public PCollectionConsumerRegistry( @@ -130,6 +147,22 @@ public PCollectionConsumerRegistry( BundleProgressReporter.Registrar bundleProgressReporterRegistrar, ProcessBundleDescriptor processBundleDescriptor, @Nullable DataSampler dataSampler) { + this( + stateTracker, + shortIdMap, + bundleProgressReporterRegistrar, + processBundleDescriptor, + dataSampler, + FusedDoFnExceptionHandling.LEGACY); + } + + public PCollectionConsumerRegistry( + ExecutionStateTracker stateTracker, + ShortIdMap shortIdMap, + BundleProgressReporter.Registrar bundleProgressReporterRegistrar, + ProcessBundleDescriptor processBundleDescriptor, + @Nullable DataSampler dataSampler, + FusedDoFnExceptionHandling fusedDoFnExceptionHandling) { this.stateTracker = stateTracker; this.shortIdMap = shortIdMap; this.bundleProgressReporterRegistrar = bundleProgressReporterRegistrar; @@ -142,6 +175,7 @@ public PCollectionConsumerRegistry( .putAllWindowingStrategies(processBundleDescriptor.getWindowingStrategiesMap()) .build()); this.dataSampler = dataSampler; + this.fusedDoFnExceptionHandling = fusedDoFnExceptionHandling; } /** @@ -274,6 +308,51 @@ public ExecutionState getProcessingExecutionState( }); } + /** + * Records a downstream exception when configured to detect exceptions swallowed by a fused + * upstream DoFn. Work cancellation is not a downstream processing failure. + */ + private void recordDownstreamException(Exception exception) { + if (fusedDoFnExceptionHandling != FusedDoFnExceptionHandling.LEGACY + && !WorkCancelledException.isWorkCancelledException(exception)) { + firstDownstreamException.compareAndSet(null, exception); + } + } + + /** Prevents additional downstream work after the first failure in fail mode. */ + private void failFastIfDownstreamFailed() throws Exception { + if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.FAIL) { + Exception exception = firstDownstreamException.get(); + if (exception != null) { + throw exception; + } + } + } + + /** + * Applies the configured policy after all bundle finish functions have run. This must be called + * before outbound data is finalized and before the bundle processor is returned to its cache. + */ + public void finishBundle() throws Exception { + Exception exception = firstDownstreamException.get(); + if (exception == null) { + return; + } + if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.WARN) { + LOG.warn( + "A downstream exception was caught by fused user code. The bundle completed, but " + + "downstream processing may have been skipped.", + exception); + } else if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.FAIL) { + throw exception; + } + } + + /** Clears per-bundle failure state when a processor is successfully returned to the cache. */ + public void reset() { + firstDownstreamException.set(null); + } + private static void logAndRethrow( Exception e, ExecutionState executionState, @@ -360,6 +439,7 @@ public MetricTrackingFnDataReceiver( @Override public void accept(WindowedValue input) throws Exception { + failFastIfDownstreamFailed(); // Increment the counter for each window the element occurs in. this.elementCountCounter.inc(input.getWindows().size()); // TODO(https://github.com/apache/beam/issues/20730): Consider updating size per window when @@ -377,6 +457,7 @@ public void accept(WindowedValue input) throws Exception { try (ExecutionState.ActiveState a = executionState.scopedActivate()) { this.delegate.accept(input); } catch (Exception e) { + recordDownstreamException(e); logAndRethrow( e, executionState, executionStateTracker, ptransformId, outputSampler, elementSample); } @@ -464,6 +545,7 @@ public void accept(WindowedValue input) throws Exception { try (ExecutionState.ActiveState a = state.scopedActivate()) { consumerAndMetadata.getConsumer().accept(input); } catch (Exception e) { + recordDownstreamException(e); logAndRethrow( e, state, From a51f5ae4a8314460e21e13b9ff8976984b28da63 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:02:27 +0530 Subject: [PATCH 02/37] Apply fused DoFn exception policy at bundle completion --- .../harness/control/ProcessBundleHandler.java | 31 ++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index 9fd1f0f0d68b..815c1290c1d5 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/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; @@ -567,6 +568,10 @@ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest re finishFunction.run(); } + // Surface downstream exceptions that user code may have caught while processing + // fused outputs. This must happen before outbound data is finalized. + bundleProcessor.getpCollectionConsumerRegistry().finishBundle(); + // If bundleProcessor has not flushed any elements, embed them in response. embedOutboundElementsIfApplicable(response, bundleProcessor); @@ -618,6 +623,26 @@ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest re } } + 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( ProcessBundleResponse.Builder response, BundleProcessor bundleProcessor) { if (bundleProcessor.getOutboundAggregators().isEmpty()) { @@ -803,13 +828,16 @@ 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 processedPTransformIds = new HashSet<>(); PTransformFunctionRegistry startFunctionRegistry = @@ -817,6 +845,7 @@ private BundleProcessor createBundleProcessor( PTransformFunctionRegistry finishFunctionRegistry = new PTransformFunctionRegistry(shortIds, stateTracker, Urns.FINISH_BUNDLE_MSECS); List resetFunctions = new ArrayList<>(); + resetFunctions.add(pCollectionConsumerRegistry::reset); List tearDownFunctions = new ArrayList<>(); // Build a multimap of PCollection ids to PTransform ids which consume said PCollections From 680b40453b9f676e7c8bafb11d2e047b9fd9b31f Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:02:29 +0530 Subject: [PATCH 03/37] Define fused DoFn exception handling experiment --- .../java/org/apache/beam/sdk/options/ExperimentalOptions.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ExperimentalOptions.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ExperimentalOptions.java index 366768a96368..0a71ddb50183 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ExperimentalOptions.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ExperimentalOptions.java @@ -35,6 +35,8 @@ public interface ExperimentalOptions extends PipelineOptions { String STATE_SAMPLING_PERIOD_MILLIS = "state_sampling_period_millis"; + 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 " From 4e9046f18838440eef5656dfb20bc6c27536a712 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:02:46 +0530 Subject: [PATCH 04/37] Test fused DoFn exception handling policies --- .../data/PCollectionConsumerRegistryTest.java | 108 ++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java index 804fd695e342..a1d929ea6f02 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java @@ -221,6 +221,114 @@ public void singleConsumerException() throws Exception { wrapperConsumer.accept(valueInGlobalWindow("elem")); } + @Test + public void failModeRemembersSwallowedExceptionAndStopsLaterOutputs() throws Exception { + final String pTransformId = "pTransformId"; + final String message = "downstream failure"; + ShortIdMap shortIds = new ShortIdMap(); + BundleProgressReporter.InMemory reporterAndRegistrar = new BundleProgressReporter.InMemory(); + PCollectionConsumerRegistry consumers = + new PCollectionConsumerRegistry( + sampler.create(), + shortIds, + reporterAndRegistrar, + TEST_DESCRIPTOR, + null, + PCollectionConsumerRegistry.FusedDoFnExceptionHandling.FAIL); + FnDataReceiver> consumer = mock(FnDataReceiver.class); + consumers.register(P_COLLECTION_A, pTransformId, pTransformId + "Name", consumer); + FnDataReceiver> wrapperConsumer = + (FnDataReceiver>) + (FnDataReceiver) consumers.getMultiplexingConsumer(P_COLLECTION_A); + doThrow(new Exception(message)).when(consumer).accept(any()); + + try { + wrapperConsumer.accept(valueInGlobalWindow("first")); + } catch (Exception expected) { + assertEquals(message, expected.getMessage()); + } + + // The upstream DoFn may catch the first exception and try to emit again. Fail mode must not + // invoke the downstream consumer again after its first failure. + try { + wrapperConsumer.accept(valueInGlobalWindow("second")); + } catch (Exception expected) { + assertEquals(message, expected.getMessage()); + } + verify(consumer, times(1)).accept(any()); + + try { + consumers.finishBundle(); + throw new AssertionError("Expected the recorded downstream exception to fail the bundle"); + } catch (Exception expected) { + assertEquals(message, expected.getMessage()); + } + } + + @Test + public void warnModeRemembersExceptionButAllowsLaterOutputs() throws Exception { + final String pTransformId = "pTransformId"; + ShortIdMap shortIds = new ShortIdMap(); + BundleProgressReporter.InMemory reporterAndRegistrar = new BundleProgressReporter.InMemory(); + PCollectionConsumerRegistry consumers = + new PCollectionConsumerRegistry( + sampler.create(), + shortIds, + reporterAndRegistrar, + TEST_DESCRIPTOR, + null, + PCollectionConsumerRegistry.FusedDoFnExceptionHandling.WARN); + FnDataReceiver> consumer = mock(FnDataReceiver.class); + consumers.register(P_COLLECTION_A, pTransformId, pTransformId + "Name", consumer); + FnDataReceiver> wrapperConsumer = + (FnDataReceiver>) + (FnDataReceiver) consumers.getMultiplexingConsumer(P_COLLECTION_A); + doThrow(new Exception("first output failed")).doNothing().when(consumer).accept(any()); + + try { + wrapperConsumer.accept(valueInGlobalWindow("first")); + } catch (Exception expected) { + assertEquals("first output failed", expected.getMessage()); + } + wrapperConsumer.accept(valueInGlobalWindow("second")); + verify(consumer, times(2)).accept(any()); + + // Warning mode reports at successful bundle completion but does not fail the bundle. + consumers.finishBundle(); + } + + @Test + public void resetClearsRecordedDownstreamException() throws Exception { + final String pTransformId = "pTransformId"; + ShortIdMap shortIds = new ShortIdMap(); + BundleProgressReporter.InMemory reporterAndRegistrar = new BundleProgressReporter.InMemory(); + PCollectionConsumerRegistry consumers = + new PCollectionConsumerRegistry( + sampler.create(), + shortIds, + reporterAndRegistrar, + TEST_DESCRIPTOR, + null, + PCollectionConsumerRegistry.FusedDoFnExceptionHandling.FAIL); + FnDataReceiver> consumer = mock(FnDataReceiver.class); + consumers.register(P_COLLECTION_A, pTransformId, pTransformId + "Name", consumer); + FnDataReceiver> wrapperConsumer = + (FnDataReceiver>) + (FnDataReceiver) consumers.getMultiplexingConsumer(P_COLLECTION_A); + doThrow(new Exception("first bundle failure")).when(consumer).accept(any()); + + try { + wrapperConsumer.accept(valueInGlobalWindow("first")); + } catch (Exception expected) { + assertEquals("first bundle failure", expected.getMessage()); + } + consumers.reset(); + doNothing().when(consumer).accept(any()); + wrapperConsumer.accept(valueInGlobalWindow("next bundle")); + consumers.finishBundle(); + verify(consumer, times(2)).accept(any()); + } + /** Test that the counter increments even when there are no consumers of the PCollection. */ @Test public void noConsumers() throws Exception { From 0b9b7bd7b1c3f306f1d8f8c2f410c9c24673e77c Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:03:16 +0530 Subject: [PATCH 05/37] Add fused two-DoFn bundle regression tests --- .../control/ProcessBundleHandlerTest.java | 223 ++++++++++++++++++ 1 file changed, 223 insertions(+) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index c03f82726740..52053e97c2c7 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -114,6 +114,7 @@ import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.fn.data.BeamFnDataOutboundAggregator; +import org.apache.beam.sdk.fn.data.FnDataReceiver; import org.apache.beam.sdk.fn.data.CloseableFnDataReceiver; import org.apache.beam.sdk.fn.data.DataEndpoint; import org.apache.beam.sdk.fn.data.TimerEndpoint; @@ -123,6 +124,7 @@ import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.metrics.MetricsEnvironment; +import org.apache.beam.sdk.options.ExperimentalOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.state.TimeDomain; import org.apache.beam.sdk.state.TimerMap; @@ -147,6 +149,7 @@ import org.apache.beam.sdk.values.CausedByDrain; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; @@ -935,6 +938,226 @@ public void testPTransformStartExceptionsArePropagated() { assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L"), empty()); } + private static final AtomicInteger FUSED_DOWNSTREAM_INVOCATIONS = new AtomicInteger(); + private static final AtomicInteger FUSED_UPSTREAM_CAUGHT_EXCEPTIONS = new AtomicInteger(); + + private static final class CatchingFusedDoFn extends DoFn { + private static final TupleTag MAIN_OUTPUT_TAG = new TupleTag<>("mainOutput"); + + @ProcessElement + public void processElement(ProcessContext context) { + try { + context.output(context.element()); + } catch (RuntimeException exception) { + FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.incrementAndGet(); + } + } + } + + private static final class ThrowingFusedDoFn extends DoFn { + private static final TupleTag MAIN_OUTPUT_TAG = new TupleTag<>("mainOutput"); + + @ProcessElement + public void processElement(ProcessContext context) { + FUSED_DOWNSTREAM_INVOCATIONS.incrementAndGet(); + throw new IllegalStateException("fused downstream failure"); + } + } + + private static PTransform fusedDoFnTransform( + String inputPCollection, + String outputPCollection, + DoFnWithExecutionInformation doFnWithExecutionInformation) { + RunnerApi.FunctionSpec doFnSpec = + RunnerApi.FunctionSpec.newBuilder() + .setUrn(ParDoTranslation.CUSTOM_JAVA_DO_FN_URN) + .setPayload( + ByteString.copyFrom( + SerializableUtils.serializeToByteArray(doFnWithExecutionInformation))) + .build(); + ParDoPayload payload = ParDoPayload.newBuilder().setDoFn(doFnSpec).build(); + return PTransform.newBuilder() + .setSpec( + FunctionSpec.newBuilder() + .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) + .setPayload(payload.toByteString())) + .putInputs("input", inputPCollection) + .putOutputs("output", outputPCollection) + .build(); + } + + private ProcessBundleHandler setupHandlerForFusedExceptionTest(String exceptionHandling) + throws Exception { + TupleTag upstreamTag = CatchingFusedDoFn.MAIN_OUTPUT_TAG; + TupleTag downstreamTag = ThrowingFusedDoFn.MAIN_OUTPUT_TAG; + DoFnWithExecutionInformation upstream = + DoFnWithExecutionInformation.of( + new CatchingFusedDoFn(), + upstreamTag, + Collections.emptyMap(), + DoFnSchemaInformation.create()); + DoFnWithExecutionInformation downstream = + DoFnWithExecutionInformation.of( + new ThrowingFusedDoFn(), + downstreamTag, + Collections.emptyMap(), + DoFnSchemaInformation.create()); + + ProcessBundleDescriptor descriptor = + ProcessBundleDescriptor.newBuilder() + .putTransforms( + "source", + PTransform.newBuilder() + .setSpec(FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()) + .putOutputs("output", "pc-source") + .build()) + .putTransforms("upstream", fusedDoFnTransform("pc-source", "pc-upstream", upstream)) + .putTransforms( + "downstream", fusedDoFnTransform("pc-upstream", "pc-downstream", downstream)) + .putPcollections( + "pc-source", + PCollection.newBuilder() + .setWindowingStrategyId("window-strategy") + .setCoderId("string-coder") + .setIsBounded(IsBounded.Enum.BOUNDED) + .build()) + .putPcollections( + "pc-upstream", + PCollection.newBuilder() + .setWindowingStrategyId("window-strategy") + .setCoderId("string-coder") + .setIsBounded(IsBounded.Enum.BOUNDED) + .build()) + .putPcollections( + "pc-downstream", + PCollection.newBuilder() + .setWindowingStrategyId("window-strategy") + .setCoderId("string-coder") + .setIsBounded(IsBounded.Enum.BOUNDED) + .build()) + .putWindowingStrategies( + "window-strategy", + WindowingStrategy.newBuilder() + .setWindowCoderId("window-coder") + .setWindowFn( + FunctionSpec.newBuilder().setUrn("beam:window_fn:global_windows:v1")) + .setOutputTime(OutputTime.Enum.END_OF_WINDOW) + .setAccumulationMode(AccumulationMode.Enum.ACCUMULATING) + .setTrigger(Trigger.newBuilder().setAlways(Always.getDefaultInstance())) + .setClosingBehavior(ClosingBehavior.Enum.EMIT_ALWAYS) + .setOnTimeBehavior(OnTimeBehavior.Enum.FIRE_ALWAYS) + .build()) + .putCoders("string-coder", CoderTranslation.toProto(StringUtf8Coder.of()).getCoder()) + .putCoders( + "window-coder", + Coder.newBuilder() + .setSpec( + FunctionSpec.newBuilder() + .setUrn(ModelCoders.GLOBAL_WINDOW_CODER_URN) + .build()) + .build()) + .build(); + + PipelineOptions options = PipelineOptionsFactory.create(); + if (exceptionHandling != null) { + options + .as(ExperimentalOptions.class) + .setExperiments( + Collections.singletonList( + ExperimentalOptions.FUSED_DOFN_EXCEPTION_HANDLING + "=" + exceptionHandling)); + } + Map factories = Maps.newHashMap(REGISTERED_RUNNER_FACTORIES); + factories.put( + DATA_INPUT_URN, + context -> + context.addIncomingDataEndpoint( + ApiServiceDescriptor.getDefaultInstance(), + StringUtf8Coder.of(), + (WindowedValue input) -> + ((FnDataReceiver>) context.getPCollectionConsumer("pc-source")) + .accept(input))); + + return new ProcessBundleHandler( + options, + Collections.emptySet(), + unused -> descriptor, + beamFnDataClient, + null /* beamFnStateGrpcClientCache */, + null /* finalizeBundleHandler */, + new ShortIdMap(), + executionStateSampler, + factories, + Caches.noop(), + new BundleProcessorCache(Duration.ZERO), + null /* dataSampler */); + } + + @Test + public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() throws Exception { + FUSED_DOWNSTREAM_INVOCATIONS.set(0); + FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); + ProcessBundleHandler handler = setupHandlerForFusedExceptionTest("fail"); + ByteStringOutputStream encoded = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("element", encoded); + Elements elements = + Elements.newBuilder() + .addData( + Data.newBuilder() + .setInstructionId("fused-fail") + .setTransformId("source") + .setData(encoded.toByteString()) + .setIsLast(true)) + .build(); + + Exception failure = + assertThrows( + Exception.class, + () -> + handler.processBundle( + InstructionRequest.newBuilder() + .setInstructionId("fused-fail") + .setProcessBundle( + ProcessBundleRequest.newBuilder() + .setProcessBundleDescriptorId("fused-descriptor") + .setElements(elements)) + .build())); + assertTrue(failure.toString(), failure.getMessage().contains("fused downstream failure")); + assertEquals(1, FUSED_DOWNSTREAM_INVOCATIONS.get()); + assertEquals(1, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); + assertThat( + handler.bundleProcessorCache.getCachedBundleProcessors().get("fused-descriptor"), empty()); + } + + @Test + public void fusedDownstreamExceptionPreservesLegacyBehaviorByDefault() throws Exception { + FUSED_DOWNSTREAM_INVOCATIONS.set(0); + FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); + ProcessBundleHandler handler = setupHandlerForFusedExceptionTest(null); + ByteStringOutputStream encoded = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("element", encoded); + Elements elements = + Elements.newBuilder() + .addData( + Data.newBuilder() + .setInstructionId("fused-legacy") + .setTransformId("source") + .setData(encoded.toByteString()) + .setIsLast(true)) + .build(); + + handler.processBundle( + InstructionRequest.newBuilder() + .setInstructionId("fused-legacy") + .setProcessBundle( + ProcessBundleRequest.newBuilder() + .setProcessBundleDescriptorId("fused-descriptor") + .setElements(elements)) + .build()); + assertEquals(1, FUSED_DOWNSTREAM_INVOCATIONS.get()); + assertEquals(1, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); + handler.shutdown(); + } + private static final class SimpleDoFn extends DoFn, String> { private static final TupleTag MAIN_OUTPUT_TAG = new TupleTag<>("mainOutput"); private static final String TIMER_FAMILY_ID = "timer_family"; From 71bd2bcf26e9d3c053cae6ca132bcce916476b02 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:03:27 +0530 Subject: [PATCH 06/37] Fix imports in fused DoFn regression test --- .../apache/beam/fn/harness/control/ProcessBundleHandlerTest.java | 1 + 1 file changed, 1 insertion(+) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index 52053e97c2c7..3ff6acd7ca5d 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -125,6 +125,7 @@ import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.metrics.MetricsEnvironment; import org.apache.beam.sdk.options.ExperimentalOptions; +import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.state.TimeDomain; import org.apache.beam.sdk.state.TimerMap; From 1239f9bb80779dee7a9e4dbbb2a122a241c06733 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:03:45 +0530 Subject: [PATCH 07/37] Correct fused regression source data receiver typing --- .../beam/fn/harness/control/ProcessBundleHandlerTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index 3ff6acd7ca5d..1f9ae44ddce8 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -1074,9 +1074,9 @@ private ProcessBundleHandler setupHandlerForFusedExceptionTest(String exceptionH context.addIncomingDataEndpoint( ApiServiceDescriptor.getDefaultInstance(), StringUtf8Coder.of(), - (WindowedValue input) -> + (String input) -> ((FnDataReceiver>) context.getPCollectionConsumer("pc-source")) - .accept(input))); + .accept(WindowedValue.valueInGlobalWindow(input)))); return new ProcessBundleHandler( options, From 54fabd1bfa07649e2ecad01d8be94fd4174bfe48 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:04:00 +0530 Subject: [PATCH 08/37] Format fused test source receiver --- .../beam/fn/harness/control/ProcessBundleHandlerTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index 1f9ae44ddce8..dadc5defecea 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -1075,7 +1075,8 @@ private ProcessBundleHandler setupHandlerForFusedExceptionTest(String exceptionH ApiServiceDescriptor.getDefaultInstance(), StringUtf8Coder.of(), (String input) -> - ((FnDataReceiver>) context.getPCollectionConsumer("pc-source")) + ((FnDataReceiver>) + context.getPCollectionConsumer("pc-source")) .accept(WindowedValue.valueInGlobalWindow(input)))); return new ProcessBundleHandler( From ad327c36d75537ab49e77207e16163dfd58e9b04 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:04:57 +0530 Subject: [PATCH 09/37] Fix Java test imports --- .../beam/fn/harness/control/ProcessBundleHandlerTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index dadc5defecea..0e35ab1c760d 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -114,8 +114,8 @@ import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.fn.data.BeamFnDataOutboundAggregator; -import org.apache.beam.sdk.fn.data.FnDataReceiver; import org.apache.beam.sdk.fn.data.CloseableFnDataReceiver; +import org.apache.beam.sdk.fn.data.FnDataReceiver; import org.apache.beam.sdk.fn.data.DataEndpoint; import org.apache.beam.sdk.fn.data.TimerEndpoint; import org.apache.beam.sdk.fn.test.TestExecutors; From 57a45eda861088b4f046bf17717bcd4f707554c6 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:05:00 +0530 Subject: [PATCH 10/37] Fix Java test imports --- .../beam/fn/harness/data/PCollectionConsumerRegistryTest.java | 1 + 1 file changed, 1 insertion(+) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java index a1d929ea6f02..305fce9ecda8 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java @@ -27,6 +27,7 @@ import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.any; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; From 84ff929f8f980dc7176a68f69d0993ac7fd29c49 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:05:06 +0530 Subject: [PATCH 11/37] Use WindowedValues helper in fused test --- .../beam/fn/harness/control/ProcessBundleHandlerTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index 0e35ab1c760d..5dcd916d5cf7 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -151,6 +151,7 @@ import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; @@ -1077,7 +1078,7 @@ private ProcessBundleHandler setupHandlerForFusedExceptionTest(String exceptionH (String input) -> ((FnDataReceiver>) context.getPCollectionConsumer("pc-source")) - .accept(WindowedValue.valueInGlobalWindow(input)))); + .accept(WindowedValues.valueInGlobalWindow(input)))); return new ProcessBundleHandler( options, From 6714eb0843737a36c933dac69b3255fb7d9593a0 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:05:18 +0530 Subject: [PATCH 12/37] Use typed PCollection consumer in fused test --- .../beam/fn/harness/control/ProcessBundleHandlerTest.java | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index 5dcd916d5cf7..6c81079fbfd7 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -115,7 +115,6 @@ import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.fn.data.BeamFnDataOutboundAggregator; import org.apache.beam.sdk.fn.data.CloseableFnDataReceiver; -import org.apache.beam.sdk.fn.data.FnDataReceiver; import org.apache.beam.sdk.fn.data.DataEndpoint; import org.apache.beam.sdk.fn.data.TimerEndpoint; import org.apache.beam.sdk.fn.test.TestExecutors; @@ -1076,8 +1075,8 @@ private ProcessBundleHandler setupHandlerForFusedExceptionTest(String exceptionH ApiServiceDescriptor.getDefaultInstance(), StringUtf8Coder.of(), (String input) -> - ((FnDataReceiver>) - context.getPCollectionConsumer("pc-source")) + context + .>getPCollectionConsumer("pc-source") .accept(WindowedValues.valueInGlobalWindow(input)))); return new ProcessBundleHandler( From 253ad2ba13fe79881410ccc9b96cc3f980a668d0 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:05:50 +0530 Subject: [PATCH 13/37] Match fused DoFn output mapping to its main output tag --- .../beam/fn/harness/control/ProcessBundleHandlerTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index 6c81079fbfd7..a2a7f0f8470c 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -983,7 +983,7 @@ private static PTransform fusedDoFnTransform( .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload(payload.toByteString())) .putInputs("input", inputPCollection) - .putOutputs("output", outputPCollection) + .putOutputs("mainOutput", outputPCollection) .build(); } From 741fdb3e010f24e4d6b09b8e0d9f74689a487291 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:06:19 +0530 Subject: [PATCH 14/37] Document fused DoFn exception handling --- .../beam/fn/harness/data/PCollectionConsumerRegistry.java | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java index a947754968bd..9b468c65ee80 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java @@ -62,8 +62,10 @@ /** * The {@code PCollectionConsumerRegistry} is used to maintain a collection of consuming * FnDataReceiver for each pCollectionId. Registering with this class allows inserting an element - * count counter for every pCollection. A combined MultiplexingConsumer (Wrapped with an - * ElementCountFnDataReceiver) is returned by calling getMultiplexingConsumer. + * count counter for every pCollection. A combined MultiplexingConsumer (wrapped with metrics + * tracking) is returned by calling getMultiplexingConsumer. When configured, this class also + * records downstream exceptions so that exceptions swallowed by fused user code can be surfaced + * when the bundle finishes. */ @SuppressWarnings({ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) From fbcb6c82c1243aa73069ad70eae022eb7c1b4948 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:06:21 +0530 Subject: [PATCH 15/37] Document fused DoFn exception handling --- .../java/org/apache/beam/sdk/options/ExperimentalOptions.java | 1 + 1 file changed, 1 insertion(+) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ExperimentalOptions.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ExperimentalOptions.java index 0a71ddb50183..34965f34f201 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ExperimentalOptions.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ExperimentalOptions.java @@ -35,6 +35,7 @@ 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( From 7668378f6225adeb6d5c2e58ba3551e1fb63c30a Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:06:24 +0530 Subject: [PATCH 16/37] Document fused DoFn exception handling --- .../apache/beam/fn/harness/control/ProcessBundleHandlerTest.java | 1 + 1 file changed, 1 insertion(+) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index a2a7f0f8470c..e02774f53f72 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -1128,6 +1128,7 @@ public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() thro assertEquals(1, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); assertThat( handler.bundleProcessorCache.getCachedBundleProcessors().get("fused-descriptor"), empty()); + handler.shutdown(); } @Test From 4751c88cb6d09971eed24c7cc821bca374bb0e74 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:06:44 +0530 Subject: [PATCH 17/37] Exercise fail-fast behavior across multiple fused inputs --- .../harness/control/ProcessBundleHandlerTest.java | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index e02774f53f72..1852cc200271 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -1099,15 +1099,22 @@ public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() thro FUSED_DOWNSTREAM_INVOCATIONS.set(0); FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); ProcessBundleHandler handler = setupHandlerForFusedExceptionTest("fail"); - ByteStringOutputStream encoded = new ByteStringOutputStream(); - StringUtf8Coder.of().encode("element", encoded); + ByteStringOutputStream firstEncodedElement = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("first", firstEncodedElement); + ByteStringOutputStream secondEncodedElement = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("second", secondEncodedElement); Elements elements = Elements.newBuilder() .addData( Data.newBuilder() .setInstructionId("fused-fail") .setTransformId("source") - .setData(encoded.toByteString()) + .setData(firstEncodedElement.toByteString())) + .addData( + Data.newBuilder() + .setInstructionId("fused-fail") + .setTransformId("source") + .setData(secondEncodedElement.toByteString()) .setIsLast(true)) .build(); @@ -1125,7 +1132,7 @@ public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() thro .build())); assertTrue(failure.toString(), failure.getMessage().contains("fused downstream failure")); assertEquals(1, FUSED_DOWNSTREAM_INVOCATIONS.get()); - assertEquals(1, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); + assertEquals(2, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); assertThat( handler.bundleProcessorCache.getCachedBundleProcessors().get("fused-descriptor"), empty()); handler.shutdown(); From 8a302c310c11eb3e458dcdcd34f9a3fe92c667b7 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:07:39 +0530 Subject: [PATCH 18/37] Separate fused failure and warning completion actions --- .../data/PCollectionConsumerRegistry.java | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java index 9b468c65ee80..79b26eaf32cc 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java @@ -332,21 +332,26 @@ private void failFastIfDownstreamFailed() throws Exception { } /** - * Applies the configured policy after all bundle finish functions have run. This must be called - * before outbound data is finalized and before the bundle processor is returned to its cache. + * Fails the bundle when configured to do so. This must run after finish functions but before + * outbound data is finalized. */ - public void finishBundle() throws Exception { + public void failBundleIfDownstreamFailed() throws Exception { Exception exception = firstDownstreamException.get(); - if (exception == null) { - return; + if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.FAIL && exception != null) { + throw exception; } - if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.WARN) { + } + + /** + * Logs a swallowed downstream exception only after the rest of bundle processing has succeeded. + */ + public void logWarningIfDownstreamFailed() { + Exception exception = firstDownstreamException.get(); + if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.WARN && exception != null) { LOG.warn( "A downstream exception was caught by fused user code. The bundle completed, but " + "downstream processing may have been skipped.", exception); - } else if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.FAIL) { - throw exception; } } From 922457a0418b4e68c348509cebbb43687d01eea6 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:07:42 +0530 Subject: [PATCH 19/37] Log fused warnings only after successful bundle work --- .../apache/beam/fn/harness/control/ProcessBundleHandler.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index 815c1290c1d5..c97b2665bab7 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java @@ -570,7 +570,7 @@ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest re // Surface downstream exceptions that user code may have caught while processing // fused outputs. This must happen before outbound data is finalized. - bundleProcessor.getpCollectionConsumerRegistry().finishBundle(); + bundleProcessor.getpCollectionConsumerRegistry().failBundleIfDownstreamFailed(); // If bundleProcessor has not flushed any elements, embed them in response. embedOutboundElementsIfApplicable(response, bundleProcessor); @@ -603,6 +603,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); From 7a834957dfe36970966970ebed98e77f94c3f721 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:07:44 +0530 Subject: [PATCH 20/37] Update registry tests for completion policies --- .../fn/harness/data/PCollectionConsumerRegistryTest.java | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java index 305fce9ecda8..6adf1543f72d 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java @@ -259,7 +259,7 @@ public void failModeRemembersSwallowedExceptionAndStopsLaterOutputs() throws Exc verify(consumer, times(1)).accept(any()); try { - consumers.finishBundle(); + consumers.failBundleIfDownstreamFailed(); throw new AssertionError("Expected the recorded downstream exception to fail the bundle"); } catch (Exception expected) { assertEquals(message, expected.getMessage()); @@ -295,7 +295,8 @@ public void warnModeRemembersExceptionButAllowsLaterOutputs() throws Exception { verify(consumer, times(2)).accept(any()); // Warning mode reports at successful bundle completion but does not fail the bundle. - consumers.finishBundle(); + consumers.failBundleIfDownstreamFailed(); + consumers.logWarningIfDownstreamFailed(); } @Test @@ -326,7 +327,8 @@ public void resetClearsRecordedDownstreamException() throws Exception { consumers.reset(); doNothing().when(consumer).accept(any()); wrapperConsumer.accept(valueInGlobalWindow("next bundle")); - consumers.finishBundle(); + consumers.failBundleIfDownstreamFailed(); + consumers.logWarningIfDownstreamFailed(); verify(consumer, times(2)).accept(any()); } From b8032590e8b244bbfa16ccc40c83cd10c9dddc94 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 10:09:06 +0530 Subject: [PATCH 21/37] Clarify failure tracker reset lifecycle --- .../beam/fn/harness/data/PCollectionConsumerRegistry.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java index 79b26eaf32cc..ee6dd0d7dab3 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java @@ -355,7 +355,7 @@ public void logWarningIfDownstreamFailed() { } } - /** Clears per-bundle failure state when a processor is successfully returned to the cache. */ + /** Clears per-bundle failure state during processor reset before it can be reused. */ public void reset() { firstDownstreamException.set(null); } From 2c0f96935dc4d0566b7db5c4819ec3b458e3d66d Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 11:51:32 +0530 Subject: [PATCH 22/37] Fix Error Prone MissingFail warnings in registry tests --- .../data/PCollectionConsumerRegistryTest.java | 42 +++++++------------ 1 file changed, 16 insertions(+), 26 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java index 6adf1543f72d..4793a887be31 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java @@ -22,6 +22,7 @@ import static org.hamcrest.Matchers.containsInAnyOrder; import static org.hamcrest.Matchers.is; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; @@ -243,27 +244,20 @@ public void failModeRemembersSwallowedExceptionAndStopsLaterOutputs() throws Exc (FnDataReceiver) consumers.getMultiplexingConsumer(P_COLLECTION_A); doThrow(new Exception(message)).when(consumer).accept(any()); - try { - wrapperConsumer.accept(valueInGlobalWindow("first")); - } catch (Exception expected) { - assertEquals(message, expected.getMessage()); - } + Exception firstFailure = + assertThrows(Exception.class, () -> wrapperConsumer.accept(valueInGlobalWindow("first"))); + assertEquals(message, firstFailure.getMessage()); // The upstream DoFn may catch the first exception and try to emit again. Fail mode must not // invoke the downstream consumer again after its first failure. - try { - wrapperConsumer.accept(valueInGlobalWindow("second")); - } catch (Exception expected) { - assertEquals(message, expected.getMessage()); - } + Exception secondFailure = + assertThrows(Exception.class, () -> wrapperConsumer.accept(valueInGlobalWindow("second"))); + assertEquals(message, secondFailure.getMessage()); verify(consumer, times(1)).accept(any()); - try { - consumers.failBundleIfDownstreamFailed(); - throw new AssertionError("Expected the recorded downstream exception to fail the bundle"); - } catch (Exception expected) { - assertEquals(message, expected.getMessage()); - } + Exception bundleFailure = + assertThrows(Exception.class, consumers::failBundleIfDownstreamFailed); + assertEquals(message, bundleFailure.getMessage()); } @Test @@ -286,11 +280,9 @@ public void warnModeRemembersExceptionButAllowsLaterOutputs() throws Exception { (FnDataReceiver) consumers.getMultiplexingConsumer(P_COLLECTION_A); doThrow(new Exception("first output failed")).doNothing().when(consumer).accept(any()); - try { - wrapperConsumer.accept(valueInGlobalWindow("first")); - } catch (Exception expected) { - assertEquals("first output failed", expected.getMessage()); - } + Exception firstFailure = + assertThrows(Exception.class, () -> wrapperConsumer.accept(valueInGlobalWindow("first"))); + assertEquals("first output failed", firstFailure.getMessage()); wrapperConsumer.accept(valueInGlobalWindow("second")); verify(consumer, times(2)).accept(any()); @@ -319,11 +311,9 @@ public void resetClearsRecordedDownstreamException() throws Exception { (FnDataReceiver) consumers.getMultiplexingConsumer(P_COLLECTION_A); doThrow(new Exception("first bundle failure")).when(consumer).accept(any()); - try { - wrapperConsumer.accept(valueInGlobalWindow("first")); - } catch (Exception expected) { - assertEquals("first bundle failure", expected.getMessage()); - } + Exception firstFailure = + assertThrows(Exception.class, () -> wrapperConsumer.accept(valueInGlobalWindow("first"))); + assertEquals("first bundle failure", firstFailure.getMessage()); consumers.reset(); doNothing().when(consumer).accept(any()); wrapperConsumer.accept(valueInGlobalWindow("next bundle")); From b356d64320f7117288fe8903466ee29c095ff847 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 12:04:58 +0530 Subject: [PATCH 23/37] Fix fused exception test input stream --- .../control/ProcessBundleHandlerTest.java | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index 1852cc200271..af613a99a33a 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -1099,22 +1099,19 @@ public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() thro FUSED_DOWNSTREAM_INVOCATIONS.set(0); FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); ProcessBundleHandler handler = setupHandlerForFusedExceptionTest("fail"); - ByteStringOutputStream firstEncodedElement = new ByteStringOutputStream(); - StringUtf8Coder.of().encode("first", firstEncodedElement); - ByteStringOutputStream secondEncodedElement = new ByteStringOutputStream(); - StringUtf8Coder.of().encode("second", secondEncodedElement); + // A single data block represents the complete encoded stream for this input endpoint. + // Encode both elements into that stream so the test exercises two distinct processElement + // invocations rather than relying on multiple embedded data blocks for one endpoint. + ByteStringOutputStream encodedElements = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("first", encodedElements); + StringUtf8Coder.of().encode("second", encodedElements); Elements elements = Elements.newBuilder() .addData( Data.newBuilder() .setInstructionId("fused-fail") .setTransformId("source") - .setData(firstEncodedElement.toByteString())) - .addData( - Data.newBuilder() - .setInstructionId("fused-fail") - .setTransformId("source") - .setData(secondEncodedElement.toByteString()) + .setData(encodedElements.toByteString()) .setIsLast(true)) .build(); From b1c33cef9db627046603e0da256f6261ae8ac475 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 12:09:01 +0530 Subject: [PATCH 24/37] Limit fused fail-fast to the downstream transform that failed --- .../data/PCollectionConsumerRegistry.java | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java index ee6dd0d7dab3..dcdab05edb0c 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java @@ -24,6 +24,7 @@ import java.util.List; import java.util.Map; import java.util.Random; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; import javax.annotation.Nullable; import org.apache.beam.fn.harness.HandlesSplits; @@ -82,6 +83,8 @@ public enum FusedDoFnExceptionHandling { private final FusedDoFnExceptionHandling fusedDoFnExceptionHandling; private final AtomicReference firstDownstreamException = new AtomicReference<>(); + private final Map downstreamExceptionsByTransform = + new ConcurrentHashMap<>(); /** Stores metadata about each consumer so that the appropriate metrics tracking can occur. */ @AutoValue @@ -314,17 +317,18 @@ public ExecutionState getProcessingExecutionState( * Records a downstream exception when configured to detect exceptions swallowed by a fused * upstream DoFn. Work cancellation is not a downstream processing failure. */ - private void recordDownstreamException(Exception exception) { + private void recordDownstreamException(String ptransformId, Exception exception) { if (fusedDoFnExceptionHandling != FusedDoFnExceptionHandling.LEGACY && !WorkCancelledException.isWorkCancelledException(exception)) { firstDownstreamException.compareAndSet(null, exception); + downstreamExceptionsByTransform.putIfAbsent(ptransformId, exception); } } - /** Prevents additional downstream work after the first failure in fail mode. */ - private void failFastIfDownstreamFailed() throws Exception { + /** Prevents retrying a downstream transform that has already failed in fail mode. */ + private void failFastIfDownstreamFailed(String ptransformId) throws Exception { if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.FAIL) { - Exception exception = firstDownstreamException.get(); + Exception exception = downstreamExceptionsByTransform.get(ptransformId); if (exception != null) { throw exception; } @@ -358,6 +362,7 @@ public void logWarningIfDownstreamFailed() { /** Clears per-bundle failure state during processor reset before it can be reused. */ public void reset() { firstDownstreamException.set(null); + downstreamExceptionsByTransform.clear(); } private static void logAndRethrow( @@ -446,7 +451,7 @@ public MetricTrackingFnDataReceiver( @Override public void accept(WindowedValue input) throws Exception { - failFastIfDownstreamFailed(); + failFastIfDownstreamFailed(ptransformId); // Increment the counter for each window the element occurs in. this.elementCountCounter.inc(input.getWindows().size()); // TODO(https://github.com/apache/beam/issues/20730): Consider updating size per window when @@ -464,7 +469,7 @@ public void accept(WindowedValue input) throws Exception { try (ExecutionState.ActiveState a = executionState.scopedActivate()) { this.delegate.accept(input); } catch (Exception e) { - recordDownstreamException(e); + recordDownstreamException(ptransformId, e); logAndRethrow( e, executionState, executionStateTracker, ptransformId, outputSampler, elementSample); } @@ -550,9 +555,10 @@ public void accept(WindowedValue input) throws Exception { ConsumerAndMetadata consumerAndMetadata = consumerAndMetadatas.get(i); ExecutionState state = consumerAndMetadata.getExecutionState(); try (ExecutionState.ActiveState a = state.scopedActivate()) { + failFastIfDownstreamFailed(consumerAndMetadata.getPTransformId()); consumerAndMetadata.getConsumer().accept(input); } catch (Exception e) { - recordDownstreamException(e); + recordDownstreamException(consumerAndMetadata.getPTransformId(), e); logAndRethrow( e, state, From 1bef43e8ecb15e834f45724aed18b30148a2ed00 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 12:21:10 +0530 Subject: [PATCH 25/37] Avoid relogging previously recorded fused failures --- .../beam/fn/harness/data/PCollectionConsumerRegistry.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java index dcdab05edb0c..4d2363934967 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java @@ -554,8 +554,10 @@ public void accept(WindowedValue input) throws Exception { for (int size = consumerAndMetadatas.size(), i = 0; i < size; ++i) { ConsumerAndMetadata consumerAndMetadata = consumerAndMetadatas.get(i); ExecutionState state = consumerAndMetadata.getExecutionState(); + // The original failure was already logged when the consumer first threw. Avoid logging + // the same failure again for every later element rejected by fail-fast handling. + failFastIfDownstreamFailed(consumerAndMetadata.getPTransformId()); try (ExecutionState.ActiveState a = state.scopedActivate()) { - failFastIfDownstreamFailed(consumerAndMetadata.getPTransformId()); consumerAndMetadata.getConsumer().accept(input); } catch (Exception e) { recordDownstreamException(consumerAndMetadata.getPTransformId(), e); From f3044058e24177eff85de9dbafdb29f4f5f6ac43 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 12:21:24 +0530 Subject: [PATCH 26/37] Always shut down handlers in fused exception tests --- .../control/ProcessBundleHandlerTest.java | 121 +++++++++--------- 1 file changed, 64 insertions(+), 57 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index af613a99a33a..f2855c0406fc 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -1099,40 +1099,44 @@ public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() thro FUSED_DOWNSTREAM_INVOCATIONS.set(0); FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); ProcessBundleHandler handler = setupHandlerForFusedExceptionTest("fail"); - // A single data block represents the complete encoded stream for this input endpoint. - // Encode both elements into that stream so the test exercises two distinct processElement - // invocations rather than relying on multiple embedded data blocks for one endpoint. - ByteStringOutputStream encodedElements = new ByteStringOutputStream(); - StringUtf8Coder.of().encode("first", encodedElements); - StringUtf8Coder.of().encode("second", encodedElements); - Elements elements = - Elements.newBuilder() - .addData( - Data.newBuilder() - .setInstructionId("fused-fail") - .setTransformId("source") - .setData(encodedElements.toByteString()) - .setIsLast(true)) - .build(); - - Exception failure = - assertThrows( - Exception.class, - () -> - handler.processBundle( - InstructionRequest.newBuilder() - .setInstructionId("fused-fail") - .setProcessBundle( - ProcessBundleRequest.newBuilder() - .setProcessBundleDescriptorId("fused-descriptor") - .setElements(elements)) - .build())); - assertTrue(failure.toString(), failure.getMessage().contains("fused downstream failure")); - assertEquals(1, FUSED_DOWNSTREAM_INVOCATIONS.get()); - assertEquals(2, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); - assertThat( - handler.bundleProcessorCache.getCachedBundleProcessors().get("fused-descriptor"), empty()); - handler.shutdown(); + try { + // A single data block represents the complete encoded stream for this input endpoint. + // Encode both elements into that stream so the test exercises two distinct processElement + // invocations rather than relying on multiple embedded data blocks for one endpoint. + ByteStringOutputStream encodedElements = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("first", encodedElements); + StringUtf8Coder.of().encode("second", encodedElements); + Elements elements = + Elements.newBuilder() + .addData( + Data.newBuilder() + .setInstructionId("fused-fail") + .setTransformId("source") + .setData(encodedElements.toByteString()) + .setIsLast(true)) + .build(); + + Exception failure = + assertThrows( + Exception.class, + () -> + handler.processBundle( + InstructionRequest.newBuilder() + .setInstructionId("fused-fail") + .setProcessBundle( + ProcessBundleRequest.newBuilder() + .setProcessBundleDescriptorId("fused-descriptor") + .setElements(elements)) + .build())); + assertTrue(failure.toString(), failure.getMessage().contains("fused downstream failure")); + assertEquals(1, FUSED_DOWNSTREAM_INVOCATIONS.get()); + assertEquals(2, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); + assertThat( + handler.bundleProcessorCache.getCachedBundleProcessors().get("fused-descriptor"), + empty()); + } finally { + handler.shutdown(); + } } @Test @@ -1140,29 +1144,32 @@ public void fusedDownstreamExceptionPreservesLegacyBehaviorByDefault() throws Ex FUSED_DOWNSTREAM_INVOCATIONS.set(0); FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); ProcessBundleHandler handler = setupHandlerForFusedExceptionTest(null); - ByteStringOutputStream encoded = new ByteStringOutputStream(); - StringUtf8Coder.of().encode("element", encoded); - Elements elements = - Elements.newBuilder() - .addData( - Data.newBuilder() - .setInstructionId("fused-legacy") - .setTransformId("source") - .setData(encoded.toByteString()) - .setIsLast(true)) - .build(); - - handler.processBundle( - InstructionRequest.newBuilder() - .setInstructionId("fused-legacy") - .setProcessBundle( - ProcessBundleRequest.newBuilder() - .setProcessBundleDescriptorId("fused-descriptor") - .setElements(elements)) - .build()); - assertEquals(1, FUSED_DOWNSTREAM_INVOCATIONS.get()); - assertEquals(1, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); - handler.shutdown(); + try { + ByteStringOutputStream encoded = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("element", encoded); + Elements elements = + Elements.newBuilder() + .addData( + Data.newBuilder() + .setInstructionId("fused-legacy") + .setTransformId("source") + .setData(encoded.toByteString()) + .setIsLast(true)) + .build(); + + handler.processBundle( + InstructionRequest.newBuilder() + .setInstructionId("fused-legacy") + .setProcessBundle( + ProcessBundleRequest.newBuilder() + .setProcessBundleDescriptorId("fused-descriptor") + .setElements(elements)) + .build()); + assertEquals(1, FUSED_DOWNSTREAM_INVOCATIONS.get()); + assertEquals(1, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); + } finally { + handler.shutdown(); + } } private static final class SimpleDoFn extends DoFn, String> { From 761f862bca05b6756e127b306474d5a020ec672d Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 12:21:50 +0530 Subject: [PATCH 27/37] Cover successful bundle completion in warning mode --- .../control/ProcessBundleHandlerTest.java | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index f2855c0406fc..d8e1760e4ced 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -1139,6 +1139,44 @@ public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() thro } } + @Test + public void fusedDownstreamExceptionWarnModeAllowsBundleToComplete() throws Exception { + FUSED_DOWNSTREAM_INVOCATIONS.set(0); + FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); + ProcessBundleHandler handler = setupHandlerForFusedExceptionTest("warn"); + try { + ByteStringOutputStream encoded = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("element", encoded); + Elements elements = + Elements.newBuilder() + .addData( + Data.newBuilder() + .setInstructionId("fused-warn") + .setTransformId("source") + .setData(encoded.toByteString()) + .setIsLast(true)) + .build(); + + handler.processBundle( + InstructionRequest.newBuilder() + .setInstructionId("fused-warn") + .setProcessBundle( + ProcessBundleRequest.newBuilder() + .setProcessBundleDescriptorId("fused-descriptor") + .setElements(elements)) + .build()); + assertEquals(1, FUSED_DOWNSTREAM_INVOCATIONS.get()); + assertEquals(1, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); + assertFalse( + handler.bundleProcessorCache + .getCachedBundleProcessors() + .get("fused-descriptor") + .isEmpty()); + } finally { + handler.shutdown(); + } + } + @Test public void fusedDownstreamExceptionPreservesLegacyBehaviorByDefault() throws Exception { FUSED_DOWNSTREAM_INVOCATIONS.set(0); From 73d5bd0bc98254d5008a184cf756d5ddbf2ab7d9 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 13:27:55 +0530 Subject: [PATCH 28/37] Handle Error subclasses and isolate failed fused consumers --- .../data/PCollectionConsumerRegistry.java | 97 ++++++++++++++----- 1 file changed, 71 insertions(+), 26 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java index 4d2363934967..3a1ec3b7501a 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java @@ -82,8 +82,8 @@ public enum FusedDoFnExceptionHandling { } private final FusedDoFnExceptionHandling fusedDoFnExceptionHandling; - private final AtomicReference firstDownstreamException = new AtomicReference<>(); - private final Map downstreamExceptionsByTransform = + private final AtomicReference firstDownstreamException = new AtomicReference<>(); + private final Map downstreamExceptionsByTransform = new ConcurrentHashMap<>(); /** Stores metadata about each consumer so that the appropriate metrics tracking can occur. */ @@ -317,7 +317,7 @@ public ExecutionState getProcessingExecutionState( * Records a downstream exception when configured to detect exceptions swallowed by a fused * upstream DoFn. Work cancellation is not a downstream processing failure. */ - private void recordDownstreamException(String ptransformId, Exception exception) { + private void recordDownstreamException(String ptransformId, Throwable exception) { if (fusedDoFnExceptionHandling != FusedDoFnExceptionHandling.LEGACY && !WorkCancelledException.isWorkCancelledException(exception)) { firstDownstreamException.compareAndSet(null, exception); @@ -325,13 +325,29 @@ private void recordDownstreamException(String ptransformId, Exception exception) } } + /** Re-throws a recorded downstream failure without changing its original type. */ + private static void rethrowDownstreamFailure(Throwable failure) throws Exception { + if (failure instanceof Exception) { + throw (Exception) failure; + } + if (failure instanceof Error) { + throw (Error) failure; + } + throw new RuntimeException(failure); + } + + /** Returns a previously recorded failure for the transform, if any. */ + private @Nullable Throwable getDownstreamFailure(String ptransformId) { + return fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.FAIL + ? downstreamExceptionsByTransform.get(ptransformId) + : null; + } + /** Prevents retrying a downstream transform that has already failed in fail mode. */ private void failFastIfDownstreamFailed(String ptransformId) throws Exception { - if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.FAIL) { - Exception exception = downstreamExceptionsByTransform.get(ptransformId); - if (exception != null) { - throw exception; - } + Throwable failure = getDownstreamFailure(ptransformId); + if (failure != null) { + rethrowDownstreamFailure(failure); } } @@ -340,9 +356,9 @@ private void failFastIfDownstreamFailed(String ptransformId) throws Exception { * outbound data is finalized. */ public void failBundleIfDownstreamFailed() throws Exception { - Exception exception = firstDownstreamException.get(); - if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.FAIL && exception != null) { - throw exception; + Throwable failure = firstDownstreamException.get(); + if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.FAIL && failure != null) { + rethrowDownstreamFailure(failure); } } @@ -350,12 +366,12 @@ public void failBundleIfDownstreamFailed() throws Exception { * Logs a swallowed downstream exception only after the rest of bundle processing has succeeded. */ public void logWarningIfDownstreamFailed() { - Exception exception = firstDownstreamException.get(); - if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.WARN && exception != null) { + Throwable failure = firstDownstreamException.get(); + if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.WARN && failure != null) { LOG.warn( "A downstream exception was caught by fused user code. The bundle completed, but " + "downstream processing may have been skipped.", - exception); + failure); } } @@ -472,6 +488,9 @@ public void accept(WindowedValue input) throws Exception { recordDownstreamException(ptransformId, e); logAndRethrow( e, executionState, executionStateTracker, ptransformId, outputSampler, elementSample); + } catch (Error e) { + recordDownstreamException(ptransformId, e); + throw e; } this.sampledByteSizeDistribution.finishLazyUpdate(); } @@ -551,25 +570,51 @@ public void accept(WindowedValue input) throws Exception { // Process Bundle Execution time metric and also ensure user counters can get an appropriate // metrics container. We specifically don't use a for-each loop since it creates an iterator // on a hot path. + Throwable firstFailure = null; for (int size = consumerAndMetadatas.size(), i = 0; i < size; ++i) { ConsumerAndMetadata consumerAndMetadata = consumerAndMetadatas.get(i); + String ptransformId = consumerAndMetadata.getPTransformId(); ExecutionState state = consumerAndMetadata.getExecutionState(); - // The original failure was already logged when the consumer first threw. Avoid logging - // the same failure again for every later element rejected by fail-fast handling. - failFastIfDownstreamFailed(consumerAndMetadata.getPTransformId()); + + // In fail mode, a failed consumer is skipped on later elements, but healthy sibling + // consumers must still receive those elements. + Throwable previousFailure = getDownstreamFailure(ptransformId); + if (previousFailure != null) { + if (firstFailure == null) { + firstFailure = previousFailure; + } + continue; + } + try (ExecutionState.ActiveState a = state.scopedActivate()) { consumerAndMetadata.getConsumer().accept(input); } catch (Exception e) { - recordDownstreamException(consumerAndMetadata.getPTransformId(), e); - logAndRethrow( - e, - state, - consumerAndMetadata.getExecutionStateTracker(), - consumerAndMetadata.getPTransformId(), - outputSampler, - elementSample); + recordDownstreamException(ptransformId, e); + if (fusedDoFnExceptionHandling != FusedDoFnExceptionHandling.FAIL) { + logAndRethrow( + e, + state, + consumerAndMetadata.getExecutionStateTracker(), + ptransformId, + outputSampler, + elementSample); + } + if (firstFailure == null) { + firstFailure = e; + } + } catch (Error e) { + recordDownstreamException(ptransformId, e); + if (fusedDoFnExceptionHandling != FusedDoFnExceptionHandling.FAIL) { + throw e; + } + if (firstFailure == null) { + firstFailure = e; + } } - this.sampledByteSizeDistribution.finishLazyUpdate(); + } + this.sampledByteSizeDistribution.finishLazyUpdate(); + if (firstFailure != null) { + rethrowDownstreamFailure(firstFailure); } } } From 21549c95d207c867bbd3bb6bcca2c86020693e71 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 13:27:57 +0530 Subject: [PATCH 29/37] Fail fused bundles before finishing failed DoFns --- .../fn/harness/control/ProcessBundleHandler.java | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index c97b2665bab7..2182c5621b07 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java @@ -561,15 +561,22 @@ 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 downstream exceptions that user code may have caught while processing - // fused outputs. This must happen before outbound data is finalized. + // 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. From 97fcdf1f37a6e60ce2523e54ed800b9b5dea801a Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 13:28:00 +0530 Subject: [PATCH 30/37] Test fused consumer isolation and Error propagation --- .../data/PCollectionConsumerRegistryTest.java | 65 +++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java index 4793a887be31..e0ec07d58f3c 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java @@ -260,6 +260,71 @@ public void failModeRemembersSwallowedExceptionAndStopsLaterOutputs() throws Exc assertEquals(message, bundleFailure.getMessage()); } + @Test + public void failModeContinuesDeliveringToHealthySiblingConsumers() throws Exception { + ShortIdMap shortIds = new ShortIdMap(); + BundleProgressReporter.InMemory reporterAndRegistrar = new BundleProgressReporter.InMemory(); + PCollectionConsumerRegistry consumers = + new PCollectionConsumerRegistry( + sampler.create(), + shortIds, + reporterAndRegistrar, + TEST_DESCRIPTOR, + null, + PCollectionConsumerRegistry.FusedDoFnExceptionHandling.FAIL); + FnDataReceiver> failedConsumer = mock(FnDataReceiver.class); + FnDataReceiver> healthyConsumer = mock(FnDataReceiver.class); + consumers.register(P_COLLECTION_A, "failed", "failed", failedConsumer); + consumers.register(P_COLLECTION_A, "healthy", "healthy", healthyConsumer); + FnDataReceiver> wrapperConsumer = + (FnDataReceiver>) + (FnDataReceiver) consumers.getMultiplexingConsumer(P_COLLECTION_A); + doThrow(new Exception("downstream failure")).when(failedConsumer).accept(any()); + + Exception firstFailure = + assertThrows( + Exception.class, () -> wrapperConsumer.accept(valueInGlobalWindow("first"))); + assertEquals("downstream failure", firstFailure.getMessage()); + verify(failedConsumer, times(1)).accept(any()); + verify(healthyConsumer, times(1)).accept(any()); + + Exception secondFailure = + assertThrows( + Exception.class, () -> wrapperConsumer.accept(valueInGlobalWindow("second"))); + assertEquals("downstream failure", secondFailure.getMessage()); + // The failed consumer is not invoked again, but its healthy sibling continues processing. + verify(failedConsumer, times(1)).accept(any()); + verify(healthyConsumer, times(2)).accept(any()); + } + + @Test + public void failModeRecordsAndSurfacesErrors() throws Exception { + ShortIdMap shortIds = new ShortIdMap(); + BundleProgressReporter.InMemory reporterAndRegistrar = new BundleProgressReporter.InMemory(); + PCollectionConsumerRegistry consumers = + new PCollectionConsumerRegistry( + sampler.create(), + shortIds, + reporterAndRegistrar, + TEST_DESCRIPTOR, + null, + PCollectionConsumerRegistry.FusedDoFnExceptionHandling.FAIL); + FnDataReceiver> consumer = mock(FnDataReceiver.class); + consumers.register(P_COLLECTION_A, "failed", "failed", consumer); + FnDataReceiver> wrapperConsumer = + (FnDataReceiver>) + (FnDataReceiver) consumers.getMultiplexingConsumer(P_COLLECTION_A); + doThrow(new AssertionError("downstream error")).when(consumer).accept(any()); + + AssertionError outputFailure = + assertThrows( + AssertionError.class, () -> wrapperConsumer.accept(valueInGlobalWindow("element"))); + assertEquals("downstream error", outputFailure.getMessage()); + AssertionError bundleFailure = + assertThrows(AssertionError.class, consumers::failBundleIfDownstreamFailed); + assertEquals("downstream error", bundleFailure.getMessage()); + } + @Test public void warnModeRemembersExceptionButAllowsLaterOutputs() throws Exception { final String pTransformId = "pTransformId"; From 82221e273259c6adf5729ae96a184d8b254491cf Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 13:28:02 +0530 Subject: [PATCH 31/37] Test finish-bundle suppression after fused failures --- .../harness/control/ProcessBundleHandlerTest.java | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index d8e1760e4ced..3a4a482e5f0f 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -941,6 +941,7 @@ public void testPTransformStartExceptionsArePropagated() { private static final AtomicInteger FUSED_DOWNSTREAM_INVOCATIONS = new AtomicInteger(); private static final AtomicInteger FUSED_UPSTREAM_CAUGHT_EXCEPTIONS = new AtomicInteger(); + private static final AtomicInteger FUSED_DOWNSTREAM_FINISH_INVOCATIONS = new AtomicInteger(); private static final class CatchingFusedDoFn extends DoFn { private static final TupleTag MAIN_OUTPUT_TAG = new TupleTag<>("mainOutput"); @@ -963,6 +964,11 @@ public void processElement(ProcessContext context) { FUSED_DOWNSTREAM_INVOCATIONS.incrementAndGet(); throw new IllegalStateException("fused downstream failure"); } + + @FinishBundle + public void finishBundle() { + FUSED_DOWNSTREAM_FINISH_INVOCATIONS.incrementAndGet(); + } } private static PTransform fusedDoFnTransform( @@ -1098,6 +1104,7 @@ private ProcessBundleHandler setupHandlerForFusedExceptionTest(String exceptionH public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() throws Exception { FUSED_DOWNSTREAM_INVOCATIONS.set(0); FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); + FUSED_DOWNSTREAM_FINISH_INVOCATIONS.set(0); ProcessBundleHandler handler = setupHandlerForFusedExceptionTest("fail"); try { // A single data block represents the complete encoded stream for this input endpoint. @@ -1131,6 +1138,10 @@ public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() thro assertTrue(failure.toString(), failure.getMessage().contains("fused downstream failure")); assertEquals(1, FUSED_DOWNSTREAM_INVOCATIONS.get()); assertEquals(2, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); + assertEquals( + "A downstream failure during @ProcessElement must prevent @FinishBundle.", + 0, + FUSED_DOWNSTREAM_FINISH_INVOCATIONS.get()); assertThat( handler.bundleProcessorCache.getCachedBundleProcessors().get("fused-descriptor"), empty()); @@ -1143,6 +1154,7 @@ public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() thro public void fusedDownstreamExceptionWarnModeAllowsBundleToComplete() throws Exception { FUSED_DOWNSTREAM_INVOCATIONS.set(0); FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); + FUSED_DOWNSTREAM_FINISH_INVOCATIONS.set(0); ProcessBundleHandler handler = setupHandlerForFusedExceptionTest("warn"); try { ByteStringOutputStream encoded = new ByteStringOutputStream(); @@ -1167,6 +1179,7 @@ public void fusedDownstreamExceptionWarnModeAllowsBundleToComplete() throws Exce .build()); assertEquals(1, FUSED_DOWNSTREAM_INVOCATIONS.get()); assertEquals(1, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); + assertEquals(1, FUSED_DOWNSTREAM_FINISH_INVOCATIONS.get()); assertFalse( handler.bundleProcessorCache .getCachedBundleProcessors() @@ -1181,6 +1194,7 @@ public void fusedDownstreamExceptionWarnModeAllowsBundleToComplete() throws Exce public void fusedDownstreamExceptionPreservesLegacyBehaviorByDefault() throws Exception { FUSED_DOWNSTREAM_INVOCATIONS.set(0); FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); + FUSED_DOWNSTREAM_FINISH_INVOCATIONS.set(0); ProcessBundleHandler handler = setupHandlerForFusedExceptionTest(null); try { ByteStringOutputStream encoded = new ByteStringOutputStream(); @@ -1205,6 +1219,7 @@ public void fusedDownstreamExceptionPreservesLegacyBehaviorByDefault() throws Ex .build()); assertEquals(1, FUSED_DOWNSTREAM_INVOCATIONS.get()); assertEquals(1, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); + assertEquals(1, FUSED_DOWNSTREAM_FINISH_INVOCATIONS.get()); } finally { handler.shutdown(); } From 7c3e48933cd50878c5e833f0af55892df600de2e Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 13:28:39 +0530 Subject: [PATCH 32/37] Preserve cancellation semantics in fused fail mode --- .../beam/fn/harness/data/PCollectionConsumerRegistry.java | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java index 3a1ec3b7501a..728a7dd78e57 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java @@ -352,8 +352,9 @@ private void failFastIfDownstreamFailed(String ptransformId) throws Exception { } /** - * Fails the bundle when configured to do so. This must run after finish functions but before - * outbound data is finalized. + * Fails the bundle when configured to do so. Call before finish functions to avoid flushing + * state after a processing failure, and after each finish function to detect swallowed failures + * raised during bundle finalization. This must run before outbound data is finalized. */ public void failBundleIfDownstreamFailed() throws Exception { Throwable failure = firstDownstreamException.get(); @@ -590,7 +591,8 @@ public void accept(WindowedValue input) throws Exception { consumerAndMetadata.getConsumer().accept(input); } catch (Exception e) { recordDownstreamException(ptransformId, e); - if (fusedDoFnExceptionHandling != FusedDoFnExceptionHandling.FAIL) { + if (fusedDoFnExceptionHandling != FusedDoFnExceptionHandling.FAIL + || WorkCancelledException.isWorkCancelledException(e)) { logAndRethrow( e, state, From cf50026e2bfd396627eaf5ca7a87ed4417ae0fba Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 13:28:41 +0530 Subject: [PATCH 33/37] Reset fused exception test counters before each test --- .../harness/control/ProcessBundleHandlerTest.java | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index 3a4a482e5f0f..3a2e65b4c4c9 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -158,6 +158,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles; import org.joda.time.Instant; +import org.junit.Before; import org.junit.After; import org.junit.Before; import org.junit.Rule; @@ -943,6 +944,10 @@ public void testPTransformStartExceptionsArePropagated() { private static final AtomicInteger FUSED_UPSTREAM_CAUGHT_EXCEPTIONS = new AtomicInteger(); private static final AtomicInteger FUSED_DOWNSTREAM_FINISH_INVOCATIONS = new AtomicInteger(); + @Before + public void resetFusedExceptionTestCounters() { + } + private static final class CatchingFusedDoFn extends DoFn { private static final TupleTag MAIN_OUTPUT_TAG = new TupleTag<>("mainOutput"); @@ -1102,9 +1107,6 @@ private ProcessBundleHandler setupHandlerForFusedExceptionTest(String exceptionH @Test public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() throws Exception { - FUSED_DOWNSTREAM_INVOCATIONS.set(0); - FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); - FUSED_DOWNSTREAM_FINISH_INVOCATIONS.set(0); ProcessBundleHandler handler = setupHandlerForFusedExceptionTest("fail"); try { // A single data block represents the complete encoded stream for this input endpoint. @@ -1152,9 +1154,6 @@ public void fusedDownstreamExceptionCaughtByUpstreamFailsBundleInFailMode() thro @Test public void fusedDownstreamExceptionWarnModeAllowsBundleToComplete() throws Exception { - FUSED_DOWNSTREAM_INVOCATIONS.set(0); - FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); - FUSED_DOWNSTREAM_FINISH_INVOCATIONS.set(0); ProcessBundleHandler handler = setupHandlerForFusedExceptionTest("warn"); try { ByteStringOutputStream encoded = new ByteStringOutputStream(); @@ -1192,9 +1191,6 @@ public void fusedDownstreamExceptionWarnModeAllowsBundleToComplete() throws Exce @Test public void fusedDownstreamExceptionPreservesLegacyBehaviorByDefault() throws Exception { - FUSED_DOWNSTREAM_INVOCATIONS.set(0); - FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); - FUSED_DOWNSTREAM_FINISH_INVOCATIONS.set(0); ProcessBundleHandler handler = setupHandlerForFusedExceptionTest(null); try { ByteStringOutputStream encoded = new ByteStringOutputStream(); From 7ac28dcc3e6c56415785ac36cad86fd4f90e3652 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 13:29:29 +0530 Subject: [PATCH 34/37] Clean up bundle processors when swallowed errors fail bundles --- .../fn/harness/control/ProcessBundleHandler.java | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index 2182c5621b07..2bba229da62e 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java @@ -630,6 +630,19 @@ 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; } } From d16dffefe0e08876f1f4137b989ab30ce5ae9678 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 13:30:13 +0530 Subject: [PATCH 35/37] Keep metrics consistent when all fused consumers have failed --- .../data/PCollectionConsumerRegistry.java | 31 ++++++++++++------- 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java index 728a7dd78e57..9ff757d9ac64 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java @@ -556,22 +556,13 @@ public MultiplexingMetricTrackingFnDataReceiver( @Override public void accept(WindowedValue input) throws Exception { - // Increment the counter for each window the element occurs in. - this.elementCountCounter.inc(input.getWindows().size()); - // TODO(https://github.com/apache/beam/issues/20730): Consider updating size per window - // when we have window optimization. - this.sampledByteSizeDistribution.tryUpdate(input.getValue(), coder); - - ElementSample elementSample = null; - if (outputSampler != null) { - elementSample = outputSampler.sample(input); - } - // Use the ExecutionStateTracker and enter an appropriate state to track the // Process Bundle Execution time metric and also ensure user counters can get an appropriate // metrics container. We specifically don't use a for-each loop since it creates an iterator // on a hot path. Throwable firstFailure = null; + boolean metricsStarted = false; + ElementSample elementSample = null; for (int size = consumerAndMetadatas.size(), i = 0; i < size; ++i) { ConsumerAndMetadata consumerAndMetadata = consumerAndMetadatas.get(i); String ptransformId = consumerAndMetadata.getPTransformId(); @@ -587,6 +578,20 @@ public void accept(WindowedValue input) throws Exception { continue; } + // Start PCollection-level metrics only when at least one consumer is eligible to receive + // this element. This keeps metrics consistent with the single-consumer fail-fast path, + // which rejects an element before incrementing the counter. + if (!metricsStarted) { + this.elementCountCounter.inc(input.getWindows().size()); + // TODO(https://github.com/apache/beam/issues/20730): Consider updating size per window + // when we have window optimization. + this.sampledByteSizeDistribution.tryUpdate(input.getValue(), coder); + if (outputSampler != null) { + elementSample = outputSampler.sample(input); + } + metricsStarted = true; + } + try (ExecutionState.ActiveState a = state.scopedActivate()) { consumerAndMetadata.getConsumer().accept(input); } catch (Exception e) { @@ -614,7 +619,9 @@ public void accept(WindowedValue input) throws Exception { } } } - this.sampledByteSizeDistribution.finishLazyUpdate(); + if (metricsStarted) { + this.sampledByteSizeDistribution.finishLazyUpdate(); + } if (firstFailure != null) { rethrowDownstreamFailure(firstFailure); } From e648eeafda8eca1b527c74d898f614f14b74f646 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 13:30:29 +0530 Subject: [PATCH 36/37] Preserve metrics for PCollections without consumers --- .../fn/harness/data/PCollectionConsumerRegistry.java | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java index 9ff757d9ac64..32b6941593fa 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java @@ -563,6 +563,15 @@ public void accept(WindowedValue input) throws Exception { Throwable firstFailure = null; boolean metricsStarted = false; ElementSample elementSample = null; + // Preserve the existing metric behavior for PCollections that have no consumers. + if (consumerAndMetadatas.isEmpty()) { + this.elementCountCounter.inc(input.getWindows().size()); + this.sampledByteSizeDistribution.tryUpdate(input.getValue(), coder); + if (outputSampler != null) { + elementSample = outputSampler.sample(input); + } + metricsStarted = true; + } for (int size = consumerAndMetadatas.size(), i = 0; i < size; ++i) { ConsumerAndMetadata consumerAndMetadata = consumerAndMetadatas.get(i); String ptransformId = consumerAndMetadata.getPTransformId(); From 928e9f53a6819ec34f616ef7ce08c75853350da3 Mon Sep 17 00:00:00 2001 From: Sanjeev Kumar Date: Fri, 13 Nov 2026 13:45:30 +0530 Subject: [PATCH 37/37] Reset fused DoFn test counters between tests --- .../beam/fn/harness/control/ProcessBundleHandlerTest.java | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index 3a2e65b4c4c9..c55222118fb3 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -158,7 +158,6 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles; import org.joda.time.Instant; -import org.junit.Before; import org.junit.After; import org.junit.Before; import org.junit.Rule; @@ -946,6 +945,9 @@ public void testPTransformStartExceptionsArePropagated() { @Before public void resetFusedExceptionTestCounters() { + FUSED_DOWNSTREAM_INVOCATIONS.set(0); + FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.set(0); + FUSED_DOWNSTREAM_FINISH_INVOCATIONS.set(0); } private static final class CatchingFusedDoFn extends DoFn { @@ -1180,7 +1182,8 @@ public void fusedDownstreamExceptionWarnModeAllowsBundleToComplete() throws Exce assertEquals(1, FUSED_UPSTREAM_CAUGHT_EXCEPTIONS.get()); assertEquals(1, FUSED_DOWNSTREAM_FINISH_INVOCATIONS.get()); assertFalse( - handler.bundleProcessorCache + handler + .bundleProcessorCache .getCachedBundleProcessors() .get("fused-descriptor") .isEmpty());