From 3dbfca58ce07946543188b4bd88b5322015b56e6 Mon Sep 17 00:00:00 2001 From: Thales Pereira <31625914+thcp@users.noreply.github.com> Date: Mon, 31 Aug 2026 18:29:42 +0100 Subject: [PATCH] fix(sse): stop connection slots leaking when a client disconnects early claim_sse_slot() runs in the handler and release_sse_slot() lived in the stream's finally. An async generator that is never started never runs its finally, so a client that disconnected before the response body began held a slot for the life of the process: StreamingResponse raises inside stream_response on its first send(), before __anext__ is ever called, and the generator body never executes. Two hundred of those and every job-progress stream and the queue stream answer 503 with nothing actually connected, until a restart. It does not need malice either -- a flaky network or a rapidly reloaded page leaks slots the same way, just slower. Claiming cannot simply move inside the generator: by the time it runs the response headers have gone out and there is no status code left, so hitting the cap could no longer answer 503. SseSlot keeps the claim in the handler and makes release idempotent, with __del__ as the backstop for the never-started case -- collecting the generator collects the closure holding the slot. _held is assigned before the claim because __del__ runs on a half-built object too: a refused claim would otherwise raise AttributeError out of __del__ rather than releasing nothing. The test for that caught it. queue.py takes the same guard, so both streams that share the budget also share the fix. Verified: removing the __del__ backstop fails the never-started test. Refs #513 --- app/api/events.py | 43 +++++++++++++++- app/api/queue.py | 6 +-- tests/test_sse_slot_budget.py | 92 +++++++++++++++++++++++++++++++++++ 3 files changed, 136 insertions(+), 5 deletions(-) create mode 100644 tests/test_sse_slot_budget.py diff --git a/app/api/events.py b/app/api/events.py index f63b6370..dedb0b84 100644 --- a/app/api/events.py +++ b/app/api/events.py @@ -38,6 +38,45 @@ def release_sse_slot() -> None: _sse_active -= 1 +class SseSlot: + """One held connection slot, released exactly once. + + Claiming has to happen in the handler so that hitting the cap can still be + answered with a 503 -- once the generator is running the response headers + have gone out and there is no status code left to send. + + That is what leaked slots: release lives in the stream's `finally`, and an + async generator that is never started never runs its `finally`. If the + client disconnects before the body begins, StreamingResponse raises inside + `stream_response` on its first `send()` -- before `__anext__` is ever + called -- so the generator body never executes and the slot was held + forever. 200 of those and every progress stream 503s with nothing actually + connected, until the process restarts (#513). + + The stream releases on its way out as before; `__del__` is the backstop for + the never-started case, where collecting the generator collects the closure + holding this. Release is idempotent so the two cannot double-count. + """ + + __slots__ = ("_held",) + + def __init__(self) -> None: + # Set first: claim_sse_slot raises at the cap, and __del__ still runs on + # a half-built object. Without this it would raise AttributeError from + # __del__ instead of releasing nothing. + self._held = False + claim_sse_slot() # may raise 503; nothing is held if it does + self._held = True + + def release(self) -> None: + if self._held: + self._held = False + release_sse_slot() + + def __del__(self) -> None: + self.release() + + @router.get("/jobs/{job_id}/events") async def job_events(job_id: str) -> StreamingResponse: """Server-Sent Events stream of job state updates. Closes when the job @@ -47,7 +86,7 @@ async def job_events(job_id: str) -> StreamingResponse: job = registry_get(job_id) if job is None: raise HTTPException(status_code=404, detail="job not found") - claim_sse_slot() + slot = SseSlot() async def stream() -> AsyncIterator[str]: try: @@ -82,7 +121,7 @@ async def stream() -> AsyncIterator[str]: keepalive_at = 0 await asyncio.sleep(0.2) finally: - release_sse_slot() + slot.release() return StreamingResponse( stream(), diff --git a/app/api/queue.py b/app/api/queue.py index 6687185a..0c2c65e5 100644 --- a/app/api/queue.py +++ b/app/api/queue.py @@ -22,7 +22,7 @@ from fastapi.responses import StreamingResponse from pydantic import BaseModel -from app.api.events import _MAX_SSE_SECONDS, claim_sse_slot, release_sse_slot +from app.api.events import _MAX_SSE_SECONDS, SseSlot from app.core.config import JOB_ID_RE, JOBS_DIR, MAX_PENDING_UPLOAD_JOBS, MAX_PENDING_URL_JOBS from app.core.registry import get as registry_get from app.core.registry import pending_count as registry_pending_count @@ -130,7 +130,7 @@ async def queue_events() -> StreamingResponse: outlives any individual job and is expected to stay open for the session, so only the 4 h ceiling ends it. """ - claim_sse_slot() + slot = SseSlot() async def stream() -> AsyncIterator[str]: try: @@ -156,7 +156,7 @@ async def stream() -> AsyncIterator[str]: keepalive_at = 0 await asyncio.sleep(0.25) finally: - release_sse_slot() + slot.release() return StreamingResponse( stream(), diff --git a/tests/test_sse_slot_budget.py b/tests/test_sse_slot_budget.py new file mode 100644 index 00000000..eb7a1613 --- /dev/null +++ b/tests/test_sse_slot_budget.py @@ -0,0 +1,92 @@ +"""The shared SSE connection budget must not leak (#513). + +claim_sse_slot() runs in the handler so hitting the cap can still answer 503. +release lives in the stream's finally -- and an async generator that is never +started never runs its finally. A client that disconnects before the response +body begins therefore held a slot forever: 200 of those and every progress +stream 503s with nothing actually connected, until the process restarts. +""" + +from __future__ import annotations + +import gc + +import pytest +from fastapi import HTTPException + +from app.api import events as _events + + +@pytest.fixture(autouse=True) +def _zero_budget(): + _events._sse_active = 0 + yield + _events._sse_active = 0 + + +def test_a_slot_is_held_then_released(): + slot = _events.SseSlot() + assert _events._sse_active == 1 + + slot.release() + assert _events._sse_active == 0 + + +def test_release_is_idempotent(): + # The stream's finally and the __del__ backstop can both fire; counting + # twice would free a slot that is still in use. + slot = _events.SseSlot() + slot.release() + slot.release() + + assert _events._sse_active == 0 + + +def test_a_slot_dropped_without_release_is_reclaimed(): + # The leak: the generator is created, so the slot is claimed, but never + # iterated, so its finally never runs. Collecting it must free the slot. + def _never_started(): + slot = _events.SseSlot() + + async def stream(): + try: + yield "data: x\n\n" + finally: + slot.release() + + return stream() # created, never iterated + + gen = _never_started() + assert _events._sse_active == 1 + + del gen + gc.collect() + + assert _events._sse_active == 0, "a slot leaked for the life of the process" + + +def test_the_cap_still_answers_503(): + held = [_events.SseSlot() for _ in range(_events._MAX_SSE_CONNECTIONS)] + assert _events._sse_active == _events._MAX_SSE_CONNECTIONS + + with pytest.raises(HTTPException) as excinfo: + _events.SseSlot() + assert excinfo.value.status_code == 503 + + for slot in held: + slot.release() + assert _events._sse_active == 0 + + +def test_a_refused_claim_holds_nothing(): + # If __init__ raises, no slot was taken -- so the failed attempt must not + # decrement on collection either. + held = [_events.SseSlot() for _ in range(_events._MAX_SSE_CONNECTIONS)] + with pytest.raises(HTTPException): + _events.SseSlot() + + gc.collect() + assert _events._sse_active == _events._MAX_SSE_CONNECTIONS + + for slot in held: + slot.release()