From 13a149743603ece59898222b6eb221bac4aed095 Mon Sep 17 00:00:00 2001 From: adavyas Date: Tue, 14 Jul 2026 11:09:25 -0400 Subject: [PATCH] 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