scwhittle commented on code in PR #40721:
URL: https://github.com/apache/beam/pull/40721#discussion_r4376149724
##########
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:
can we remove inboundDrained variable? Seems like we can close the
inboundObserver unconditionally as it is idempotent.
It might also be clearer to move to closing it on line 554 so it is closer
to where we are either awaiting for input or not.
--
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]
