prodriguezdefino commented on code in PR #40606:
URL: https://github.com/apache/beam/pull/40606#discussion_r4300042616
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java:
##########
@@ -52,23 +54,150 @@
* location with a specific outbound observer.
*/
public class BeamFnDataGrpcMultiplexer implements AutoCloseable {
+ // The inbound stream uses manual gRPC flow control, so onNext never blocks.
A blocked onNext
+ // also blocks the onReady callback of the outbound stream, because gRPC
runs both on one
+ // executor.
+ //
+ // Number of inbound messages that the stream asks for when no output waits.
+ private static final int STEADY_STATE_INBOUND_WINDOW = 16;
+ // Largest number of inbound messages for each active bundle while output
waits.
+ private static final int MAX_INBOUND_WINDOW = 100;
+ // Inbound messages that the stream asked for and that no bundle has
finished.
+ private final AtomicInteger inFlightInbound = new
AtomicInteger(STEADY_STATE_INBOUND_WINDOW);
+
+ /**
+ * Asks for one more inbound message while output waits on flow control, up
to {@code
+ * MAX_INBOUND_WINDOW} for each active bundle. The extra message lets a
bundle that waits for
+ * input continue, so the runner can read the output.
+ */
+ public void onStalledPoll() {
+ synchronized (inboundObserver) {
+ int maxWindow = Math.max(1, receivers.size()) * MAX_INBOUND_WINDOW;
+ if (inFlightInbound.get() < maxWindow) {
+ inFlightInbound.incrementAndGet();
+ inboundObserver.request(1);
+ }
+ }
+ }
+
private static final Logger LOG =
LoggerFactory.getLogger(BeamFnDataGrpcMultiplexer.class);
private static final Duration POISONED_INSTRUCTION_ID_CACHE_TIMEOUT =
Duration.ofMinutes(20);
private final Endpoints.@Nullable ApiServiceDescriptor apiServiceDescriptor;
- private final StreamObserver<BeamFnApi.Elements> inboundObserver;
+ private final InboundObserver inboundObserver;
private final StreamObserver<BeamFnApi.Elements> outboundObserver;
- private final ConcurrentHashMap<
- /* instructionId= */ String,
- CompletableFuture<CloseableFnDataReceiver<BeamFnApi.Elements>>>
- receivers;
- private final Cache</* instructionId= */ String, /* unused= */ Boolean>
poisonedInstructionIds;
- private static class PoisonedException extends RuntimeException {
- private PoisonedException() {
- super("Instruction poisoned");
+ private static final class PendingElement {
+ BeamFnApi.Elements elements;
+ Runnable callback;
+
+ PendingElement(BeamFnApi.Elements elements, Runnable callback) {
+ this.elements = elements;
+ this.callback = callback;
+ }
+ }
+
+ private static final class ReceiverHolder {
+ @Nullable CloseableFnDataReceiver<BeamFnApi.Elements> receiver;
+
+ ArrayDeque<PendingElement> pending = new ArrayDeque<>();
+ boolean poisoned;
+
+ private Runnable wrapCallback(Runnable callback) {
+ return () -> {
+ drainPending();
+ callback.run();
+ };
+ }
+
+ synchronized void drainPending() {
+ if (poisoned || !(receiver instanceof BeamFnDataInboundObserver)) {
+ return;
+ }
+ BeamFnDataInboundObserver inbound = (BeamFnDataInboundObserver) receiver;
+ while (!pending.isEmpty()) {
+ PendingElement next = pending.peekFirst();
+ if (next == null) {
+ break;
+ }
+ try {
+ if (inbound.tryAcceptWithCallback(next.elements,
wrapCallback(next.callback))) {
+ pending.removeFirst();
+ } else {
+ break;
+ }
+ } catch (Exception e) {
+ pending.removeFirst();
+ next.callback.run();
+ }
+ }
+ }
+
+ synchronized boolean accept(BeamFnApi.Elements elements, Runnable
callback) throws Exception {
+ if (poisoned) {
+ return false;
+ }
+ if (receiver != null) {
+ if (receiver instanceof BeamFnDataInboundObserver) {
+ BeamFnDataInboundObserver inbound = (BeamFnDataInboundObserver)
receiver;
+ drainPending();
+ if (!pending.isEmpty()
+ || !inbound.tryAcceptWithCallback(elements,
wrapCallback(callback))) {
+ pending.addLast(new PendingElement(elements, callback));
+ }
+ } else {
+ receiver.accept(elements);
+ callback.run();
+ }
+ return true;
+ } else {
+ pending.addLast(new PendingElement(elements, callback));
+ return true;
+ }
+ }
+
+ synchronized boolean complete(CloseableFnDataReceiver<BeamFnApi.Elements>
receiver)
+ throws Exception {
+ if (this.receiver != null) {
+ return false;
+ }
+ this.receiver = receiver;
+ if (receiver instanceof BeamFnDataInboundObserver) {
+ drainPending();
Review Comment:
Fixed in b7532bb. `complete()` runs the callback of every buffered element,
including the one the receiver failed on, and clears the buffer even if the
receiver throws. `registerConsumer` rethrows runtime exceptions as they are, so
duplicate registration throws `IllegalArgumentException` directly again.
Covered by `testPermitsReleasedWhenReceiverThrowsOnBufferedElements` and the
duplicate-registration check in
`testReceiverHolderBuffersBeforeConsumerRegisteredAndDrains`.
--
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]
