This is an automated email from the ASF dual-hosted git repository.
scwhittle pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new e0790a97367 [Java fn-execution] Drain bundle output before waiting for
the ProcessBundleResponse (#40721)
e0790a97367 is described below
commit e0790a9736763469e3fdb05689702d7449d90814
Author: Pablo Villanueva <[email protected]>
AuthorDate: Tue Oct 27 15:48:28 2026 +0100
[Java fn-execution] Drain bundle output before waiting for the
ProcessBundleResponse (#40721)
* [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.
---
CHANGES.md | 1 +
.../fnexecution/control/SdkHarnessClient.java | 89 +++++++---
.../fnexecution/control/SdkHarnessClientTest.java | 182 +++++++++++++++++++++
3 files changed, 247 insertions(+), 25 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 04ac961fcd9..aa9f5aadb04 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -139,6 +139,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)).
* (Python) DataFrame API: `drop_duplicates`, `duplicated` and `nunique` no
longer fail on rows with NaN values or on categorical columns with unused
categories, and `nunique(dropna=False)` now counts NaN
([#40570](https://github.com/apache/beam/issues/40570)).
* (Java) BigQuery Storage Write API `TableRow` conversion now honors
`allowMissingRequiredFields` when building protocol-buffer messages, so callers
that pass `true` with required descriptors get partial messages instead of a
conversion-time `UninitializedMessageException`
([#40189](https://github.com/apache/beam/issues/40189)).
* (Go) Fixed generic registration (for example `register.Function1x1` or
`register.DoFn1x1`) panicking on functions and methods that take or return an
anonymous struct. Unnamed struct types are now only registered with the schema
registry, since they have no type key for the runtime type registry
([#21789](https://github.com/apache/beam/issues/21789)).
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 704d298a195..b540c1dba76 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,22 +516,68 @@ public class SdkHarnessClient implements AutoCloseable {
} 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.
+ 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);
+ }
+ });
+ try {
+ inboundObserver.awaitCompletion();
+ fnApiDataService.unregisterReceiver(bundleId);
+ } catch (Exception e) {
+ 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) {
BeamFnApi.ProcessBundleResponse completedResponse =
MoreFutures.get(response);
- if (completedResponse.hasElements()) {
-
beamFnDataInboundObserver.get().accept(completedResponse.getElements());
- }
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
@@ -546,6 +592,15 @@ public class SdkHarnessClient implements AutoCloseable {
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,22 +614,6 @@ public class SdkHarnessClient implements AutoCloseable {
exception.addSuppressed(e);
}
}
- if (beamFnDataInboundObserver.isPresent()) {
- try {
- if (exception == null) {
- beamFnDataInboundObserver.get().awaitCompletion();
- fnApiDataService.unregisterReceiver(bundleId);
- } else {
- beamFnDataInboundObserver.get().close();
- }
- } catch (Exception e) {
- if (exception == null) {
- exception = e;
- } else {
- exception.addSuppressed(e);
- }
- }
- }
if (exception != null) {
throw exception;
}
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 7b19376426c..5dc6afe2808 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.Collections;
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,186 @@ public class SdkHarnessClientTest {
verifyNoMoreInteractions(mockCheckpointHandler, mockSplitHandler);
}
+ /** Output of {@code count} byte arrays for {@code bundleId}, one per data
message. */
+ private static List<Elements> outputMessages(String bundleId, int count)
throws Exception {
+ List<Elements> 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<InstructionResponse> processBundleResponseFuture = new
CompletableFuture<>();
+ when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
+ .thenReturn(processBundleResponseFuture);
+ FullWindowedValueCoder<String> 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<byte[]> 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<InstructionResponse> processBundleResponseFuture = new
CompletableFuture<>();
+ when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
+ .thenReturn(processBundleResponseFuture);
+ FullWindowedValueCoder<String> 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<byte[]> 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());
+ }
+
+ @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<InstructionResponse> processBundleResponseFuture = new
CompletableFuture<>();
+ when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
+ .thenReturn(processBundleResponseFuture);
+ FullWindowedValueCoder<String> 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<String, String> {
@ProcessElement
public void processElement(ProcessContext context) {
