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()