From 76dc6c3c7981de3de2c256093a14c0716967f4f9 Mon Sep 17 00:00:00 2001 From: adavyas Date: Tue, 14 Jul 2026 11:48:01 -0400 Subject: [PATCH] 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",