From 9d92655570df82eb65566374b3f7a74273776018 Mon Sep 17 00:00:00 2001 From: Pablo Rodriguez Defino Date: Tue, 13 Oct 2026 01:08:21 -0700 Subject: [PATCH 1/5] Prevent a Fn data plane deadlock in BeamFnDataGrpcMultiplexer When a bundle queue was full, onNext of the inbound data stream blocked. gRPC runs onNext and the onReady callback of the outbound stream on the same executor. So the outbound stream could not become ready again, and the bundle, the gRPC thread and the runner waited for each other. The inbound stream uses manual gRPC flow control and gives messages to the bundle queue without blocking. A message that does not fit waits in a buffer. The stream asks for one more message each time a bundle finishes one. When output waits 500 ms on flow control, the stream asks for one extra message, up to 100 for each active bundle. Each bundle queue keeps its limit of 100 messages. DirectStreamObserver checks isReady before each message and checks again every 50 ms while it waits. --- .../apache/beam/sdk/fn/CancellableQueue.java | 52 +- .../fn/data/BeamFnDataGrpcMultiplexer.java | 337 ++++++++---- .../fn/data/BeamFnDataInboundObserver.java | 36 +- .../fn/stream/BufferingStreamObserver.java | 46 +- .../sdk/fn/stream/DirectStreamObserver.java | 49 +- .../ForwardingClientResponseObserver.java | 15 +- .../fn/stream/OutboundObserverFactory.java | 37 +- .../sdk/fn/stream/StalledPollHandler.java | 26 + .../beam/sdk/fn/CancellableQueueTest.java | 108 ++++ .../data/BeamFnDataGrpcMultiplexerTest.java | 482 ++++++++++++++++++ .../data/BeamFnDataInboundObserverTest.java | 80 +++ .../stream/BufferingStreamObserverTest.java | 29 ++ .../fn/stream/DirectStreamObserverTest.java | 39 +- .../ForwardingClientResponseObserverTest.java | 34 ++ .../stream/OutboundObserverFactoryTest.java | 44 ++ 15 files changed, 1280 insertions(+), 134 deletions(-) create mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/StalledPollHandler.java diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/CancellableQueue.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/CancellableQueue.java index e67423a8a044..d65356b75372 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/CancellableQueue.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/CancellableQueue.java @@ -19,9 +19,13 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Consumer; import javax.annotation.concurrent.ThreadSafe; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; @@ -79,6 +83,25 @@ public void put(T t) throws Exception, InterruptedException { } } + public boolean tryPut(T t) throws Exception { + lock.lock(); + try { + if (cancellationException != null) { + throw cancellationException; + } + if (count >= capacity) { + return false; + } + elements[addIndex] = t; + addIndex = (addIndex + 1) % elements.length; + count += 1; + notEmpty.signal(); + return true; + } finally { + lock.unlock(); + } + } + /** * Takes an element from this queue. Will block until the queue is not full or is cancelled. * @@ -116,27 +139,50 @@ public T take() throws Exception, InterruptedException { * queue clears the exception. */ public void cancel(Exception exception) { + cancel(exception, null); + } + + /** + * Like {@link #cancel(Exception)}, and passes each element that the queue drops to {@code + * onDrop}. + * + *

{@code onDrop} runs after the queue releases its lock, so it may take other locks, including + * locks that a thread holds while it calls {@link #tryPut}. + */ + public void cancel(Exception exception, @Nullable Consumer onDrop) { checkNotNull(exception); + List dropped = Collections.emptyList(); lock.lock(); try { if (cancellationException == null) { cancellationException = exception; - clearElementsLocked(); + dropped = clearElementsLocked(onDrop != null); } notEmpty.signalAll(); notFull.signalAll(); } finally { lock.unlock(); } + if (onDrop != null) { + for (T element : dropped) { + onDrop.accept(element); + } + } } - private void clearElementsLocked() { + @SuppressWarnings({"cast"}) + private List clearElementsLocked(boolean collect) { + List dropped = collect ? new ArrayList<>(count) : Collections.emptyList(); for (int i = takeIndex; count > 0; i = (i + 1) % elements.length) { + if (collect && elements[i] != null) { + dropped.add((T) elements[i]); + } elements[i] = null; --count; } addIndex = 0; takeIndex = 0; + return dropped; } /** Enables the queue to be re-used after it has been cancelled. */ @@ -144,7 +190,7 @@ public void reset() { lock.lock(); try { cancellationException = null; - clearElementsLocked(); + clearElementsLocked(false); } finally { lock.unlock(); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java index e8afd1b7e6cd..c76afba849c0 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java @@ -18,17 +18,19 @@ package org.apache.beam.sdk.fn.data; import java.time.Duration; +import java.util.ArrayDeque; import java.util.HashSet; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.pipeline.v1.Endpoints; import org.apache.beam.sdk.fn.stream.OutboundObserverFactory; +import org.apache.beam.sdk.fn.stream.StalledPollHandler; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Status; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientCallStreamObserver; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientResponseObserver; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; @@ -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 inboundObserver; + private final InboundObserver inboundObserver; private final StreamObserver outboundObserver; - private final ConcurrentHashMap< - /* instructionId= */ String, - CompletableFuture>> - receivers; - private final Cache 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 receiver; + + ArrayDeque 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 receiver) + throws Exception { + if (this.receiver != null) { + return false; + } + this.receiver = receiver; + if (receiver instanceof BeamFnDataInboundObserver) { + drainPending(); + } else { + for (PendingElement p : pending) { + receiver.accept(p.elements); + p.callback.run(); + } + pending.clear(); + } + return true; + } + + synchronized void poison() { + poisoned = true; + for (PendingElement p : pending) { + p.callback.run(); + } + pending.clear(); + } + + synchronized boolean isDone() { + return receiver != null || poisoned; + } + + synchronized void close() throws Exception { + if (receiver != null) { + receiver.close(); + } + for (PendingElement p : pending) { + p.callback.run(); + } + pending.clear(); } } - ; + + private final ConcurrentMap receivers; + private final Cache poisonedInstructionIds; public BeamFnDataGrpcMultiplexer( Endpoints.@Nullable ApiServiceDescriptor apiServiceDescriptor, @@ -93,6 +222,20 @@ public String toString() { .toString(); } + /** + * Called when a bundle finishes one inbound message. Asks for more messages until {@code + * STEADY_STATE_INBOUND_WINDOW} messages are open again. + */ + public void requestInbound(int count) { + inFlightInbound.decrementAndGet(); + synchronized (inboundObserver) { + while (inFlightInbound.get() < STEADY_STATE_INBOUND_WINDOW) { + inFlightInbound.incrementAndGet(); + inboundObserver.request(1); + } + } + } + public StreamObserver getInboundObserver() { return inboundObserver; } @@ -118,26 +261,39 @@ public void registerConsumer( instructionId, (unused, existing) -> { if (existing != null) { - if (!existing.complete(receiver)) { - throw new IllegalArgumentException("Instruction id was registered twice"); + try { + if (!existing.complete(receiver)) { + throw new IllegalArgumentException("Instruction id was registered twice"); + } + } catch (Exception e) { + throw new RuntimeException(e); } + return existing; } if (poisonedInstructionIds.getIfPresent(instructionId) != null) { throw new IllegalArgumentException("Instruction id was poisoned"); } - return CompletableFuture.completedFuture(receiver); + ReceiverHolder h = new ReceiverHolder(); + try { + h.complete(receiver); + } catch (Exception e) { + throw new RuntimeException(e); + } + return h; }); } /** Unregisters a previously registered consumer. */ public void unregisterConsumer(String instructionId) { - @Nullable CompletableFuture> receiverFuture = - receivers.remove(instructionId); - if (receiverFuture != null && !receiverFuture.isDone()) { - // The future must have been inserted by the inbound observer since registerConsumer completes - // the future. - throw new IllegalArgumentException("Unregistering consumer which was not registered."); + @Nullable ReceiverHolder receiverFuture = receivers.remove(instructionId); + if (receiverFuture != null) { + if (!receiverFuture.isDone()) { + // The inbound observer inserted this holder before registerConsumer completed it. + receiverFuture.poison(); + throw new IllegalArgumentException("Unregistering consumer which was not registered."); + } + receiverFuture.poison(); } } @@ -149,18 +305,14 @@ public void unregisterConsumer(String instructionId) { */ public void poisonInstructionId(String instructionId) { poisonedInstructionIds.put(instructionId, Boolean.TRUE); - @Nullable CompletableFuture> receiverFuture = - receivers.remove(instructionId); + @Nullable ReceiverHolder receiverFuture = receivers.remove(instructionId); if (receiverFuture != null) { - // Completing exceptionally has no effect if the future was already notified. In that case - // whatever registered the receiver needs to handle cancelling it. - receiverFuture.completeExceptionally(new PoisonedException()); - if (!receiverFuture.isCompletedExceptionally()) { - try { - receiverFuture.get().close(); - } catch (Exception e) { - LOG.warn("Unexpected error closing existing observer"); - } + // Release pending callbacks and close any registered receiver. + receiverFuture.poison(); + try { + receiverFuture.close(); + } catch (Exception e) { + LOG.warn("Unexpected error closing existing observer"); } } } @@ -173,20 +325,17 @@ boolean hasConsumer(String instructionId) { @Override public void close() throws Exception { Exception exception = null; - for (CompletableFuture> receiver : - ImmutableList.copyOf(receivers.values())) { + for (ReceiverHolder receiver : ImmutableList.copyOf(receivers.values())) { // Cancel any observer waiting for the client to complete. If the receiver has already been // completed or cancelled, this call will be ignored. - receiver.cancel(true); - if (!receiver.isCompletedExceptionally()) { - try { - receiver.get().close(); - } catch (Exception e) { - if (exception == null) { - exception = e; - } else { - exception.addSuppressed(e); - } + receiver.poison(); + try { + receiver.close(); + } catch (Exception e) { + if (exception == null) { + exception = e; + } else { + exception.addSuppressed(e); } } } @@ -203,14 +352,36 @@ public void close() throws Exception { * A multiplexing {@link StreamObserver} that selects the inbound {@link Consumer} to pass the * elements to. * - *

The inbound observer blocks until the {@link Consumer} is bound allowing for the sending - * harness to initiate transmitting data without needing for the receiving harness to signal that - * it is ready to consume that data. + *

Data for an instruction with no {@link Consumer} yet waits in a buffer until the consumer is + * registered, so the sending harness can send data before the receiving harness is ready. The + * observer does not block: when a bundle queue is full, the message waits in the buffer. */ - private final class InboundObserver implements StreamObserver { + private final class InboundObserver + implements ClientResponseObserver, + StalledPollHandler { + @Override + public void onStalledPoll() { + BeamFnDataGrpcMultiplexer.this.onStalledPoll(); + } + + private @Nullable ClientCallStreamObserver requestStream; + + @Override + public void beforeStart(ClientCallStreamObserver requestStream) { + this.requestStream = requestStream; + requestStream.disableAutoRequestWithInitial(STEADY_STATE_INBOUND_WINDOW); + } + + public void request(int count) { + if (requestStream != null) { + requestStream.request(count); + } + } + @SuppressWarnings("LabelledBreakTarget") @Override public void onNext(BeamFnApi.Elements value) { + Runnable callback = () -> BeamFnDataGrpcMultiplexer.this.requestInbound(1); // Have a fast path to handle the common case and provide a short circuit to exit if we detect // multiple instruction ids. SINGLE_INSTRUCTION_ID: @@ -233,9 +404,12 @@ public void onNext(BeamFnApi.Elements value) { } } if (instructionId == null) { + callback.run(); return; } - forwardToConsumerForInstructionId(instructionId, value); + if (!forwardToConsumerForInstructionId(instructionId, value, callback)) { + callback.run(); + } return; } @@ -247,6 +421,14 @@ public void onNext(BeamFnApi.Elements value) { for (BeamFnApi.Elements.Timers timers : value.getTimersList()) { instructionIds.add(timers.getInstructionId()); } + AtomicInteger countdown = new AtomicInteger(instructionIds.size()); + Runnable multiCallback = + () -> { + if (countdown.decrementAndGet() == 0) { + callback.run(); + } + }; + for (String instructionId : instructionIds) { BeamFnApi.Elements.Builder builder = BeamFnApi.Elements.newBuilder(); for (BeamFnApi.Elements.Data data : value.getDataList()) { @@ -259,61 +441,34 @@ public void onNext(BeamFnApi.Elements value) { builder.addTimers(timers); } } - forwardToConsumerForInstructionId(instructionId, builder.build()); + if (!forwardToConsumerForInstructionId(instructionId, builder.build(), multiCallback)) { + multiCallback.run(); + } } } - private void forwardToConsumerForInstructionId(String instructionId, BeamFnApi.Elements value) { - CloseableFnDataReceiver consumer; + private boolean forwardToConsumerForInstructionId( + String instructionId, BeamFnApi.Elements value, Runnable callback) { try { - CompletableFuture> consumerFuture = + ReceiverHolder consumerFuture = receivers.computeIfAbsent( instructionId, (unused) -> { if (poisonedInstructionIds.getIfPresent(instructionId) != null) { - throw new PoisonedException(); + ReceiverHolder h = new ReceiverHolder(); + h.poison(); + return h; } LOG.debug( "Received data for instruction {} without consumer ready. " - + "Waiting for consumer to be registered.", + + "Buffering elements until consumer is registered.", instructionId); - return new CompletableFuture<>(); + return new ReceiverHolder(); }); - // The consumer may not be registered until the bundle processor is fully constructed so we - // conservatively set - // a high timeout. Poisoning will prevent this for occurring for consumers that will not be - // registered. - consumer = consumerFuture.get(3, TimeUnit.HOURS); - /* - * TODO: On failure we should fail any bundles that were impacted eagerly - * instead of relying on the Runner harness to do all the failure handling. - */ - } catch (TimeoutException e) { - LOG.error( - "Timed out waiting to observe consumer data stream for instruction {}", - instructionId, - e); - outboundObserver.onError(e); - return; - } catch (ExecutionException | InterruptedException | PoisonedException e) { - if (e instanceof PoisonedException || e.getCause() instanceof PoisonedException) { - LOG.debug("Received data for poisoned instruction {}. Dropping input.", instructionId); - return; - } - LOG.error( - "Client interrupted during handling of data for instruction {}", instructionId, e); - outboundObserver.onError(e); - return; - } catch (RuntimeException e) { - LOG.error("Client failed to handle data for instruction {}", instructionId, e); - outboundObserver.onError(e); - return; - } - - try { - consumer.accept(value); + return consumerFuture.accept(value, callback); } catch (Exception e) { poisonInstructionId(instructionId); + return false; } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java index 1b23c35c7faf..bf43a67e978f 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java @@ -66,7 +66,18 @@ private static class EndpointStatus { private final Map>> transformIdToDataEndpoint; private final Map>>> transformIdToTimerFamilyIdToTimerEndpoint; - private final CancellableQueue queue; + + static class ElementWithCallback { + final BeamFnApi.Elements elements; + final Runnable callback; + + ElementWithCallback(BeamFnApi.Elements elements, Runnable callback) { + this.elements = elements; + this.callback = callback; + } + } + + private final CancellableQueue queue; // We use a custom exception for closing to avoid the expense of stack trace generation. @SuppressWarnings("StaticAssignmentOfThrowable") @@ -107,7 +118,16 @@ private BeamFnDataInboundObserver( @Override public void accept(BeamFnApi.Elements elements) throws Exception { - queue.put(elements); + acceptWithCallback(elements, () -> {}); + } + + public void acceptWithCallback(BeamFnApi.Elements elements, Runnable callback) throws Exception { + queue.put(new ElementWithCallback(elements, callback)); + } + + public boolean tryAcceptWithCallback(BeamFnApi.Elements elements, Runnable callback) + throws Exception { + return queue.tryPut(new ElementWithCallback(elements, callback)); } @Override @@ -117,7 +137,7 @@ public void flush() throws Exception { @Override public void close() throws Exception { - queue.cancel(CloseException.INSTANCE); + queue.cancel(CloseException.INSTANCE, wrapper -> wrapper.callback.run()); } public boolean isConsumingReceivedData() { @@ -177,19 +197,23 @@ public void awaitCompletion() throws Exception { // collection. Iterator dataIterator; Iterator timersIterator; + ElementWithCallback wrapper; { - BeamFnApi.Elements elements = queue.take(); + wrapper = queue.take(); + BeamFnApi.Elements elements = wrapper.elements; dataIterator = createDiscardingIterator(elements.getDataList()); timersIterator = createDiscardingIterator(elements.getTimersList()); } // The SDK is no longer blocked on receiving more data from the Runner. consumingReceivedData.set(true); - if (multiplexElements(dataIterator, timersIterator)) { + boolean isDone = multiplexElements(dataIterator, timersIterator); + wrapper.callback.run(); + if (isDone) { return; } } } catch (Exception e) { - queue.cancel(e); + queue.cancel(e, wrapper -> wrapper.callback.run()); throw e; } finally { consumingReceivedData.set(false); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserver.java index b54018d7095a..e48190c9c9a7 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserver.java @@ -20,12 +20,15 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.Phaser; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import javax.annotation.concurrent.ThreadSafe; import org.apache.beam.sdk.fn.CancellableQueue; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.CallStreamObserver; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A thread safe {@link StreamObserver} which uses a bounded queue to pass elements to a processing @@ -55,22 +58,40 @@ public OnErrorException(@NonNull Throwable throwable) { } private static final Object POISON_PILL = new Object(); + private static final int STALL_POLL_INTERVAL_MS = 50; + private static final int STALL_NOTIFY_INTERVAL_MS = 500; private final CancellableQueue queue; private final Phaser phaser; private final CallStreamObserver outboundObserver; + private final @Nullable StalledPollHandler stalledPollHandler; private final Future queueDrainer; private final int bufferSize; - @SuppressWarnings("methodref.receiver.bound") public BufferingStreamObserver( Phaser phaser, CallStreamObserver outboundObserver, ExecutorService executor, int bufferSize) { + this(phaser, outboundObserver, executor, bufferSize, null); + } + + /** + * Like {@link #BufferingStreamObserver(Phaser, CallStreamObserver, ExecutorService, int)}. While + * the drainer waits for the outbound stream to become ready, it calls {@code stalledPollHandler} + * every 500 ms. + */ + @SuppressWarnings("methodref.receiver.bound") + public BufferingStreamObserver( + Phaser phaser, + CallStreamObserver outboundObserver, + ExecutorService executor, + int bufferSize, + @Nullable StalledPollHandler stalledPollHandler) { this.phaser = phaser; this.bufferSize = bufferSize; this.queue = new CancellableQueue<>(bufferSize); this.outboundObserver = outboundObserver; + this.stalledPollHandler = stalledPollHandler; this.queueDrainer = executor.submit(this::drainQueue); } @@ -87,7 +108,7 @@ private void drainQueue() { return; } } - phaser.awaitAdvance(currentPhase); + awaitReady(currentPhase); } } catch (OnErrorException e) { outboundObserver.onError(e.getCause()); @@ -97,6 +118,27 @@ private void drainQueue() { } } + /** Waits for the phase to advance. With a stall handler, also polls readiness and notifies it. */ + private void awaitReady(int phase) throws InterruptedException { + StalledPollHandler handler = stalledPollHandler; + if (handler == null) { + phaser.awaitAdvance(phase); + return; + } + int totalMsWaited = 0; + while (phase >= 0 && !outboundObserver.isReady()) { + try { + phaser.awaitAdvanceInterruptibly(phase, STALL_POLL_INTERVAL_MS, TimeUnit.MILLISECONDS); + return; + } catch (TimeoutException e) { + totalMsWaited += STALL_POLL_INTERVAL_MS; + if (totalMsWaited % STALL_NOTIFY_INTERVAL_MS == 0) { + handler.onStalledPoll(); + } + } + } + } + @Override public void onNext(T value) { try { diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java index bafacfed1f62..ca5ff6f5ed03 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java @@ -23,6 +23,7 @@ import javax.annotation.concurrent.ThreadSafe; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.CallStreamObserver; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver; +import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -39,7 +40,7 @@ @ThreadSafe public final class DirectStreamObserver implements StreamObserver { private static final Logger LOG = LoggerFactory.getLogger(DirectStreamObserver.class); - private static final int DEFAULT_MAX_MESSAGES_BEFORE_CHECK = 100; + private static final int DEFAULT_MAX_MESSAGES_BEFORE_CHECK = 0; private final Phaser phaser; private final CallStreamObserver outboundObserver; @@ -55,15 +56,34 @@ public final class DirectStreamObserver implements StreamObserver { private final Object lock = new Object(); private int numMessages; + private final @Nullable StalledPollHandler stalledPollHandler; + public DirectStreamObserver(Phaser phaser, CallStreamObserver outboundObserver) { - this(phaser, outboundObserver, DEFAULT_MAX_MESSAGES_BEFORE_CHECK); + this(phaser, outboundObserver, null); + } + + public DirectStreamObserver( + Phaser phaser, + CallStreamObserver outboundObserver, + @Nullable StalledPollHandler stalledPollHandler) { + + this(phaser, outboundObserver, DEFAULT_MAX_MESSAGES_BEFORE_CHECK, stalledPollHandler); } DirectStreamObserver( Phaser phaser, CallStreamObserver outboundObserver, int maxMessagesBeforeCheck) { + this(phaser, outboundObserver, maxMessagesBeforeCheck, null); + } + + DirectStreamObserver( + Phaser phaser, + CallStreamObserver outboundObserver, + int maxMessagesBeforeCheck, + @Nullable StalledPollHandler stalledPollHandler) { this.phaser = phaser; this.outboundObserver = outboundObserver; this.maxMessagesBeforeCheck = maxMessagesBeforeCheck; + this.stalledPollHandler = stalledPollHandler; } @Override @@ -71,8 +91,8 @@ public void onNext(T value) { synchronized (lock) { if (numMessages >= maxMessagesBeforeCheck) { numMessages = 0; - int waitSeconds = 1; - int totalSecondsWaited = 0; + int waitMs = 50; + int totalMsWaited = 0; int phase = phaser.getPhase(); // Record the initial phase in case we are in the inbound gRPC thread where the phase won't // advance. @@ -80,35 +100,36 @@ public void onNext(T value) { // A negative phase indicates that the phaser is terminated. while (phase >= 0 && !outboundObserver.isReady()) { try { - phase = phaser.awaitAdvanceInterruptibly(phase, waitSeconds, TimeUnit.SECONDS); + phase = phaser.awaitAdvanceInterruptibly(phase, waitMs, TimeUnit.MILLISECONDS); } catch (TimeoutException e) { - totalSecondsWaited += waitSeconds; - // Double the backoff for re-evaluating the isReady bit up to a maximum of once per - // minute. This bounds the waiting if the onReady callback is not called as expected. - waitSeconds = Math.min(waitSeconds * 2, 60); + totalMsWaited += waitMs; + if (stalledPollHandler != null && totalMsWaited >= 500 && totalMsWaited % 500 == 0) { + stalledPollHandler.onStalledPoll(); + } + waitMs = 50; } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new RuntimeException(e); } } - if (totalSecondsWaited > 0) { + if (totalMsWaited > 0) { // If the phase didn't change, this means that the installed onReady callback had not // been invoked. if (initialPhase == phase) { LOG.info( "Output channel stalled for {}s, outbound thread {}. OnReady notification was " + "not invoked, ensure the inbound gRPC thread is not used for output.", - totalSecondsWaited, + totalMsWaited / 1000, Thread.currentThread().getName()); - } else if (totalSecondsWaited > 60) { + } else if (totalMsWaited > 60000) { LOG.warn( "Output channel stalled for {}s, outbound thread {}.", - totalSecondsWaited, + totalMsWaited / 1000, Thread.currentThread().getName()); } else { LOG.debug( "Output channel stalled for {}s, outbound thread {}.", - totalSecondsWaited, + totalMsWaited / 1000, Thread.currentThread().getName()); } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java index d690108f8a2e..4e99870b1a85 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java @@ -31,7 +31,7 @@ *

This is as thread-safe as the underlying stream observer that is being wrapped. */ public final class ForwardingClientResponseObserver - implements ClientResponseObserver { + implements ClientResponseObserver, StalledPollHandler { public static ForwardingClientResponseObserver create( StreamObserver inbound, Runnable onReadyHandler) { return new ForwardingClientResponseObserver<>(inbound, onReadyHandler, () -> {}); @@ -73,5 +73,18 @@ public void onCompleted() { @Override public void beforeStart(ClientCallStreamObserver stream) { stream.setOnReadyHandler(onReadyHandler); + if (inboundObserver instanceof ClientResponseObserver) { + @SuppressWarnings("unchecked") + ClientResponseObserver clientResponseObserver = + (ClientResponseObserver) inboundObserver; + clientResponseObserver.beforeStart(stream); + } + } + + @Override + public void onStalledPoll() { + if (inboundObserver instanceof StalledPollHandler) { + ((StalledPollHandler) inboundObserver).onStalledPoll(); + } } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactory.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactory.java index 7acbe2e573a0..fc518fe40283 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactory.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactory.java @@ -89,13 +89,17 @@ public StreamObserver outboundObserverFor( BasicFactory baseOutboundObserverFactory, StreamObserver inboundObserver) { AdvancingPhaser phaser = new AdvancingPhaser(1); - inboundObserver = - ForwardingClientResponseObserver.create( - inboundObserver, phaser::arrive, phaser::forceTermination); CallStreamObserver outboundObserver = (CallStreamObserver) - baseOutboundObserverFactory.outboundObserverFor(inboundObserver); - return new DirectStreamObserver<>(phaser, outboundObserver); + baseOutboundObserverFactory.outboundObserverFor( + ForwardingClientResponseObserver.create( + inboundObserver, phaser::arrive, phaser::forceTermination)); + return new DirectStreamObserver<>( + phaser, + outboundObserver, + inboundObserver instanceof StalledPollHandler + ? (StalledPollHandler) inboundObserver + : null); } } @@ -109,7 +113,12 @@ public StreamObserver outboundObserverFor( (CallStreamObserver) baseOutboundObserverFactory.outboundObserverFor(inboundObserver); outboundObserver.setOnReadyHandler(phaser::arrive); - return new DirectStreamObserver<>(phaser, outboundObserver); + return new DirectStreamObserver<>( + phaser, + outboundObserver, + inboundObserver instanceof StalledPollHandler + ? (StalledPollHandler) inboundObserver + : null); } } @@ -128,13 +137,19 @@ public StreamObserver outboundObserverFor( BasicFactory baseOutboundObserverFactory, StreamObserver inboundObserver) { AdvancingPhaser phaser = new AdvancingPhaser(1); - inboundObserver = - ForwardingClientResponseObserver.create( - inboundObserver, phaser::arrive, phaser::forceTermination); CallStreamObserver outboundObserver = (CallStreamObserver) - baseOutboundObserverFactory.outboundObserverFor(inboundObserver); - return new BufferingStreamObserver<>(phaser, outboundObserver, executorService, bufferSize); + baseOutboundObserverFactory.outboundObserverFor( + ForwardingClientResponseObserver.create( + inboundObserver, phaser::arrive, phaser::forceTermination)); + return new BufferingStreamObserver<>( + phaser, + outboundObserver, + executorService, + bufferSize, + inboundObserver instanceof StalledPollHandler + ? (StalledPollHandler) inboundObserver + : null); } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/StalledPollHandler.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/StalledPollHandler.java new file mode 100644 index 000000000000..8d88982874cc --- /dev/null +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/StalledPollHandler.java @@ -0,0 +1,26 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.fn.stream; + +/** + * Receives a call each time an outbound stream waits 500 ms more on gRPC flow control. {@link + * DirectStreamObserver} calls it on the inbound observer of the same stream. + */ +public interface StalledPollHandler { + void onStalledPoll(); +} diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/CancellableQueueTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/CancellableQueueTest.java index aa85e7a2cec6..2780240a4075 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/CancellableQueueTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/CancellableQueueTest.java @@ -20,6 +20,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.empty; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; @@ -28,11 +29,14 @@ import java.util.HashSet; import java.util.List; import java.util.Set; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import org.apache.beam.sdk.fn.test.TestExecutors; import org.apache.beam.sdk.fn.test.TestExecutors.TestExecutorService; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; import org.junit.Rule; import org.junit.Test; @@ -247,4 +251,108 @@ public void testMemoryReferenceOnTake() throws Exception { assertTrue(weakReference1.get() == null); assertTrue(weakReference2.get() == null); } + + @Test + public void testTryPut() throws Exception { + CancellableQueue queue = new CancellableQueue<>(2); + assertTrue(queue.tryPut("A")); + assertTrue(queue.tryPut("B")); + assertFalse(queue.tryPut("C")); + + assertEquals("A", queue.take()); + assertTrue(queue.tryPut("C")); + assertEquals("B", queue.take()); + assertEquals("C", queue.take()); + } + + @Test + public void testTryPutThrowsWhenCancelled() { + CancellableQueue queue = new CancellableQueue<>(2); + queue.cancel(new Exception("Cancelled")); + assertThrows(Exception.class, () -> queue.tryPut("A")); + } + + @Test + public void testCancelWithOnDrop() throws Exception { + CancellableQueue queue = new CancellableQueue<>(5); + queue.put("A"); + queue.put("B"); + queue.put("C"); + + List dropped = new ArrayList<>(); + queue.cancel(new Exception("Cancelled"), dropped::add); + + assertEquals(3, dropped.size()); + assertEquals("A", dropped.get(0)); + assertEquals("B", dropped.get(1)); + assertEquals("C", dropped.get(2)); + assertThrows(Exception.class, queue::take); + } + + @Test(timeout = 10_000) + public void testOnDropRunsWithoutQueueLock() throws Exception { + CancellableQueue queue = new CancellableQueue<>(5); + queue.put("A"); + List otherThreadGotLock = new ArrayList<>(); + queue.cancel( + new Exception("Cancelled"), + dropped -> { + // Another thread must be able to take the queue lock while onDrop runs. + Future other = executor.submit(() -> queue.tryPut("B")); + try { + other.get(5, TimeUnit.SECONDS); + } catch (ExecutionException e) { + // tryPut takes the lock and then throws the cancellation exception. + otherThreadGotLock.add(true); + } catch (Exception e) { + otherThreadGotLock.add(false); + } + }); + assertEquals(ImmutableList.of(true), otherThreadGotLock); + } + + /** + * A producer holds its own lock while it calls {@link CancellableQueue#tryPut}, and a consumer + * cancels the queue with an {@code onDrop} that takes the producer's lock. This is how {@code + * BeamFnDataGrpcMultiplexer} and {@code BeamFnDataInboundObserver} use the queue. + */ + @Test(timeout = 30_000) + public void testCancelWithOnDropDoesNotDeadlockWithLockedProducer() throws Exception { + for (int run = 0; run < 200; run++) { + CancellableQueue queue = new CancellableQueue<>(1); + Object producerLock = new Object(); + queue.put("A"); + CountDownLatch start = new CountDownLatch(1); + Future producer = + executor.submit( + () -> { + start.await(); + while (true) { + synchronized (producerLock) { + try { + queue.tryPut("B"); + } catch (Exception cancelled) { + return null; + } + } + } + }); + Future canceller = + executor.submit( + () -> { + start.await(); + queue.cancel( + new Exception("Cancelled"), + dropped -> { + synchronized (producerLock) { + // The producer's lock is free once the producer waits for the queue. + } + }); + return null; + }); + start.countDown(); + canceller.get(5, TimeUnit.SECONDS); + producer.get(5, TimeUnit.SECONDS); + } + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java index 6ee06d0a36a4..81bd4e3638da 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java @@ -19,23 +19,42 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.core.StringContains.containsString; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicBoolean; import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.pipeline.v1.Endpoints; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.fn.stream.OutboundObserverFactory; +import org.apache.beam.sdk.fn.stream.StalledPollHandler; import org.apache.beam.sdk.fn.test.TestExecutors; import org.apache.beam.sdk.fn.test.TestExecutors.TestExecutorService; import org.apache.beam.sdk.fn.test.TestStreams; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.ByteStringOutputStream; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientCallStreamObserver; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientResponseObserver; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.hamcrest.MatcherAssert; import org.hamcrest.Matchers; @@ -363,4 +382,467 @@ public void accept(BeamFnApi.Elements input) throws Exception { Iterables.getOnlyElement(errorWasReturned).getMessage(), containsString("Multiplexer hanging up")); } + + @Test + public void testFlowControlManualRequestAndStalledPoll() { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + @SuppressWarnings("unchecked") + ClientCallStreamObserver mockCallStream = + mock(ClientCallStreamObserver.class); + + @SuppressWarnings("unchecked") + ClientResponseObserver clientResponseObserver = + (ClientResponseObserver) + multiplexer.getInboundObserver(); + + clientResponseObserver.beforeStart(mockCallStream); + verify(mockCallStream).disableAutoRequestWithInitial(16); + + // Call onStalledPoll directly on multiplexer + multiplexer.onStalledPoll(); + verify(mockCallStream, times(1)).request(1); + + // Call onStalledPoll on InboundObserver + StalledPollHandler stalledPollHandler = (StalledPollHandler) multiplexer.getInboundObserver(); + stalledPollHandler.onStalledPoll(); + verify(mockCallStream, times(2)).request(1); + + // inFlightInbound started at 16, incremented twice to 18. + // requestInbound decrements inFlightInbound. + // When decremented below 16 (STEADY_STATE_INBOUND_WINDOW), it requests more. + multiplexer.requestInbound(1); // inFlight becomes 17 (>= 16, no request) + verify(mockCallStream, times(2)).request(1); + + multiplexer.requestInbound(1); // inFlight becomes 16 (>= 16, no request) + verify(mockCallStream, times(2)).request(1); + + multiplexer.requestInbound(1); // inFlight becomes 15 (< 16, requests 1 and inFlight becomes 16) + verify(mockCallStream, times(3)).request(1); + } + + @Test + public void testEmptyElementsHandled() { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + // Should not throw and should safely exit + multiplexer.getInboundObserver().onNext(BeamFnApi.Elements.getDefaultInstance()); + } + + @Test + public void testReceiverHolderBuffersBeforeConsumerRegisteredAndDrains() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + String instructionId = "unregistered_instruction"; + BeamFnApi.Elements elements = + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId("t1") + .setData(ByteString.copyFromUtf8("hello"))) + .build(); + + // Send elements before consumer is registered - should buffer and not block + multiplexer.getInboundObserver().onNext(elements); + + // Now register a generic receiver + List received = new ArrayList<>(); + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) { + received.add(input); + } + + @Override + public void flush() {} + + @Override + public void close() {} + }); + + assertThat(received, Matchers.contains(elements)); + + // Registering twice throws IllegalArgumentException wrapped in RuntimeException + RuntimeException e = + assertThrows( + RuntimeException.class, + () -> + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) {} + + @Override + public void flush() {} + + @Override + public void close() {} + })); + assertThat(e.getCause(), Matchers.instanceOf(IllegalArgumentException.class)); + } + + @Test + public void testReceiverHolderBuffersAndDrainsWithBeamFnDataInboundObserver() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + String instructionId = "inbound_obs_instruction"; + Coder> coder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + ByteStringOutputStream out = new ByteStringOutputStream(); + coder.encode(WindowedValues.valueInGlobalWindow("payload"), out); + + BeamFnApi.Elements elements = + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId("t1") + .setData(out.toByteString())) + .build(); + BeamFnApi.Elements terminal = + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId("t1") + .setIsLast(true)) + .build(); + + // Buffer elements before registration + multiplexer.getInboundObserver().onNext(elements); + + List> receivedValues = new ArrayList<>(); + BeamFnDataInboundObserver inboundObserver = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList(DataEndpoint.create("t1", coder, receivedValues::add)), + Collections.emptyList()); + + multiplexer.registerConsumer(instructionId, inboundObserver); + multiplexer.getInboundObserver().onNext(terminal); + + inboundObserver.awaitCompletion(); + assertThat(receivedValues, Matchers.contains(WindowedValues.valueInGlobalWindow("payload"))); + } + + @Test + public void testUnregisterConsumerNotRegisteredThrows() { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + String instructionId = "unregistered_instruction"; + // Send data to create ReceiverHolder in receivers without registering consumer + multiplexer + .getInboundObserver() + .onNext( + BeamFnApi.Elements.newBuilder() + .addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(instructionId)) + .build()); + + assertThrows( + IllegalArgumentException.class, () -> multiplexer.unregisterConsumer(instructionId)); + } + + @Test + public void testPoisonInstructionId() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + String instructionId = "poison_instruction"; + AtomicBoolean closed = new AtomicBoolean(); + + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) {} + + @Override + public void flush() {} + + @Override + public void close() { + closed.set(true); + } + }); + + multiplexer.poisonInstructionId(instructionId); + assertTrue(closed.get()); + + // Registering consumer for poisoned instruction fails + assertThrows( + IllegalArgumentException.class, + () -> + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) {} + + @Override + public void flush() {} + + @Override + public void close() {} + })); + + // Subsequent elements for poisoned instruction are dropped + multiplexer + .getInboundObserver() + .onNext( + BeamFnApi.Elements.newBuilder() + .addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(instructionId)) + .build()); + } + + @Test + public void testPoisonInstructionIdWhenPendingElements() { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + String instructionId = "pending_poison_instruction"; + // Buffer elements + multiplexer + .getInboundObserver() + .onNext( + BeamFnApi.Elements.newBuilder() + .addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(instructionId)) + .build()); + + // Poison while pending + multiplexer.poisonInstructionId(instructionId); + + // Registering now throws + assertThrows( + IllegalArgumentException.class, + () -> + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) {} + + @Override + public void flush() {} + + @Override + public void close() {} + })); + } + + @Test + public void testToString() { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + String str = multiplexer.toString(); + assertThat(str, containsString("apiServiceDescriptor")); + assertThat(str, containsString("consumers")); + } + + @Test + public void testConsumerAcceptThrowsExceptionPoisonsInstruction() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + String instructionId = "throw_instruction"; + AtomicBoolean closed = new AtomicBoolean(); + + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) throws Exception { + throw new RuntimeException("Accept failure"); + } + + @Override + public void flush() {} + + @Override + public void close() { + closed.set(true); + } + }); + + multiplexer + .getInboundObserver() + .onNext( + BeamFnApi.Elements.newBuilder() + .addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(instructionId)) + .build()); + + assertTrue(closed.get()); + } + + @Test + public void testBeamFnDataInboundObserverBufferingWhenFull() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> + TestStreams.withOnNext((BeamFnApi.Elements val) -> {}) + .withOnError(err -> {}) + .build()); + + String instructionId = "full_queue_instruction"; + BeamFnDataInboundObserver mockInbound = mock(BeamFnDataInboundObserver.class); + when(mockInbound.tryAcceptWithCallback(any(), any())).thenReturn(false).thenReturn(true); + + multiplexer.registerConsumer(instructionId, mockInbound); + + BeamFnApi.Elements elements = + BeamFnApi.Elements.newBuilder() + .addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(instructionId)) + .build(); + + // First call returns false -> added to pending + multiplexer.getInboundObserver().onNext(elements); + + // Closing drains / clears pending + multiplexer.close(); + } + + @Test + public void testBeamFnDataInboundObserverDrainPendingExceptionHandled() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + String instructionId = "drain_exception_instruction"; + BeamFnApi.Elements elements = + BeamFnApi.Elements.newBuilder() + .addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(instructionId)) + .build(); + + // Buffer elements first + multiplexer.getInboundObserver().onNext(elements); + + BeamFnDataInboundObserver mockInbound = mock(BeamFnDataInboundObserver.class); + when(mockInbound.tryAcceptWithCallback(any(), any())) + .thenThrow(new RuntimeException("queue error")); + + // Registering calls drainPending(), which catches the exception and removes from pending + multiplexer.registerConsumer(instructionId, mockInbound); + } + + /** + * The gRPC thread delivers data for one bundle while the bundle thread drains the pending buffer + * from its callbacks, and the consumer registers in the middle of the stream. Every element must + * arrive once and in order. + */ + @Test(timeout = 60_000) + public void testConcurrentHandoffKeepsOrder() throws Exception { + int numElements = 3_000; + for (int run = 0; run < 3; run++) { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + String instructionId = "concurrent_" + run; + Coder> coder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + + List received = new ArrayList<>(); + BeamFnDataInboundObserver observer = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList( + DataEndpoint.create( + "t1", + coder, + value -> { + received.add(value.getValue()); + if (received.size() % 50 == 0) { + // A slow consumer lets the bundle queue fill, so elements wait in the + // pending buffer and the bundle thread drains them. + Thread.sleep(1); + } + })), + Collections.emptyList()); + + CountDownLatch halfway = new CountDownLatch(1); + Future grpcThread = + executor.submit( + () -> { + for (int i = 0; i < numElements; i++) { + ByteStringOutputStream out = new ByteStringOutputStream(); + coder.encode(WindowedValues.valueInGlobalWindow(Integer.toString(i)), out); + multiplexer + .getInboundObserver() + .onNext( + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId("t1") + .setData(out.toByteString())) + .build()); + if (i == 200) { + halfway.countDown(); + } + } + multiplexer + .getInboundObserver() + .onNext( + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId("t1") + .setIsLast(true)) + .build()); + return null; + }); + + halfway.await(); + multiplexer.registerConsumer(instructionId, observer); + observer.awaitCompletion(); + grpcThread.get(); + + List expected = new ArrayList<>(); + for (int i = 0; i < numElements; i++) { + expected.add(Integer.toString(i)); + } + assertEquals(expected, received); + multiplexer.unregisterConsumer(instructionId); + } + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java index f915f9285367..24ebdf067b82 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java @@ -19,8 +19,10 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.contains; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; import java.util.ArrayList; import java.util.Arrays; @@ -218,6 +220,84 @@ public void testBadProducerDataFailureVisibleToAwaitCompletionCallerAndProducer( future.get(); } + @Test + public void testAcceptWithCallbackInvokedOnProcessing() throws Exception { + Collection> values = new ArrayList<>(); + BeamFnDataInboundObserver observer = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList(DataEndpoint.create(TRANSFORM_ID, CODER, values::add)), + Collections.emptyList()); + + AtomicBoolean callback1Ran = new AtomicBoolean(); + AtomicBoolean callback2Ran = new AtomicBoolean(); + AtomicBoolean callbackLastRan = new AtomicBoolean(); + + Future future = + executor.submit( + () -> { + observer.acceptWithCallback(dataWith("ABC"), () -> callback1Ran.set(true)); + assertTrue( + observer.tryAcceptWithCallback(dataWith("DEF"), () -> callback2Ran.set(true))); + observer.acceptWithCallback(lastData(), () -> callbackLastRan.set(true)); + return null; + }); + + observer.awaitCompletion(); + future.get(); + + assertTrue(callback1Ran.get()); + assertTrue(callback2Ran.get()); + assertTrue(callbackLastRan.get()); + assertThat( + values, + contains( + WindowedValues.valueInGlobalWindow("ABC"), WindowedValues.valueInGlobalWindow("DEF"))); + } + + @Test + public void testCallbacksInvokedOnClose() throws Exception { + BeamFnDataInboundObserver observer = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList(DataEndpoint.create(TRANSFORM_ID, CODER, val -> {})), + Collections.emptyList()); + + AtomicBoolean callback1 = new AtomicBoolean(); + AtomicBoolean callback2 = new AtomicBoolean(); + + observer.acceptWithCallback(dataWith("ABC"), () -> callback1.set(true)); + observer.tryAcceptWithCallback(dataWith("DEF"), () -> callback2.set(true)); + + observer.close(); + + assertTrue(callback1.get()); + assertTrue(callback2.get()); + } + + @Test + public void testCallbacksInvokedOnConsumerException() throws Exception { + BeamFnDataInboundObserver observer = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList( + DataEndpoint.create( + TRANSFORM_ID, + CODER, + val -> { + throw new RuntimeException("consumer error"); + })), + Collections.emptyList()); + + AtomicBoolean callback1 = new AtomicBoolean(); + AtomicBoolean callback2 = new AtomicBoolean(); + + observer.acceptWithCallback(dataWith("ABC"), () -> callback1.set(true)); + observer.acceptWithCallback(dataWith("DEF"), () -> callback2.set(true)); + + assertThrows(RuntimeException.class, observer::awaitCompletion); + + assertFalse(callback1.get()); + assertTrue(callback2.get()); + } + private BeamFnApi.Elements dataWith(String... values) throws Exception { ByteStringOutputStream output = new ByteStringOutputStream(); for (String value : values) { diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserverTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserverTest.java index a294ba0bedc0..63f2f205336e 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserverTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserverTest.java @@ -28,6 +28,7 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import org.apache.beam.sdk.fn.test.TestExecutors; import org.apache.beam.sdk.fn.test.TestExecutors.TestExecutorService; import org.apache.beam.sdk.fn.test.TestStreams; @@ -129,4 +130,32 @@ public void testIsReadyIsHonored() throws Exception { } streamObserver.onCompleted(); } + + @Test(timeout = 10_000) + public void testStalledPollHandlerInvokedWhileNotReady() throws Exception { + AdvancingPhaser phaser = new AdvancingPhaser(1); + AtomicBoolean isReady = new AtomicBoolean(false); + AtomicInteger stalls = new AtomicInteger(); + List values = new ArrayList<>(); + BufferingStreamObserver streamObserver = + new BufferingStreamObserver<>( + phaser, + TestStreams.withOnNext((String t) -> values.add(t)).withIsReady(isReady::get).build(), + executor, + 3, + stalls::incrementAndGet); + + streamObserver.onNext("A"); + while (stalls.get() == 0) { + Uninterruptibles.sleepUninterruptibly(50, TimeUnit.MILLISECONDS); + } + // The handler runs every 500 ms of waiting, not on each 50 ms poll. + Uninterruptibles.sleepUninterruptibly(300, TimeUnit.MILLISECONDS); + assertTrue(stalls.get() <= 2); + + // The drainer polls readiness, so it sends without an onReady callback. + isReady.set(true); + streamObserver.onCompleted(); + assertEquals(ImmutableList.of("A"), values); + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java index a8a1385749b8..404b355c7a2a 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java @@ -57,7 +57,7 @@ public void testThreadSafety() throws Exception { AdvancingPhaser phaser = new AdvancingPhaser(1); final AtomicBoolean isCriticalSectionShared = new AtomicBoolean(); final DirectStreamObserver streamObserver = - new DirectStreamObserver<>( + new DirectStreamObserver( phaser, TestStreams.withOnNext( (String t) -> { @@ -101,7 +101,7 @@ public void testIsReadyIsHonored() throws Exception { AdvancingPhaser phaser = new AdvancingPhaser(1); final AtomicBoolean elementsAllowed = new AtomicBoolean(); final DirectStreamObserver streamObserver = - new DirectStreamObserver<>( + new DirectStreamObserver( phaser, TestStreams.withOnNext((String t) -> assertTrue(elementsAllowed.get())) .withIsReady(elementsAllowed::get) @@ -137,7 +137,7 @@ public void testIsReadyIsHonoredTermination() throws Exception { AdvancingPhaser phaser = new AdvancingPhaser(1); final AtomicBoolean elementsAllowed = new AtomicBoolean(); final DirectStreamObserver streamObserver = - new DirectStreamObserver<>( + new DirectStreamObserver( phaser, TestStreams.withOnNext( (String t) -> { @@ -191,7 +191,7 @@ public void testIsReadyCheckDoesntBlockIfPhaserCallbackNeverHappens() throws Exc // Note that we never advance the phaser in this test. final AtomicBoolean elementsAllowed = new AtomicBoolean(); final DirectStreamObserver streamObserver = - new DirectStreamObserver<>( + new DirectStreamObserver( new AdvancingPhaser(1), TestStreams.withOnNext((String t) -> assertTrue(elementsAllowed.get())) .withIsReady(elementsAllowed::get) @@ -237,7 +237,7 @@ public void testMessageCheckInterval() throws Exception { .build(); final DirectStreamObserver streamObserver = - new DirectStreamObserver<>(new AdvancingPhaser(1), bucketingObserver, 10); + new DirectStreamObserver(new AdvancingPhaser(1), bucketingObserver, 10); List prefixes = ImmutableList.of("0", "1", "2", "3", "4"); List> results = new ArrayList<>(); @@ -274,7 +274,7 @@ public void testPhaserTermination() throws Exception { final AtomicInteger index = new AtomicInteger(); ArrayListMultimap values = ArrayListMultimap.create(); final DirectStreamObserver streamObserver = - new DirectStreamObserver<>( + new DirectStreamObserver( new AdvancingPhaser(1), TestStreams.withOnNext((String t) -> assertTrue(values.put(index.get(), t))) .withIsReady( @@ -314,4 +314,31 @@ public void testPhaserTermination() throws Exception { prefixesIndex[prefix] += 1; } } + + @Test + public void testStalledPollHandlerInvokedWhenStalled() throws Exception { + AdvancingPhaser phaser = new AdvancingPhaser(1); + final AtomicBoolean isReady = new AtomicBoolean(false); + final AtomicInteger stallCount = new AtomicInteger(0); + StalledPollHandler handler = stallCount::incrementAndGet; + + final DirectStreamObserver streamObserver = + new DirectStreamObserver<>( + phaser, + TestStreams.withOnNext((String t) -> {}).withIsReady(isReady::get).build(), + handler); + + Future future = executor.submit(() -> streamObserver.onNext("test")); + + while (stallCount.get() == 0) { + Uninterruptibles.sleepUninterruptibly(50, TimeUnit.MILLISECONDS); + } + assertTrue(stallCount.get() >= 1); + + isReady.set(true); + phaser.arrive(); + + future.get(5, TimeUnit.SECONDS); + streamObserver.onCompleted(); + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java index fc0ebf3ca4e4..9cf208f46259 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java @@ -89,4 +89,38 @@ public void testCallsAreForwardedAndOnReadyHandlerBoundError() { verifyNoMoreInteractions(delegateObserver, callStreamObserver); verify(onDoneHandler).run(); } + + private interface StalledResponseObserver + extends ClientResponseObserver, StalledPollHandler {} + + @Test + public void testBeforeStartForwardsToClientResponseObserver() { + @SuppressWarnings("unchecked") + ClientResponseObserver delegateObserver = mock(ClientResponseObserver.class); + @SuppressWarnings("unchecked") + ClientCallStreamObserver callStreamObserver = mock(ClientCallStreamObserver.class); + Runnable onReadyHandler = () -> {}; + ClientResponseObserver observer = + ForwardingClientResponseObserver.create(delegateObserver, onReadyHandler); + observer.beforeStart(callStreamObserver); + verify(callStreamObserver).setOnReadyHandler(onReadyHandler); + verify(delegateObserver).beforeStart(callStreamObserver); + } + + @Test + public void testOnStalledPollForwardsToStalledPollHandler() { + @SuppressWarnings("unchecked") + StalledResponseObserver delegateObserver = mock(StalledResponseObserver.class); + ForwardingClientResponseObserver observer = + ForwardingClientResponseObserver.create(delegateObserver, () -> {}); + observer.onStalledPoll(); + verify(delegateObserver).onStalledPoll(); + + @SuppressWarnings("unchecked") + StreamObserver plainObserver = mock(StreamObserver.class); + ForwardingClientResponseObserver observer2 = + ForwardingClientResponseObserver.create(plainObserver, () -> {}); + observer2.onStalledPoll(); + verifyNoMoreInteractions(plainObserver); + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java index ee523dd83412..5ad89087e5a8 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java @@ -20,6 +20,10 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.instanceOf; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import java.util.concurrent.Executors; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.CallStreamObserver; @@ -37,6 +41,8 @@ public class OutboundObserverFactoryTest { @Mock private StreamObserver mockRequestObserver; @Mock private CallStreamObserver mockResponseObserver; + private interface StalledStreamObserver extends StreamObserver, StalledPollHandler {} + @Before public void setUp() { MockitoAnnotations.initMocks(this); @@ -50,6 +56,36 @@ public void testDefaultInstantiation() { assertThat(observer, instanceOf(DirectStreamObserver.class)); } + @Test + public void testClientDirectWithStalledPollHandler() { + @SuppressWarnings("unchecked") + StalledStreamObserver stalledObserver = mock(StalledStreamObserver.class); + StreamObserver observer = + OutboundObserverFactory.clientDirect() + .outboundObserverFor(this.fakeFactory(), stalledObserver); + assertThat(observer, instanceOf(DirectStreamObserver.class)); + } + + @Test + public void testServerDirectInstantiation() { + StreamObserver observer = + OutboundObserverFactory.serverDirect() + .outboundObserverFor(inbound -> mockResponseObserver, mockRequestObserver); + assertThat(observer, instanceOf(DirectStreamObserver.class)); + verify(mockResponseObserver).setOnReadyHandler(any()); + } + + @Test + public void testServerDirectWithStalledPollHandler() { + @SuppressWarnings("unchecked") + StalledStreamObserver stalledObserver = mock(StalledStreamObserver.class); + StreamObserver observer = + OutboundObserverFactory.serverDirect() + .outboundObserverFor(inbound -> mockResponseObserver, stalledObserver); + assertThat(observer, instanceOf(DirectStreamObserver.class)); + verify(mockResponseObserver).setOnReadyHandler(any()); + } + @Test public void testBufferedStreamInstantiation() { StreamObserver observer = @@ -67,6 +103,14 @@ public void testBufferedStreamWithLimitInstantiation() { assertEquals(1, ((BufferingStreamObserver) observer).getBufferSize()); } + @Test + public void testTrivial() { + StreamObserver observer = + OutboundObserverFactory.trivial() + .outboundObserverFor(inbound -> mockResponseObserver, mockRequestObserver); + assertSame(mockResponseObserver, observer); + } + private OutboundObserverFactory.BasicFactory fakeFactory() { return inboundObserver -> { assertThat(inboundObserver, instanceOf(ForwardingClientResponseObserver.class)); From b7532bb089ac3f422bda8d8e71a6abcc0ebbc835 Mon Sep 17 00:00:00 2001 From: Pablo Rodriguez Defino Date: Fri, 16 Oct 2026 15:47:52 -0700 Subject: [PATCH 2/5] Release inbound permits on failure and poisoned paths Fixes from review of the Fn data plane change: - BeamFnDataInboundObserver runs the callback of the element that was being processed when the consumer throws. queue.take() had removed it, so the cancel did not reach it, and each failed bundle kept one inbound permit for the life of the multiplexer. - Late data for a poisoned instruction no longer adds a ReceiverHolder to receivers. The entry was never removed, grew the stall window, and let a later registerConsumer succeed on a poisoned holder. - ReceiverHolder.complete refuses a poisoned holder and runs every buffered callback even if the receiver throws. - registerConsumer throws IllegalArgumentException unwrapped again. - requestInbound releases count permits instead of one. Tests count request(1) calls on a mocked stream for a failed bundle, late data for a poisoned instruction, and a receiver that throws on buffered elements, and restore the expected behavior in three earlier tests. --- .../fn/data/BeamFnDataGrpcMultiplexer.java | 77 +++++--- .../fn/data/BeamFnDataInboundObserver.java | 8 + .../data/BeamFnDataGrpcMultiplexerTest.java | 184 +++++++++++++++--- .../data/BeamFnDataInboundObserverTest.java | 4 +- 4 files changed, 222 insertions(+), 51 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java index c76afba849c0..346c0688cf0f 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java @@ -155,20 +155,35 @@ synchronized boolean accept(BeamFnApi.Elements elements, Runnable callback) thro } } + /** + * Sets the receiver and hands it the buffered elements. Returns false if the holder already has + * a receiver or is poisoned. Every buffered callback runs, even if the receiver throws. + */ synchronized boolean complete(CloseableFnDataReceiver receiver) throws Exception { - if (this.receiver != null) { + if (this.receiver != null || poisoned) { return false; } this.receiver = receiver; if (receiver instanceof BeamFnDataInboundObserver) { drainPending(); } else { - for (PendingElement p : pending) { - receiver.accept(p.elements); - p.callback.run(); + try { + while (!pending.isEmpty()) { + PendingElement p = pending.removeFirst(); + try { + receiver.accept(p.elements); + } finally { + p.callback.run(); + } + } + } finally { + // If the receiver threw, release the callbacks of the elements it did not see. + for (PendingElement p : pending) { + p.callback.run(); + } + pending.clear(); } - pending.clear(); } return true; } @@ -185,6 +200,10 @@ synchronized boolean isDone() { return receiver != null || poisoned; } + synchronized boolean isPoisoned() { + return poisoned; + } + synchronized void close() throws Exception { if (receiver != null) { receiver.close(); @@ -227,7 +246,7 @@ public String toString() { * STEADY_STATE_INBOUND_WINDOW} messages are open again. */ public void requestInbound(int count) { - inFlightInbound.decrementAndGet(); + inFlightInbound.addAndGet(-count); synchronized (inboundObserver) { while (inFlightInbound.get() < STEADY_STATE_INBOUND_WINDOW) { inFlightInbound.incrementAndGet(); @@ -261,29 +280,35 @@ public void registerConsumer( instructionId, (unused, existing) -> { if (existing != null) { - try { - if (!existing.complete(receiver)) { - throw new IllegalArgumentException("Instruction id was registered twice"); - } - } catch (Exception e) { - throw new RuntimeException(e); + if (existing.isPoisoned()) { + throw new IllegalArgumentException("Instruction id was poisoned"); + } + if (!complete(existing, receiver)) { + throw new IllegalArgumentException("Instruction id was registered twice"); } - return existing; } if (poisonedInstructionIds.getIfPresent(instructionId) != null) { throw new IllegalArgumentException("Instruction id was poisoned"); } ReceiverHolder h = new ReceiverHolder(); - try { - h.complete(receiver); - } catch (Exception e) { - throw new RuntimeException(e); - } + complete(h, receiver); return h; }); } + /** {@link ReceiverHolder#complete} that rethrows runtime exceptions as they are. */ + private static boolean complete( + ReceiverHolder holder, CloseableFnDataReceiver receiver) { + try { + return holder.complete(receiver); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + /** Unregisters a previously registered consumer. */ public void unregisterConsumer(String instructionId) { @Nullable ReceiverHolder receiverFuture = receivers.remove(instructionId); @@ -447,17 +472,21 @@ public void onNext(BeamFnApi.Elements value) { } } + // computeIfAbsent documents that a null result inserts no mapping; the nullness checker's + // model of the function does not allow null. The poison check stays inside computeIfAbsent so + // it cannot race with poisonInstructionId. + @SuppressWarnings("return") private boolean forwardToConsumerForInstructionId( String instructionId, BeamFnApi.Elements value, Runnable callback) { try { - ReceiverHolder consumerFuture = + // computeIfAbsent inserts nothing when the function returns null, so data for a poisoned + // instruction leaves no entry in receivers. + @Nullable ReceiverHolder consumerFuture = receivers.computeIfAbsent( instructionId, (unused) -> { if (poisonedInstructionIds.getIfPresent(instructionId) != null) { - ReceiverHolder h = new ReceiverHolder(); - h.poison(); - return h; + return null; } LOG.debug( "Received data for instruction {} without consumer ready. " @@ -465,6 +494,10 @@ private boolean forwardToConsumerForInstructionId( instructionId); return new ReceiverHolder(); }); + if (consumerFuture == null) { + LOG.debug("Received data for poisoned instruction {}. Dropping input.", instructionId); + return false; + } return consumerFuture.accept(value, callback); } catch (Exception e) { poisonInstructionId(instructionId); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java index bf43a67e978f..e4930debe959 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java @@ -186,6 +186,9 @@ private static Iterator createDiscardingIterator(List list) { *

Erroneous elements passed from the producer will be visible to the caller of this method. */ public void awaitCompletion() throws Exception { + // The callback of the element that is being processed. queue.take() has removed the element, + // so a cancel does not reach it, and it must still run if processing fails. + @Nullable Runnable inProgressCallback = null; try { while (true) { // The SDK indicates it has consumed all the received data before it attempts to take @@ -200,6 +203,7 @@ public void awaitCompletion() throws Exception { ElementWithCallback wrapper; { wrapper = queue.take(); + inProgressCallback = wrapper.callback; BeamFnApi.Elements elements = wrapper.elements; dataIterator = createDiscardingIterator(elements.getDataList()); timersIterator = createDiscardingIterator(elements.getTimersList()); @@ -207,6 +211,7 @@ public void awaitCompletion() throws Exception { // The SDK is no longer blocked on receiving more data from the Runner. consumingReceivedData.set(true); boolean isDone = multiplexElements(dataIterator, timersIterator); + inProgressCallback = null; wrapper.callback.run(); if (isDone) { return; @@ -214,6 +219,9 @@ public void awaitCompletion() throws Exception { } } catch (Exception e) { queue.cancel(e, wrapper -> wrapper.callback.run()); + if (inProgressCallback != null) { + inProgressCallback.run(); + } throw e; } finally { consumingReceivedData.set(false); diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java index 81bd4e3638da..83ef535b9f15 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java @@ -477,24 +477,22 @@ public void close() {} assertThat(received, Matchers.contains(elements)); - // Registering twice throws IllegalArgumentException wrapped in RuntimeException - RuntimeException e = - assertThrows( - RuntimeException.class, - () -> - multiplexer.registerConsumer( - instructionId, - new CloseableFnDataReceiver() { - @Override - public void accept(BeamFnApi.Elements input) {} - - @Override - public void flush() {} - - @Override - public void close() {} - })); - assertThat(e.getCause(), Matchers.instanceOf(IllegalArgumentException.class)); + // Registering twice throws IllegalArgumentException. + assertThrows( + IllegalArgumentException.class, + () -> + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) {} + + @Override + public void flush() {} + + @Override + public void close() {} + })); } @Test @@ -593,8 +591,18 @@ public void close() { multiplexer.poisonInstructionId(instructionId); assertTrue(closed.get()); + assertFalse(multiplexer.hasConsumer(instructionId)); + + // Late elements for the poisoned instruction are dropped and leave no entry behind. + multiplexer + .getInboundObserver() + .onNext( + BeamFnApi.Elements.newBuilder() + .addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(instructionId)) + .build()); + assertFalse(multiplexer.hasConsumer(instructionId)); - // Registering consumer for poisoned instruction fails + // Registering a consumer for the poisoned instruction still fails after the late data. assertThrows( IllegalArgumentException.class, () -> @@ -610,14 +618,7 @@ public void flush() {} @Override public void close() {} })); - - // Subsequent elements for poisoned instruction are dropped - multiplexer - .getInboundObserver() - .onNext( - BeamFnApi.Elements.newBuilder() - .addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(instructionId)) - .build()); + assertFalse(multiplexer.hasConsumer(instructionId)); } @Test @@ -845,4 +846,133 @@ public void testConcurrentHandoffKeepsOrder() throws Exception { multiplexer.unregisterConsumer(instructionId); } } + + /** A multiplexer whose inbound stream is a mock, so a test can count {@code request(1)} calls. */ + @SuppressWarnings("unchecked") + private static ClientCallStreamObserver startWithMockStream( + BeamFnDataGrpcMultiplexer multiplexer) { + ClientCallStreamObserver stream = mock(ClientCallStreamObserver.class); + ((ClientResponseObserver) + multiplexer.getInboundObserver()) + .beforeStart(stream); + return stream; + } + + private static BeamFnApi.Elements dataFor(String instructionId, String transformId) + throws Exception { + Coder> coder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + ByteStringOutputStream out = new ByteStringOutputStream(); + coder.encode(WindowedValues.valueInGlobalWindow("v"), out); + return BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId(transformId) + .setData(out.toByteString())) + .build(); + } + + private static BeamFnDataGrpcMultiplexer newMultiplexer() { + return new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + } + + /** + * Each inbound message releases its flow-control permit once, which asks gRPC for one more + * message. A bundle whose consumer throws must release the permit of the message it was + * processing as well as of the messages still queued. + */ + @Test + public void testPermitsReleasedWhenBundleFails() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = newMultiplexer(); + ClientCallStreamObserver stream = startWithMockStream(multiplexer); + String instructionId = "failing_bundle"; + Coder> coder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + BeamFnDataInboundObserver observer = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList( + DataEndpoint.create( + "t1", + coder, + value -> { + throw new IllegalStateException("consumer failed"); + })), + Collections.emptyList()); + multiplexer.registerConsumer(instructionId, observer); + + for (int i = 0; i < 3; i++) { + multiplexer.getInboundObserver().onNext(dataFor(instructionId, "t1")); + } + assertThrows(IllegalStateException.class, observer::awaitCompletion); + multiplexer.poisonInstructionId(instructionId); + + verify(stream, times(3)).request(1); + assertFalse(multiplexer.hasConsumer(instructionId)); + } + + @Test + public void testPermitsReleasedForLateDataOfPoisonedInstruction() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = newMultiplexer(); + ClientCallStreamObserver stream = startWithMockStream(multiplexer); + String instructionId = "poisoned_late_data"; + multiplexer.poisonInstructionId(instructionId); + + for (int i = 0; i < 3; i++) { + multiplexer.getInboundObserver().onNext(dataFor(instructionId, "t1")); + } + + verify(stream, times(3)).request(1); + assertFalse(multiplexer.hasConsumer(instructionId)); + } + + @Test + public void testPermitsReleasedWhenReceiverThrowsOnBufferedElements() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = newMultiplexer(); + ClientCallStreamObserver stream = startWithMockStream(multiplexer); + String instructionId = "throwing_receiver"; + for (int i = 0; i < 3; i++) { + multiplexer.getInboundObserver().onNext(dataFor(instructionId, "t1")); + } + verify(stream, times(0)).request(1); + + // The receiver's own exception reaches the caller, not a wrapped one. + assertThrows( + IllegalStateException.class, + () -> + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) { + throw new IllegalStateException("receiver failed"); + } + + @Override + public void flush() {} + + @Override + public void close() {} + })); + + verify(stream, times(3)).request(1); + } + + @Test + public void testRequestInboundReleasesCount() { + BeamFnDataGrpcMultiplexer multiplexer = newMultiplexer(); + ClientCallStreamObserver stream = startWithMockStream(multiplexer); + // Three stalled polls open three messages beyond the steady-state window. + multiplexer.onStalledPoll(); + multiplexer.onStalledPoll(); + multiplexer.onStalledPoll(); + verify(stream, times(3)).request(1); + + // Releasing five brings the window two below steady state, so it asks for two more. + multiplexer.requestInbound(5); + verify(stream, times(5)).request(1); + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java index 24ebdf067b82..4b5bb0f0500c 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java @@ -19,7 +19,6 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.contains; -import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; @@ -294,7 +293,8 @@ public void testCallbacksInvokedOnConsumerException() throws Exception { assertThrows(RuntimeException.class, observer::awaitCompletion); - assertFalse(callback1.get()); + // callback1 belongs to the element that was being processed when the consumer threw. + assertTrue(callback1.get()); assertTrue(callback2.get()); } From 6f9d4442a1e3858c68fd8f8941bba2682a01d275 Mon Sep 17 00:00:00 2001 From: Pablo Rodriguez Defino Date: Fri, 16 Oct 2026 17:13:38 -0700 Subject: [PATCH 3/5] Poison on failed registration, keep elements collectable, quiet short stalls Fixes from the next review of the Fn data plane change: - registerConsumer poisons the instruction when the receiver throws on a buffered element. The exception aborted compute(), which left the holder in receivers with its receiver open. Poisoning removes it and closes the receiver, as a failure on the inbound gRPC thread does. A duplicate or poisoned registration still only throws. - awaitCompletion keeps the queued element inside its scope block again, so the Elements proto can be collected while it is processed, and releases the callback of that element in finally, so an Error also releases it. - DirectStreamObserver logs a stall only after a second of waiting, as with the earlier one-second poll. Shorter waits logged "0s" and a false "OnReady notification was not invoked" when isReady turned true just before the onReady callback ran. The stall hook still runs every 500 ms. --- .../fn/data/BeamFnDataGrpcMultiplexer.java | 63 ++++++++++------ .../fn/data/BeamFnDataInboundObserver.java | 21 +++--- .../sdk/fn/stream/DirectStreamObserver.java | 4 +- .../data/BeamFnDataGrpcMultiplexerTest.java | 72 ++++++++++++++++++- .../data/BeamFnDataInboundObserverTest.java | 26 +++++++ .../fn/stream/DirectStreamObserverTest.java | 29 ++++++++ 6 files changed, 184 insertions(+), 31 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java index 346c0688cf0f..4a893a661438 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java @@ -276,36 +276,59 @@ public StreamObserver getOutboundObserver() { */ public void registerConsumer( String instructionId, CloseableFnDataReceiver receiver) { - receivers.compute( - instructionId, - (unused, existing) -> { - if (existing != null) { - if (existing.isPoisoned()) { - throw new IllegalArgumentException("Instruction id was poisoned"); + try { + receivers.compute( + instructionId, + (unused, existing) -> { + if (existing != null) { + if (existing.isPoisoned()) { + throw new IllegalArgumentException("Instruction id was poisoned"); + } + if (!complete(existing, receiver)) { + throw new IllegalArgumentException("Instruction id was registered twice"); + } + return existing; } - if (!complete(existing, receiver)) { - throw new IllegalArgumentException("Instruction id was registered twice"); + if (poisonedInstructionIds.getIfPresent(instructionId) != null) { + throw new IllegalArgumentException("Instruction id was poisoned"); } - return existing; - } - if (poisonedInstructionIds.getIfPresent(instructionId) != null) { - throw new IllegalArgumentException("Instruction id was poisoned"); - } - ReceiverHolder h = new ReceiverHolder(); - complete(h, receiver); - return h; - }); + ReceiverHolder h = new ReceiverHolder(); + complete(h, receiver); + return h; + }); + } catch (ReceiverFailedException e) { + // The receiver threw on a buffered element. compute() left the holder in receivers, so + // poison the instruction, which removes it and closes the receiver, as a failure on the + // inbound gRPC thread does. + poisonInstructionId(instructionId); + throw e.getCause(); + } + } + + /** Thrown by {@link #complete} when the receiver fails, as opposed to a rejected registration. */ + private static final class ReceiverFailedException extends RuntimeException { + private final RuntimeException cause; + + ReceiverFailedException(RuntimeException cause) { + super(cause); + this.cause = cause; + } + + @Override + public synchronized RuntimeException getCause() { + return cause; + } } - /** {@link ReceiverHolder#complete} that rethrows runtime exceptions as they are. */ + /** {@link ReceiverHolder#complete} that reports a failure of the receiver. */ private static boolean complete( ReceiverHolder holder, CloseableFnDataReceiver receiver) { try { return holder.complete(receiver); } catch (RuntimeException e) { - throw e; + throw new ReceiverFailedException(e); } catch (Exception e) { - throw new RuntimeException(e); + throw new ReceiverFailedException(new RuntimeException(e)); } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java index e4930debe959..a3ebb0c0e4af 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java @@ -187,7 +187,7 @@ private static Iterator createDiscardingIterator(List list) { */ public void awaitCompletion() throws Exception { // The callback of the element that is being processed. queue.take() has removed the element, - // so a cancel does not reach it, and it must still run if processing fails. + // so a cancel does not reach it, and it must still run if processing fails in any way. @Nullable Runnable inProgressCallback = null; try { while (true) { @@ -200,9 +200,8 @@ public void awaitCompletion() throws Exception { // collection. Iterator dataIterator; Iterator timersIterator; - ElementWithCallback wrapper; { - wrapper = queue.take(); + ElementWithCallback wrapper = queue.take(); inProgressCallback = wrapper.callback; BeamFnApi.Elements elements = wrapper.elements; dataIterator = createDiscardingIterator(elements.getDataList()); @@ -211,21 +210,25 @@ public void awaitCompletion() throws Exception { // The SDK is no longer blocked on receiving more data from the Runner. consumingReceivedData.set(true); boolean isDone = multiplexElements(dataIterator, timersIterator); + Runnable callback = inProgressCallback; inProgressCallback = null; - wrapper.callback.run(); + callback.run(); if (isDone) { return; } } } catch (Exception e) { queue.cancel(e, wrapper -> wrapper.callback.run()); - if (inProgressCallback != null) { - inProgressCallback.run(); - } throw e; } finally { - consumingReceivedData.set(false); - close(); + try { + consumingReceivedData.set(false); + close(); + } finally { + if (inProgressCallback != null) { + inProgressCallback.run(); + } + } } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java index ca5ff6f5ed03..21f66e0699cd 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java @@ -112,7 +112,9 @@ public void onNext(T value) { throw new RuntimeException(e); } } - if (totalMsWaited > 0) { + // Log only waits of a second or more, as with the earlier one-second poll. A shorter wait + // is normal backpressure, and isReady() can turn true just before onReady runs. + if (totalMsWaited >= 1000) { // If the phase didn't change, this means that the installed onReady callback had not // been invoked. if (initialPhase == phase) { diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java index 83ef535b9f15..12e52ab16a60 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java @@ -940,6 +940,7 @@ public void testPermitsReleasedWhenReceiverThrowsOnBufferedElements() throws Exc verify(stream, times(0)).request(1); // The receiver's own exception reaches the caller, not a wrapped one. + AtomicBoolean closed = new AtomicBoolean(); assertThrows( IllegalStateException.class, () -> @@ -955,10 +956,79 @@ public void accept(BeamFnApi.Elements input) { public void flush() {} @Override - public void close() {} + public void close() { + closed.set(true); + } })); verify(stream, times(3)).request(1); + // As for a failure on the inbound gRPC thread, the instruction is poisoned: its entry is + // removed, the receiver is closed, and a later registration fails. + assertFalse(multiplexer.hasConsumer(instructionId)); + assertTrue(closed.get()); + IllegalArgumentException poisoned = + assertThrows( + IllegalArgumentException.class, + () -> + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) {} + + @Override + public void flush() {} + + @Override + public void close() {} + })); + assertThat(poisoned.getMessage(), containsString("poisoned")); + } + + @Test + public void testDuplicateRegistrationDoesNotPoison() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = newMultiplexer(); + String instructionId = "duplicate"; + AtomicBoolean closed = new AtomicBoolean(); + List received = new ArrayList<>(); + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) { + received.add(input); + } + + @Override + public void flush() {} + + @Override + public void close() { + closed.set(true); + } + }); + assertThrows( + IllegalArgumentException.class, + () -> + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) {} + + @Override + public void flush() {} + + @Override + public void close() {} + })); + + // The first registration keeps working. + assertTrue(multiplexer.hasConsumer(instructionId)); + assertFalse(closed.get()); + BeamFnApi.Elements elements = dataFor(instructionId, "t1"); + multiplexer.getInboundObserver().onNext(elements); + assertThat(received, Matchers.contains(elements)); } @Test diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java index 4b5bb0f0500c..27d895270e1b 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java @@ -298,6 +298,32 @@ public void testCallbacksInvokedOnConsumerException() throws Exception { assertTrue(callback2.get()); } + @Test + public void testCallbacksInvokedOnConsumerError() throws Exception { + BeamFnDataInboundObserver observer = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList( + DataEndpoint.create( + TRANSFORM_ID, + CODER, + val -> { + throw new AssertionError("consumer error"); + })), + Collections.emptyList()); + + AtomicBoolean callback1 = new AtomicBoolean(); + AtomicBoolean callback2 = new AtomicBoolean(); + + observer.acceptWithCallback(dataWith("ABC"), () -> callback1.set(true)); + observer.acceptWithCallback(dataWith("DEF"), () -> callback2.set(true)); + + assertThrows(AssertionError.class, observer::awaitCompletion); + + // An Error is not an Exception, and both callbacks must still run. + assertTrue(callback1.get()); + assertTrue(callback2.get()); + } + private BeamFnApi.Elements dataWith(String... values) throws Exception { ByteStringOutputStream output = new ByteStringOutputStream(); for (String value : values) { diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java index 404b355c7a2a..789d982e4333 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java @@ -36,6 +36,7 @@ import org.apache.beam.sdk.fn.test.TestExecutors; import org.apache.beam.sdk.fn.test.TestExecutors.TestExecutorService; import org.apache.beam.sdk.fn.test.TestStreams; +import org.apache.beam.sdk.testing.ExpectedLogs; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.CallStreamObserver; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ArrayListMultimap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; @@ -50,6 +51,7 @@ @RunWith(JUnit4.class) public class DirectStreamObserverTest { @Rule public TestExecutorService executor = TestExecutors.from(Executors.newCachedThreadPool()); + @Rule public ExpectedLogs expectedLogs = ExpectedLogs.none(DirectStreamObserver.class); @Test public void testThreadSafety() throws Exception { @@ -341,4 +343,31 @@ public void testStalledPollHandlerInvokedWhenStalled() throws Exception { future.get(5, TimeUnit.SECONDS); streamObserver.onCompleted(); } + + /** Sends one element on a stream that is not ready for about {@code notReadyMs}. */ + private void sendWithStall(long notReadyMs) throws Exception { + AtomicBoolean isReady = new AtomicBoolean(false); + DirectStreamObserver streamObserver = + new DirectStreamObserver<>( + new AdvancingPhaser(1), + TestStreams.withOnNext((String t) -> {}).withIsReady(isReady::get).build(), + 0); + Future future = executor.submit(() -> streamObserver.onNext("A")); + Uninterruptibles.sleepUninterruptibly(notReadyMs, TimeUnit.MILLISECONDS); + // The phaser never advances, as when onReady has not run yet. + isReady.set(true); + future.get(10, TimeUnit.SECONDS); + } + + @Test + public void testShortStallIsNotLogged() throws Exception { + sendWithStall(300); + expectedLogs.verifyNotLogged("Output channel stalled"); + } + + @Test + public void testStallOfASecondIsLogged() throws Exception { + sendWithStall(1_300); + expectedLogs.verifyInfo("Output channel stalled for 1s"); + } } From 9388a8a89b7d5c1f8642e274a555091421493209 Mon Sep 17 00:00:00 2001 From: Pablo Rodriguez Defino Date: Fri, 16 Oct 2026 17:34:17 -0700 Subject: [PATCH 4/5] Make the DirectStreamObserver stall log tests independent of timing DirectStreamObserver counts 50 ms for each wait that times out, not the wall time. The tests kept the stream not ready for a wall time, so on a slow machine fewer waits timed out, the counted wait stayed under a second, and testStallOfASecondIsLogged failed. The stream now becomes ready after a number of readiness polls, so the counted wait is exact. --- .../fn/stream/DirectStreamObserverTest.java | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java index 789d982e4333..7d4a345c8ca5 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java @@ -344,30 +344,36 @@ public void testStalledPollHandlerInvokedWhenStalled() throws Exception { streamObserver.onCompleted(); } - /** Sends one element on a stream that is not ready for about {@code notReadyMs}. */ - private void sendWithStall(long notReadyMs) throws Exception { - AtomicBoolean isReady = new AtomicBoolean(false); + /** + * Sends one element on a stream that becomes ready after {@code timeouts} waits time out. + * + *

The observer counts 50 ms for each wait that times out, whatever the wall time, and polls + * {@code isReady()} once before it waits and once after each timeout. Counting polls keeps the + * counted wait exact on a slow machine. The phaser never advances, as when onReady has not run. + */ + private void sendWithStall(int timeouts) throws Exception { + AtomicInteger polls = new AtomicInteger(); DirectStreamObserver streamObserver = new DirectStreamObserver<>( new AdvancingPhaser(1), - TestStreams.withOnNext((String t) -> {}).withIsReady(isReady::get).build(), + TestStreams.withOnNext((String t) -> {}) + .withIsReady(() -> polls.getAndIncrement() >= timeouts) + .build(), 0); - Future future = executor.submit(() -> streamObserver.onNext("A")); - Uninterruptibles.sleepUninterruptibly(notReadyMs, TimeUnit.MILLISECONDS); - // The phaser never advances, as when onReady has not run yet. - isReady.set(true); - future.get(10, TimeUnit.SECONDS); + executor.submit(() -> streamObserver.onNext("A")).get(30, TimeUnit.SECONDS); } @Test public void testShortStallIsNotLogged() throws Exception { - sendWithStall(300); + // 6 timeouts count 300 ms. + sendWithStall(6); expectedLogs.verifyNotLogged("Output channel stalled"); } @Test public void testStallOfASecondIsLogged() throws Exception { - sendWithStall(1_300); + // 22 timeouts count 1,100 ms. + sendWithStall(22); expectedLogs.verifyInfo("Output channel stalled for 1s"); } } From 658b801ca6ff583d1f6c67e5b6c459bae67ef2f4 Mon Sep 17 00:00:00 2001 From: Pablo Rodriguez Defino Date: Sun, 25 Oct 2026 11:36:45 -0700 Subject: [PATCH 5/5] Check isReady() every 100 messages again in DirectStreamObserver Checking before each message cost about 17% on ProcessBundleBenchmark.testLargeBundle. Adds a test that data for an unregistered instruction does not block onReady. --- .../sdk/fn/stream/DirectStreamObserver.java | 2 +- .../InboundUnregisteredConsumerReproTest.java | 176 ++++++++++++++++++ .../fn/stream/DirectStreamObserverTest.java | 1 + 3 files changed, 178 insertions(+), 1 deletion(-) create mode 100644 sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/InboundUnregisteredConsumerReproTest.java diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java index 21f66e0699cd..3d3b58377496 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java @@ -40,7 +40,7 @@ @ThreadSafe public final class DirectStreamObserver implements StreamObserver { private static final Logger LOG = LoggerFactory.getLogger(DirectStreamObserver.class); - private static final int DEFAULT_MAX_MESSAGES_BEFORE_CHECK = 0; + private static final int DEFAULT_MAX_MESSAGES_BEFORE_CHECK = 100; private final Phaser phaser; private final CallStreamObserver outboundObserver; diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/InboundUnregisteredConsumerReproTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/InboundUnregisteredConsumerReproTest.java new file mode 100644 index 000000000000..8bfb3e68d4f2 --- /dev/null +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/InboundUnregisteredConsumerReproTest.java @@ -0,0 +1,176 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.fn.data; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; +import org.apache.beam.model.pipeline.v1.Endpoints; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.fn.stream.OutboundObserverFactory; +import org.apache.beam.sdk.util.ByteStringOutputStream; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientCallStreamObserver; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientResponseObserver; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Test; + +/** + * Data for an instruction whose consumer is not registered yet must not block the gRPC inbound + * thread. gRPC runs inbound onNext and the outbound onReady handler on one serializing executor, so + * a blocked onNext also stops every bundle that waits on onReady. + */ +public class InboundUnregisteredConsumerReproTest { + private static final Endpoints.ApiServiceDescriptor DESCRIPTOR = + Endpoints.ApiServiceDescriptor.newBuilder().setUrl("test").build(); + + @Test + public void testUnregisteredConsumerDoesNotBlockOnReady() throws Exception { + AtomicReference<@Nullable Runnable> onReadyHandler = new AtomicReference<>(); + AtomicBoolean ready = new AtomicBoolean(false); + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inbound -> { + ClientCallStreamObserver call = + new ClientCallStreamObserver() { + @Override + public boolean isReady() { + return ready.get(); + } + + @Override + public void setOnReadyHandler(Runnable handler) { + onReadyHandler.set(handler); + } + + @Override + public void disableAutoInboundFlowControl() {} + + @Override + public void disableAutoRequestWithInitial(int request) {} + + @Override + public void request(int count) {} + + @Override + public void setMessageCompression(boolean enable) {} + + @Override + public void cancel(@Nullable String message, @Nullable Throwable cause) {} + + @Override + public void onNext(BeamFnApi.Elements value) {} + + @Override + public void onError(Throwable t) {} + + @Override + public void onCompleted() {} + }; + if (inbound instanceof ClientResponseObserver) { + @SuppressWarnings("unchecked") + ClientResponseObserver r = + (ClientResponseObserver) inbound; + r.beforeStart(call); + } + return call; + }); + + // Stands in for gRPC's per-call SerializingExecutor. + ExecutorService grpcExecutor = Executors.newSingleThreadExecutor(); + try { + String instructionId = "not-yet-registered"; + ByteStringOutputStream encoded = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("elem", encoded); + BeamFnApi.Elements data = + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId("t") + .setData(encoded.toByteString())) + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId("t") + .setIsLast(true)) + .build(); + + // 1. The runner sends data before the bundle processor registers its consumer. + Future onNextFuture = + grpcExecutor.submit(() -> multiplexer.getInboundObserver().onNext(data)); + + // 2. The outbound stream becomes ready; gRPC queues onReady on the same executor. + ready.set(true); + CountDownLatch onReadyRan = new CountDownLatch(1); + Future onReadyFuture = + grpcExecutor.submit( + () -> { + Runnable handler = onReadyHandler.get(); + if (handler != null) { + handler.run(); + } + onReadyRan.countDown(); + }); + + boolean onReadyRanInTime = onReadyRan.await(2, TimeUnit.SECONDS); + + // 3. The consumer registers; the data must still arrive. + List received = Collections.synchronizedList(new ArrayList<>()); + BeamFnDataInboundObserver consumer = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList( + DataEndpoint.create("t", StringUtf8Coder.of(), received::add)), + Collections.emptyList()); + multiplexer.registerConsumer(instructionId, consumer); + ExecutorService bundle = Executors.newSingleThreadExecutor(); + try { + Future done = + bundle.submit( + () -> { + consumer.awaitCompletion(); + return null; + }); + done.get(5, TimeUnit.SECONDS); + } finally { + bundle.shutdownNow(); + } + onNextFuture.get(5, TimeUnit.SECONDS); + onReadyFuture.get(5, TimeUnit.SECONDS); + assertEquals(Collections.singletonList("elem"), received); + + assertTrue( + "onReady was stuck behind an inbound onNext waiting for an unregistered consumer", + onReadyRanInTime); + } finally { + grpcExecutor.shutdownNow(); + } + } +} diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java index 7d4a345c8ca5..f1d1d5803ae1 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/DirectStreamObserverTest.java @@ -328,6 +328,7 @@ public void testStalledPollHandlerInvokedWhenStalled() throws Exception { new DirectStreamObserver<>( phaser, TestStreams.withOnNext((String t) -> {}).withIsReady(isReady::get).build(), + 0, handler); Future future = executor.submit(() -> streamObserver.onNext("test"));