Repository navigation
[Java SDK harness] Prevent a Fn data plane deadlock in BeamFnDataGrpcMultiplexer - #40606
prodriguezdefino wants to merge 4 commits into
Conversation
20d81c9 to
2bdb35a
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #40606 +/- ##
============================================
+ Coverage 56.16% 58.70% +2.53%
- Complexity 2288 14015 +11727
============================================
Files 1124 2594 +1470
Lines 178336 272493 +94157
Branches 1489 11337 +9848
============================================
+ Hits 100170 159959 +59789
- Misses 75645 106501 +30856
- Partials 2521 6033 +3512
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ 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.
| consumingReceivedData.set(true); | ||
| if (multiplexElements(dataIterator, timersIterator)) { | ||
| boolean isDone = multiplexElements(dataIterator, timersIterator); | ||
| wrapper.callback.run(); |
There was a problem hiding this 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).
There was a problem hiding this comment.
Now awaitCompletion keeps the callback of the element it took from the queue and runs it after cancelling the queue, so a failed bundle releases the permit of that element as well as the queued ones. BeamFnDataGrpcMultiplexerTest.testPermitsReleasedWhenBundleFails sends three messages to a bundle whose consumer throws and checks for three request(1) calls on a mocked stream.
|
|
||
| assertThrows(RuntimeException.class, observer::awaitCompletion); | ||
|
|
||
| assertFalse(callback1.get()); |
There was a problem hiding this 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.
| assertFalse(callback1.get()); | |
| assertTrue(callback1.get()); |
There was a problem hiding this comment.
Agreed. I'd changed this to match the bug. It asserts assertTrue(callback1.get()) now.
| @@ -233,9 +404,12 @@ public void onNext(BeamFnApi.Elements value) { | |||
| } | |||
| } | |||
| if (instructionId == null) { | |||
There was a problem hiding this 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:
- A permanent memory leak in
receivers(defeating the 20-minute TTL onpoisonedInstructionIds) and permanent inflation ofreceivers.size(), which increasesmaxWindowinonStalledPoll()byMAX_INBOUND_WINDOWper poisoned instruction. - If
registerConsumeris called afterward for thatinstructionId, it seesexisting != nullbefore checkingpoisonedInstructionIdsand callsexisting.complete(receiver), which succeeds becausecomplete()only checksthis.receiver != nulland notthis.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().
There was a problem hiding this comment.
Now the computeIfAbsent function returns null for a poisoned ID, so late data inserts nothing and is dropped, and its permit is released. ReceiverHolder.complete() also returns false for a poisoned holder, and registerConsumer reports it as poisoned. testPoisonInstructionId now sends late data after poisoning, checks that no entry remains, and checks that a later registration throws IllegalArgumentException. testPermitsReleasedForLateDataOfPoisonedInstruction checks the permits.
| } | ||
| this.receiver = receiver; | ||
| if (receiver instanceof BeamFnDataInboundObserver) { | ||
| drainPending(); |
There was a problem hiding this 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.
There was a problem hiding this comment.
Fixed in b7532bb. complete() runs the callback of every buffered element, including the one the receiver failed on, and clears the buffer even if the receiver throws. registerConsumer rethrows runtime exceptions as they are, so duplicate registration throws IllegalArgumentException directly again. Covered by testPermitsReleasedWhenReceiverThrowsOnBufferedElements and the duplicate-registration check in testReceiverHolderBuffersBeforeConsumerRegisteredAndDrains.
| * STEADY_STATE_INBOUND_WINDOW} messages are open again. | ||
| */ | ||
| public void requestInbound(int count) { | ||
| inFlightInbound.decrementAndGet(); |
There was a problem hiding this 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).
| inFlightInbound.decrementAndGet(); | |
| inFlightInbound.addAndGet(-count); |
There was a problem hiding this comment.
Fixed in b7532bb with addAndGet(-count). testRequestInboundReleasesCount opens three extra messages with onStalledPoll(), then releases five and checks that two more are requested.
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.
| if (existing.isPoisoned()) { | ||
| throw new IllegalArgumentException("Instruction id was poisoned"); | ||
| } | ||
| if (!complete(existing, receiver)) { |
There was a problem hiding this comment.
🟠 High: When elements arrive before registerConsumer is called, receivers already contains a ReceiverHolder (existing) with buffered elements. If complete(existing, receiver) throws an exception because receiver.accept(p.elements) fails while draining pending (as in testPermitsReleasedWhenReceiverThrowsOnBufferedElements), existing.complete(receiver) has already set this.receiver = receiver (line 167), and the thrown exception aborts receivers.compute(...).
Per ConcurrentHashMap.compute semantics, throwing from the remapping function leaves the existing entry in receivers. As a result, hasConsumer(instructionId) remains true, instructionId is not poisoned, receiver.close() is not called (unlike when receiver.accept throws during onNext in forwardToConsumerForInstructionId), and receivers.size() stays permanently inflated by the leaked entry (which also inflates maxWindow in onStalledPoll()).
Consider catching failures from holder.complete(receiver) and calling poisonInstructionId(instructionId) outside receivers.compute before rethrowing the exception.
There was a problem hiding this comment.
Fixed in 6f9d444.
When the receiver throws on a buffered element, registerConsumer poisons the instruction after compute() returns, which removes the holder and closes the receiver, as a failure on the inbound gRPC thread does. A duplicate or poisoned registration still only throws, without poisoning.
testPermitsReleasedWhenReceiverThrowsOnBufferedElementsnow also checks that the entry is gone, that the receiver was closed, and that a later registration fails as poisoned.testDuplicateRegistrationDoesNotPoisonchecks that a rejected duplicate leaves the first registration working.
|
|
||
| // We use discarding iterators and don't reference the elements longer than necessary to | ||
| // avoid pinning possibly large data pages and making them ineligible for garbage | ||
| // collection. |
There was a problem hiding this comment.
🟡 Medium: Two issues here:
- The
{ ... }block on lines 201–207 exists specifically so no local variable in the outerwhilescope retains a reference toBeamFnApi.ElementswhilemultiplexElements(dataIterator, timersIterator)runs (see the comment on lines 195–198 andcreateDiscardingIterator). DeclaringElementWithCallback wrapperoutside{ ... }and callingwrapper.callback.run()on line 215 keepswrapper(and itsfinal BeamFnApi.Elements elementsfield) strongly reachable for the entire duration ofmultiplexElements. SinceinProgressCallback = wrapper.callbackis already captured inside{ ... },wrappercan remain scoped inside{ ... }and lines 214–215 can invokeinProgressCallbackafter clearing it. if (inProgressCallback != null) { inProgressCallback.run(); }is only executed incatch (Exception e)(lines 222–224). IfmultiplexElementsthrows anError(e.g.,AssertionError),finallycancels the queue viaclose()but leaks theinProgressCallbackpermit. Moving theinProgressCallbackcleanup intofinally(afterclose()) covers allThrowables.
There was a problem hiding this comment.
Both fixed in 6f9d444.
The queued element is declared inside the scope block again, so nothing outside it holds the Elements proto while multiplexElements runs, the callback is kept separately.
The callback of the in-progress element is now released in finally, after close(), so it covers every throwable.
testCallbacksInvokedOnConsumerErrorhas a consumer throw an AssertionError and checks that both callbacks run.
| } | ||
| } | ||
| if (totalSecondsWaited > 0) { | ||
| if (totalMsWaited > 0) { |
There was a problem hiding this comment.
🟡 Medium: Because waitMs was reduced from 1s to 50ms and DEFAULT_MAX_MESSAGES_BEFORE_CHECK from 100 to 0, totalMsWaited > 0 now triggers after a single 50ms poll. For any wait between 50ms and 950ms, totalMsWaited / 1000 evaluates to 0 ("Output channel stalled for 0s"). Furthermore, under normal backpressure, outboundObserver.isReady() can become true during a 50ms poll slightly before the gRPC executor dispatches the onReady callback (phaser.arrive()), exiting the loop with totalMsWaited == 50 and initialPhase == phase and logging a false-positive LOG.info warning. Requiring totalMsWaited >= 1000 preserves the previous 1-second minimum threshold for stall logging and avoids logging 0s.
| if (totalMsWaited > 0) { | |
| if (totalMsWaited >= 1000) { |
There was a problem hiding this comment.
Fixed in 6f9d444 with totalMsWaited >= 1000, as suggested.
The stall hook still runs every 500 ms.
testShortStallIsNotLoggedchecks that a 300 ms stall logs nothing,- and
testStallOfASecondIsLoggedchecks that a 1.3 s stall still logs "stalled for 1s".
… 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.
1 participant
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.