PabloPauling commented on code in PR #40721:
URL: https://github.com/apache/beam/pull/40721#discussion_r4376836165
##########
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;
Review Comment:
Done in 900eb56. Removed inboundDrained and now close the inbound observer
unconditionally right after the drain, since close() is idempotent. This also
covers the case where sending the input failed, so the close at the end is gone.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
