From 13a149743603ece59898222b6eb221bac4aed095 Mon Sep 17 00:00:00 2001 From: adavyas Date: Tue, 14 Jul 2026 11:09:25 -0400 Subject: [PATCH 1/6] Harden DB session cleanup against task cancellation (DEV-1861) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cancelled client-facing requests (search, dialectic chat/stream) could abandon an open transaction: when a request task is cancelled, the CancelledError can be re-delivered onto the ROLLBACK/close awaits in the session-cleanup `finally`, so ROLLBACK never reaches Postgres and close() never runs. Supavisor (transaction mode) does not reset the orphaned backend on disconnect, leaving it `idle in transaction` indefinitely. These zombies accrued ~5/hour, exhausting the connection pool and pinning the xmin horizon. Fix, centralized in the three session entry points (get_db, get_read_db, tracked_db) so every subsystem inherits it: - `_run_to_completion`: run cleanup as a detached task and await it through asyncio.shield in a loop, so a (re-delivered) cancellation cannot interrupt ROLLBACK/close. Survives a cancel storm. - `_finalize_session`: rollback then always close (inner try/finally); invalidate the connection on a broken-mid-protocol InterfaceError/ OperationalError/DBAPIError so a dead connection is not returned to the pool; never let a cleanup error mask the original. Note: anyio.CancelScope(shield=True) — the originally-proposed primitive — does NOT defer a native asyncio.Task.cancel() (verified empirically under both asyncio and anyio loops), which is how Starlette and the deriver's uvloop cancel; hence the stdlib approach. No new dependency. Tests: cancellation, re-delivery-during-cleanup, cancel-storm, broken-connection->invalidate, and DI-path teardown. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/dependencies.py | 108 ++++++++++++++++++----- tests/test_dependencies.py | 175 ++++++++++++++++++++++++++++++++++++- 2 files changed, 258 insertions(+), 25 deletions(-) diff --git a/src/dependencies.py b/src/dependencies.py index 060186b4..d07e45b8 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -1,11 +1,90 @@ +import asyncio +import logging import uuid +from collections.abc import Coroutine from contextlib import asynccontextmanager +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__) + + +async def _run_to_completion(coro: Coroutine[Any, Any, None]) -> None: + """Run ``coro`` to completion even if the current task is cancelled. + + This is the cancellation shield for DB-session cleanup (DEV-1861). A client + disconnect cancels the request task, and the cancellation can be re-delivered + onto the ROLLBACK/close awaits — leaving the backend parked 'idle in + transaction' with an open BEGIN that Supavisor v2 does not reset. + + NOTE: ``anyio.CancelScope(shield=True)`` does NOT cover this: it only defers + anyio-scoped cancellation, not a native ``asyncio.Task.cancel()`` (which is + how the deriver's uvloop and the ASGI server actually cancel). We instead run + cleanup as a detached task and keep awaiting it through ``asyncio.shield`` — + which never cancels the inner task — re-raising the cancellation only after + cleanup has fully finished. Survives repeated re-delivery (a cancel storm). + """ + fut = asyncio.ensure_future(coro) + cancelled = False + while not fut.done(): + try: + await asyncio.shield(fut) + except asyncio.CancelledError: + # Our task was cancelled; the shielded cleanup keeps running. Note it + # and keep waiting so ROLLBACK/close actually reach Postgres. + cancelled = True + exc = fut.exception() + if exc is not None: # _finalize_session swallows its own errors; be safe. + logger.exception("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. + + Runs via ``_run_to_completion`` (see callers) so it cannot be interrupted by + task cancellation. Structurally, close() is ALSO guaranteed by an inner + ``finally`` even if the rollback await is cut off — a second line of defense + the shield backs up: if ROLLBACK is interrupted before reaching Postgres and + close() never runs, the pooled connection is orphaned with an open BEGIN and + the backend parks 'idle in transaction' indefinitely (DEV-1861). Supavisor + v2 does NOT reset such orphaned transactions on client disconnect in + transaction-pooling mode. (close() returning the connection to the pool + triggers a DBAPI-level ROLLBACK on reset, so the transaction is closed even + if the explicit rollback below was interrupted.) + + Each step is guarded independently: a failed ROLLBACK (e.g. the connection + broke mid-protocol) must not prevent close(), and a broken connection is + invalidated so the pool discards it rather than handing out a poisoned + connection on the next checkout. ROLLBACK is unconditional (a wire-level + no-op under read_only/AUTOCOMMIT, and cheap if the lazy session never + checked out a connection). + """ + try: + await db.rollback() + except (InterfaceError, OperationalError, DBAPIError) as e: + # The connection is broken mid-protocol; the ROLLBACK never reached + # Postgres. Invalidate so the pool discards the dead connection instead + # of returning it. Do not let this mask the original error path. + logger.warning("cleanup ROLLBACK failed; invalidating connection: %s", e) + try: + await db.invalidate() + except Exception: + logger.exception("failed to invalidate broken connection during cleanup") + except Exception: + logger.exception("unexpected error during cleanup ROLLBACK") + finally: + try: + await db.close() + except Exception: + logger.exception("failed to close session during cleanup") + async def get_db(): """FastAPI Dependency Generator for Database. @@ -18,18 +97,10 @@ 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() + # Cleanup must survive a re-delivered cancellation on the rollback/close + # awaits, or the connection leaks with an open BEGIN (DEV-1861). + await _run_to_completion(_finalize_session(db)) async def get_read_db(): @@ -47,10 +118,7 @@ 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 @@ -77,14 +145,10 @@ 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() + # Run cleanup to completion despite (re-delivered) cancellation — see + # get_db / _run_to_completion. + await _run_to_completion(_finalize_session(db)) if token: # Only reset if we set it request_context.reset(token) diff --git a/tests/test_dependencies.py b/tests/test_dependencies.py index 1b5c219b..6fad04f2 100644 --- a/tests/test_dependencies.py +++ b/tests/test_dependencies.py @@ -1,7 +1,9 @@ +import asyncio import uuid from typing import Any 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,151 @@ 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. + + +@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. + 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() + for _ in range(5): # cancel storm straddling the cleanup window + task.cancel() + await asyncio.sleep(0.05) + with pytest.raises(asyncio.CancelledError): + await task + + assert fake_db.rollback_calls == 1 + assert fake_db.close_calls == 1 # cleanup completed despite the storm From c2001407691811996f42f050292556f3560387cc Mon Sep 17 00:00:00 2001 From: adavyas Date: Tue, 14 Jul 2026 11:36:51 -0400 Subject: [PATCH 2/6] Trim docstrings/comments in dependencies.py to be proportionate Keep the load-bearing "why" (the anyio-shield caveat; close-in-finally as the backstop) and drop the restated Supavisor/idle-in-transaction prose that the PR/ticket already carry. Also collapse the try/except-pass cleanup guards to contextlib.suppress. No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/dependencies.py | 68 ++++++++++++--------------------------------- 1 file changed, 18 insertions(+), 50 deletions(-) diff --git a/src/dependencies.py b/src/dependencies.py index d07e45b8..9277edf8 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -1,8 +1,8 @@ import asyncio +import contextlib import logging import uuid from collections.abc import Coroutine -from contextlib import asynccontextmanager from typing import Any from fastapi import Depends @@ -15,19 +15,13 @@ logger = logging.getLogger(__name__) async def _run_to_completion(coro: Coroutine[Any, Any, None]) -> None: - """Run ``coro`` to completion even if the current task is cancelled. + """Run ``coro`` to completion even if the current task is cancelled (DEV-1861). - This is the cancellation shield for DB-session cleanup (DEV-1861). A client - disconnect cancels the request task, and the cancellation can be re-delivered - onto the ROLLBACK/close awaits — leaving the backend parked 'idle in - transaction' with an open BEGIN that Supavisor v2 does not reset. - - NOTE: ``anyio.CancelScope(shield=True)`` does NOT cover this: it only defers - anyio-scoped cancellation, not a native ``asyncio.Task.cancel()`` (which is - how the deriver's uvloop and the ASGI server actually cancel). We instead run - cleanup as a detached task and keep awaiting it through ``asyncio.shield`` — - which never cancels the inner task — re-raising the cancellation only after - cleanup has fully finished. Survives repeated re-delivery (a cancel storm). + Runs the coro as a detached task and keeps awaiting it through + ``asyncio.shield`` (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 @@ -35,11 +29,8 @@ async def _run_to_completion(coro: Coroutine[Any, Any, None]) -> None: try: await asyncio.shield(fut) except asyncio.CancelledError: - # Our task was cancelled; the shielded cleanup keeps running. Note it - # and keep waiting so ROLLBACK/close actually reach Postgres. - cancelled = True - exc = fut.exception() - if exc is not None: # _finalize_session swallows its own errors; be safe. + cancelled = True # keep waiting so ROLLBACK/close reach Postgres + if (exc := fut.exception()) is not None: # cleanup swallows its own errors logger.exception("session cleanup task failed", exc_info=exc) if cancelled: raise asyncio.CancelledError() @@ -48,42 +39,22 @@ async def _run_to_completion(coro: Coroutine[Any, Any, None]) -> None: async def _finalize_session(db: AsyncSession) -> None: """Roll back and close a session, releasing its pooled connection. - Runs via ``_run_to_completion`` (see callers) so it cannot be interrupted by - task cancellation. Structurally, close() is ALSO guaranteed by an inner - ``finally`` even if the rollback await is cut off — a second line of defense - the shield backs up: if ROLLBACK is interrupted before reaching Postgres and - close() never runs, the pooled connection is orphaned with an open BEGIN and - the backend parks 'idle in transaction' indefinitely (DEV-1861). Supavisor - v2 does NOT reset such orphaned transactions on client disconnect in - transaction-pooling mode. (close() returning the connection to the pool - triggers a DBAPI-level ROLLBACK on reset, so the transaction is closed even - if the explicit rollback below was interrupted.) - - Each step is guarded independently: a failed ROLLBACK (e.g. the connection - broke mid-protocol) must not prevent close(), and a broken connection is - invalidated so the pool discards it rather than handing out a poisoned - connection on the next checkout. ROLLBACK is unconditional (a wire-level - no-op under read_only/AUTOCOMMIT, and cheap if the lazy session never - checked out a 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: - # The connection is broken mid-protocol; the ROLLBACK never reached - # Postgres. Invalidate so the pool discards the dead connection instead - # of returning it. Do not let this mask the original error path. logger.warning("cleanup ROLLBACK failed; invalidating connection: %s", e) - try: + with contextlib.suppress(Exception): await db.invalidate() - except Exception: - logger.exception("failed to invalidate broken connection during cleanup") except Exception: logger.exception("unexpected error during cleanup ROLLBACK") finally: - try: + with contextlib.suppress(Exception): await db.close() - except Exception: - logger.exception("failed to close session during cleanup") async def get_db(): @@ -98,8 +69,7 @@ async def get_db(): try: yield db finally: - # Cleanup must survive a re-delivered cancellation on the rollback/close - # awaits, or the connection leaks with an open BEGIN (DEV-1861). + # Shielded so a cancelled request can't leak an open BEGIN (DEV-1861). await _run_to_completion(_finalize_session(db)) @@ -121,7 +91,7 @@ async def get_read_db(): 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. @@ -146,9 +116,7 @@ async def tracked_db(operation_name: str | None = None, *, read_only: bool = Fal try: yield db finally: - # Run cleanup to completion despite (re-delivered) cancellation — see - # get_db / _run_to_completion. - await _run_to_completion(_finalize_session(db)) + await _run_to_completion(_finalize_session(db)) # shielded — see get_db if token: # Only reset if we set it request_context.reset(token) From 76dc6c3c7981de3de2c256093a14c0716967f4f9 Mon Sep 17 00:00:00 2001 From: adavyas Date: Tue, 14 Jul 2026 11:48:01 -0400 Subject: [PATCH 3/6] Address review: bound cleanup, always reset context, add real-DB cancel test Review findings on the cancellation-shield fix: - Bound _run_to_completion with _CLEANUP_TIMEOUT_SECONDS. Shielded cleanup is uninterruptible by design, so a wedged connection (dead socket, no libpq timeout) would otherwise pin the task forever; past the bound we abandon cleanup rather than hang. Switched the wait to asyncio.wait (shields fut, supports a timeout) and guarded fut.exception()/fut.cancelled() so loop shutdown can't raise from the tail. - tracked_db: reset request_context in its own finally so a cancelled cleanup path can't leak the task-scoped contextvar into a reused long-lived task (deriver). - Tests: reset-context-on-cancellation, abandon-wedged-cleanup, and a real-DB wire cancellation test asserting no `idle in transaction` backend lingers after a cancelled in-flight write transaction (pg_stat_activity). Full suite: 21 passed live (incl. 4 pre-existing + 4 new real-DB tests). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/dependencies.py | 46 +++++++++++++++----- tests/test_dependencies.py | 87 ++++++++++++++++++++++++++++++++++++++ uv.lock | 4 -- 3 files changed, 122 insertions(+), 15 deletions(-) diff --git a/src/dependencies.py b/src/dependencies.py index 9277edf8..40f90a64 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -13,25 +13,45 @@ 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 keeps awaiting it through - ``asyncio.shield`` (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. + 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 while not fut.done(): try: - await asyncio.shield(fut) + # wait() shields fut from our cancellation but still bounds the wait. + await asyncio.wait({fut}, timeout=_CLEANUP_TIMEOUT_SECONDS) except asyncio.CancelledError: cancelled = True # keep waiting so ROLLBACK/close reach Postgres - if (exc := fut.exception()) is not None: # cleanup swallows its own errors - logger.exception("session cleanup task failed", exc_info=exc) + continue + if not fut.done(): + # 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 + 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() @@ -116,9 +136,13 @@ async def tracked_db(operation_name: str | None = None, *, read_only: bool = Fal try: yield db finally: - await _run_to_completion(_finalize_session(db)) # shielded — see get_db - 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 6fad04f2..427d7cf3 100644 --- a/tests/test_dependencies.py +++ b/tests/test_dependencies.py @@ -477,3 +477,90 @@ async def test_tracked_db_survives_cancellation_storm_during_cleanup( 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()), timeout=5) + assert started.is_set() + + +@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 + + 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" diff --git a/uv.lock b/uv.lock index a71283bb..f0763d94 100644 --- a/uv.lock +++ b/uv.lock @@ -7,10 +7,6 @@ resolution-markers = [ "python_full_version < '3.13'", ] -[options] -exclude-newer = "2026-06-10T19:45:16.447512Z" -exclude-newer-span = "P5D" - [manifest] members = [ "honcho", From 3bf7bd4b579ed29196e617a2f1e005a611fae944 Mon Sep 17 00:00:00 2001 From: adavyas Date: Tue, 14 Jul 2026 11:48:30 -0400 Subject: [PATCH 4/6] Restore uv.lock (unrelated exclude-newer change swept in by git add -A) Co-Authored-By: Claude Opus 4.8 (1M context) --- uv.lock | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/uv.lock b/uv.lock index f0763d94..a71283bb 100644 --- a/uv.lock +++ b/uv.lock @@ -7,6 +7,10 @@ resolution-markers = [ "python_full_version < '3.13'", ] +[options] +exclude-newer = "2026-06-10T19:45:16.447512Z" +exclude-newer-span = "P5D" + [manifest] members = [ "honcho", From 6d517bbd3b21529764c976d655ed273aa73f4637 Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:12:20 -0400 Subject: [PATCH 5/6] fix: make DB cleanup abandon-deadline monotonic across cancel re-delivery _run_to_completion passed _CLEANUP_TIMEOUT_SECONDS to each asyncio.wait() call, so every re-delivered CancelledError re-armed a fresh 10s window and the deadline never elapsed. anyio re-delivers cancellation on every event loop iteration (CancelScope._deliver_cancellation reschedules itself via call_soon until the task leaves the scope), and Starlette's BaseHTTPMiddleware -- registered for every request in main.py -- wraps the request in an anyio task group. So on the API path the abandon valve was dead exactly where it was needed: a wedged connection pinned the request task, hot-spinning the loop, for as long as the storm lasted. Compute one monotonic deadline before the loop and derive each wait timeout from the remaining time. Tests: storm helper now re-cancels at anyio's call_soon cadence instead of 5 sparse cancels (sparse cancels leave quiet ticks where a per-wait timeout re-arms and cleanup finishes, which is why the existing storm test passed either way). New test_run_to_completion_deadline_holds_under_cancel_storm fails on the old loop ("not abandoned after 5.0s") and passes with the fix. Co-Authored-By: Claude Opus 5 (1M context) --- src/dependencies.py | 23 ++++++++------ tests/test_dependencies.py | 65 +++++++++++++++++++++++++++++++++++--- 2 files changed, 74 insertions(+), 14 deletions(-) diff --git a/src/dependencies.py b/src/dependencies.py index 40f90a64..07c87a59 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -32,22 +32,27 @@ async def _run_to_completion(coro: Coroutine[Any, Any, None]) -> None: """ 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(): - try: - # wait() shields fut from our cancellation but still bounds the wait. - await asyncio.wait({fut}, timeout=_CLEANUP_TIMEOUT_SECONDS) - except asyncio.CancelledError: - cancelled = True # keep waiting so ROLLBACK/close reach Postgres - continue - if 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", + "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: diff --git a/tests/test_dependencies.py b/tests/test_dependencies.py index 427d7cf3..442931a3 100644 --- a/tests/test_dependencies.py +++ b/tests/test_dependencies.py @@ -1,6 +1,6 @@ import asyncio import uuid -from typing import Any +from typing import Any, final import pytest from sqlalchemy.exc import InterfaceError @@ -339,6 +339,25 @@ async def test_read_only_session_works_with_tracing_checkout_hook() -> None: # 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, @@ -452,6 +471,7 @@ async def test_tracked_db_survives_cancellation_storm_during_cleanup( # 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 @@ -469,9 +489,7 @@ async def test_tracked_db_survives_cancellation_storm_during_cleanup( task = asyncio.ensure_future(worker()) await entered.wait() - for _ in range(5): # cancel storm straddling the cleanup window - task.cancel() - await asyncio.sleep(0.05) + storm_until_done(task) with pytest.raises(asyncio.CancelledError): await task @@ -515,10 +533,47 @@ async def test_run_to_completion_abandons_wedged_cleanup( 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()), timeout=5) + 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 From cde84654bf71bb4f458d19ca3ca4b3e1bbe5d8d6 Mon Sep 17 00:00:00 2001 From: adavyas Date: Thu, 30 Jul 2026 13:01:22 -0400 Subject: [PATCH 6/6] test: terminate leaked backend if the real-DB cancellation test fails MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On assertion failure the backend under test is parked 'idle in transaction' — exactly the bug this PR fixes — and would poison the shared test pool for the rest of the session. Terminate it in a finally regardless of pass/fail. (CodeRabbit review, tests/test_dependencies.py) Co-Authored-By: Claude Fable 5 --- tests/test_dependencies.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/tests/test_dependencies.py b/tests/test_dependencies.py index 442931a3..a4d59adb 100644 --- a/tests/test_dependencies.py +++ b/tests/test_dependencies.py @@ -612,10 +612,17 @@ async def test_tracked_db_no_idle_in_transaction_after_real_cancellation() -> No with pytest.raises(asyncio.CancelledError): await task - 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) + try: state = await backend_state(pid) - assert state != "idle in transaction", f"backend left {state!r} after cancel" + 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})