RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Harden DB session cleanup against task cancellation (DEV-1861) 

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) <noreply@anthropic.com>
This commit is contained in:
adavyas 2026-07-14 11:09:25 -04:00
parent a2adeb9f45
commit 13a1497436
2 changed files with 258 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

108
src/dependencies.py
View File

@ -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)

175
tests/test_dependencies.py
View File

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