fix(server): let subscriber taps evict on full instead of wedging dispatch#1137
fix(server): let subscriber taps evict on full instead of wedging dispatch#1137astrogilda wants to merge 2 commits into
Conversation
…patch One undrained sink stalls EventQueueSource._dispatch_loop for every sink: the dispatcher's gather awaits a blocking put on each sink, so a single full queue whose consumer was abandoned blocks the gather forever. The incoming queue fills behind it, producers wedge in enqueue_event, and under sustained load the task's event flow never recovers (see the task dumps and repro on the tracking issue). Split the sink semantics. The default sink is flow control: its consumer drives task state, so a blocking put remains correct backpressure. Tapped subscriber sinks are broadcast observers whose remote consumers can be abandoned, so tap() now accepts evict_on_full: when such a sink is full at delivery time it is force-closed and detached, dispatch to the remaining sinks continues, and blocked producers recover. The full() pre-check is race-free because the dispatcher is the only writer to a sink queue. tap() defaults to evict_on_full=False, preserving the documented blocking contract for existing callers; ActiveTask's subscriber tap opts in explicitly, which fixes the production wedge without changing the public primitive's default behavior. Regression tests cover eviction with continued dispatch, producer recovery after eviction, preserved backpressure on default taps, and the subscriber tap opting in; the eviction tests fail without the fix.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an evict_on_full mechanism for event queue sinks to prevent abandoned remote subscribers from blocking the dispatcher and wedging the entire event pipeline. Sinks configured with this option are force-closed and detached when full, while default sinks retain their standard backpressure behavior. The changes are well-covered by new unit tests. The review feedback highlights a potential race condition in _deliver_to_sink where a gracefully closing sink might be prematurely force-closed if it is full, and suggests adding a sink.is_closed() check to avoid this issue.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if sink._evict_on_full and sink.queue.full(): # noqa: SLF001 | ||
| logger.warning( | ||
| 'Evicting event queue sink %r: its queue is full and the ' | ||
| 'consumer is not draining it. Closing the sink so dispatch ' | ||
| 'to the remaining sinks continues.', | ||
| sink, | ||
| ) | ||
| await sink.close(immediate=True) | ||
| return | ||
| await sink._put_internal(event) # noqa: SLF001 |
There was a problem hiding this comment.
There is a potential race condition when a sink is being closed gracefully (i.e., sink.close(immediate=False)). During a graceful close, the sink is marked as closed (_is_closed = True) and removed from self._sinks, but it may still exist in the active_sinks snapshot currently being processed by _dispatch_loop.
If the queue is full and evict_on_full is True, the dispatcher will see sink.queue.full() == True and call await sink.close(immediate=True). This will immediately upgrade the graceful close to an immediate close, discarding any remaining events that the consumer was in the middle of draining.
Checking sink.is_closed() at the beginning of _deliver_to_sink prevents this race condition and ensures that already-closed or gracefully-closing sinks are ignored by the dispatcher.
| if sink._evict_on_full and sink.queue.full(): # noqa: SLF001 | |
| logger.warning( | |
| 'Evicting event queue sink %r: its queue is full and the ' | |
| 'consumer is not draining it. Closing the sink so dispatch ' | |
| 'to the remaining sinks continues.', | |
| sink, | |
| ) | |
| await sink.close(immediate=True) | |
| return | |
| await sink._put_internal(event) # noqa: SLF001 | |
| if sink.is_closed(): | |
| return | |
| if sink._evict_on_full and sink.queue.full(): # noqa: SLF001 | |
| logger.warning( | |
| 'Evicting event queue sink %r: its queue is full and the ' | |
| 'consumer is not draining it. Closing the sink so dispatch ' | |
| 'to the remaining sinks continues.', | |
| sink, | |
| ) | |
| await sink.close(immediate=True) | |
| return | |
| await sink._put_internal(event) # noqa: SLF001 |
There was a problem hiding this comment.
Good catch, this race is real. Fixed in 9830251: _deliver_to_sink now checks sink.is_closed() up front and just skips the sink, so a graceful close(immediate=False) that's still draining can't get upgraded to an immediate close when the dispatcher reaches it from a stale snapshot with a full queue. I kept the check best-effort rather than taking the sink lock; if a close lands right after the check the put is still fine, since _put_internal already swallows QueueShutDown. Regression test is test_graceful_close_not_upgraded_to_immediate_by_eviction (fails without the skip).
🧪 Code Coverage (vs
|
| Base | PR | Delta | |
|---|---|---|---|
| src/a2a/server/events/event_queue_v2.py | 91.71% | 91.75% | 🟢 +0.04% |
| src/a2a/utils/telemetry.py | 91.47% | 90.70% | 🔴 -0.78% |
| Total | 93.00% | 92.99% | 🔴 -0.01% |
Generated by coverage-comment.yml
A sink mid-graceful-close (close(immediate=False), consumer still draining) can remain in the dispatcher's active_sinks snapshot. With evict_on_full=True and a full queue, delivering to it would upgrade the close to immediate and discard the events the consumer was draining. Skip closed sinks before the evict-on-full check; a close landing after the check is harmless because the put already tolerates a shut-down queue.
1 participant
Fixes #1136 (the runtime issue split out of #1101).
Problem
One undrained sink stalls
EventQueueSource._dispatch_loopfor every sink. The dispatcher'sasyncio.gatherawaits a blocking put on each sink, so a single full queue (default cap 1024) whose consumer went away blocks the gather forever. The incoming queue then fills, producers wedge inenqueue_event, and the task's event flow never recovers; see the task dumps and repro in #1101.Change: split the sink semantics
The two kinds of sink want different fullness behavior.
The default sink is flow control. Its consumer drives task state, so a full queue that back-pressures the dispatcher (and, transitively, the producer) is correct and stays exactly as it is — a slow task store slows the pipeline down rather than losing events.
Tapped subscriber sinks are broadcast observers. Their consumers are remote and can be abandoned, most simply by a non-blocking
message/sendwhose HTTP response has already returned.tap()therefore gainsevict_on_full: bool = False. When an evict-on-full sink is full at delivery time, the dispatcher force-closes it (close(immediate=True), warning logged) and it detaches through the existingremove_sinkpath; dispatch to the remaining sinks continues and any producer parked on the incoming queue recovers. The consumer of an evicted sink seesQueueShutDownon its next dequeue, same as any closed queue — no new exception type, no API surface beyond the keyword.The
full()pre-check is race-free: the dispatcher is the only writer to a sink queue and consumers only read, so a queue observed non-full cannot become full before the put.tap()defaults toFalse, so the documented blocking contract is unchanged for existing callers.ActiveTask's subscriber tap opts in explicitly, which is the one production site that wedges — the fix reaches existing deployments without changing the public primitive's default.Tests
test_evict_on_full_sink_is_evicted_and_dispatch_continues: an abandoned capacity-1 evict sink is closed and detached; the default sink still receives every event.test_producer_unblocked_after_evict_on_full: producers blocked on a full incoming queue recover once the sink is evicted (the [Bug]: DefaultRequestHandlerV2 ActiveTask producer can remain pending during EventQueue subscriber cleanup #1101 runtime shape in miniature).test_default_tap_keeps_backpressure: a default tap still blocks dispatch when full (the existing flow-control contract, pinned).test_subscriber_taps_are_evict_on_full:ActiveTask.subscribecreates evict-on-full sinks.The two eviction tests fail without the fix. Full suite: 1755 passed, 90 skipped, 3 xfailed, 1 xpassed. The 4
tests/integration/cross_versionsetup errors are environmental (the fixture spawns versioned servers with uv) and occur identically without this change.Out of scope
The deeper root cause is subscriber-sink lifecycle: a sink tapped for a consumer that has gone away should be closed by its owner (connection teardown), not discovered full later. That needs its own look at the request-handler layer; this PR keeps the dispatch layer from letting any such sink take the task down in the meantime.