diff --git a/app/db/session.py b/app/db/session.py index d7d6ae7f9d..a5caced988 100644 --- a/app/db/session.py +++ b/app/db/session.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import inspect import logging import os import sqlite3 @@ -8,12 +9,13 @@ from contextlib import asynccontextmanager from enum import StrEnum from pathlib import Path -from typing import TYPE_CHECKING, AsyncIterator, Awaitable, Callable, Protocol, TypeVar +from typing import TYPE_CHECKING, Any, AsyncIterator, Awaitable, Callable, Protocol, TypeVar import anyio from anyio import to_thread from sqlalchemy import event, text -from sqlalchemy.engine import Engine +from sqlalchemy import util as sqlalchemy_util +from sqlalchemy.engine import Connection, Engine from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.pool import NullPool @@ -55,6 +57,26 @@ "begin exclusive", ) _SQLITE_WATCHDOG_STATEMENT_PREVIEW_CHARS = 300 +# Hard deadline for the shielded rollback/close teardown on SQLite (part 2 of +# the issue #1682 plan). A teardown wedged behind a stuck aiosqlite worker +# keeps holding the single writer slot, so it must be reclaimed well before +# other writers exhaust their busy timeout and surface "database is locked"; +# one-sixth of the busy timeout (5s) matches the leader-gated shielded-drain +# grace. Abandoning the wedged await alone would NOT release the lock — the +# aiosqlite worker thread still holds it — so on timeout the reclaim below +# interrupts the driver and invalidates the connection, disposing the worker. +_SQLITE_TEARDOWN_TIMEOUT_SECONDS = _SQLITE_BUSY_TIMEOUT_SECONDS / 6 +# Session.info marker set once a teardown step was abandoned as wedged: the +# session must never be driven by another coroutine again (the abandoned +# greenlet may still resume), and the deferred cleanup takes over. +_SQLITE_TEARDOWN_WEDGED_INFO_KEY = "sqlite_teardown_wedged" +# Abandoned wedged teardown tasks and the deferred bookkeeping closes they +# schedule on late completion, owned until completion so shutdown (close_db) +# drains them instead of closing the event loop over pending tasks. The +# bookkeeping closes are bounded by _SQLITE_TEARDOWN_TIMEOUT_SECONDS; an +# abandoned teardown may outlive its reclaim (the interrupt is best-effort), +# so the close_db drain is explicitly bounded as a whole. +_wedged_teardown_cleanup_tasks: set[asyncio.Task[Any]] = set() # PostgreSQL pool checkout timeout and connection recycle window. Fixed # application constants (issue #1340): recycle keeps pooled connections @@ -350,20 +372,320 @@ async def _shielded(awaitable: Awaitable[object]) -> None: raise +async def _shielded_bounded(awaitable: Awaitable[object], timeout: float) -> asyncio.Task[object] | None: + """Shield ``awaitable`` from the caller's cancellation, waiting at most ``timeout``. + + Returns ``None`` when the awaitable finished inside the bound (re-raising + its exception like ``_shielded``); returns the still-running task when the + deadline passed — the caller must treat the underlying connection as + wedged and reclaim it, because the abandoned await does not release + anything the aiosqlite worker thread holds (issue #1682). + """ + task = asyncio.ensure_future(awaitable) + waiter = asyncio.ensure_future(asyncio.wait({task}, timeout=timeout)) + while not waiter.done(): + try: + await asyncio.shield(waiter) + except asyncio.CancelledError: + # Teardown runs in ``finally`` blocks: the bound, not the caller's + # cancellation, decides abandonment. ``asyncio.wait`` cannot + # outlive its timeout, so this drain stays bounded. + continue + if task.done(): + task.result() + return None + return task + + +def _sqlite_uri_mode_active(url: Any) -> bool: + """Whether the pysqlite/aiosqlite dialect will connect in URI mode. + + Mirrors the dialect's ``create_connect_args``: URI mode is enabled only + when the URL query carries a ``uri`` value that coerces to true. + """ + query = getattr(url, "query", None) + if query is None: + return False + try: + value = query.get("uri") + except Exception: + return False + values = value if isinstance(value, (tuple, list)) else (value,) + for item in values: + if item is None: + continue + try: + if sqlalchemy_util.asbool(str(item)): + return True + except Exception: + # An unrecognized value would fail at connect time anyway; + # classify conservatively as file-backed (bounded). + continue + return False + + +def _session_teardown_bound_seconds(session: AsyncSession) -> float | None: + """Teardown deadline for this session, or None for the unbounded path. + + Only file-backed SQLite gets a bound: its single writer slot turns a + wedged teardown into a database-wide write stall (issue #1682). + PostgreSQL teardown semantics are deliberately untouched, and in-memory + SQLite shares one StaticPool connection with the whole process — the + reclaim's invalidation would destroy the entire database (the + database-backends spec requires preserving shared in-memory state), and + with a single shared connection there is no cross-connection writer + contention to starve in the first place. + """ + try: + bind = session.get_bind() + except Exception: + return None + if getattr(getattr(bind, "dialect", None), "name", None) != "sqlite": + return None + url = getattr(bind, "url", None) + if url is not None: + database = getattr(url, "database", None) + if not database: + return None + database_text = str(database) + if database_text == ":memory:": + return None + # SQLite URI forms (``sqlite:///file:name?mode=memory&cache=shared&uri=true``) + # are in-memory only when the pysqlite/aiosqlite dialect actually + # passes the database string to the driver as a URI, which it does + # only when the URL query carries a truthy ``uri`` — and SQLite itself + # parses a filename as a URI only when it starts with ``file:``. + # Without ``uri=true``, ``file:name?mode=memory`` is a *file-backed* + # database whose filename literally contains those characters, so it + # must keep the bounded teardown. + if _sqlite_uri_mode_active(url) and database_text.startswith("file:"): + # ``mode=memory`` normally rides the parsed URL's query; it only + # appears inside ``url.database`` when the URL escaped the query + # into the database portion. + if ":memory:" in database_text or "mode=memory" in database_text: + return None + query = getattr(url, "query", None) + if query is not None: + try: + mode = query.get("mode") + except Exception: + mode = None + modes = mode if isinstance(mode, (tuple, list)) else (mode,) + if any(str(value).lower() == "memory" for value in modes if value is not None): + return None + return _SQLITE_TEARDOWN_TIMEOUT_SECONDS + + +def _session_is_teardown_wedged(session: AsyncSession) -> bool: + try: + return bool(session.info.get(_SQLITE_TEARDOWN_WEDGED_INFO_KEY)) + except Exception: + return False + + +def _session_sync_connections(session: AsyncSession) -> tuple[Connection, ...]: + """Best-effort snapshot of the sync Connections held by the session's transaction. + + Captured before a teardown attempt so a wedged rollback can be attributed + and its connection reclaimed. Diagnostics only — never raises. + """ + try: + transaction = session.sync_session.get_transaction() + if transaction is None: + return () + connections = getattr(transaction, "_connections", None) + if not isinstance(connections, dict): + return () + # The transaction tracks each Connection under two keys (the + # Connection itself and its Engine); deduplicate by identity. + unique: dict[int, Connection] = {} + for value in connections.values(): + if isinstance(value, tuple) and value and isinstance(value[0], Connection): + unique[id(value[0])] = value[0] + return tuple(unique.values()) + except Exception: + return () + + +def _sqlite_watchdog_identifiers(connection: Connection) -> str: + """Render the long-write watchdog's identifiers for the wedged connection. + + Invalidation prevents the connection from ever reaching the watchdog's + deferred report (next begin / pool checkin), so the reclaim log carries + the same attribution instead. + """ + try: + info = connection.info + started_at = info.get("sqlite_write_started_at") + first_statement = info.get("sqlite_first_write_statement") + last_statement = info.get("sqlite_last_write_statement") + task_name = info.get("sqlite_write_task") + if started_at is None: + # The watchdog's commit/rollback listener already moved the + # identifiers into the deferred report — the wedge happened inside + # the transaction-ending call itself, exactly the issue #1682 + # shape. + pending = info.get("sqlite_write_pending_report") + if isinstance(pending, tuple) and len(pending) == 5: + started_at, _, first_statement, last_statement, task_name = pending + held = f"{time.monotonic() - started_at:.1f}" if isinstance(started_at, float) else "unknown" + return ( + f"write_held_seconds={held} write_task={task_name!r} " + f"first_statement={first_statement!r} last_statement={last_statement!r}" + ) + except Exception: + return "write_held_seconds=unknown" + + +async def _reclaim_wedged_sqlite_session( + session: AsyncSession, + abandoned: asyncio.Task[object], + connections: tuple[Connection, ...], + *, + phase: str, +) -> None: + """Release what a wedged SQLite teardown still holds and fence the session. + + Abandoning the wedged rollback/close is not enough: the aiosqlite worker + thread keeps holding the write lock (issue #1682). Interrupting the driver + aborts the C-level call the worker is stuck in, and invalidating the + connection terminates it at the pool — aiosqlite's ``stop()`` queues a + hard close of the underlying ``sqlite3`` connection, which releases the + writer slot and disposes the worker thread — so leader election and every + other writer recover instead of stalling behind the wedge. The invalidated + connection can never be handed out again. + """ + try: + session.info[_SQLITE_TEARDOWN_WEDGED_INFO_KEY] = True + except Exception: + logger.exception("Failed to fence a wedged SQLite session during teardown reclaim") + # Own the abandoned teardown before this coroutine's first await: if + # close_db runs concurrently with the reclaim, it must already see the + # pending task in the registry instead of returning while the rollback is + # still pending. The completion callbacks are attached only after the + # connection is invalidated below, so the deferred bookkeeping close can + # never touch a live connection. + _wedged_teardown_cleanup_tasks.add(abandoned) + for connection in connections: + logger.warning( + "sqlite_wedged_teardown phase=%s bound_seconds=%.1f %s — interrupting and invalidating the " + "connection so the writer slot is released instead of stalling every writer (issue #1682)", + phase, + _SQLITE_TEARDOWN_TIMEOUT_SECONDS, + _sqlite_watchdog_identifiers(connection), + ) + try: + driver = connection.connection.driver_connection + if driver is not None: + # aiosqlite's ``interrupt`` runs sqlite3_interrupt inline on + # this task — it never enters the (wedged) worker queue. In + # the pinned aiosqlite (0.22.x) it is a coroutine function; + # await the result only when it is awaitable so a driver that + # makes ``interrupt`` synchronous keeps working. + result = driver.interrupt() + if inspect.isawaitable(result): + await result + except Exception: + logger.warning( + "Interrupting a wedged SQLite connection failed — the invalidation below still " + "reclaims the writer slot, but the stuck statement may run to completion first", + exc_info=True, + ) + try: + connection.invalidate() + except Exception: + logger.debug("Invalidating a wedged SQLite connection failed", exc_info=True) + if not connections: + logger.warning( + "sqlite_wedged_teardown phase=%s bound_seconds=%.1f — no held connection to reclaim; " + "abandoning the wedged %s (issue #1682)", + phase, + _SQLITE_TEARDOWN_TIMEOUT_SECONDS, + phase, + ) + # The abandoned teardown is owned until completion (registered above, so + # close_db drains it and shutdown waits for — or boundedly abandons — the + # reclaimed rollback/close instead of returning while it is still + # pending). The discard callback is registered first so that when the task + # completes during the drain, deregistration happens before + # _finish_abandoned_teardown registers the follow-up bookkeeping close. + abandoned.add_done_callback(_wedged_teardown_cleanup_tasks.discard) + abandoned.add_done_callback(lambda task: _finish_abandoned_teardown(session, task, phase=phase)) + + +def _finish_abandoned_teardown(session: AsyncSession, task: asyncio.Task[object], *, phase: str) -> None: + if not task.cancelled(): + # The wedged teardown resuming into an interrupted/invalidated + # connection is expected to error; consume it so the abandoned task + # never logs "exception was never retrieved". + task.exception() + logger.info("Wedged SQLite teardown finished late phase=%s", phase) + if phase != "rollback": + return + + # The session was abandoned before ``close`` ran. Now that no other + # coroutine can be driving it, close it for bookkeeping — the connection + # is already invalidated, so this cannot touch the database. + async def _close_late() -> None: + try: + await asyncio.wait_for(session.close(), timeout=_SQLITE_TEARDOWN_TIMEOUT_SECONDS) + except BaseException: + logger.debug("Late close of a wedged SQLite session failed", exc_info=True) + + try: + cleanup_task = asyncio.get_running_loop().create_task(_close_late()) + except RuntimeError: + # Event loop already gone (shutdown); the invalidated connection was + # closed at the pool, nothing is leaked. + return + # Own the task until completion: close_db drains it so shutdown cannot + # skip the promised bookkeeping close or leave a pending-task warning. + _wedged_teardown_cleanup_tasks.add(cleanup_task) + cleanup_task.add_done_callback(_wedged_teardown_cleanup_tasks.discard) + + async def _safe_rollback(session: AsyncSession) -> None: if not session.in_transaction(): return + if _session_is_teardown_wedged(session): + # A previous bounded teardown abandoned a wedged rollback; the + # abandoned greenlet may still resume, so never drive this session + # concurrently. The reclaim already released the connection. + return + bound = _session_teardown_bound_seconds(session) + if bound is None: + try: + await _shielded(session.rollback()) + except BaseException: + return + return + held_connections = _session_sync_connections(session) try: - await _shielded(session.rollback()) + abandoned = await _shielded_bounded(session.rollback(), bound) except BaseException: return + if abandoned is not None: + await _reclaim_wedged_sqlite_session(session, abandoned, held_connections, phase="rollback") async def _safe_close(session: AsyncSession) -> None: + if _session_is_teardown_wedged(session): + # Deferred cleanup owns the session now; see _finish_abandoned_teardown. + return + bound = _session_teardown_bound_seconds(session) + if bound is None: + try: + await _shielded(session.close()) + except BaseException: + return + return + held_connections = _session_sync_connections(session) try: - await _shielded(session.close()) + abandoned = await _shielded_bounded(session.close(), bound) except BaseException: return + if abandoned is not None: + await _reclaim_wedged_sqlite_session(session, abandoned, held_connections, phase="close") async def close_session(session: AsyncSession) -> None: @@ -620,6 +942,30 @@ async def init_db() -> None: async def close_db() -> None: + if _wedged_teardown_cleanup_tasks: + # Abandoned wedged teardowns plus their deferred bookkeeping closes. + # Drain until the registry is stable — an abandoned teardown that + # completes during the drain schedules its bookkeeping close only + # after any one-time snapshot — and bound the whole drain so a + # teardown still wedged despite the reclaim (the interrupt is + # best-effort) cannot wedge shutdown too: one deadline covers the + # abandoned teardown and the bounded close it chains. + loop = asyncio.get_running_loop() + deadline = loop.time() + 2 * _SQLITE_TEARDOWN_TIMEOUT_SECONDS + while _wedged_teardown_cleanup_tasks: + remaining = deadline - loop.time() + if remaining <= 0: + logger.warning( + "close_db abandoned %d still-pending wedged-teardown task(s) after the bounded " + "drain; their connections were already reclaimed (issue #1682)", + len(_wedged_teardown_cleanup_tasks), + ) + break + await asyncio.wait(tuple(_wedged_teardown_cleanup_tasks), timeout=remaining) + # Completion callbacks (deregistration and scheduling of the + # deferred bookkeeping close) run via call_soon; yield once so + # the registry reflects them before the next stability check. + await asyncio.sleep(0) await engine.dispose() if _background_engine is not None: await _background_engine.dispose() diff --git a/openspec/changes/bound-sqlite-wedged-teardown/proposal.md b/openspec/changes/bound-sqlite-wedged-teardown/proposal.md new file mode 100644 index 0000000000..f6a1e4dfd6 --- /dev/null +++ b/openspec/changes/bound-sqlite-wedged-teardown/proposal.md @@ -0,0 +1,21 @@ +## Why + +Issue #1682, part 2 of the plan. Session teardown shields rollback/close unboundedly (`app/db/session.py`), so a wedged teardown pins SQLite's single writer slot with nothing to reclaim it: every writer — including the `scheduler_leader` INSERT that would re-establish leadership — surfaces `database is locked` until the wedge spontaneously resolves (~17 minutes in the report). Part 1 (`report-sqlite-long-write-holders`) made the holder attributable; the teardown itself must now be bounded. Crucially, abandoning the wedged await alone releases nothing — the aiosqlite worker thread still holds the lock — so the bound must come with reclaiming the connection. + +## What Changes + +- The shielded rollback/close teardown gets a hard deadline on file-backed SQLite (one-sixth of the busy timeout, 5s — reclaimed well before other writers exhaust their 30s busy timeout). PostgreSQL teardown semantics are untouched, and in-memory SQLite keeps the unbounded path: its one shared StaticPool connection is the whole database (invalidation would destroy it) and cannot starve other writers. +- A teardown that misses the deadline is reclaimed, not merely abandoned: the driver connection is interrupted (aborting the C-level call the aiosqlite worker is stuck in) and the connection is invalidated — terminating it at the pool disposes the worker and hard-closes the underlying `sqlite3` connection, which releases the writer slot and guarantees the connection is never handed out again. +- The reclaim report carries part 1's watchdog identifiers (held duration, owning task, first/last write statements), including when the watchdog had already deferred them into its pending report because the wedge is inside the transaction-ending call itself — invalidation would otherwise suppress that deferred report. +- A wedged session is fenced: later teardown attempts return immediately instead of driving the session concurrently with the abandoned work, and once the abandoned teardown finishes late the session is closed for bookkeeping — a deferred task owned until completion and drained at `close_db`, never fire-and-forget. +- No new settings: the deadline derives from the existing busy timeout. + +## Capabilities + +### New Capabilities + +None. + +### Modified Capabilities + +- `database-backends`: a wedged SQLite session teardown is bounded and its connection reclaimed so the writer slot is released. diff --git a/openspec/changes/bound-sqlite-wedged-teardown/specs/database-backends/spec.md b/openspec/changes/bound-sqlite-wedged-teardown/specs/database-backends/spec.md new file mode 100644 index 0000000000..b234160df7 --- /dev/null +++ b/openspec/changes/bound-sqlite-wedged-teardown/specs/database-backends/spec.md @@ -0,0 +1,40 @@ +## ADDED Requirements + +### Requirement: Wedged SQLite session teardown is bounded and reclaimed + +Session teardown (rollback and close) on file-backed SQLite MUST complete within a hard deadline derived from the busy timeout while remaining shielded from the caller's cancellation. A teardown that misses the deadline MUST NOT be merely abandoned — the aiosqlite worker thread would keep holding the writer slot — it MUST be reclaimed: the driver connection is interrupted to abort the call the worker is stuck in, and the connection is invalidated so the worker is disposed, the underlying `sqlite3` connection is hard-closed releasing the writer slot, and the connection can never be handed out again. The reclaim MUST be reported with the long-write watchdog's identifiers where available (held duration, owning task, first and last write statements), including identifiers the watchdog already deferred into its pending report. A session whose teardown was abandoned MUST be fenced from further teardown attempts, the abandoned work finishing late MUST NOT surface unretrieved errors, and the deferred bookkeeping close MUST be owned until completion (drained at database shutdown, never fire-and-forget). PostgreSQL teardown semantics MUST remain unchanged, and in-memory SQLite — whose single shared connection is the entire database and cannot starve other writers — MUST keep the unbounded teardown and never be reclaimed. + +#### Scenario: A wedged rollback no longer starves every other writer + +- **GIVEN** a session holding an open SQLite write transaction whose rollback wedges during teardown +- **WHEN** the teardown deadline passes +- **THEN** teardown returns, the connection is interrupted and invalidated, and another writer — such as the leader-election `scheduler_leader` INSERT — acquires the writer slot immediately instead of surfacing `database is locked` + +#### Scenario: The reclaim is attributed with the watchdog's identifiers + +- **GIVEN** a wedged teardown whose transaction ran write statements tracked by the long-write watchdog +- **WHEN** the connection is reclaimed +- **THEN** the report names the held duration, owning task, and first/last write statements, even though invalidation prevents the watchdog's own deferred report from firing + +#### Scenario: A wedged session cannot be driven concurrently + +- **GIVEN** a session whose teardown was abandoned as wedged +- **WHEN** teardown is attempted again +- **THEN** it returns immediately, and the session is closed for bookkeeping only after the abandoned teardown finishes late + +#### Scenario: PostgreSQL teardown is untouched + +- **GIVEN** a session bound to a non-SQLite dialect +- **WHEN** its rollback or close outlives the SQLite deadline +- **THEN** the teardown still awaits completion unboundedly and no connection is reclaimed + +#### Scenario: The shared in-memory SQLite connection is never reclaimed + +- **GIVEN** a session bound to an in-memory SQLite database, whose one shared connection is the entire database +- **WHEN** its teardown outlives the deadline +- **THEN** the teardown still awaits completion unboundedly and the connection is never invalidated, preserving schema and data for later sessions + +#### Scenario: The bound never abandons healthy teardown + +- **WHEN** rollback and close complete within the deadline +- **THEN** teardown behaves exactly as before, including re-raising the completed call's exception to the existing swallow points diff --git a/openspec/changes/bound-sqlite-wedged-teardown/tasks.md b/openspec/changes/bound-sqlite-wedged-teardown/tasks.md new file mode 100644 index 0000000000..0ccedd5515 --- /dev/null +++ b/openspec/changes/bound-sqlite-wedged-teardown/tasks.md @@ -0,0 +1,14 @@ +## 1. Bounded teardown and reclaim + +- [x] 1.1 Bound the shielded rollback/close teardown for file-backed SQLite sessions with a deadline derived from the busy timeout, preserving the shield against caller cancellation; PostgreSQL and in-memory SQLite (one shared StaticPool connection is the whole database — reclaim would destroy it, and it cannot starve other writers) keep the unbounded path +- [x] 1.2 On a missed deadline, interrupt the driver connection and invalidate it so the aiosqlite worker is disposed, the writer slot is released, and the connection can never be handed out again; report the reclaim with the long-write watchdog's identifiers (including ones already deferred into its pending report) +- [x] 1.3 Fence the wedged session against further teardown, consume the abandoned task's late failure, and close the session for bookkeeping once the abandoned teardown finishes; own the deferred close until completion and drain it at close_db so shutdown cannot abandon it + +## 2. Tests + +- [x] 2.1 A wedged sqlite rollback: close_session returns within the bound (fails on pre-fix unbounded teardown), the driver is interrupted, the connection invalidated, the reclaim log carries the watchdog identifiers, and an independent writer succeeds immediately while the wedge is still pending +- [x] 2.2 A wedged session is fenced from further teardown; the abandoned teardown finishing late is observed and followed by the bookkeeping close +- [x] 2.3 The bounded shield completes fast work, abandons at the deadline without cancelling, and absorbs caller cancellation like the unbounded shield +- [x] 2.4 Non-sqlite sessions keep the unbounded teardown: a slow rollback/close beyond the sqlite bound still runs to completion and is never reclaimed +- [x] 2.5 A wedged close without a transaction is bounded and fenced too +- [x] 2.6 In-memory SQLite keeps the unbounded teardown: a slow teardown is never reclaimed, the shared connection is never invalidated, and schema/data survive for later sessions diff --git a/tests/unit/test_db_session.py b/tests/unit/test_db_session.py index e75bf781db..895520b70b 100644 --- a/tests/unit/test_db_session.py +++ b/tests/unit/test_db_session.py @@ -13,6 +13,7 @@ import pytest from sqlalchemy import event as sa_event from sqlalchemy import text as sa_text +from sqlalchemy.engine import make_url from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.pool import NullPool @@ -1148,3 +1149,516 @@ def failing_commit(self) -> None: assert "outcome=commit " not in records[0].getMessage() finally: await engine.dispose() + + +@pytest.mark.asyncio +async def test_shielded_bounded_returns_none_when_the_awaitable_finishes_in_time() -> None: + async def _fast() -> str: + return "done" + + assert await session_module._shielded_bounded(_fast(), 1.0) is None + + async def _boom() -> None: + raise RuntimeError("teardown failed") + + with pytest.raises(RuntimeError, match="teardown failed"): + await session_module._shielded_bounded(_boom(), 1.0) + + +@pytest.mark.asyncio +async def test_shielded_bounded_abandons_a_wedged_awaitable_at_the_deadline() -> None: + release = asyncio.Event() + + async def _wedged() -> None: + await release.wait() + + abandoned = await session_module._shielded_bounded(_wedged(), 0.05) + assert abandoned is not None + assert not abandoned.done(), "the wedged awaitable must be left running, not cancelled" + release.set() + await abandoned + + +@pytest.mark.asyncio +async def test_shielded_bounded_absorbs_caller_cancellation_like_the_unbounded_shield() -> None: + """Teardown runs in ``finally`` blocks: the bound, not the caller's + cancellation, must decide abandonment (matching ``_shielded`` + the + swallow in ``_safe_rollback``/``_safe_close``).""" + started = asyncio.Event() + release = asyncio.Event() + finished: list[bool] = [] + + async def _work() -> None: + started.set() + await release.wait() + finished.append(True) + + async def _caller() -> asyncio.Task[object] | None: + return await session_module._shielded_bounded(_work(), 5.0) + + caller = asyncio.ensure_future(_caller()) + await started.wait() + caller.cancel() + await asyncio.sleep(0.05) + assert not caller.done(), "cancellation must not abandon the shielded teardown" + release.set() + assert await caller is None + assert finished, "the shielded work must run to completion despite the cancellation" + + +@pytest.mark.asyncio +async def test_close_session_reclaims_a_wedged_sqlite_rollback_so_other_writers_recover( + tmp_path, monkeypatch, caplog +) -> None: + """Issue #1682 part 2: a wedged rollback used to be awaited forever while + the aiosqlite worker kept the single writer slot — a self-sustaining + 'database is locked' stall that starved leader election itself. The + teardown must be bounded, and the bound alone is not enough: the wedged + connection must be interrupted and invalidated so the writer slot is + actually released and the connection is never handed out again.""" + monkeypatch.setattr(session_module, "_SQLITE_TEARDOWN_TIMEOUT_SECONDS", 0.2) + db_path = tmp_path / "wedged-rollback.db" + engine = create_async_engine( + f"sqlite+aiosqlite:///{db_path}", + poolclass=NullPool, + connect_args={"timeout": 5.0}, + ) + session_module._configure_sqlite_engine(engine.sync_engine, enable_wal=True) + # An independent writer with a short busy timeout: if the reclaim fails to + # release the writer slot, its INSERT surfaces 'database is locked' fast. + other_writer = create_async_engine( + f"sqlite+aiosqlite:///{db_path}", + poolclass=NullPool, + connect_args={"timeout": 1.0}, + ) + release_wedge = asyncio.Event() + try: + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + caplog.clear() + + factory = async_sessionmaker(engine, expire_on_commit=False) + session = factory() + # Take the writer slot with an uncommitted write. + await session.execute(sa_text("DELETE FROM accounts")) + + held = session_module._session_sync_connections(session) + assert held, "the open write transaction must expose its sync connection" + sync_connection = held[0] + driver = sync_connection.connection.driver_connection + assert driver is not None + + # Wedge this connection's rollback (a stuck aiosqlite worker queues + # the teardown behind itself exactly like this) and spy on interrupt. + original_rollback = driver.rollback + interrupted = asyncio.Event() + original_interrupt = driver.interrupt + + async def _wedged_rollback() -> None: + await release_wedge.wait() + await original_rollback() + + # Delegate without changing the installed driver's shape: the reclaim + # awaits ``interrupt()``'s result only when it is awaitable, so the + # spy hands back exactly what the real aiosqlite method returns and + # the production awaitable-handling is exercised against the installed + # contract (a coroutine in the pinned aiosqlite) instead of a stand-in. + def _spying_interrupt() -> object: + interrupted.set() + return original_interrupt() + + driver.rollback = _wedged_rollback + driver.interrupt = _spying_interrupt + + with caplog.at_level(logging.INFO, logger=session_module.__name__): + close_task = asyncio.ensure_future(session_module.close_session(session)) + done, _ = await asyncio.wait({close_task}, timeout=2.0) + # RED on the pre-fix teardown: the shielded rollback was awaited + # unboundedly, so close_session never returned. + assert done, "close_session must be bounded when the sqlite rollback wedges" + + assert interrupted.is_set(), "the wedged driver must be interrupted to unstick its worker" + assert sync_connection.invalidated, "the wedged connection must be invalidated, never reused" + assert session.info.get(session_module._SQLITE_TEARDOWN_WEDGED_INFO_KEY) is True + + reclaim_logs = [ + record + for record in caplog.records + if record.levelno == logging.WARNING and "sqlite_wedged_teardown" in record.getMessage() + ] + assert reclaim_logs, "the reclaim must be reported with the watchdog's identifiers" + message = reclaim_logs[0].getMessage() + assert "phase=rollback" in message + assert "DELETE FROM accounts" in message, "part 1 watchdog identifiers must attribute the holder" + + # The stall must not be self-sustaining: with the wedged rollback + # still pending, another writer takes the slot immediately. + async with other_writer.begin() as writer: + await writer.execute(sa_text("DELETE FROM accounts")) + + # A wedged session is fenced: further teardown returns immediately + # instead of driving the session concurrently with the abandoned + # greenlet. + await asyncio.wait_for(session_module.close_session(session), timeout=1.0) + + # Late completion: once the wedge resolves, the abandoned teardown + # finishes and the session is closed for bookkeeping. + release_wedge.set() + for _ in range(100): + if any("finished late" in record.getMessage() for record in caplog.records): + break + await asyncio.sleep(0.02) + assert any("finished late" in record.getMessage() for record in caplog.records), ( + "the abandoned teardown must be observed finishing late" + ) + # The deferred bookkeeping close is owned until completion (drained + # by close_db on shutdown), never fire-and-forget. + pending_cleanup = tuple(session_module._wedged_teardown_cleanup_tasks) + if pending_cleanup: + await asyncio.wait_for(asyncio.gather(*pending_cleanup, return_exceptions=True), timeout=2.0) + assert not session_module._wedged_teardown_cleanup_tasks, ( + "the deferred close must deregister itself once it completes" + ) + finally: + release_wedge.set() + await asyncio.sleep(0.05) + await engine.dispose() + await other_writer.dispose() + + +@pytest.mark.asyncio +async def test_close_session_keeps_the_unbounded_shield_for_non_sqlite_sessions(monkeypatch) -> None: + """PostgreSQL teardown semantics are untouched: a slow rollback/close far + beyond the SQLite bound is still awaited to completion, never reclaimed.""" + monkeypatch.setattr(session_module, "_SQLITE_TEARDOWN_TIMEOUT_SECONDS", 0.01) + + class _FakeDialect: + name = "postgresql" + + class _FakeBind: + dialect = _FakeDialect() + + class _FakeSession: + def __init__(self) -> None: + self.info: dict[str, object] = {} + self.rolled_back = False + self.closed = False + + def get_bind(self) -> _FakeBind: + return _FakeBind() + + def in_transaction(self) -> bool: + return not self.rolled_back + + async def rollback(self) -> None: + await asyncio.sleep(0.1) + self.rolled_back = True + + async def close(self) -> None: + await asyncio.sleep(0.1) + self.closed = True + + fake = _FakeSession() + await session_module.close_session(cast(session_module.AsyncSession, fake)) + + assert fake.rolled_back, "the slow PostgreSQL rollback must be awaited to completion" + assert fake.closed, "the slow PostgreSQL close must be awaited to completion" + assert session_module._SQLITE_TEARDOWN_WEDGED_INFO_KEY not in fake.info + + +@pytest.mark.asyncio +async def test_close_session_bounds_a_wedged_sqlite_close_without_a_transaction(monkeypatch, caplog) -> None: + """The close step can wedge on its own (connection release goes through + the same aiosqlite worker); it must be bounded and fenced too.""" + monkeypatch.setattr(session_module, "_SQLITE_TEARDOWN_TIMEOUT_SECONDS", 0.05) + release = asyncio.Event() + + class _FakeDialect: + name = "sqlite" + + class _FakeUrl: + database = "/tmp/wedged-close.db" + query: dict[str, str] = {} + + class _FakeBind: + dialect = _FakeDialect() + url = _FakeUrl() + + class _FakeSyncSession: + def get_transaction(self) -> None: + return None + + class _FakeSession: + def __init__(self) -> None: + self.info: dict[str, object] = {} + self.sync_session = _FakeSyncSession() + + def get_bind(self) -> _FakeBind: + return _FakeBind() + + def in_transaction(self) -> bool: + return False + + async def close(self) -> None: + await release.wait() + + fake = _FakeSession() + try: + with caplog.at_level(logging.WARNING, logger=session_module.__name__): + await asyncio.wait_for(session_module.close_session(cast(session_module.AsyncSession, fake)), timeout=2.0) + + assert fake.info.get(session_module._SQLITE_TEARDOWN_WEDGED_INFO_KEY) is True + messages = [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "sqlite_wedged_teardown" in record.getMessage() + ] + assert messages + assert "phase=close" in messages[0] + finally: + release.set() + await asyncio.sleep(0.05) + + +@pytest.mark.asyncio +async def test_close_session_never_reclaims_the_shared_in_memory_sqlite_connection(monkeypatch) -> None: + """In-memory SQLite shares one StaticPool connection with the whole + process: invalidating it would destroy the entire database (the + database-backends spec preserves shared in-memory state), and a single + shared connection cannot starve other writers. The teardown must keep the + unbounded shield there.""" + monkeypatch.setattr(session_module, "_SQLITE_TEARDOWN_TIMEOUT_SECONDS", 0.01) + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + try: + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + factory = async_sessionmaker(engine, expire_on_commit=False) + + session = factory() + await session.execute(sa_text("DELETE FROM accounts")) + assert session_module._session_teardown_bound_seconds(session) is None + + held = session_module._session_sync_connections(session) + assert held + driver = held[0].connection.driver_connection + assert driver is not None + original_rollback = driver.rollback + + async def _slow_rollback() -> None: + await asyncio.sleep(0.1) + await original_rollback() + + driver.rollback = _slow_rollback + await session_module.close_session(session) + driver.rollback = original_rollback + + assert not held[0].invalidated, "the shared in-memory connection must never be invalidated" + assert session_module._SQLITE_TEARDOWN_WEDGED_INFO_KEY not in session.info + + # The database survives: schema and connection are intact. + verify = factory() + (await verify.execute(sa_text("SELECT count(*) FROM accounts"))).scalar_one() + await session_module.close_session(verify) + finally: + await engine.dispose() + + +@pytest.mark.parametrize( + "url_text", + [ + "sqlite+aiosqlite:///:memory:", + "sqlite+aiosqlite://", + "sqlite+aiosqlite:///file:shared?mode=memory&cache=shared&uri=true", + "sqlite+aiosqlite:///file::memory:?cache=shared&uri=true", + ], +) +def test_session_teardown_bound_skips_every_in_memory_sqlite_url_form(url_text: str) -> None: + """Every in-memory SQLite URL form must keep the unbounded teardown: the + SQLite URI forms carry ``mode=memory`` in the parsed URL's query, not in + ``url.database``, and a shared in-memory database reclaimed by invalidation + would be destroyed for the whole process. URI forms count only with + ``uri=true`` — that is what makes the dialect pass the string as a URI.""" + + class _FakeDialect: + name = "sqlite" + + class _FakeBind: + dialect = _FakeDialect() + url = make_url(url_text) + + class _FakeSession: + def get_bind(self) -> _FakeBind: + return _FakeBind() + + fake = _FakeSession() + assert session_module._session_teardown_bound_seconds(cast(session_module.AsyncSession, fake)) is None + + +@pytest.mark.parametrize( + "url_text", + [ + "sqlite+aiosqlite:////data/store.db", + "sqlite+aiosqlite:///file:/data/store.db?uri=true", + # Without ``uri=true`` the dialect never enables SQLite URI mode: this + # connects to a file literally named ``file:shared`` and must keep the + # bounded teardown despite carrying ``mode=memory`` in the query. + "sqlite+aiosqlite:///file:shared?mode=memory&cache=shared", + ], +) +def test_session_teardown_bound_applies_to_file_backed_sqlite_url_forms(url_text: str) -> None: + """File-backed SQLite (plain path, ``file:`` URI without ``mode=memory``, + or a ``mode=memory`` query without ``uri=true``) is exactly the + wedge-prone single-writer case and must stay bounded.""" + + class _FakeDialect: + name = "sqlite" + + class _FakeBind: + dialect = _FakeDialect() + url = make_url(url_text) + + class _FakeSession: + def get_bind(self) -> _FakeBind: + return _FakeBind() + + fake = _FakeSession() + assert ( + session_module._session_teardown_bound_seconds(cast(session_module.AsyncSession, fake)) + == session_module._SQLITE_TEARDOWN_TIMEOUT_SECONDS + ) + + +@pytest.mark.asyncio +async def test_reclaim_interrupts_the_real_aiosqlite_driver_without_a_spy(tmp_path, caplog) -> None: + """The reclaim invokes the driver's real ``interrupt()`` and awaits the + result only when it is awaitable. Exercise the production path against the + installed aiosqlite with no stand-in, so a driver signature change + surfaces as a failure here instead of being swallowed by the reclaim's + broad except (the failure is logged, and asserted absent).""" + engine = create_async_engine( + f"sqlite+aiosqlite:///{tmp_path / 'interrupt-contract.db'}", + poolclass=NullPool, + ) + try: + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + factory = async_sessionmaker(engine, expire_on_commit=False) + session = factory() + await session.execute(sa_text("DELETE FROM accounts")) + held = session_module._session_sync_connections(session) + assert held, "the open write transaction must expose its sync connection" + + async def _already_finished_teardown() -> None: + return None + + abandoned = asyncio.ensure_future(_already_finished_teardown()) + with caplog.at_level(logging.DEBUG, logger=session_module.__name__): + await session_module._reclaim_wedged_sqlite_session(session, abandoned, held, phase="rollback") + + assert not any( + "Interrupting a wedged SQLite connection failed" in record.getMessage() for record in caplog.records + ), "the installed aiosqlite interrupt() contract must be handled without error" + assert held[0].invalidated, "the reclaim must still invalidate the connection" + + # Drain the bookkeeping the reclaim registered so no task outlives + # the test (mirrors the close_db drain). + for _ in range(100): + pending = tuple(session_module._wedged_teardown_cleanup_tasks) + if not pending: + break + await asyncio.wait_for(asyncio.gather(*pending, return_exceptions=True), timeout=2.0) + await asyncio.sleep(0) + assert not session_module._wedged_teardown_cleanup_tasks + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_close_db_drains_a_pending_reclaimed_rollback_and_its_bookkeeping_close( + tmp_path, monkeypatch, caplog +) -> None: + """A rollback reclaimed as wedged can still be pending when close_db runs. + The abandoned task is registered in the teardown registry immediately, so + close_db must wait for it — and for the bookkeeping close it schedules only + after any one-time snapshot — instead of returning while the event loop + still has pending teardown tasks.""" + monkeypatch.setattr(session_module, "_SQLITE_TEARDOWN_TIMEOUT_SECONDS", 0.5) + db_path = tmp_path / "close-db-drain.db" + engine = create_async_engine( + f"sqlite+aiosqlite:///{db_path}", + poolclass=NullPool, + connect_args={"timeout": 5.0}, + ) + session_module._configure_sqlite_engine(engine.sync_engine, enable_wal=True) + release_wedge = asyncio.Event() + try: + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + + factory = async_sessionmaker(engine, expire_on_commit=False) + session = factory() + await session.execute(sa_text("DELETE FROM accounts")) + + held = session_module._session_sync_connections(session) + assert held + driver = held[0].connection.driver_connection + assert driver is not None + original_rollback = driver.rollback + + async def _wedged_rollback() -> None: + await release_wedge.wait() + await original_rollback() + + driver.rollback = _wedged_rollback + + with caplog.at_level(logging.INFO, logger=session_module.__name__): + await asyncio.wait_for(session_module.close_session(session), timeout=5.0) + abandoned_pending = [task for task in session_module._wedged_teardown_cleanup_tasks if not task.done()] + # RED pre-fix: the reclaim only registered the deferred bookkeeping + # close (which does not exist yet), never the abandoned rollback. + assert abandoned_pending, "the reclaimed rollback must be registered while still pending" + + async def _release_soon() -> None: + await asyncio.sleep(0.05) + release_wedge.set() + + releaser = asyncio.ensure_future(_release_soon()) + await asyncio.wait_for(session_module.close_db(), timeout=5.0) + # RED pre-fix: close_db saw an empty registry and returned + # immediately, before the wedge was even released. + assert release_wedge.is_set(), "close_db must drain the pending reclaimed rollback" + assert all(task.done() for task in abandoned_pending), ( + "close_db must wait for the abandoned rollback itself" + ) + assert not session_module._wedged_teardown_cleanup_tasks, ( + "close_db must also drain the bookkeeping close scheduled after its first snapshot" + ) + await releaser + + assert any("finished late" in record.getMessage() for record in caplog.records) + finally: + release_wedge.set() + await asyncio.sleep(0.05) + await engine.dispose() + + +@pytest.mark.asyncio +async def test_close_db_bounds_the_wedged_teardown_drain(monkeypatch, caplog) -> None: + """A teardown that stays wedged despite the reclaim (the interrupt is + best-effort) must not wedge shutdown too: the registry drain is explicitly + bounded and abandons whatever remains after the deadline.""" + monkeypatch.setattr(session_module, "_SQLITE_TEARDOWN_TIMEOUT_SECONDS", 0.05) + never = asyncio.Event() + stuck: asyncio.Task[bool] = asyncio.ensure_future(never.wait()) + session_module._wedged_teardown_cleanup_tasks.add(stuck) + try: + with caplog.at_level(logging.WARNING, logger=session_module.__name__): + await asyncio.wait_for(session_module.close_db(), timeout=2.0) + assert any("still-pending wedged-teardown" in record.getMessage() for record in caplog.records), ( + "the bounded drain must report what it abandoned" + ) + assert stuck in session_module._wedged_teardown_cleanup_tasks + finally: + session_module._wedged_teardown_cleanup_tasks.discard(stuck) + never.set() + await stuck