github-actions[bot] commented on code in PR #40606:
URL: https://github.com/apache/beam/pull/40606#discussion_r4299667571
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java:
##########
@@ -177,19 +197,23 @@ public void awaitCompletion() throws Exception {
// collection.
Iterator<Elements.Data> dataIterator;
Iterator<Elements.Timers> 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();
Review Comment:
**🔴 Critical**: If `multiplexElements(dataIterator, timersIterator)` throws
an exception on line 209, `wrapper` has already been removed from `queue` by
`queue.take()`, so neither line 210 (`wrapper.callback.run()`) nor
`queue.cancel(e, wrapper -> wrapper.callback.run())` in the `catch` block will
invoke `wrapper.callback`. In `BeamFnDataGrpcMultiplexer`, skipping this
callback permanently leaks an `inFlightInbound` flow-control permit (and
prevents `multiCallback`'s `countdown` from reaching `0` for multi-instruction
`Elements` messages). Repeated bundle failures on a shared multiplexer will
eventually exhaust `STEADY_STATE_INBOUND_WINDOW` / `maxWindow` and deadlock
inbound data flow.
Ensure `wrapper.callback.run()` is always executed (for example, in a
`finally` block after cancelling the queue on exception).
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java:
##########
@@ -233,9 +404,12 @@ public void onNext(BeamFnApi.Elements value) {
}
}
if (instructionId == null) {
Review Comment:
**🟠High**: When late data arrives for an instruction that was already
poisoned and removed from `receivers`, returning a new poisoned
`ReceiverHolder` from `receivers.computeIfAbsent(...)` inserts it into the
unbounded `receivers` map where it is never removed. This causes:
1. A permanent memory leak in `receivers` (defeating the 20-minute TTL on
`poisonedInstructionIds`) and permanent inflation of `receivers.size()`, which
increases `maxWindow` in `onStalledPoll()` by `MAX_INBOUND_WINDOW` per poisoned
instruction.
2. If `registerConsumer` is called afterward for that `instructionId`, it
sees `existing != null` before checking `poisonedInstructionIds` and calls
`existing.complete(receiver)`, which succeeds because `complete()` only checks
`this.receiver != null` and not `this.poisoned`.
Return `null` from `computeIfAbsent` when
`poisonedInstructionIds.getIfPresent(instructionId) != null` (so no mapping is
inserted into `receivers`), guard `consumerFuture != null` before calling
`accept`, and check `if (this.receiver != null || this.poisoned) { return
false; }` in `ReceiverHolder.complete()`.
##########
sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserverTest.java:
##########
@@ -218,6 +220,84 @@ public void
testBadProducerDataFailureVisibleToAwaitCompletionCallerAndProducer(
future.get();
}
+ @Test
+ public void testAcceptWithCallbackInvokedOnProcessing() throws Exception {
+ Collection<WindowedValue<String>> 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());
Review Comment:
**🟠High**: Corresponding to the permit leak in
`BeamFnDataInboundObserver.awaitCompletion()`, the callback for the element
being processed when the consumer throws (`callback1`) must also be invoked so
that `BeamFnDataGrpcMultiplexer.inFlightInbound` is decremented.
```suggestion
assertTrue(callback1.get());
```
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java:
##########
@@ -52,23 +54,150 @@
* location with a specific outbound observer.
*/
public class BeamFnDataGrpcMultiplexer implements AutoCloseable {
+ // The inbound stream uses manual gRPC flow control, so onNext never blocks.
A blocked onNext
+ // also blocks the onReady callback of the outbound stream, because gRPC
runs both on one
+ // executor.
+ //
+ // Number of inbound messages that the stream asks for when no output waits.
+ private static final int STEADY_STATE_INBOUND_WINDOW = 16;
+ // Largest number of inbound messages for each active bundle while output
waits.
+ private static final int MAX_INBOUND_WINDOW = 100;
+ // Inbound messages that the stream asked for and that no bundle has
finished.
+ private final AtomicInteger inFlightInbound = new
AtomicInteger(STEADY_STATE_INBOUND_WINDOW);
+
+ /**
+ * Asks for one more inbound message while output waits on flow control, up
to {@code
+ * MAX_INBOUND_WINDOW} for each active bundle. The extra message lets a
bundle that waits for
+ * input continue, so the runner can read the output.
+ */
+ public void onStalledPoll() {
+ synchronized (inboundObserver) {
+ int maxWindow = Math.max(1, receivers.size()) * MAX_INBOUND_WINDOW;
+ if (inFlightInbound.get() < maxWindow) {
+ inFlightInbound.incrementAndGet();
+ inboundObserver.request(1);
+ }
+ }
+ }
+
private static final Logger LOG =
LoggerFactory.getLogger(BeamFnDataGrpcMultiplexer.class);
private static final Duration POISONED_INSTRUCTION_ID_CACHE_TIMEOUT =
Duration.ofMinutes(20);
private final Endpoints.@Nullable ApiServiceDescriptor apiServiceDescriptor;
- private final StreamObserver<BeamFnApi.Elements> inboundObserver;
+ private final InboundObserver inboundObserver;
private final StreamObserver<BeamFnApi.Elements> outboundObserver;
- private final ConcurrentHashMap<
- /* instructionId= */ String,
- CompletableFuture<CloseableFnDataReceiver<BeamFnApi.Elements>>>
- receivers;
- private final Cache</* instructionId= */ String, /* unused= */ Boolean>
poisonedInstructionIds;
- private static class PoisonedException extends RuntimeException {
- private PoisonedException() {
- super("Instruction poisoned");
+ private static final class PendingElement {
+ BeamFnApi.Elements elements;
+ Runnable callback;
+
+ PendingElement(BeamFnApi.Elements elements, Runnable callback) {
+ this.elements = elements;
+ this.callback = callback;
+ }
+ }
+
+ private static final class ReceiverHolder {
+ @Nullable CloseableFnDataReceiver<BeamFnApi.Elements> receiver;
+
+ ArrayDeque<PendingElement> pending = new ArrayDeque<>();
+ boolean poisoned;
+
+ private Runnable wrapCallback(Runnable callback) {
+ return () -> {
+ drainPending();
+ callback.run();
+ };
+ }
+
+ synchronized void drainPending() {
+ if (poisoned || !(receiver instanceof BeamFnDataInboundObserver)) {
+ return;
+ }
+ BeamFnDataInboundObserver inbound = (BeamFnDataInboundObserver) receiver;
+ while (!pending.isEmpty()) {
+ PendingElement next = pending.peekFirst();
+ if (next == null) {
+ break;
+ }
+ try {
+ if (inbound.tryAcceptWithCallback(next.elements,
wrapCallback(next.callback))) {
+ pending.removeFirst();
+ } else {
+ break;
+ }
+ } catch (Exception e) {
+ pending.removeFirst();
+ next.callback.run();
+ }
+ }
+ }
+
+ synchronized boolean accept(BeamFnApi.Elements elements, Runnable
callback) throws Exception {
+ if (poisoned) {
+ return false;
+ }
+ if (receiver != null) {
+ if (receiver instanceof BeamFnDataInboundObserver) {
+ BeamFnDataInboundObserver inbound = (BeamFnDataInboundObserver)
receiver;
+ drainPending();
+ if (!pending.isEmpty()
+ || !inbound.tryAcceptWithCallback(elements,
wrapCallback(callback))) {
+ pending.addLast(new PendingElement(elements, callback));
+ }
+ } else {
+ receiver.accept(elements);
+ callback.run();
+ }
+ return true;
+ } else {
+ pending.addLast(new PendingElement(elements, callback));
+ return true;
+ }
+ }
+
+ synchronized boolean complete(CloseableFnDataReceiver<BeamFnApi.Elements>
receiver)
+ throws Exception {
+ if (this.receiver != null) {
+ return false;
+ }
+ this.receiver = receiver;
+ if (receiver instanceof BeamFnDataInboundObserver) {
+ drainPending();
Review Comment:
**🟡 Medium**: In `ReceiverHolder.complete()`, if a
non-`BeamFnDataInboundObserver` receiver throws an exception during
`receiver.accept(p.elements)`, the remaining elements in `pending` (as well as
the failing element) never have `p.callback.run()` called and `pending` is not
cleared, leaking `inFlightInbound` permits. Also, in `registerConsumer` (lines
251-257), throwing `IllegalArgumentException` inside `try { ... } catch
(Exception e)` causes duplicate registration to throw
`RuntimeException(IllegalArgumentException)` instead of throwing
`IllegalArgumentException` directly as before.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java:
##########
@@ -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();
Review Comment:
**🟡 Medium**: `requestInbound(int count)` accepts a `count` parameter but
ignores it and always decrements `inFlightInbound` by 1. Use
`inFlightInbound.addAndGet(-count)` (or remove the parameter if it is only ever
intended to release a single message).
```suggestion
inFlightInbound.addAndGet(-count);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
