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..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,9 @@ public interface ExperimentalOptions extends PipelineOptions { String STATE_SAMPLING_PERIOD_MILLIS = "state_sampling_period_millis"; + /** Experiment controlling exception handling at fused Java DoFn boundaries. */ + String FUSED_DOFN_EXCEPTION_HANDLING = "fused_dofn_exception_handling"; + @Description( "Apache Beam provides a number of experimental features that can " + "be enabled with this flag. If executing against a managed service, please contact the " 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..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 @@ -79,6 +79,7 @@ import org.apache.beam.sdk.metrics.MetricsContainer; import org.apache.beam.sdk.metrics.MetricsEnvironment; import org.apache.beam.sdk.metrics.MetricsEnvironment.MetricsEnvironmentState; +import org.apache.beam.sdk.options.ExperimentalOptions; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.SdkHarnessOptions; import org.apache.beam.sdk.transforms.DoFn.BundleFinalizer; @@ -560,13 +561,24 @@ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest re bundleProcessor.getInboundEndpointApiServiceDescriptors()); } - // Need to reverse this since we want to call finish in topological order. + // Do not run @FinishBundle after a downstream failure was swallowed during element + // processing. Finish functions can flush state or write to external sinks. + bundleProcessor.getpCollectionConsumerRegistry().failBundleIfDownstreamFailed(); + + // Need to reverse this since we want to call finish in topological order. Check after + // each function as well, so a swallowed failure during @FinishBundle prevents later + // finish functions from running. for (ThrowingRunnable finishFunction : Lists.reverse(finishFunctionRegistry.getFunctions())) { LOG.debug("Finishing function {}", finishFunction); finishFunction.run(); + bundleProcessor.getpCollectionConsumerRegistry().failBundleIfDownstreamFailed(); } + // Surface a failure raised by the final @FinishBundle function before outbound data is + // finalized. + bundleProcessor.getpCollectionConsumerRegistry().failBundleIfDownstreamFailed(); + // If bundleProcessor has not flushed any elements, embed them in response. embedOutboundElementsIfApplicable(response, bundleProcessor); @@ -598,6 +610,9 @@ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest re } } + // A warning is appropriate only once all bundle work has succeeded. + bundleProcessor.getpCollectionConsumerRegistry().logWarningIfDownstreamFailed(); + // Mark the bundle processor as re-usable. bundleProcessorCache.release( request.getProcessBundle().getProcessBundleDescriptorId(), bundleProcessor); @@ -615,7 +630,40 @@ public BeamFnApi.InstructionResponse.Builder processBundle(InstructionRequest re // Ensure that if more data arrives for the instruction it is discarded. beamFnDataClient.poisonInstructionId(instructionId); throw e; + } catch (Error e) { + LOG.debug( + "Error processing bundle {} with bundleProcessor for {} after error", + instructionId, + request.getProcessBundle().getProcessBundleDescriptorId(), + e); + if (bundleProcessor != null) { + // Make sure we clean up from the active set of bundle processors. + bundleProcessorCache.discard(bundleProcessor); + } + // Ensure that if more data arrives for the instruction it is discarded. + beamFnDataClient.poisonInstructionId(instructionId); + throw e; + } + } + + private static PCollectionConsumerRegistry.FusedDoFnExceptionHandling + getFusedDoFnExceptionHandling(PipelineOptions options) { + String value = + ExperimentalOptions.getExperimentValue( + options, ExperimentalOptions.FUSED_DOFN_EXCEPTION_HANDLING); + if (value == null || value.equalsIgnoreCase("legacy")) { + return PCollectionConsumerRegistry.FusedDoFnExceptionHandling.LEGACY; + } + if (value.equalsIgnoreCase("warn")) { + return PCollectionConsumerRegistry.FusedDoFnExceptionHandling.WARN; + } + if (value.equalsIgnoreCase("fail")) { + return PCollectionConsumerRegistry.FusedDoFnExceptionHandling.FAIL; } + throw new IllegalArgumentException( + String.format( + "Unsupported value '%s' for experiment %s. Expected legacy, warn, or fail.", + value, ExperimentalOptions.FUSED_DOFN_EXCEPTION_HANDLING)); } private void embedOutboundElementsIfApplicable( @@ -803,13 +851,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 +868,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 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..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 @@ -24,6 +24,8 @@ 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; import org.apache.beam.fn.harness.control.BundleProgressReporter; @@ -61,8 +63,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) @@ -70,6 +74,18 @@ }) 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<>(); + private final Map downstreamExceptionsByTransform = + new ConcurrentHashMap<>(); + /** Stores metadata about each consumer so that the appropriate metrics tracking can occur. */ @AutoValue @AutoValue.CopyAnnotations @@ -121,7 +137,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 +152,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 +180,7 @@ public PCollectionConsumerRegistry( .putAllWindowingStrategies(processBundleDescriptor.getWindowingStrategiesMap()) .build()); this.dataSampler = dataSampler; + this.fusedDoFnExceptionHandling = fusedDoFnExceptionHandling; } /** @@ -274,6 +313,75 @@ 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, Throwable exception) { + if (fusedDoFnExceptionHandling != FusedDoFnExceptionHandling.LEGACY + && !WorkCancelledException.isWorkCancelledException(exception)) { + firstDownstreamException.compareAndSet(null, exception); + downstreamExceptionsByTransform.putIfAbsent(ptransformId, 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 { + Throwable failure = getDownstreamFailure(ptransformId); + if (failure != null) { + rethrowDownstreamFailure(failure); + } + } + + /** + * 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(); + if (fusedDoFnExceptionHandling == FusedDoFnExceptionHandling.FAIL && failure != null) { + rethrowDownstreamFailure(failure); + } + } + + /** + * Logs a swallowed downstream exception only after the rest of bundle processing has succeeded. + */ + public void logWarningIfDownstreamFailed() { + 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.", + failure); + } + } + + /** 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( Exception e, ExecutionState executionState, @@ -360,6 +468,7 @@ public MetricTrackingFnDataReceiver( @Override public void accept(WindowedValue input) throws Exception { + 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 @@ -377,8 +486,12 @@ public void accept(WindowedValue input) throws Exception { try (ExecutionState.ActiveState a = executionState.scopedActivate()) { this.delegate.accept(input); } catch (Exception e) { + recordDownstreamException(ptransformId, e); logAndRethrow( e, executionState, executionStateTracker, ptransformId, outputSampler, elementSample); + } catch (Error e) { + recordDownstreamException(ptransformId, e); + throw e; } this.sampledByteSizeDistribution.finishLazyUpdate(); } @@ -443,37 +556,84 @@ 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; + // 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(); ExecutionState state = consumerAndMetadata.getExecutionState(); + + // 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; + } + + // 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) { - logAndRethrow( - e, - state, - consumerAndMetadata.getExecutionStateTracker(), - consumerAndMetadata.getPTransformId(), - outputSampler, - elementSample); + recordDownstreamException(ptransformId, e); + if (fusedDoFnExceptionHandling != FusedDoFnExceptionHandling.FAIL + || WorkCancelledException.isWorkCancelledException(e)) { + 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; + } } + } + if (metricsStarted) { this.sampledByteSizeDistribution.finishLazyUpdate(); } + if (firstFailure != null) { + rethrowDownstreamFailure(firstFailure); + } } } 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..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 @@ -123,6 +123,8 @@ 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.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.state.TimeDomain; import org.apache.beam.sdk.state.TimerMap; @@ -147,6 +149,8 @@ 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.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; @@ -935,6 +939,291 @@ 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 AtomicInteger FUSED_DOWNSTREAM_FINISH_INVOCATIONS = new AtomicInteger(); + + @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 { + 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"); + } + + @FinishBundle + public void finishBundle() { + FUSED_DOWNSTREAM_FINISH_INVOCATIONS.incrementAndGet(); + } + } + + 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("mainOutput", 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(), + (String input) -> + context + .>getPCollectionConsumer("pc-source") + .accept(WindowedValues.valueInGlobalWindow(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 { + ProcessBundleHandler handler = setupHandlerForFusedExceptionTest("fail"); + 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()); + assertEquals( + "A downstream failure during @ProcessElement must prevent @FinishBundle.", + 0, + FUSED_DOWNSTREAM_FINISH_INVOCATIONS.get()); + assertThat( + handler.bundleProcessorCache.getCachedBundleProcessors().get("fused-descriptor"), + empty()); + } finally { + handler.shutdown(); + } + } + + @Test + public void fusedDownstreamExceptionWarnModeAllowsBundleToComplete() throws Exception { + 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()); + assertEquals(1, FUSED_DOWNSTREAM_FINISH_INVOCATIONS.get()); + assertFalse( + handler + .bundleProcessorCache + .getCachedBundleProcessors() + .get("fused-descriptor") + .isEmpty()); + } finally { + handler.shutdown(); + } + } + + @Test + public void fusedDownstreamExceptionPreservesLegacyBehaviorByDefault() throws Exception { + ProcessBundleHandler handler = setupHandlerForFusedExceptionTest(null); + 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()); + assertEquals(1, FUSED_DOWNSTREAM_FINISH_INVOCATIONS.get()); + } finally { + 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"; 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..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 @@ -22,11 +22,13 @@ 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; 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; @@ -221,6 +223,170 @@ 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()); + + 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. + Exception secondFailure = + assertThrows(Exception.class, () -> wrapperConsumer.accept(valueInGlobalWindow("second"))); + assertEquals(message, secondFailure.getMessage()); + verify(consumer, times(1)).accept(any()); + + Exception bundleFailure = + assertThrows(Exception.class, consumers::failBundleIfDownstreamFailed); + 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"; + 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()); + + 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()); + + // Warning mode reports at successful bundle completion but does not fail the bundle. + consumers.failBundleIfDownstreamFailed(); + consumers.logWarningIfDownstreamFailed(); + } + + @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()); + + 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")); + consumers.failBundleIfDownstreamFailed(); + consumers.logWarningIfDownstreamFailed(); + 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 {