Repository navigation
[Java SDK harness] Prevent a Fn data plane deadlock in BeamFnDataGrpcMultiplexer - #40606
prodriguezdefino wants to merge 5 commits into
Conversation
20d81c9 to
2bdb35a
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #40606 +/- ##
=========================================
Coverage 56.16% 56.16%
Complexity 2288 2288
=========================================
Files 1124 1124
Lines 178336 178336
Branches 1489 1489
=========================================
Hits 100170 100170
Misses 75645 75645
Partials 2521 2521 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
2bdb35a to
8de9191
Compare
|
Assigning reviewers: R: @ahmedabu98 for label java. 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 Available commands:
The PR bot will only process comments in the main thread (not review comments). |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
📋 Review Summary
This PR introduces per-bundle inbound queuing in BeamFnDataInboundObserver and adds a StalledPollHandler hook across DirectStreamObserver and ForwardingClientResponseObserver so that when outbound gRPC writes block on flow control, the bundle thread can cooperatively drain pending inbound elements for other bundles from BeamFnDataGrpcMultiplexer. The architectural approach effectively addresses cross-bundle head-of-line blocking when inbound and outbound data share a single gRPC stream.
🔍 General Feedback
- Thread-Safety & Concurrency: There is a critical synchronization gap when
BeamFnDataGrpcMultiplexer.tryDrainPendingForOtherBundlesinvokesBeamFnDataInboundObserver.acceptconcurrently from a stalled bundle's worker thread while the gRPC event loop thread is also callingacceptfor the same bundle. - Deadlock Prevention in Buffered Mode: The stall-callback mechanism is only wired into
DirectStreamObserver(OutboundObserverFactory.Direct), leavingBufferingStreamObserver(beam_fn_api_buffered_stream) vulnerable to the same deadlock when its boundedCancellableQueuefills up.
💡 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/BeamFnDataInboundObserver.java(line 122): 🔴 Critical:acceptis no longer exclusively called on the serialized gRPC inbound thread:BeamFnDataGrpcMultiplexer.tryDrainPendingForOtherBundlescan invokeconsumer.accept(element)on a stalled bundle worker thread concurrently with the gRPC thread callingconsumer.accept(nextElement)for the sameinstructionId(once the consumer has been registered inconsumers).
Because accept is not synchronized (and Status fields like numMessages and byteCounter are plain non-volatile primitives), concurrent calls to accept on the same BeamFnDataInboundObserver instance can:
- Reorder elements within the same bundle (e.g., if
isLast=truearrives on the gRPC thread whiledrainPendingis mid-flight,queue.close()can be called before earlier elements are enqueued intoqueue, causing those elements to be silently dropped or processed out of order). - Cause data races on
status.numMessagesandstatus.byteCounter.
Consider synchronizing accept (or guarding enqueueing + status updates under a lock) and ensuring BeamFnDataGrpcMultiplexer.registerConsumer holds a per-instruction lock while draining pending elements before exposing the consumer to concurrent callers.
public synchronized void accept(BeamFnApi.Elements t) throws Exception {
sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java(line 146): 🟠 High: InregisterConsumer,consumers.put(instructionId, consumer)happens beforedrainPending(instructionId, consumer). If an inbound message arrives on the gRPC thread immediately afterconsumers.put(...),InboundObserver.onNextwill bypasspendingElementsand deliver the new message directly toconsumer.accept(val)beforedrainPendinghas finished delivering the older buffered messages, resulting in out-of-order element delivery (or early stream termination if the new message hasisLast=true).
Additionally, tryDrainPendingForOtherBundles and registerConsumer can both call drainPending concurrently for the same instructionId, causing two threads to poll from the same ConcurrentLinkedQueue and deliver elements out of order.
Synchronizing on the Queue<Elements> (or a per-instruction lock) across registerConsumer, drainPending, and the fallback path in onNext would preserve strict FIFO ordering.
sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactory.java(line 140): 🟡 Medium:OutboundObserverFactory.Bufferedignores the newForwardingClientResponseObserver.create(inboundObserver, inboundObserver::onCompleted)overload that forwardsStalledPollHandler, andBufferingStreamObserverdoes not invokeonStalledPollwhenqueue.put(value)blocks on a full buffer while the background drainer thread is stalled inDirectStreamObserver.onNext. When thebeam_fn_api_buffered_streamexperiment is enabled, a bundle thread that fills the outboundBufferingStreamObserverqueue will block inqueue.putwithout ever callingonStalledPoll(), reintroducing the cross-bundle deadlock.
When a bundle queue was full, onNext of the inbound data stream blocked. gRPC runs onNext and the onReady callback of the outbound stream on the same executor. So the outbound stream could not become ready again, and the bundle, the gRPC thread and the runner waited for each other. The inbound stream uses manual gRPC flow control and gives messages to the bundle queue without blocking. A message that does not fit waits in a buffer. The stream asks for one more message each time a bundle finishes one. When output waits 500 ms on flow control, the stream asks for one extra message, up to 100 for each active bundle. Each bundle queue keeps its limit of 100 messages. DirectStreamObserver checks isReady before each message and checks again every 50 ms while it waits.
8de9191 to
9d92655
Compare
|
Thanks for the review. I pushed 9d92655, which addresses the buffered-mode point and fixes a separate deadlock I found while checking the other two. Concurrent accept and ordering during registerConsumer (🔴 / 🟠). I don't think these apply to the code in this PR. There is no
The bundle thread does hand off buffered elements from its completion callback, so I added a test for that rather than relying on reading the code:
Lock-order deadlock in CancellableQueue.cancel. I found this while checking the paths above.
Buffered mode (🟡). Good catch. The main deadlock was already fixed there, because inbound
|
|
/gemini review |
There was a problem hiding this comment.
📋 Review Summary
This PR replaces blocking inbound gRPC data multiplexing with manual gRPC flow control, non-blocking bounded queuing (CancellableQueue.tryPut), and a StalledPollHandler mechanism to request additional inbound messages when outbound streams stall on flow control. While the overall design effectively breaks the single-executor gRPC inbound/outbound deadlock, there are a few resource/permit leaks on error and poisoned-instruction paths that could cause inFlightInbound permits or receivers entries to leak over time.
🔍 General Feedback
- Inbound flow-control permit accounting on exceptions: Every inbound
Elementsmessage must reliably trigger its completion callback even when consumer processing throws an exception; otherwiseinFlightInboundpermanently drifts upward and can eventually starve the data stream. - Poisoned instruction lifecycle:
receivers.computeIfAbsentshould not insert newReceiverHolderentries for already-poisoned instruction IDs, as nothing removes them afterward.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
Fixes from review of the Fn data plane change: - BeamFnDataInboundObserver runs the callback of the element that was being processed when the consumer throws. queue.take() had removed it, so the cancel did not reach it, and each failed bundle kept one inbound permit for the life of the multiplexer. - Late data for a poisoned instruction no longer adds a ReceiverHolder to receivers. The entry was never removed, grew the stall window, and let a later registerConsumer succeed on a poisoned holder. - ReceiverHolder.complete refuses a poisoned holder and runs every buffered callback even if the receiver throws. - registerConsumer throws IllegalArgumentException unwrapped again. - requestInbound releases count permits instead of one. Tests count request(1) calls on a mocked stream for a failed bundle, late data for a poisoned instruction, and a receiver that throws on buffered elements, and restore the expected behavior in three earlier tests.
|
/gemini review |
There was a problem hiding this comment.
📋 Review Summary
This PR resolves a bidirectional gRPC flow-control deadlock in the Java SDK harness by replacing blocking consumer lookup and queue insertion on the inbound gRPC thread with manual gRPC flow control (disableAutoRequestWithInitial / request(1)), non-blocking buffering in BeamFnDataGrpcMultiplexer, and a StalledPollHandler callback from outbound stream observers. The overall design directly addresses the shared client-call executor deadlock and is well-tested, with a few edge cases around failed consumer registration, memory retention during bundle element multiplexing, and sub-second stall logging detailed below.
🔍 General Feedback
- When a
CloseableFnDataReceiverthrows while draining pre-registration buffered elements inBeamFnDataGrpcMultiplexer.registerConsumer, the exception abortsConcurrentHashMap.compute, leaving the half-completedReceiverHolderregistered inreceiversinstead of poisoning and removing it. - In
BeamFnDataInboundObserver.awaitCompletion(), keepingElementWithCallback wrapperin the outer loop scope acrossmultiplexElements(...)retains a strong reference toBeamFnApi.Elements, defeating the inner scope block intended to allow GC of largeElementsprotos during processing. - Lowering
DirectStreamObserver's poll interval from 1s to 50ms while keepingif (totalMsWaited > 0)causes sub-second waits (50ms..950ms) to log"Output channel stalled for 0s"(and can triggerLOG.infowhenisReady()flips totruebeforeonReadyHandlerruns).
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
… stalls Fixes from the next review of the Fn data plane change: - registerConsumer poisons the instruction when the receiver throws on a buffered element. The exception aborted compute(), which left the holder in receivers with its receiver open. Poisoning removes it and closes the receiver, as a failure on the inbound gRPC thread does. A duplicate or poisoned registration still only throws. - awaitCompletion keeps the queued element inside its scope block again, so the Elements proto can be collected while it is processed, and releases the callback of that element in finally, so an Error also releases it. - DirectStreamObserver logs a stall only after a second of waiting, as with the earlier one-second poll. Shorter waits logged "0s" and a false "OnReady notification was not invoked" when isReady turned true just before the onReady callback ran. The stall hook still runs every 500 ms.
|
/gemini review |
There was a problem hiding this comment.
📋 Review Summary
This PR resolves a Fn data plane deadlock in BeamFnDataGrpcMultiplexer by switching the inbound stream to manual gRPC flow control (disableAutoRequestWithInitial(16)) with non-blocking buffering per instruction via ReceiverHolder and elastic window expansion via StalledPollHandler when outbound streams stall on flow control. All previously identified issues—including permit leaks on bundle or registration failure, late data handling for poisoned instructions, GC reference scoping in BeamFnDataInboundObserver, and stall log thresholds—have been cleanly addressed with comprehensive unit test coverage.
🔍 General Feedback
- The synchronization and lock-ordering design between
ReceiverHolderandCancellableQueue(releasing the queue lock before invokingonDropcallbacks) properly avoids ABBA deadlocks between the gRPC inbound thread and bundle processing threads. - Permit accounting across single-instruction and multi-instruction
Elementsbatches, early buffering, cancellation, and error paths is well-tested and sound. - The changes look ready to merge.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
DirectStreamObserver counts 50 ms for each wait that times out, not the wall time. The tests kept the stream not ready for a wall time, so on a slow machine fewer waits timed out, the counted wait stayed under a second, and testStallOfASecondIsLogged failed. The stream now becomes ready after a number of readiness polls, so the counted wait is exact.
|
/gemini review |
There was a problem hiding this comment.
📋 Review Summary
This PR resolves a bidirectional gRPC flow-control deadlock in BeamFnDataGrpcMultiplexer by switching the inbound stream to manual flow control, decoupling inbound onNext from blocking on full bundle queues or unregistered consumers, and incrementally expanding the inbound window via StalledPollHandler when outbound streams stall. All previously identified issues—including permit accounting on bundle failures, poisoned instruction cleanup, Elements GC reachability during processing, and deterministic stall logging tests—have been cleanly addressed.
🔍 General Feedback
- Concurrency & Lock Ordering:
CancellableQueue.cancel(..., onDrop)safely invokes drop callbacks outside the queue lock, andReceiverHoldersynchronizes state transitions and pending queue draining consistently. - Flow-Control Accounting: Inbound permits are reliably released across normal processing, consumer exceptions/errors, unregistered/poisoned instructions, and multi-instruction
Elementsbatches. - Test Coverage: The test suite thoroughly covers concurrent producer/consumer handoffs, lock-ordering edge cases, failure permit recovery, and deterministic stall logging.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
|
Reminder, please take a look at this pr: @ahmedabu98 |
|
Hey! I reviewed this change and noticed quite a bit of complexity around synchronization/locking was added and I worried that it might affect bundle processing throughput. I found a notable throughput regression that seems to be caused by cross-thread I've addressed this in #40788 by using a simple state machine and non-blocking Benchmarks
./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"This change shows a ~17% throughput regression on
|
Checking before each message cost about 17% on ProcessBundleBenchmark.testLargeBundle. Adds a test that data for an unregistered instruction does not block onReady.
|
@sjvanrossum Thank you for the review. I can see the same decrease in throughput. One change causes it:
I also tried a lock-free counter that calls JIC #40788 does not fix all of the deadlock. Data can arrive before
|
I overlooked the easiest thing! Yeah, that looks a lot better now. I'll run this benchmark on a few different machine types for an extended period of time when I'm at the office tomorrow.
|
|
@sjvanrossum if your PR covers the case I care about (not having both delivery threads blocking each other when a data producer is faster than the data consumer from the harness perspective), and you feel your code is easier to maintain than what I proposed here; feel free to close this PR once yours reach main. I can test my particular case then, and open another fix if #40788 does not covers it. |
When a bundle's input queue is full,
BeamFnDataGrpcMultiplexerblocks inonNext.gRPC runs
onNextand the outboundonReadycallback on the same executor, so a bundle thread that waits inDirectStreamObserverforonReadyonly continues when its wait times out (up to 60 s). The runner logs "Ignoring data element byte limit … to avoid stuckness" and the job crawls or stalls. We hit this on Dataflow with a Java step fused between two steps of another SDK.The inbound data stream now uses manual gRPC flow control, so
onNextnever blocks. A message that does not fit in the bundle queue waits in a buffer, and the stream asks for more messages as bundles finish them. While output waits on flow control, it asks for one more message every 500 ms, up to 100 for each active bundle.DirectStreamObserverchecksisReady()before each message.Related to #20028.