[sjvanrossum]

Fixes a bidirectional flow-control deadlock on the Java SDK Fn API data plane stream (as described in #40606) without introducing secondary per-instruction queues or per-element allocations.

On a client-side gRPC stream (OutboundObserverFactory.DirectClient / BufferingStreamObserver), gRPC uses a single per-call SerializingExecutor to execute both InboundObserver.onNext (inbound elements) and ClientCallStreamObserver.onReady (which calls AdvancingPhaser.arrive() to unblock outbound DirectStreamObserver.onNext calls).

When a bundle processing thread in BeamFnDataInboundObserver.awaitCompletion() emits elements back over the same data plane stream and blocks in DirectStreamObserver.onNext waiting for the outbound stream to become ready (AdvancingPhaser), it stops draining its CancellableQueue (capacity 100). Once 100 inbound Elements messages accumulate in that CancellableQueue, the next InboundObserver.onNext blocks in CancellableQueue.put() directly on the gRPC SerializingExecutor thread. Because the SerializingExecutor thread is blocked in put(), the pending onReady callback queued behind it on the same SerializingExecutor can never run to unblock the consumer thread, resulting in a permanent deadlock.

Benchmarks

ProcessBundleBenchmark executed on Apple M5 Pro.

./gradlew :sdks:java:harness:jmh:jmh --args="org.apache.beam.fn.harness.jmh.ProcessBundleBenchmark -p elementsEmbedding=false -f 1 -wi 3 -i 5 -w 2s -r 2s -foe=true"

Benchmark	Base	This PR	vs base
testLargeBundle	14,991.4 ± 496.3 ops/s	14,931.6 ± 119.8 ops/s	-0.4%
testTinyBundle	35,573.1 ± 1,610.7 ops/s	37,160.9 ± 223.5 ops/s	+4.5%
testStateWithCaching	6,748.1 ± 167.0 ops/s	6,768.4 ± 38.9 ops/s	+0.3%
testStateWithoutCaching	3,936.6 ± 216.9 ops/s	4,113.7 ± 96.3 ops/s	+4.5%

[github-actions]

📋 Review Summary

This PR introduces manual gRPC inbound flow control on BeamFnDataGrpcMultiplexer, requesting a new Elements message only after a bundle consumer finishes processing, drops, or rejects the current message. While the overall approach is much simpler than PR #40606, there are several critical flow-control bugs and edge cases—including permanent stalls when messages contain multiple instruction IDs, broken flow control when BeamFnDataInboundObserver is reused across bundles via reset(), and an order-dependent callback wiring issue in registerConsumer—that must be addressed before merging.

🔍 General Feedback

• Multi-instruction Elements messages cause permanent gRPC stalls: Under disableAutoInboundFlowControl(), gRPC provides only 1 initial permit. When a single Elements message is split across $N > 1$ instructions in InboundObserver.onNext, only 1 permit was consumed on the wire, yet $N$ callbacks are created, which either stalls when $N-1$ consumers wait for is_last or inflates the in-flight window by $N-1$ permits.
• BeamFnDataInboundObserver reuse across bundles (reset()): ProcessBundleHandler caches BeamFnDataInboundObserver instances and calls reset() between bundles before calling registerConsumer() again, which breaks the onElementsConsumed lifecycle as currently written.
• Wait-with-timeout before poisoning unregistered consumers: Replacing consumerFuture.get(3, TimeUnit.HOURS) with an immediate ExecutionException if registerConsumer has not completed within the computeIfAbsent block (or immediately thereafter) changes the contract when data arrives slightly before bundle registration finishes.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

💬 Inline Findings

• sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java (line 236): 🔴 Critical: When an Elements message contains empty data and timers lists (instructionId == null), onNext returns without calling request(1). Because disableAutoInboundFlowControl() is enabled and this message consumed the single in-flight gRPC delivery permit, failing to call request(1) here permanently stalls the inbound gRPC stream. Call request(1) before returning.

        if (instructionId == null) {
          request(1);
          return;
        }

• sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java (line 258): 🔴 Critical: In the multi-instruction slow path, a single inbound gRPC Elements message (which consumed only 1 gRPC permit) is split and forwarded to instructionIds.size() different consumers:

1. If those consumers are BeamFnDataInboundObserver instances that already have onElementsConsumed = this::request wired, each of the $N$ bundles will call request(1) when it processes its split piece—inflating the gRPC flow-control window by $N - 1$ extra permits per multi-instruction message.
2. Worse, if the consumers are not yet registered when onNext runs, the first loop iteration blocks in consumerFuture.get(3, TimeUnit.HOURS) and if the bundle needs subsequent messages before completing, or if registerConsumer was called before beforeStart, permits are either never requested or requested $N$ times.

Instead of relying on per-consumer callbacks firing $N$ times for 1 gRPC message, either gate the request(1) for a split message behind an AtomicInteger(instructionIds.size()) countdown (as in #40606) or ensure exactly 1 permit is replenished per onNext invocation.

• sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java (line 122): 🔴 Critical: inboundObserver.requestObserver is only initialized when gRPC invokes beforeStart() on the client call. If registerConsumer() is called before beforeStart() runs (for example, in unit tests, or if a consumer is registered while the channel/stub is still starting), inboundObserver.requestObserver != null is false, so setOnElementsConsumed(inboundObserver::request) is skipped and never wired even after beforeStart() later sets requestObserver.

Additionally, inboundObserver.requestObserver is read on the thread calling registerConsumer and written on the gRPC thread in beforeStart without volatile or synchronization.

Always wire ((BeamFnDataInboundObserver) receiver).setOnElementsConsumed(inboundObserver::request) unconditionally in registerConsumer, make requestObserver volatile, and have inboundObserver.request(int) check ClientCallStreamObserver<BeamFnApi.Elements> observer = this.requestObserver; if (observer != null) { observer.request(count); }.

    if (receiver instanceof BeamFnDataInboundObserver) {

• sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java (line 157): 🟠 High: BeamFnDataInboundObserver instances are cached and reused across bundles by ProcessBundleHandler (which calls dataReceiver.reset() before reusing the observer for a new instruction). Two issues arise:

1. onElementsConsumed is a non-volatile field written by registerConsumer (which may run on a different thread than the bundle execution thread or gRPC inbound thread) and read inside reset(), close(), and awaitCompletion().
2. When ProcessBundleHandler calls reset() at the end of a bundle (after awaitCompletion() has already closed the queue), reset() clears cancellationException. If any late data was queued or if reset() is called before reuse, onElementsConsumed should be volatile.

  private volatile @Nullable Consumer<Integer> onElementsConsumed;

• sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java (line 318): 🟠 High: When consumer is not a BeamFnDataInboundObserver (e.g., a custom CloseableFnDataReceiver or in tests), request(1) is only called if consumer.accept(value) succeeds. If consumer.accept(value) throws an exception, poisonInstructionId(instructionId) is called in the catch block, inboundObserver.requestObserver == null is false (when requestObserver != null), and request(1) is never called for the failed message—permanently leaking the only gRPC flow-control permit and stalling the stream for all subsequent bundles.

Similarly, if consumer is a BeamFnDataInboundObserver and consumer.accept(value) throws (because the bundle already failed/closed its CancellableQueue), queue.put(elements) throws cancellationException without ever enqueuing the element, so awaitCompletion()/close() never sees it and never calls onElementsConsumed.accept(1)! Thus, whenever consumer.accept(value) throws, the permit for value is leaked unless request(1) is called in the catch (Exception e) block.

      } catch (Exception e) {
        LOG.error("Failed to handle data for instruction {}", instructionId, e);
        poisonInstructionId(instructionId);
        request(1);
      }

• sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java (line 292): 🟠 High: With disableAutoInboundFlowControl() (which sets the initial gRPC window to 1), blocking the gRPC inbound thread inside consumer.accept(value) (queue.put(elements)) until the bundle thread calls queue.take() and finishes multiplexElements before calling request(1) can deadlock pipelines when:

1. A bundle has a large CancellableQueue (default capacity 10,000), so queue.put() does not block, but with an initial window of 1, gRPC will not deliver message #2 until message #1 is completely processed by multiplexElements() and onElementsConsumed.accept(1) is called.
2. If a transform buffers elements in memory across multiple Elements messages (or waits for a side input / timer / second input in the same bundle) while emitting output that fills the outbound gRPC buffer, or if message #1 is for Bundle A (whose transform is blocked on output backpressure) and message #2 on the same multiplexed stream is is_last or timer data needed by Bundle B to unblock the runner, having an initial window of only 1 (with no request(1) until multiplexElements finishes) serializes all inbound messages across all concurrent bundles on the multiplexer to 1 in-flight message at a time.

[github-actions]

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

[sjvanrossum]

This was certainly unexpected since I added this test and it passes locally...

  Gradle Test Run :sdks:java:core:test > Gradle Test Executor 6 > BeamFnDataGrpcMultiplexerTest > testBidirectionalFlowControlDoesNotDeadlock FAILED
      java.lang.AssertionError: expected:<151> but was:<107>
          at org.junit.Assert.fail(Assert.java:89)
          at org.junit.Assert.failNotEquals(Assert.java:835)
          at org.junit.Assert.assertEquals(Assert.java:647)
          at org.junit.Assert.assertEquals(Assert.java:633)
          at org.apache.beam.sdk.fn.data.BeamFnDataGrpcMultiplexerTest.testBidirectionalFlowControlDoesNotDeadlock(BeamFnDataGrpcMultiplexerTest.java:784)

Directed by Robert B. Weide

[sjvanrossum]

• sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java (line 236): 🔴 Critical: When an Elements message contains empty data and timers lists (instructionId == null), onNext returns without calling request(1). Because disableAutoInboundFlowControl() is enabled and this message consumed the single in-flight gRPC delivery permit, failing to call request(1) here permanently stalls the inbound gRPC stream. Call request(1) before returning.

        if (instructionId == null) {
          request(1);
          return;
        }

I'm guessing this was supposed to be regarding line 326 of onNext?
The finding is useless, onNext is entirely wrapped in a try { ... } finally { ... } block which executes releaseOrParkFlowControlToken().

Edit: I just realized this was a stale review. Triggering another one.

[github-actions]

Assigning reviewers:

R: @kennknowles for label java.
Backup reviewer: ahmedabu98 (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

This pull request likely touches a core component ("core" label). Please review with scrutiny.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 0.00%. Comparing base (a8d6afd) to head (af2841e).
⚠️ Report is 74 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff             @@
##           master   #40788       +/-   ##
===========================================
- Coverage   80.04%        0   -80.05%     
===========================================
  Files         490        0      -490     
  Lines       82916        0    -82916     
===========================================
- Hits        66373        0    -66373     
+ Misses      16543        0    -16543     

Flag	Coverage Δ
python	?

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[sjvanrossum]

/gemini review

[github-actions]

📋 Review Summary

This PR resolves a bidirectional gRPC flow-control deadlock on the Fn Data plane by forwarding beforeStart through ForwardingClientResponseObserver, disabling automatic inbound flow control in BeamFnDataGrpcMultiplexer, switching inbound delivery to a non-blocking CancellableQueue.add, and tracking buffered elements via a lock-free VarHandle state machine. The overall design and unit test coverage are thorough and well-structured.

🔍 General Feedback

• The bit-packed flowControlState state machine in BeamFnDataGrpcMultiplexer cleanly coordinates request(1) parking and wake-ups across the gRPC SerializingExecutor thread and bundle processing threads.
• There are a couple of thread-visibility and registration-ordering details around BeamFnDataInboundObserver.onElementsConsumed and BeamFnDataGrpcMultiplexer.registerConsumer worth tightening up.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.