From 69eb1238187d1e3983aac1a19667ecc3a7a38980 Mon Sep 17 00:00:00 2001 From: Sania Parveen Date: Sun, 27 Sep 2026 20:44:40 +0000 Subject: [PATCH 1/2] Reduce logging severity for state requests cancelled by runner in Java Fn Harness --- .../harness/control/BeamFnControlClient.java | 19 +++-- .../data/PCollectionConsumerRegistry.java | 17 ++-- .../state/BeamFnStateGrpcClientCache.java | 2 + .../harness/state/WorkCancelledException.java | 47 +++++++++++ .../control/BeamFnControlClientTest.java | 22 +++++ .../data/PCollectionConsumerRegistryTest.java | 84 +++++++++++++++++++ .../state/BeamFnStateGrpcClientCacheTest.java | 20 ++++- 7 files changed, 198 insertions(+), 13 deletions(-) create mode 100644 sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/WorkCancelledException.java diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/BeamFnControlClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/BeamFnControlClient.java index 119bdd73d942..74e7d7cf1a41 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/BeamFnControlClient.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/BeamFnControlClient.java @@ -23,6 +23,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import org.apache.beam.fn.harness.logging.BeamFnLoggingMDC; +import org.apache.beam.fn.harness.state.WorkCancelledException; import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.fnexecution.v1.BeamFnControlGrpc; import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor; @@ -151,11 +152,19 @@ public BeamFnApi.InstructionResponse delegateOnInstructionRequestType( .setInstructionId(value.getInstructionId()) .build(); } catch (Exception e) { - LOG.error( - "Exception while trying to handle {} {}", - BeamFnApi.InstructionRequest.class.getSimpleName(), - value.getInstructionId(), - e); + if (WorkCancelledException.isWorkCancelledException(e)) { + LOG.info( + "Instruction {} {} cancelled by runner", + BeamFnApi.InstructionRequest.class.getSimpleName(), + value.getInstructionId(), + e); + } else { + LOG.error( + "Exception while trying to handle {} {}", + BeamFnApi.InstructionRequest.class.getSimpleName(), + value.getInstructionId(), + e); + } return BeamFnApi.InstructionResponse.newBuilder() .setInstructionId(value.getInstructionId()) .setError(getStackTraceAsString(e)) 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 3ba8b4e76c3c..58c06bac310f 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 @@ -36,6 +36,7 @@ import org.apache.beam.fn.harness.debug.DataSampler; import org.apache.beam.fn.harness.debug.ElementSample; import org.apache.beam.fn.harness.debug.OutputSampler; +import org.apache.beam.fn.harness.state.WorkCancelledException; import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor; import org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo; import org.apache.beam.model.pipeline.v1.RunnerApi; @@ -281,14 +282,16 @@ private static void logAndRethrow( @Nullable OutputSampler outputSampler, @Nullable ElementSample elementSample) throws Exception { - ExecutionStateSampler.ExecutionStateTrackerStatus status = executionStateTracker.getStatus(); - String processBundleId = status == null ? null : status.getProcessBundleId(); - if (outputSampler != null) { - outputSampler.exception(elementSample, e, ptransformId, processBundleId); - } + if (!WorkCancelledException.isWorkCancelledException(e)) { + ExecutionStateSampler.ExecutionStateTrackerStatus status = executionStateTracker.getStatus(); + String processBundleId = status == null ? null : status.getProcessBundleId(); + if (outputSampler != null) { + outputSampler.exception(elementSample, e, ptransformId, processBundleId); + } - if (executionState.error()) { - LOG.error("Failed to process element for bundle \"{}\"", processBundleId, e); + if (executionState.error()) { + LOG.error("Failed to process element for bundle \"{}\"", processBundleId, e); + } } throw e; } diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java index 453dd31e0d42..c9adb9d7d142 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java @@ -187,6 +187,8 @@ public void onNext(StateResponse value) { } if (value.getError().isEmpty()) { responseFuture.complete(value); + } else if (value.getErrorReason() == StateResponse.ErrorReason.CANCELLED) { + responseFuture.completeExceptionally(new WorkCancelledException(value.getError())); } else { responseFuture.completeExceptionally(new IllegalStateException(value.getError())); } diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/WorkCancelledException.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/WorkCancelledException.java new file mode 100644 index 000000000000..401416abaa82 --- /dev/null +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/WorkCancelledException.java @@ -0,0 +1,47 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.fn.harness.state; + +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Indicates that the work item is no longer valid on the runner and should be cancelled without + * logging an error. + */ +public class WorkCancelledException extends IllegalStateException { + + public WorkCancelledException(String message) { + super(message); + } + + public WorkCancelledException(Throwable cause) { + super(cause); + } + + /** Returns whether an exception was caused by a {@link WorkCancelledException}. */ + public static boolean isWorkCancelledException(@Nullable Throwable t) { + @Nullable Throwable throwable = t; + while (throwable != null) { + if (throwable instanceof WorkCancelledException) { + return true; + } + throwable = throwable.getCause(); + } + return false; + } +} diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/BeamFnControlClientTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/BeamFnControlClientTest.java index 5afd2c475f67..c54c35c46e6a 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/BeamFnControlClientTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/BeamFnControlClientTest.java @@ -88,6 +88,18 @@ public class BeamFnControlClientTest { .setInstructionId("3L") .setError(getStackTraceAsString(FAILURE)) .build(); + private static final org.apache.beam.fn.harness.state.WorkCancelledException CANCELLED = + new org.apache.beam.fn.harness.state.WorkCancelledException("Work item cancelled"); + private static final BeamFnApi.InstructionRequest CANCELLED_REQUEST = + BeamFnApi.InstructionRequest.newBuilder() + .setInstructionId("4L") + .setSampleData(BeamFnApi.SampleDataRequest.getDefaultInstance()) + .build(); + private static final BeamFnApi.InstructionResponse CANCELLED_RESPONSE = + BeamFnApi.InstructionResponse.newBuilder() + .setInstructionId("4L") + .setError(getStackTraceAsString(CANCELLED)) + .build(); @Rule public TestRule restoreMDCAfterTest = new RestoreBeamFnLoggingMDC(); @@ -137,6 +149,12 @@ public StreamObserver control( assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId()); throw FAILURE; }); + handlers.put( + BeamFnApi.InstructionRequest.RequestCase.SAMPLE_DATA, + value -> { + assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId()); + throw CANCELLED; + }); ExecutorService executor = Executors.newCachedThreadPool(); BeamFnControlClient client = @@ -163,6 +181,10 @@ public StreamObserver control( outboundServerObserver.onNext(FAILURE_REQUEST); assertEquals(FAILURE_RESPONSE, values.take()); + // Ensure that WorkCancelledException is caught and translated to a failure response + outboundServerObserver.onNext(CANCELLED_REQUEST); + assertEquals(CANCELLED_RESPONSE, values.take()); + // Ensure that the server completing the stream translates to the completable future // being completed allowing for a successful shutdown of the client. outboundServerObserver.onCompleted(); 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 7ba2f8921be3..2a7ec115b3cb 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 @@ -703,6 +703,90 @@ public StreamObserver logging( } } + @Test + public void doesNotLogWorkCancelledException() throws Exception { + final String pTransformId = "pTransformId"; + final String message = "Work item cancelled"; + final String instructionId = "instruction"; + final Exception thrownException = + new RuntimeException(new org.apache.beam.fn.harness.state.WorkCancelledException(message)); + + AtomicBoolean clientClosedStream = new AtomicBoolean(); + Collection values = new ConcurrentLinkedQueue<>(); + AtomicReference> outboundServerObserver = + new AtomicReference<>(); + CallStreamObserver inboundServerObserver = + TestStreams.withOnNext( + (BeamFnApi.LogEntry.List logEntries) -> + values.addAll(logEntries.getLogEntriesList())) + .withOnCompleted( + () -> { + clientClosedStream.set(true); + outboundServerObserver.get().onCompleted(); + }) + .build(); + + Endpoints.ApiServiceDescriptor apiServiceDescriptor = + Endpoints.ApiServiceDescriptor.newBuilder() + .setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()) + .build(); + Server server = + InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()) + .addService( + new BeamFnLoggingGrpc.BeamFnLoggingImplBase() { + @Override + public StreamObserver logging( + StreamObserver outboundObserver) { + outboundServerObserver.set(outboundObserver); + return inboundServerObserver; + } + }) + .build(); + server.start(); + ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build(); + + ExecutionStateSampler sampler = + new ExecutionStateSampler(PipelineOptionsFactory.create(), System::currentTimeMillis, null); + ExecutionStateSampler.ExecutionStateTracker stateTracker = sampler.create(); + stateTracker.start("process-bundle"); + ExecutionStateSampler.ExecutionState state = + stateTracker.create("shortId", pTransformId, pTransformId, "process"); + state.activate(); + + BeamFnLoggingMDC.setInstructionId(instructionId); + BeamFnLoggingMDC.setStateTracker(stateTracker); + + try (LoggingClient ignored = + LoggingClientFactory.createAndStart( + PipelineOptionsFactory.create(), + apiServiceDescriptor, + (Endpoints.ApiServiceDescriptor descriptor) -> channel)) { + + ShortIdMap shortIds = new ShortIdMap(); + BundleProgressReporter.InMemory reporterAndRegistrar = new BundleProgressReporter.InMemory(); + PCollectionConsumerRegistry consumers = + new PCollectionConsumerRegistry( + stateTracker, shortIds, reporterAndRegistrar, TEST_DESCRIPTOR); + FnDataReceiver> consumer = mock(FnDataReceiver.class); + + consumers.register(P_COLLECTION_A, pTransformId, pTransformId + "Name", consumer); + + FnDataReceiver> wrapperConsumer = + (FnDataReceiver>) + (FnDataReceiver) consumers.getMultiplexingConsumer(P_COLLECTION_A); + + doThrow(thrownException).when(consumer).accept(any()); + expectedException.expectMessage(message); + expectedException.expect(Exception.class); + + wrapperConsumer.accept(valueInGlobalWindow("elem")); + + } finally { + assertTrue(values.isEmpty()); + server.shutdownNow(); + } + } + private static class TestElementByteSizeObservableIterable extends ElementByteSizeObservableIterable> { private List elements; diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java index 2d7de3489b28..77e9ca1b6a98 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java @@ -148,13 +148,16 @@ public void testRequestResponses() throws Exception { client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS)); CompletableFuture unsuccessfulResponse = client.handle(StateRequest.newBuilder().setInstructionId(FAIL)); + CompletableFuture cancelledResponse = + client.handle(StateRequest.newBuilder().setInstructionId("CANCELLED")); // Wait for the client to connect. StreamObserver outboundServerObserver = outboundServerObservers.take(); // Ensure the client doesn't break when sent garbage. outboundServerObserver.onNext(StateResponse.newBuilder().setId("UNKNOWN ID").build()); - // We expect to receive and handle two requests + // We expect to receive and handle three requests + handleServerRequest(outboundServerObserver, values.take()); handleServerRequest(outboundServerObserver, values.take()); handleServerRequest(outboundServerObserver, values.take()); @@ -166,6 +169,13 @@ public void testRequestResponses() throws Exception { } catch (ExecutionException e) { assertThat(e.toString(), containsString(TEST_ERROR)); } + try { + cancelledResponse.get(); + fail("Expected cancelled response"); + } catch (ExecutionException e) { + assertThat(e.toString(), containsString(TEST_ERROR)); + org.junit.Assert.assertTrue(WorkCancelledException.isWorkCancelledException(e)); + } } @Test @@ -237,6 +247,14 @@ private void handleServerRequest( outboundObserver.onNext( StateResponse.newBuilder().setId(value.getId()).setError(TEST_ERROR).build()); return; + case "CANCELLED": + outboundObserver.onNext( + StateResponse.newBuilder() + .setId(value.getId()) + .setError(TEST_ERROR) + .setErrorReason(StateResponse.ErrorReason.CANCELLED) + .build()); + return; default: outboundObserver.onNext(StateResponse.newBuilder().setId(value.getId()).build()); return; From 60eda2cbc704c379fd5b44e539b98fda00cd6ae1 Mon Sep 17 00:00:00 2001 From: Sania Parveen Date: Thu, 1 Oct 2026 05:35:19 +0000 Subject: [PATCH 2/2] Address PR review comments --- .../fn/harness/state/WorkCancelledException.java | 15 ++++++--------- .../data/PCollectionConsumerRegistryTest.java | 8 ++++---- .../state/BeamFnStateGrpcClientCacheTest.java | 3 ++- 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/WorkCancelledException.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/WorkCancelledException.java index 401416abaa82..587f5962f03a 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/WorkCancelledException.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/WorkCancelledException.java @@ -17,13 +17,15 @@ */ package org.apache.beam.fn.harness.state; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.checkerframework.checker.nullness.qual.Nullable; /** * Indicates that the work item is no longer valid on the runner and should be cancelled without * logging an error. */ -public class WorkCancelledException extends IllegalStateException { +public class WorkCancelledException extends RuntimeException { public WorkCancelledException(String message) { super(message); @@ -35,13 +37,8 @@ public WorkCancelledException(Throwable cause) { /** Returns whether an exception was caused by a {@link WorkCancelledException}. */ public static boolean isWorkCancelledException(@Nullable Throwable t) { - @Nullable Throwable throwable = t; - while (throwable != null) { - if (throwable instanceof WorkCancelledException) { - return true; - } - throwable = throwable.getCause(); - } - return false; + return t != null + && !Iterables.isEmpty( + Iterables.filter(Throwables.getCausalChain(t), WorkCancelledException.class)); } } 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 2a7ec115b3cb..804fd695e342 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 @@ -20,6 +20,7 @@ import static org.apache.beam.sdk.values.WindowedValues.valueInGlobalWindow; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.hamcrest.Matchers.is; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; @@ -50,6 +51,7 @@ import org.apache.beam.fn.harness.logging.BeamFnLoggingMDC; import org.apache.beam.fn.harness.logging.LoggingClient; import org.apache.beam.fn.harness.logging.LoggingClientFactory; +import org.apache.beam.fn.harness.state.WorkCancelledException; import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor; import org.apache.beam.model.fnexecution.v1.BeamFnLoggingGrpc; @@ -708,8 +710,7 @@ public void doesNotLogWorkCancelledException() throws Exception { final String pTransformId = "pTransformId"; final String message = "Work item cancelled"; final String instructionId = "instruction"; - final Exception thrownException = - new RuntimeException(new org.apache.beam.fn.harness.state.WorkCancelledException(message)); + final Exception thrownException = new RuntimeException(new WorkCancelledException(message)); AtomicBoolean clientClosedStream = new AtomicBoolean(); Collection values = new ConcurrentLinkedQueue<>(); @@ -776,8 +777,7 @@ public StreamObserver logging( (FnDataReceiver) consumers.getMultiplexingConsumer(P_COLLECTION_A); doThrow(thrownException).when(consumer).accept(any()); - expectedException.expectMessage(message); - expectedException.expect(Exception.class); + expectedException.expect(is(thrownException)); wrapperConsumer.accept(valueInGlobalWindow("elem")); diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java index 77e9ca1b6a98..6360c4595116 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java @@ -22,6 +22,7 @@ import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNotSame; import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import java.util.UUID; @@ -174,7 +175,7 @@ public void testRequestResponses() throws Exception { fail("Expected cancelled response"); } catch (ExecutionException e) { assertThat(e.toString(), containsString(TEST_ERROR)); - org.junit.Assert.assertTrue(WorkCancelledException.isWorkCancelledException(e)); + assertTrue(WorkCancelledException.isWorkCancelledException(e)); } }