[Java fn-execution] Drain bundle output before waiting for the ProcessBundleResponse

CHANGES.md

@@ -133,6 +133,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

runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java

@@ -516,22 +516,58 @@ 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();
 
-            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 +582,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 +604,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;

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,186 @@ public void testBundleFinalizationCallback() throws Exception {
     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) {