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] 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