From 9fce4f2fb0ecc9f4d85679ef060d4decdf316259 Mon Sep 17 00:00:00 2001 From: pabloPauling Date: Thu, 22 Oct 2026 10:36:51 +0200 Subject: [PATCH 1/3] [Java fn-execution] Drain bundle output before waiting for the ProcessBundleResponse ActiveBundle.close() waited for the ProcessBundleResponse and only then consumed the bundle's output. The inbound queue holds 100 messages, and since #38407 the Python SDK's send queue is bounded too (100 MB). The SDK does not complete a bundle until its output has been sent, so a bundle whose output exceeds both buffers deadlocks. On Flink batch one bundle is a whole partition, so this triggers at about 1 GB of output per subtask. close() now drains the inbound observer before waiting for the response. Elements embedded in the response are delivered from an async callback, and a failed response closes the observer so the drain cannot wait for an is_last that will never arrive. The response's error stays primary. Fixes #40720. Co-Authored-By: Claude Opus 5.5 --- CHANGES.md | 1 + .../fnexecution/control/SdkHarnessClient.java | 57 ++++++-- .../control/SdkHarnessClientTest.java | 126 ++++++++++++++++++ 3 files changed, 174 insertions(+), 10 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 4e2fffb89380..8477b6a142bf 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -132,6 +132,7 @@ * (Python) Setting `max_batch_weight` on a RunInference `ModelHandler` no longer fails with a `TypeError`: `BatchElements` now supports `max_batch_weight`, which limits the total `element_size_fn` of a batch in addition to its number of elements ([#40644](https://github.com/apache/beam/issues/40644)). * (Python) `BatchElements` now honors `element_size_fn` when `max_batch_duration_secs` is set, and no longer emits an empty batch before an element that exceeds the batch size on its own ([#40644](https://github.com/apache/beam/issues/40644)). * (Python) `KeyedModelHandler` and `MaybeKeyedModelHandler` now apply the `element_size_fn` and `batch_length_fn` of the wrapped `ModelHandler` to examples, instead of to `(key, example)` pairs ([#40644](https://github.com/apache/beam/issues/40644)). +* (Java) Portable runners no longer deadlock when a single bundle produces more output than the Fn API data buffers hold (about 1 GB with the Python SDK, e.g. one Flink batch partition). The runner now consumes a bundle's output before waiting for its `ProcessBundleResponse` ([#40720](https://github.com/apache/beam/issues/40720)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java index 704d298a195d..1995ebf01020 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java @@ -516,13 +516,45 @@ public void close() throws Exception { } finally { beamFnDataOutboundAggregator.discard(); } + // Consume the bundle's output BEFORE waiting for the ProcessBundleResponse. The inbound + // queue is bounded, and the SDK does not finish the bundle until its output has been sent. + // Waiting for the response first therefore deadlocks once a single bundle produces more + // output than the inbound queue and the SDK's own send buffer hold together: the SDK + // blocks sending, the response never arrives, and nothing drains the queue. Batch runners + // that hand a whole partition to the SDK as one bundle hit this with ~1 GB of output. + boolean inboundDrained = false; + Exception drainException = null; + if (exception == null && beamFnDataInboundObserver.isPresent()) { + BeamFnDataInboundObserver inboundObserver = beamFnDataInboundObserver.get(); + // Async: when the response is already complete, a synchronous callback would run on + // this thread and could block putting elements into the queue this thread must drain. + response.whenCompleteAsync( + (completedResponse, failure) -> { + try { + if (failure != null) { + // The SDK will not send the rest of the output; unblock the drain below. + inboundObserver.close(); + } else if (completedResponse.hasElements()) { + inboundObserver.accept(completedResponse.getElements()); + } + } catch (Exception e) { + // Only possible once the drain has failed and cancelled the queue, and that + // failure is reported by close(). + LOG.debug("Dropping the response elements of failed bundle {}", bundleId, e); + } + }); + inboundDrained = true; + try { + inboundObserver.awaitCompletion(); + fnApiDataService.unregisterReceiver(bundleId); + } catch (Exception e) { + drainException = e; + } + } try { // We don't have to worry about the completion stage. if (exception == null) { BeamFnApi.ProcessBundleResponse completedResponse = MoreFutures.get(response); - if (completedResponse.hasElements()) { - beamFnDataInboundObserver.get().accept(completedResponse.getElements()); - } outstandingRequests.arriveAndAwaitAdvance(); @@ -546,6 +578,15 @@ public void close() throws Exception { exception.addSuppressed(e); } } + // After the response, so that a failed bundle reports the SDK's error rather than the + // cancelled queue it caused. + if (drainException != null) { + if (exception == null) { + exception = drainException; + } else { + exception.addSuppressed(drainException); + } + } try { if (exception == null) { stateRegistration.deregister(); @@ -559,14 +600,10 @@ public void close() throws Exception { exception.addSuppressed(e); } } - if (beamFnDataInboundObserver.isPresent()) { + if (beamFnDataInboundObserver.isPresent() && !inboundDrained) { + // Sending the bundle's input failed, so its output was never consumed. try { - if (exception == null) { - beamFnDataInboundObserver.get().awaitCompletion(); - fnApiDataService.unregisterReceiver(bundleId); - } else { - beamFnDataInboundObserver.get().close(); - } + beamFnDataInboundObserver.get().close(); } catch (Exception e) { if (exception == null) { exception = e; diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java index 7b19376426cc..2daff006730a 100644 --- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java +++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java @@ -43,7 +43,9 @@ import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import org.apache.beam.model.fnexecution.v1.BeamFnApi; @@ -965,6 +967,130 @@ public void testBundleFinalizationCallback() throws Exception { verifyNoMoreInteractions(mockCheckpointHandler, mockSplitHandler); } + /** Output of {@code count} byte arrays for {@code bundleId}, one per data message. */ + private static List outputMessages(String bundleId, int count) throws Exception { + List messages = new ArrayList<>(); + for (int i = 0; i < count; i++) { + ByteString.Output data = ByteString.newOutput(); + ByteArrayCoder.of().encode(new byte[] {(byte) i}, data); + messages.add( + Elements.newBuilder() + .addData( + Elements.Data.newBuilder() + .setInstructionId(bundleId) + .setTransformId(SDK_GRPC_WRITE_TRANSFORM) + .setData(data.toByteString())) + .build()); + } + return messages; + } + + private static Elements lastOutputMessage(String bundleId) { + return Elements.newBuilder() + .addData( + Elements.Data.newBuilder() + .setInstructionId(bundleId) + .setTransformId(SDK_GRPC_WRITE_TRANSFORM) + .setIsLast(true)) + .build(); + } + + @Test + public void testOutputLargerThanTheInboundQueueDoesNotDeadlock() throws Exception { + // The SDK only completes a bundle once all of its output has been sent. If close() waits for + // the ProcessBundleResponse before consuming that output, a bundle whose output does not fit + // in the inbound queue (100 messages) never finishes. + CompletableFuture processBundleResponseFuture = new CompletableFuture<>(); + when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))) + .thenReturn(processBundleResponseFuture); + FullWindowedValueCoder coder = + FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE); + BundleProcessor processor = + sdkHarnessClient.getProcessor( + descriptor, + Collections.singletonList( + RemoteInputDestination.of( + (FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM))); + doNothing().when(dataService).registerReceiver(any(), outputReceiverCaptor.capture()); + when(dataService.createOutboundAggregator(any(), anyBoolean())) + .thenReturn(mock(BeamFnDataOutboundAggregator.class)); + + List received = Collections.synchronizedList(new ArrayList<>()); + ActiveBundle activeBundle = + processor.newBundle( + ImmutableMap.of( + SDK_GRPC_WRITE_TRANSFORM, + RemoteOutputReceiver.of(ByteArrayCoder.of(), received::add)), + BundleProgressHandler.ignored()); + String bundleId = activeBundle.getId(); + + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future sdkSendsOutputThenResponse = + executor.submit( + () -> { + for (Elements message : outputMessages(bundleId, 150)) { + outputReceiverCaptor.getValue().accept(message); + } + outputReceiverCaptor.getValue().accept(lastOutputMessage(bundleId)); + processBundleResponseFuture.complete( + InstructionResponse.newBuilder() + .setProcessBundle(ProcessBundleResponse.getDefaultInstance()) + .build()); + return null; + }); + + Future closed = + executor.submit( + () -> { + activeBundle.close(); + return null; + }); + // Without the fix this times out: the SDK blocks on message 101 and close() waits forever. + closed.get(60, TimeUnit.SECONDS); + sdkSendsOutputThenResponse.get(60, TimeUnit.SECONDS); + } finally { + executor.shutdownNow(); + } + assertEquals(150, received.size()); + verify(dataService).unregisterReceiver(activeBundle.getId()); + } + + @Test + public void testOutputEmbeddedInAnAlreadyCompletedResponseIsDelivered() throws Exception { + CompletableFuture processBundleResponseFuture = new CompletableFuture<>(); + when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))) + .thenReturn(processBundleResponseFuture); + FullWindowedValueCoder coder = + FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE); + BundleProcessor processor = + sdkHarnessClient.getProcessor( + descriptor, + Collections.singletonList( + RemoteInputDestination.of( + (FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM))); + when(dataService.createOutboundAggregator(any(), anyBoolean())) + .thenReturn(mock(BeamFnDataOutboundAggregator.class)); + + List received = Collections.synchronizedList(new ArrayList<>()); + try (ActiveBundle activeBundle = + processor.newBundle( + ImmutableMap.of( + SDK_GRPC_WRITE_TRANSFORM, + RemoteOutputReceiver.of(ByteArrayCoder.of(), received::add)), + BundleProgressHandler.ignored())) { + Elements embedded = + outputMessages(activeBundle.getId(), 1).get(0).toBuilder() + .mergeFrom(lastOutputMessage(activeBundle.getId())) + .build(); + processBundleResponseFuture.complete( + InstructionResponse.newBuilder() + .setProcessBundle(ProcessBundleResponse.newBuilder().setElements(embedded)) + .build()); + } + assertEquals(1, received.size()); + } + private static class TestFn extends DoFn { @ProcessElement public void processElement(ProcessContext context) { From 5b2c22305c0e2774735f13c03f04e2f43405e153 Mon Sep 17 00:00:00 2001 From: pabloPauling Date: Thu, 22 Oct 2026 10:45:11 +0200 Subject: [PATCH 2/3] Do not report completion for a bundle whose output drain failed If consuming the bundle's output fails while the SDK completes the ProcessBundleResponse normally, the bundle has still failed. Skip progressHandler.onCompleted, the checkpoint and the finalization request in that case, so a failed bundle does not register residuals or finalization. The state registration is already aborted. Addresses review feedback on #40721. Co-Authored-By: Claude Opus 5.5 --- .../fnexecution/control/SdkHarnessClient.java | 16 ++++-- .../control/SdkHarnessClientTest.java | 56 +++++++++++++++++++ 2 files changed, 66 insertions(+), 6 deletions(-) diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java index 1995ebf01020..da3d98118751 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java @@ -558,12 +558,16 @@ public void close() throws Exception { outstandingRequests.arriveAndAwaitAdvance(); - progressHandler.onCompleted(completedResponse); - if (completedResponse.getResidualRootsCount() > 0) { - checkpointHandler.onCheckpoint(completedResponse); - } - if (completedResponse.getRequiresFinalization()) { - finalizationHandler.requestsFinalization(bundleId); + // A bundle whose output could not be consumed has failed even if the SDK finished + // it, so it must not report completion or register residuals or finalization. + if (drainException == null) { + progressHandler.onCompleted(completedResponse); + if (completedResponse.getResidualRootsCount() > 0) { + checkpointHandler.onCheckpoint(completedResponse); + } + if (completedResponse.getRequiresFinalization()) { + finalizationHandler.requestsFinalization(bundleId); + } } } else { // TODO: [https://github.com/apache/beam/issues/18756] Handle aborting the bundle being diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java index 2daff006730a..5dc6afe2808f 100644 --- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java +++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java @@ -1091,6 +1091,62 @@ public void testOutputEmbeddedInAnAlreadyCompletedResponseIsDelivered() throws E assertEquals(1, received.size()); } + @Test + public void testFailedOutputDrainDoesNotReportCompletion() throws Exception { + // The SDK finishes the bundle normally, but consuming its output fails. The bundle has + // failed, so it must not report completion, register residuals, or request finalization. + CompletableFuture processBundleResponseFuture = new CompletableFuture<>(); + when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))) + .thenReturn(processBundleResponseFuture); + FullWindowedValueCoder coder = + FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE); + BundleProcessor processor = + sdkHarnessClient.getProcessor( + descriptor, + Collections.singletonList( + RemoteInputDestination.of( + (FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM))); + doNothing().when(dataService).registerReceiver(any(), outputReceiverCaptor.capture()); + when(dataService.createOutboundAggregator(any(), anyBoolean())) + .thenReturn(mock(BeamFnDataOutboundAggregator.class)); + + BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class); + BundleCheckpointHandler mockCheckpointHandler = mock(BundleCheckpointHandler.class); + BundleFinalizationHandler mockFinalizationHandler = mock(BundleFinalizationHandler.class); + RuntimeException receiverFailure = new RuntimeException("receiver failed"); + + ActiveBundle activeBundle = + processor.newBundle( + ImmutableMap.of( + SDK_GRPC_WRITE_TRANSFORM, + RemoteOutputReceiver.of( + ByteArrayCoder.of(), + value -> { + throw receiverFailure; + })), + Collections.emptyMap(), + (request) -> { + throw new UnsupportedOperationException(); + }, + mockProgressHandler, + mock(BundleSplitHandler.class), + mockCheckpointHandler, + mockFinalizationHandler); + outputReceiverCaptor.getValue().accept(outputMessages(activeBundle.getId(), 1).get(0)); + processBundleResponseFuture.complete( + InstructionResponse.newBuilder() + .setProcessBundle( + ProcessBundleResponse.newBuilder() + .setRequiresFinalization(true) + .addResidualRoots(DelayedBundleApplication.getDefaultInstance())) + .build()); + + Exception thrown = assertThrows(Exception.class, activeBundle::close); + assertSame(receiverFailure, thrown); + verify(mockProgressHandler, never()).onCompleted(any()); + verifyNoMoreInteractions(mockCheckpointHandler, mockFinalizationHandler); + } + private static class TestFn extends DoFn { @ProcessElement public void processElement(ProcessContext context) { From 900eb563036a596b2f360faa6e4970be3325c140 Mon Sep 17 00:00:00 2001 From: pabloPauling Date: Mon, 26 Oct 2026 12:35:51 +0100 Subject: [PATCH 3/3] Close the inbound observer unconditionally after draining close() is idempotent, so the inboundDrained flag is not needed. Co-Authored-By: Claude Opus 5.5 --- .../fnexecution/control/SdkHarnessClient.java | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java index da3d98118751..b540c1dba760 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java @@ -522,7 +522,6 @@ public void close() throws Exception { // output than the inbound queue and the SDK's own send buffer hold together: the SDK // blocks sending, the response never arrives, and nothing drains the queue. Batch runners // that hand a whole partition to the SDK as one bundle hit this with ~1 GB of output. - boolean inboundDrained = false; Exception drainException = null; if (exception == null && beamFnDataInboundObserver.isPresent()) { BeamFnDataInboundObserver inboundObserver = beamFnDataInboundObserver.get(); @@ -543,7 +542,6 @@ public void close() throws Exception { LOG.debug("Dropping the response elements of failed bundle {}", bundleId, e); } }); - inboundDrained = true; try { inboundObserver.awaitCompletion(); fnApiDataService.unregisterReceiver(bundleId); @@ -551,6 +549,18 @@ public void close() throws Exception { drainException = e; } } + if (beamFnDataInboundObserver.isPresent()) { + // Idempotent. Also releases the observer when sending the bundle's input failed. + try { + beamFnDataInboundObserver.get().close(); + } catch (Exception e) { + if (drainException == null) { + drainException = e; + } else { + drainException.addSuppressed(e); + } + } + } try { // We don't have to worry about the completion stage. if (exception == null) { @@ -604,18 +614,6 @@ public void close() throws Exception { exception.addSuppressed(e); } } - if (beamFnDataInboundObserver.isPresent() && !inboundDrained) { - // Sending the bundle's input failed, so its output was never consumed. - try { - beamFnDataInboundObserver.get().close(); - } catch (Exception e) { - if (exception == null) { - exception = e; - } else { - exception.addSuppressed(e); - } - } - } if (exception != null) { throw exception; }