diff --git a/src/dependencies.py b/src/dependencies.py index 060186b4..07c87a59 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -1,11 +1,86 @@ +import asyncio +import contextlib +import logging import uuid -from contextlib import asynccontextmanager +from collections.abc import Coroutine +from typing import Any from fastapi import Depends +from sqlalchemy.exc import DBAPIError, InterfaceError, OperationalError from sqlalchemy.ext.asyncio import AsyncSession from src.db import ReadSessionLocal, SessionLocal, request_context +logger = logging.getLogger(__name__) + +# Upper bound on how long shielded cleanup may run before we stop waiting on it. +# Healthy ROLLBACK/close take milliseconds; this only bites a wedged connection +# (e.g. a black-holed socket with no libpq timeout) so it can't pin the task +# forever. Cleanup is abandoned (not cancelled) past this — the pool/OS reclaims +# the connection — since a shielded task cannot be interrupted anyway. +_CLEANUP_TIMEOUT_SECONDS = 10.0 + + +async def _run_to_completion(coro: Coroutine[Any, Any, None]) -> None: + """Run ``coro`` to completion even if the current task is cancelled (DEV-1861). + + Runs the coro as a detached task and waits on it via ``asyncio.wait`` (which + never cancels the inner task), re-raising the cancellation only afterward. + Survives repeated re-delivery. NB: ``anyio.CancelScope(shield=True)`` does + NOT work here — it ignores a native ``asyncio.Task.cancel()``, which is how + Starlette and the deriver cancel. + """ + fut = asyncio.ensure_future(coro) + cancelled = False + loop = asyncio.get_running_loop() + # One monotonic deadline for the whole loop, not a per-wait timeout: anyio + # (Starlette's BaseHTTPMiddleware wraps every request in a task group) + # re-delivers cancellation on every event-loop tick, so a per-wait timeout + # re-arms each tick and never elapses. + deadline = loop.time() + _CLEANUP_TIMEOUT_SECONDS + while not fut.done(): + remaining = deadline - loop.time() + if remaining <= 0: + # Wedged cleanup (dead socket). Stop pinning this task; leave fut to + # error out / be reclaimed rather than block indefinitely. + logger.error( + "DB session cleanup exceeded %.0fs; abandoning to avoid pinning the task", + _CLEANUP_TIMEOUT_SECONDS, + ) + break + try: + # wait() shields fut from our cancellation but still bounds the wait. + await asyncio.wait({fut}, timeout=remaining) + except asyncio.CancelledError: + cancelled = True # keep waiting so ROLLBACK/close reach Postgres + if fut.cancelled(): + cancelled = True # cleanup itself was cancelled (close() still ran) + elif fut.done() and (exc := fut.exception()) is not None: + logger.error("session cleanup task failed", exc_info=exc) + if cancelled: + raise asyncio.CancelledError() + + +async def _finalize_session(db: AsyncSession) -> None: + """Roll back and close a session, releasing its pooled connection. + + close() is guaranteed by the inner ``finally`` even if the rollback await is + interrupted — and returning the connection to the pool triggers a DBAPI-level + ROLLBACK on reset, so the transaction closes regardless. A broken-mid-protocol + connection is invalidated so the pool discards it instead of reusing it. + """ + try: + await db.rollback() + except (InterfaceError, OperationalError, DBAPIError) as e: + logger.warning("cleanup ROLLBACK failed; invalidating connection: %s", e) + with contextlib.suppress(Exception): + await db.invalidate() + except Exception: + logger.exception("unexpected error during cleanup ROLLBACK") + finally: + with contextlib.suppress(Exception): + await db.close() + async def get_db(): """FastAPI Dependency Generator for Database. @@ -18,18 +93,9 @@ async def get_db(): db: AsyncSession = SessionLocal() try: yield db - except Exception: - await db.rollback() - raise finally: - # Always send ROLLBACK unconditionally so the wire-level transaction - # is closed before the TCP connection drops. Supavisor v2 does NOT - # clean up orphaned transactions on client disconnect in transaction- - # pooling mode, so relying on `in_transaction()` (Python-side state) - # can leave the backend pinned with an open BEGIN. (Cheap no-op if the - # lazy session never checked out a connection.) - await db.rollback() - await db.close() + # Shielded so a cancelled request can't leak an open BEGIN (DEV-1861). + await _run_to_completion(_finalize_session(db)) async def get_read_db(): @@ -47,13 +113,10 @@ async def get_read_db(): try: yield db finally: - # rollback is a wire-level no-op under AUTOCOMMIT; kept to reset any - # Python-side session state before close, mirroring get_db. - await db.rollback() - await db.close() + await _run_to_completion(_finalize_session(db)) -@asynccontextmanager +@contextlib.asynccontextmanager async def tracked_db(operation_name: str | None = None, *, read_only: bool = False): """Context manager for tracked database sessions. @@ -77,16 +140,14 @@ async def tracked_db(operation_name: str | None = None, *, read_only: bool = Fal db = (ReadSessionLocal if read_only else SessionLocal)() try: yield db - except Exception: - await db.rollback() - raise finally: - # Always send ROLLBACK unconditionally — see get_db() comment. (Under - # read_only/AUTOCOMMIT it is a wire-level no-op.) - await db.rollback() - await db.close() - if token: # Only reset if we set it - request_context.reset(token) + try: + await _run_to_completion(_finalize_session(db)) # shielded — see get_db + finally: + # Must run even when cleanup re-raises CancelledError, or a reused + # long-lived task (e.g. the deriver) leaks this contextvar. + if token: # Only reset if we set it + request_context.reset(token) db: AsyncSession = Depends(get_db) diff --git a/tests/test_dependencies.py b/tests/test_dependencies.py index 1b5c219b..a4d59adb 100644 --- a/tests/test_dependencies.py +++ b/tests/test_dependencies.py @@ -1,7 +1,9 @@ +import asyncio import uuid -from typing import Any +from typing import Any, final import pytest +from sqlalchemy.exc import InterfaceError import src.dependencies as dependencies_module from src.config import settings @@ -11,12 +13,24 @@ from src.dependencies import tracked_db as real_tracked_db class FakeSession: - def __init__(self, *, in_transaction: bool = False): + def __init__( + self, + *, + in_transaction: bool = False, + rollback_exc: BaseException | None = None, + close_exc: BaseException | None = None, + ): self._in_transaction: bool = in_transaction + # Optional failures injected into the cleanup path to exercise + # cancellation-hardening (DEV-1861): a CancelledError re-delivered on + # rollback, or a broken-connection error that must trigger invalidate(). + self._rollback_exc: BaseException | None = rollback_exc + self._close_exc: BaseException | None = close_exc self.execute_calls: list[tuple[Any, ...]] = [] self.rollback_calls: int = 0 self.close_calls: int = 0 self.connection_calls: int = 0 + self.invalidate_calls: int = 0 async def connection(self) -> None: # Tracks checkout attempts so tests can assert get_db/tracked_db stay @@ -28,9 +42,16 @@ class FakeSession: async def rollback(self) -> None: self.rollback_calls += 1 + if self._rollback_exc is not None: + raise self._rollback_exc + + async def invalidate(self) -> None: + self.invalidate_calls += 1 async def close(self) -> None: self.close_calls += 1 + if self._close_exc is not None: + raise self._close_exc def in_transaction(self) -> bool: return self._in_transaction @@ -74,7 +95,7 @@ async def test_get_db_rolls_back_and_closes_when_consumer_raises( with pytest.raises(RuntimeError, match="boom"): await dep_gen.athrow(RuntimeError("boom")) - assert fake_db.rollback_calls == 2 # once in except, once in finally + assert fake_db.rollback_calls == 1 # single rollback in shielded finally assert fake_db.close_calls == 1 @@ -133,7 +154,7 @@ async def test_tracked_db_rolls_back_on_error_and_closes( async with real_tracked_db("operation"): raise ValueError("failed operation") - assert fake_db.rollback_calls == 2 # once in except, once in finally + assert fake_db.rollback_calls == 1 # single rollback in shielded finally assert fake_db.close_calls == 1 @@ -308,3 +329,300 @@ async def test_read_only_session_works_with_tracing_checkout_hook() -> None: finally: event.remove(engine.sync_engine, "checkout", _set_application_name_on_checkout) request_context.reset(context_token) + + +# --- Cancellation-hardened cleanup (DEV-1861) -------------------------------- +# +# The regression: on a cancelled request task, cleanup must still ROLLBACK and +# close so the pooled connection is released and no backend is left parked +# 'idle in transaction'. The shielded finally in dependencies.py must survive +# both a CancelledError raised in the body AND one re-delivered during cleanup. + + +def storm_until_done(task: "asyncio.Task[None]") -> None: + """Re-cancel `task` at anyio's cadence until it finishes. + + anyio's CancelScope._deliver_cancellation re-cancels via loop.call_soon on + every event-loop tick for as long as the task is inside the scope — and + Starlette's BaseHTTPMiddleware wraps every request in a task group. A few + sparse cancel()s leave quiet ticks that let a per-wait timeout re-arm and + cleanup finish; this cadence leaves none. + """ + loop = asyncio.get_running_loop() + + def restorm() -> None: + if not task.done(): + task.cancel() + loop.call_soon(restorm) + + restorm() + + +@pytest.mark.asyncio +async def test_tracked_db_rolls_back_and_closes_on_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Body cancelled mid-transaction (the client-disconnect case): cleanup must + # still run and the CancelledError must propagate (never be swallowed). + fake_db = FakeSession() + monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) + + with pytest.raises(asyncio.CancelledError): + async with real_tracked_db("cancelled_op"): + raise asyncio.CancelledError() + + assert fake_db.rollback_calls == 1 + assert fake_db.close_calls == 1 + + +@pytest.mark.asyncio +async def test_tracked_db_closes_when_cancel_redelivered_during_rollback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # The core DEV-1861 defense: a cancellation re-delivered on the ROLLBACK + # await must NOT skip close(). We simulate re-delivery by making rollback() + # itself raise CancelledError; close() must still run so the connection is + # released rather than orphaned with an open BEGIN. + fake_db = FakeSession(rollback_exc=asyncio.CancelledError()) + monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) + + with pytest.raises(asyncio.CancelledError): + async with real_tracked_db("cancelled_cleanup_op"): + pass + + assert fake_db.rollback_calls == 1 + assert fake_db.close_calls == 1 # close ran despite the interrupted rollback + + +@pytest.mark.asyncio +async def test_tracked_db_invalidates_connection_on_broken_rollback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Mid-protocol break: ROLLBACK fails with a driver InterfaceError. The dead + # connection must be invalidated (not returned to the pool) and close() must + # still run. The broken-connection error must not surface to the caller — + # cleanup is best-effort by this point. + broken = InterfaceError("ROLLBACK", None, Exception("connection is closed")) + fake_db = FakeSession(rollback_exc=broken) + monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) + + async with real_tracked_db("broken_conn_op"): + pass + + assert fake_db.rollback_calls == 1 + assert fake_db.invalidate_calls == 1 + assert fake_db.close_calls == 1 + + +@pytest.mark.asyncio +async def test_get_db_rolls_back_and_closes_on_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Same guarantee via the FastAPI DI generator path (Starlette cancels the + # request task on client disconnect and drives teardown through athrow). + fake_db = FakeSession() + monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) + + dep_gen = real_get_db() + await anext(dep_gen) + + with pytest.raises(asyncio.CancelledError): + await dep_gen.athrow(asyncio.CancelledError()) + + assert fake_db.rollback_calls == 1 + assert fake_db.close_calls == 1 + + +@pytest.mark.asyncio +async def test_tracked_db_survives_real_task_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # End-to-end check with a genuine asyncio cancellation (not a hand-raised + # exception): cancel a task blocked inside the tracked_db body and assert + # cleanup still completed and the CancelledError still propagates. + fake_db = FakeSession() + monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) + + entered = asyncio.Event() + + async def worker() -> None: + async with real_tracked_db("blocked_op"): + entered.set() + await asyncio.sleep(3600) # block until cancelled + + task = asyncio.ensure_future(worker()) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert fake_db.rollback_calls == 1 + assert fake_db.close_calls == 1 + + +@pytest.mark.asyncio +async def test_tracked_db_survives_cancellation_storm_during_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # The strongest DEV-1861 guarantee: repeated cancellations re-delivered + # WHILE cleanup is running must not stop rollback+close from completing. + # `_run_to_completion` runs cleanup as a shielded detached task, so even a + # storm of task.cancel()s cannot orphan the connection. (This is precisely + # the case `anyio.CancelScope(shield=True)` fails to cover for a native + # task.cancel(), which is why we do not use it.) rollback() here awaits, so + # a re-delivered cancel has a real window to land mid-cleanup. + @final + class SlowRollbackSession(FakeSession): + async def rollback(self) -> None: + self.rollback_calls += 1 + await asyncio.sleep(0.2) # cancellation window during cleanup + + fake_db = SlowRollbackSession() + monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) + + entered = asyncio.Event() + + async def worker() -> None: + async with real_tracked_db("stormed_op"): + entered.set() + await asyncio.sleep(3600) + + task = asyncio.ensure_future(worker()) + await entered.wait() + storm_until_done(task) + with pytest.raises(asyncio.CancelledError): + await task + + assert fake_db.rollback_calls == 1 + assert fake_db.close_calls == 1 # cleanup completed despite the storm + + +@pytest.mark.asyncio +async def test_tracked_db_resets_request_context_on_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # The task-scoped contextvar must be reset even when cleanup re-raises + # CancelledError, else a reused long-lived task (e.g. the deriver) carries a + # stale `task:...` context into later work and corrupts tracing/attribution. + fake_db = FakeSession() + monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) + + clear_token = request_context.set(None) + try: + with pytest.raises(asyncio.CancelledError): + async with real_tracked_db("ctx_cancel_op"): + assert (request_context.get() or "").startswith("task:ctx_cancel_op:") + raise asyncio.CancelledError() + assert request_context.get() is None # reset despite the cancel + finally: + request_context.reset(clear_token) + + +@pytest.mark.asyncio +async def test_run_to_completion_abandons_wedged_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Shielded cleanup is uninterruptible by design; a wedged connection (dead + # socket, no libpq timeout) must therefore NOT pin the caller forever. + # _run_to_completion abandons cleanup past _CLEANUP_TIMEOUT_SECONDS. + monkeypatch.setattr(dependencies_module, "_CLEANUP_TIMEOUT_SECONDS", 0.2) + started = asyncio.Event() + + async def wedged() -> None: + started.set() + await asyncio.sleep(3600) # never completes + + # Returns (abandons) well under the sleep instead of hanging. + await asyncio.wait_for( + dependencies_module._run_to_completion(wedged()), # pyright: ignore[reportPrivateUsage] + timeout=5, + ) + assert started.is_set() + + +@pytest.mark.asyncio +async def test_run_to_completion_deadline_holds_under_cancel_storm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # The abandon deadline must be monotonic across re-deliveries. A per-wait + # timeout re-arms on every CancelledError, and anyio re-delivers on every + # event-loop tick (see storm_until_done), so it would never elapse — a wedged + # connection would pin the request task for as long as the storm lasts. + monkeypatch.setattr(dependencies_module, "_CLEANUP_TIMEOUT_SECONDS", 0.3) + + async def wedged() -> None: + await asyncio.sleep(3600) # never completes + + loop = asyncio.get_running_loop() + task = asyncio.ensure_future( + dependencies_module._run_to_completion(wedged()) # pyright: ignore[reportPrivateUsage] + ) + await asyncio.sleep(0) # let it reach the wait loop + storm_until_done(task) + + began = loop.time() + for _ in range(50): # bounded poll: a broken deadline hangs rather than fails + if task.done(): + break + await asyncio.sleep(0.1) + elapsed = loop.time() - began + if not task.done(): + task.cancel() + pytest.fail(f"cleanup not abandoned after {elapsed:.1f}s under a cancel storm") + + assert task.cancelled() # cancellation still propagates after abandoning + assert elapsed < 2.0, f"abandoned only after {elapsed:.1f}s (deadline was 0.3s)" + + +@pytest.mark.asyncio +async def test_tracked_db_no_idle_in_transaction_after_real_cancellation() -> None: + # Wire-level chaos test (ticket verification): cancel a REAL in-flight write + # transaction mid-flight and assert, from an independent connection, that the + # backend is not left parked 'idle in transaction'. Complements the mock + # tests above with a real pooled connection receiving ROLLBACK under cancel. + from sqlalchemy import text + + from src.db import read_engine + + entered = asyncio.Event() + pid_box: dict[str, int] = {} + + async def worker() -> None: + async with real_tracked_db("chaos_cancel") as db: + pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar() + pid_box["pid"] = int(pid) # pyright: ignore[reportArgumentType] + entered.set() + await asyncio.sleep(3600) # hold the open BEGIN until cancelled + + async def backend_state(pid: int) -> str | None: + async with read_engine.connect() as obs: + return ( + await obs.execute( + text("SELECT state FROM pg_stat_activity WHERE pid = :p"), + {"p": pid}, + ) + ).scalar() + + task = asyncio.ensure_future(worker()) + await entered.wait() + pid = pid_box["pid"] + assert await backend_state(pid) == "idle in transaction" # precondition + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + try: + state = await backend_state(pid) + for _ in range(20): # allow shielded cleanup a moment to land + if state != "idle in transaction": + break + await asyncio.sleep(0.1) + state = await backend_state(pid) + assert state != "idle in transaction", f"backend left {state!r} after cancel" + finally: + # On failure the backend is parked exactly like the bug under test; + # terminate it so it can't poison the shared pool for later tests. + if await backend_state(pid) == "idle in transaction": + async with read_engine.connect() as obs: + await obs.execute(text("SELECT pg_terminate_backend(:p)"), {"p": pid})