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..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 @@ -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,169 @@ * 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; + } + } + + /** + * 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 || poisoned) { + return false; + } + this.receiver = receiver; + if (receiver instanceof BeamFnDataInboundObserver) { + drainPending(); + } else { + 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(); + } + } + return true; + } + + synchronized void poison() { + poisoned = true; + for (PendingElement p : pending) { + p.callback.run(); + } + pending.clear(); + } + + synchronized boolean isDone() { + return receiver != null || poisoned; + } + + synchronized boolean isPoisoned() { + return 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 +241,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.addAndGet(-count); + synchronized (inboundObserver) { + while (inFlightInbound.get() < STEADY_STATE_INBOUND_WINDOW) { + inFlightInbound.incrementAndGet(); + inboundObserver.request(1); + } + } + } + public StreamObserver getInboundObserver() { return inboundObserver; } @@ -114,30 +276,72 @@ public StreamObserver getOutboundObserver() { */ public void registerConsumer( String instructionId, CloseableFnDataReceiver receiver) { - receivers.compute( - instructionId, - (unused, existing) -> { - if (existing != null) { - if (!existing.complete(receiver)) { - throw new IllegalArgumentException("Instruction id was registered twice"); + 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; } - return existing; - } - if (poisonedInstructionIds.getIfPresent(instructionId) != null) { - throw new IllegalArgumentException("Instruction id was poisoned"); - } - return CompletableFuture.completedFuture(receiver); - }); + if (poisonedInstructionIds.getIfPresent(instructionId) != null) { + throw new IllegalArgumentException("Instruction id was poisoned"); + } + 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 reports a failure of the receiver. */ + private static boolean complete( + ReceiverHolder holder, CloseableFnDataReceiver receiver) { + try { + return holder.complete(receiver); + } catch (RuntimeException e) { + throw new ReceiverFailedException(e); + } catch (Exception e) { + throw new ReceiverFailedException(new RuntimeException(e)); + } } /** 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 +353,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 +373,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 +400,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 +452,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 +469,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 +489,42 @@ 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; + // 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 { - CompletableFuture> 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) { - throw new PoisonedException(); + return null; } 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) { + if (consumerFuture == null) { LOG.debug("Received data for poisoned instruction {}. Dropping input.", instructionId); - return; + return false; } - 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..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 @@ -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() { @@ -166,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 in any way. + @Nullable Runnable inProgressCallback = null; try { while (true) { // The SDK indicates it has consumed all the received data before it attempts to take @@ -178,22 +201,34 @@ public void awaitCompletion() throws Exception { Iterator dataIterator; Iterator timersIterator; { - BeamFnApi.Elements elements = queue.take(); + ElementWithCallback wrapper = queue.take(); + inProgressCallback = wrapper.callback; + 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); + Runnable callback = inProgressCallback; + inProgressCallback = null; + callback.run(); + if (isDone) { return; } } } catch (Exception e) { - queue.cancel(e); + queue.cancel(e, wrapper -> wrapper.callback.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/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..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 @@ -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,38 @@ 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) { + // 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) { 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..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 @@ -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,667 @@ 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. + 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 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()); + 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 a consumer for the poisoned instruction still fails after the late data. + assertThrows( + IllegalArgumentException.class, + () -> + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) {} + + @Override + public void flush() {} + + @Override + public void close() {} + })); + assertFalse(multiplexer.hasConsumer(instructionId)); + } + + @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); + } + } + + /** 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. + AtomicBoolean closed = new AtomicBoolean(); + 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() { + 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 + 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 f915f9285367..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 @@ -21,6 +21,7 @@ import static org.hamcrest.Matchers.contains; 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 +219,111 @@ 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); + + // callback1 belongs to the element that was being processed when the consumer threw. + assertTrue(callback1.get()); + 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/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..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 @@ -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 { @@ -57,7 +59,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 +103,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 +139,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 +193,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 +239,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 +276,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 +316,64 @@ 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(); + } + + /** + * 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(() -> polls.getAndIncrement() >= timeouts) + .build(), + 0); + executor.submit(() -> streamObserver.onNext("A")).get(30, TimeUnit.SECONDS); + } + + @Test + public void testShortStallIsNotLogged() throws Exception { + // 6 timeouts count 300 ms. + sendWithStall(6); + expectedLogs.verifyNotLogged("Output channel stalled"); + } + + @Test + public void testStallOfASecondIsLogged() throws Exception { + // 22 timeouts count 1,100 ms. + sendWithStall(22); + expectedLogs.verifyInfo("Output channel stalled for 1s"); + } } 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));