Harden DB session cleanup against task cancellation (DEV-1861)#908
Harden DB session cleanup against task cancellation (DEV-1861)#908adavyas wants to merge 4 commits into
Conversation
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>
WalkthroughAsync database dependency teardown now uses bounded, cancellation-resistant session finalization, including rollback failure invalidation and guaranteed context reset. Tests cover cancellation, cleanup failures, timeout abandonment, and pooled PostgreSQL transaction state. ChangesDatabase cleanup hardening
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Request
participant DatabaseDependency
participant CleanupRunner
participant Session
participant PostgreSQL
Request->>DatabaseDependency: use database session
Request-->>DatabaseDependency: cancellation or request failure
DatabaseDependency->>CleanupRunner: run bounded finalization
CleanupRunner->>Session: rollback
Session->>PostgreSQL: release transaction state
CleanupRunner->>Session: invalidate on rollback error
CleanupRunner->>Session: close
DatabaseDependency->>DatabaseDependency: reset request context
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
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) <noreply@anthropic.com>
…el 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) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/dependencies.py`:
- Around line 24-56: The _run_to_completion function must enforce one monotonic
overall cleanup deadline instead of resetting the timeout after each
cancellation. Compute a deadline before the wait loop, derive each asyncio.wait
timeout from the remaining time, and abandon cleanup once the deadline is
reached while preserving cancellation re-raising and existing task-error
handling.
In `@tests/test_dependencies.py`:
- Around line 554-566: Ensure the cancellation test’s backend cleanup always
runs when the final state assertion fails: wrap the retry and assertion logic
around backend_state(pid) in failure-safe cleanup that terminates the backend
identified by pid before propagating the assertion failure. Preserve the
existing successful cleanup verification and cancellation behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 15f55013-9c39-4579-a36b-718e007ec26f
📒 Files selected for processing (2)
src/dependencies.pytests/test_dependencies.py
| 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 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: | ||
| # 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(): | ||
| # 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() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== outline src/dependencies.py ==\n'
ast-grep outline src/dependencies.py --view expanded || true
printf '\n== outline tests/test_dependencies.py ==\n'
ast-grep outline tests/test_dependencies.py --view expanded || true
printf '\n== relevant lines src/dependencies.py ==\n'
sed -n '1,140p' src/dependencies.py | cat -n
printf '\n== relevant lines tests/test_dependencies.py ==\n'
sed -n '1,260p' tests/test_dependencies.py | cat -nRepository: plastic-labs/honcho
Length of output: 20326
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "stormed_op|chaos_cancel|_run_to_completion|CLEANUP_TIMEOUT" src testsRepository: plastic-labs/honcho
Length of output: 1435
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '430,560p' tests/test_dependencies.py | cat -nRepository: plastic-labs/honcho
Length of output: 6162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '131,150p' tests/test_dependencies.py | cat -nRepository: plastic-labs/honcho
Length of output: 992
Track a monotonic deadline across cancellation retries. asyncio.wait(..., timeout=_CLEANUP_TIMEOUT_SECONDS) resets on every CancelledError, so a cancellation storm can keep _run_to_completion() alive well past the intended 10s cap. The stormed_op/chaos_cancel tests exercise cleanup under repeated cancellation, but they don’t assert a hard wall-clock bound.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/dependencies.py` around lines 24 - 56, The _run_to_completion function
must enforce one monotonic overall cleanup deadline instead of resetting the
timeout after each cancellation. Compute a deadline before the wait loop, derive
each asyncio.wait timeout from the remaining time, and abandon cleanup once the
deadline is reached while preserving cancellation re-raising and existing
task-error handling.
| 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" |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
On assertion failure, this test leaks a real orphaned backend connection.
If state never leaves idle in transaction within the 2s retry budget, the final assert fails but pid is never terminated — the backend stays parked exactly like the bug this PR fixes, potentially exhausting the shared pool for the rest of the test session. Terminate the backend as cleanup regardless of pass/fail.
🧹 Proposed fix to terminate the backend on failure
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"
+ try:
+ assert state != "idle in transaction", f"backend left {state!r} after cancel"
+ finally:
+ if state == "idle in transaction":
+ async with read_engine.connect() as obs:
+ await obs.execute(text("SELECT pg_terminate_backend(:p)"), {"p": pid})📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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" | |
| 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) | |
| try: | |
| assert state != "idle in transaction", f"backend left {state!r} after cancel" | |
| finally: | |
| if state == "idle in transaction": | |
| async with read_engine.connect() as obs: | |
| await obs.execute(text("SELECT pg_terminate_backend(:p)"), {"p": pid}) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_dependencies.py` around lines 554 - 566, Ensure the cancellation
test’s backend cleanup always runs when the final state assertion fails: wrap
the retry and assertion logic around backend_state(pid) in failure-safe cleanup
that terminates the backend identified by pid before propagating the assertion
failure. Preserve the existing successful cleanup verification and cancellation
behavior.
ulyssepence
left a comment
ulyssepence
left a comment
There was a problem hiding this comment.
I read over the code in src/dependencies.py and it looks good to me. Interesting how asyncio.wait can dodge its coroutine being canceled!
The only thing that I wasn't clear on from the ticket is how confident we are that this is the problem. Was there something like a smoking gun or was this more a best guess?
|
Finn-loop review of 3bf7bd4 CI: not configured (no required checks reported on the branch) ReviewSummary: Hardens all three DB session entry points ( 1. Must fix before merge[CI] No required CI is configured on this branch — 2. Should fix soonNone — the code changes are clean and well-tested. 3. CI check details (not required)
4. Code reviewContract compliance ✓Per the PR body (no formal Paperclip issue linked — Linear DEV-1861): the fix addresses the
Code quality
Tests ✓21 tests in 5. Safe to mergeThe code changes are correct, self-contained, and comprehensively tested. The merge is blocked by CI configuration (no required checks, stale branch-protection rule) — an admin needs to resolve that before merging. Labels to apply (no admin rights on this token):
|
4 participants
Summary
Hardens the DB session cleanup path against task cancellation — one of the causes of the
idle in transactionzombie-connection leak in DEV-1861, not all of them (see Coverage below). Apg_stat_activitysnapshot showed 101 backends stuckidle in transactionvia Supavisor (93 for 1–24h, worst 21h27m) across 34 tenant schemas, accruing ~5/hour — draining the pool, pinning the xmin horizon (vacuum starvation), and blocking DDL.Root cause
Cancellation (client disconnect, stream abort, request-scoped timeout, worker shutdown) is delivered at an
await. Session cleanup lived in afinally:except Exceptionnever catches cancellation (CancelledErroris aBaseException), so only thefinallycovered it.finally's ownawaits are cancellable: a re-deliveredCancelledErrorinterruptsROLLBACKbefore it reaches Postgres and skipsclose(), orphaning the connection mid-BEGIN.idle in transactionindefinitely.Where the re-delivery comes from. A single
Task.cancel()deliversCancelledErrorexactly once, so the oldfinallywould have completed on its own — the bug needs cancellation re-delivered during cleanup. anyio provides exactly that:CancelScope._deliver_cancellationreschedules itself vialoop.call_soonon every event-loop tick for as long as the task is inside the scope, and Starlette'sBaseHTTPMiddleware(registered for every request insrc/main.py) wraps each request in an anyio task group. So the storm is real and reachable on the whole API path.This is not the held-session-across-external-call antipattern — the dominant path (post-Turbopuffer batch fetch) correctly opens the DB session after the embedding/vector call. The leak is purely in the cleanup path.
Coverage — what this does and does not fix
On process death no cleanup code runs at all: the socket to Supavisor just drops, and Supavisor in transaction mode still does not reset the server backend. Given the observed age distribution (93 backends aged 1–24h, worst 21h27m), an unknown but plausibly large share of the snapshot is this second class. The only thing that reclaims those is the server-side
idle_in_transaction_session_timeout— see Required alongside this PR.Fix
Centralized in the three session entry points (
get_db,get_read_db,tracked_db) — every subsystem (deriver, dreamer, reconciler, webhooks, streaming) opens sessions through these, so they all inherit the fix with no call-site changes._run_to_completion— runs cleanup as a detached task and waits on it viaasyncio.wait(which never cancels the inner task), re-raising the cancellation only after cleanup finishes. Survives a cancel storm. A single monotonic deadline bounds the whole loop so a wedged connection can't pin the task._finalize_session— rollback then always close (innertry/finally);invalidate()the connection on a broken-mid-protocolInterfaceError/OperationalError/DBAPIErrorso a dead connection isn't returned to the pool; cleanup errors are caught+logged so they can't mask the original error.Two guarantees stack: structural (
close()runs even ifrollback()is cut off) + shield (survives repeated re-delivery).Why not
anyio.CancelScope(shield=True)The ticket proposed it, but it does not defer a native
asyncio.Task.cancel()— only anyio-scoped cancellation — verified empirically under bothasyncio.run()andanyio.run(). Starlette (ASGI) and the deriver's uvloop both cancel natively, so the anyio shield would be inert exactly where we need it. The stdlib approach works in both runtimes and adds no dependency.Testing
tests/test_dependencies.py— 22 tests pass against a live Postgres (ruffclean,basedpyright0 errors):CancelledError; on cancel re-delivered during cleanup; under a cancel storm; broken-connection →invalidate(); DI-generator teardown.request_contextreset even on the cancelled path (no contextvar leak in reused deriver tasks).idle in transactionbackend (asserted viapg_stat_activity).tests/test_db_resilience.py— 11 pass, no regression.Review pass (addressed)
An adversarial correctness review + a ticket-compliance audit were run. Fixes applied:
_CLEANUP_TIMEOUT_SECONDS.asyncio.wait()call, so it re-armed on every re-delivered cancel and never elapsed — dead exactly on the anyio/BaseHTTPMiddlewarepath it was written for. Now one deadline is computed before the loop. The storm test also re-cancels at anyio'scall_sooncadence; sparse cancels leave quiet ticks that let a per-wait timeout re-arm, which is why the old test passed either way.finallyso a cancelled cleanup can't leak the task-scoped contextvar.fut.cancelled()/fut.exception()so loop shutdown can't raise from the tail.Still to verify before merge
chaos_http.pywritten, needs a running server). The committed real-DB test covers the wire level; this adds the full request path.Required alongside this PR (not code — do these regardless)
The process-death class above is only reclaimed server-side, so the first item is a companion to this PR rather than optional follow-up:
ALTER ROLE postgres SET idle_in_transaction_session_timeout = '2min';— one line, covers every class including the ones no application fix can reach.pg_terminate_backendof stuck backends).idle_in_txnsampler inlock_monitoring.sql(count > 10 or max age > 10min).Follow-up (not in this PR)
tcp_user_timeouton the psycopgconnect_argsindb.py, so a black-holed socket errors out (→invalidate()) instead of relying on the abandon-timeout. Shared infra; out of this ticket's scope.Part of DEV-1861 — leaving the ticket open until the server-side timeout is applied and the 24h prod watch confirms idle-in-transaction ~0.
🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Tests