Shield pool cleanup against cancel and drain pending-close on release

Two cancel-hygiene gaps in the pool's lifecycle paths:

- acquire's broken-conn cleanup awaited conn.close() bare. An outer
  cancel landing during the close skipped it and left the user's
  transport open until GC, plus left _pool_released=False so a later
  close() did not short-circuit. Wrap the close in
  asyncio.shield + suppress(CancelledError) and move _pool_released
  into a finally so the flag flips even when close raises an
  unrecognised exception.

- _release set _pool_released=True before observing _pending_drain
  (scheduled by _invalidate's bounded wait_closed task on cancel
  mid-ROLLBACK). The next close() would short-circuit past the
  pending-drain await and the reader-task could outlive the
  connection. Drain the pending task in the finally before flipping
  _pool_released.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: antoineleclair

src/dqliteclient/pool.py

@@ -764,15 +764,34 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
                             "pool.acquire cleanup: _drain_idle failed",
                             exc_info=True,
                         )
+                    # Shield ``conn.close()`` so an outer cancel landing
+                    # mid-cleanup does not skip it — the user's
+                    # checked-out connection's transport would otherwise
+                    # stay open until GC. The ``_close_timeout`` bound
+                    # on ``wait_closed`` keeps the shielded await
+                    # bounded; ``contextlib.suppress(CancelledError)``
+                    # absorbs any nested cancel delivered after the
+                    # shield released so the original cancel still
+                    # propagates via the surrounding ``raise``.
+                    # ``_drain_idle()`` above stays bare: it protects
+                    # siblings, not the leaked conn, and shielding it
+                    # could turn outer cancel into an unbounded wait.
                     try:
-                        await conn.close()
-                    except (OSError, DqliteConnectionError):
-                        logger.debug(
-                            "pool.acquire cleanup: conn.close(%r) failed",
-                            getattr(conn, "_address", "?"),
-                            exc_info=True,
-                        )
-                    conn._pool_released = True
+                        try:
+                            with contextlib.suppress(asyncio.CancelledError):
+                                await asyncio.shield(conn.close())
+                        except (OSError, DqliteConnectionError):
+                            logger.debug(
+                                "pool.acquire cleanup: conn.close(%r) failed",
+                                getattr(conn, "_address", "?"),
+                                exc_info=True,
+                            )
+                    finally:
+                        # Always set ``_pool_released`` so a subsequent
+                        # close() short-circuits and the slot accounting
+                        # stays consistent — landed even if the shielded
+                        # close above raised an unrecognised exception.
+                        conn._pool_released = True
             finally:
                 if not returned_to_queue:
                     # ``asyncio.shield`` already prevents an outer cancel
@@ -941,6 +960,21 @@ async def _release(self, conn: DqliteConnection) -> None:
                 # takes the early-return path instead of running a
                 # redundant close against a protocol that's already
                 # None.
+                #
+                # Drain ``_pending_drain`` BEFORE setting
+                # ``_pool_released=True``: a cancel mid-ROLLBACK has
+                # ``_invalidate`` schedule a bounded ``wait_closed``
+                # drain task on the connection. ``close()`` would
+                # normally await that task at ``connection.py``'s
+                # pending-drain block, but the early-return on
+                # ``_pool_released=True`` we set below short-circuits
+                # past it. Snapshot and shield-await the drain here so
+                # the reader-task doesn't outlive the connection (no
+                # "Task was destroyed but it is pending" on shutdown).
+                pending = getattr(conn, "_pending_drain", None)
+                if pending is not None and not pending.done():
+                    with contextlib.suppress(BaseException):
+                        await asyncio.shield(pending)
                 conn._pool_released = True
                 with contextlib.suppress(asyncio.CancelledError):
                     await asyncio.shield(self._release_reservation())

tests/test_pool_acquire_close_shielded_on_cancel.py

@@ -0,0 +1,157 @@
+"""Pin: pool ``acquire``'s broken-conn cleanup shields ``conn.close()``.
+
+The exception-cleanup branch in ``acquire`` runs when the user's
+checked-out connection is broken (invalidated by execute/fetch error).
+It calls ``await conn.close()`` to release the transport. If an outer
+``asyncio.timeout`` or ``CancelScope`` fires during the cleanup, the
+unshielded close is interrupted and the user's transport leaks open
+until GC. The fix wraps the close in ``asyncio.shield`` and
+``contextlib.suppress(CancelledError)`` so the close completes
+(bounded by ``_close_timeout``) before the cancellation propagates;
+the inner ``finally`` then guarantees ``_pool_released = True``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import patch
+
+import pytest
+
+from dqliteclient import DqliteConnection
+from dqliteclient.exceptions import DqliteConnectionError
+from dqliteclient.pool import ConnectionPool
+
+
+async def _build_pool_with_breakable_conn() -> tuple[ConnectionPool, DqliteConnection]:
+    """Construct a pool whose ``_create_connection`` returns a conn
+    that LOOKS connected at acquire time but transitions to broken
+    once the user invalidates it (simulating leader flip mid-execute).
+    """
+    pool = ConnectionPool(["localhost:9001"], min_size=0, max_size=1)
+    conn = DqliteConnection("localhost:9001")
+    # Pretend the conn is fully connected so ``acquire`` does not run
+    # the entry-time ``_drain_idle`` (which would mask the cleanup
+    # branch we are testing).
+    conn._protocol = object()  # type: ignore[assignment]
+    conn._db_id = 1
+    assert conn.is_connected is True
+
+    async def _fake_create_connection() -> DqliteConnection:
+        return conn
+
+    pool._create_connection = _fake_create_connection  # type: ignore[method-assign]
+    return pool, conn
+
+
+def _break_conn(conn: DqliteConnection) -> None:
+    """Simulate a leader-flip-style invalidation: drop ``_protocol``
+    so ``conn.is_connected`` flips to False, sending the cleanup down
+    the broken-conn branch."""
+    conn._protocol = None
+    conn._db_id = None
+
+
+@pytest.mark.asyncio
+async def test_acquire_cleanup_close_completes_under_outer_cancel() -> None:
+    """Cancel during the broken-conn cleanup must NOT skip
+    ``conn.close()``: shield + suppress(CancelledError) keeps the
+    close running while the original cancel propagates via the
+    surrounding ``raise``. Pin via ``_pool_released=True`` (set in the
+    inner ``finally`` after the close) — without the fix it would
+    stay False because cancel propagated past the bare close."""
+    pool, conn = await _build_pool_with_breakable_conn()
+
+    close_completed = False
+
+    async def fake_close() -> None:
+        nonlocal close_completed
+        # Yield once so a concurrent cancel can land mid-close. The
+        # shield must keep us running through the second statement.
+        await asyncio.sleep(0)
+        # Without the shield, the cancel queued via raise CancelledError
+        # below would propagate through this await and skip the
+        # assignment — the test would catch that as close_completed=False.
+        close_completed = True
+
+    async def fake_drain_idle() -> None:
+        return
+
+    pool._drain_idle = fake_drain_idle  # type: ignore[method-assign]
+
+    with (
+        patch.object(conn, "close", new=fake_close),
+        pytest.raises((asyncio.CancelledError, ValueError)),
+    ):
+        async with pool.acquire():
+            _break_conn(conn)
+            # Schedule a cancel of the current task so it lands
+            # while ``conn.close`` is suspended on the
+            # ``await asyncio.sleep(0)`` above. The shield must
+            # absorb the cancel and let close finish.
+            current = asyncio.current_task()
+            assert current is not None
+            asyncio.get_running_loop().call_soon(current.cancel)
+            raise ValueError("user code error")
+
+    assert close_completed, (
+        "conn.close() must run to completion despite cancel — shield "
+        "around the close await keeps it running"
+    )
+    assert conn._pool_released is True, (
+        "_pool_released must be set in inner finally — the shield+suppress "
+        "around close ensures we reach the flag-set"
+    )
+
+
+@pytest.mark.asyncio
+async def test_acquire_cleanup_close_failure_still_sets_pool_released() -> None:
+    """If ``conn.close()`` raises an unhandled exception (not
+    OSError/DqliteConnectionError), the inner ``finally`` must still
+    set ``_pool_released=True`` so subsequent close() short-circuits."""
+    pool, conn = await _build_pool_with_breakable_conn()
+
+    async def fake_close() -> None:
+        # Simulate an unrecognised close failure (not OSError /
+        # DqliteConnectionError). Must propagate, but the finally
+        # still flips _pool_released.
+        raise RuntimeError("unexpected close failure")
+
+    async def fake_drain_idle() -> None:
+        return
+
+    pool._drain_idle = fake_drain_idle  # type: ignore[method-assign]
+
+    with patch.object(conn, "close", new=fake_close), pytest.raises((ValueError, RuntimeError)):
+        async with pool.acquire():
+            _break_conn(conn)
+            raise ValueError("user code error")
+
+    assert conn._pool_released is True
+
+
+@pytest.mark.asyncio
+async def test_drain_idle_failure_does_not_skip_close() -> None:
+    """Pin existing behaviour: the bare ``_drain_idle()`` failure path
+    is intentionally narrow (DqliteConnectionError / OSError swallowed
+    via _POOL_CLEANUP_EXCEPTIONS); the close still runs after."""
+    pool, conn = await _build_pool_with_breakable_conn()
+
+    close_call_count = 0
+
+    async def fake_close() -> None:
+        nonlocal close_call_count
+        close_call_count += 1
+
+    async def fake_drain_idle() -> None:
+        raise DqliteConnectionError("simulated drain failure")
+
+    pool._drain_idle = fake_drain_idle  # type: ignore[method-assign]
+
+    with patch.object(conn, "close", new=fake_close), pytest.raises(ValueError):
+        async with pool.acquire():
+            _break_conn(conn)
+            raise ValueError("user code error")
+
+    assert close_call_count == 1
+    assert conn._pool_released is True

tests/test_pool_release_pending_drain_awaited_on_cancel.py

@@ -0,0 +1,159 @@
+"""Pin: ``_release``'s finally drains ``_pending_drain`` before
+setting ``_pool_released=True``.
+
+Cancel mid-ROLLBACK (during ``_release._reset_connection``) lands in
+``_run_protocol``'s cancellation branch, which schedules a bounded
+``wait_closed`` drain task on ``_pending_drain``. ``close()`` would
+normally await that drain at its pending-drain block, but the
+``_pool_released=True`` flag set by ``_release``'s finally short-
+circuits past it. The reader-task then outlives the connection,
+producing "Task was destroyed but it is pending" warnings on
+shutdown under cancel-heavy workloads.
+
+The fix snapshots and shield-awaits the drain BEFORE setting
+``_pool_released=True``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import patch
+
+import pytest
+
+from dqliteclient import DqliteConnection
+from dqliteclient.pool import ConnectionPool
+
+
+@pytest.mark.asyncio
+async def test_release_drains_pending_before_setting_pool_released() -> None:
+    """Pin: the drain task completes BEFORE ``_pool_released`` flips.
+
+    Use a release_reservation stub that captures the drain's done-state
+    at the moment it runs. ``_release_reservation`` is awaited AFTER
+    ``_pool_released = True`` in the finally clause, so capturing
+    ``drain.done()`` from inside the stub tells us whether the fix's
+    pre-flag drain ran.
+    """
+    pool = ConnectionPool(["localhost:9001"], min_size=0, max_size=1)
+    pool._size = 1  # reservation already taken
+    conn = DqliteConnection("localhost:9001")
+    conn._protocol = object()  # type: ignore[assignment]
+    conn._db_id = 1
+    conn._in_transaction = True  # force the ROLLBACK branch in _reset_connection
+
+    # Drain task that requires multiple loop turns to complete — so a
+    # single yield from _release_reservation cannot accidentally drain
+    # it. The fix's explicit ``await asyncio.shield(pending)`` is the
+    # only mechanism that drains it before _pool_released flips.
+    async def slow_drain() -> None:
+        for _ in range(5):
+            await asyncio.sleep(0)
+
+    conn._pending_drain = asyncio.create_task(slow_drain())
+
+    async def fake_reset(c: DqliteConnection) -> bool:
+        # Take the close-and-drop branch (mimics
+        # cancel-during-ROLLBACK). The pending drain is what we are
+        # pinning the order on.
+        return False
+
+    pool._reset_connection = fake_reset  # type: ignore[method-assign]
+
+    # Capture the drain task's done-state at the moment
+    # _release_reservation runs. With the fix, the drain has been
+    # explicitly awaited before this point. Without the fix, the
+    # drain is still pending when this fires.
+    drain_done_when_release_reservation_ran: list[bool] = []
+
+    async def observing_release_reservation() -> None:
+        drain_done_when_release_reservation_ran.append(
+            conn._pending_drain is not None and conn._pending_drain.done()
+        )
+
+    pool._release_reservation = observing_release_reservation  # type: ignore[method-assign]
+
+    async def fake_close() -> None:
+        return
+
+    with patch.object(conn, "close", new=fake_close):
+        await pool._release(conn)
+
+    assert drain_done_when_release_reservation_ran == [True], (
+        "drain task must have completed before _release_reservation ran "
+        "(i.e. before _pool_released was set in the finally)"
+    )
+    assert conn._pool_released is True
+    assert conn._pending_drain.done()
+
+
+@pytest.mark.asyncio
+async def test_release_with_no_pending_drain_path_unchanged() -> None:
+    """Negative pin: when there is no ``_pending_drain``, the existing
+    behaviour is unchanged — the drain block is a no-op."""
+    pool = ConnectionPool(["localhost:9001"], min_size=0, max_size=1)
+    pool._size = 1
+    conn = DqliteConnection("localhost:9001")
+    conn._protocol = object()  # type: ignore[assignment]
+    conn._db_id = 1
+    conn._in_transaction = True
+    assert conn._pending_drain is None
+
+    async def fake_reset(c: DqliteConnection) -> bool:
+        return False
+
+    pool._reset_connection = fake_reset  # type: ignore[method-assign]
+
+    async def noop() -> None:
+        return
+
+    pool._release_reservation = noop  # type: ignore[method-assign]
+
+    async def fake_close() -> None:
+        return
+
+    with patch.object(conn, "close", new=fake_close):
+        await pool._release(conn)
+
+    assert conn._pool_released is True
+
+
+@pytest.mark.asyncio
+async def test_release_pending_drain_failure_does_not_block_release() -> None:
+    """If the drain task itself raises, ``_release`` must still
+    complete (set ``_pool_released``, release the reservation). The
+    fix uses ``contextlib.suppress(BaseException)`` around the
+    shielded await so the drain's exception is observed and discarded
+    — ``_release`` is a cleanup path; the original cause is what the
+    user cares about."""
+    pool = ConnectionPool(["localhost:9001"], min_size=0, max_size=1)
+    pool._size = 1
+    conn = DqliteConnection("localhost:9001")
+    conn._protocol = object()  # type: ignore[assignment]
+    conn._db_id = 1
+    conn._in_transaction = True
+
+    async def failing_drain() -> None:
+        raise RuntimeError("drain failed")
+
+    conn._pending_drain = asyncio.create_task(failing_drain())
+
+    async def fake_reset(c: DqliteConnection) -> bool:
+        return False
+
+    pool._reset_connection = fake_reset  # type: ignore[method-assign]
+
+    async def noop() -> None:
+        return
+
+    pool._release_reservation = noop  # type: ignore[method-assign]
+
+    async def fake_close() -> None:
+        return
+
+    with patch.object(conn, "close", new=fake_close):
+        await pool._release(conn)
+
+    assert conn._pool_released is True
+    # The drain task is done — the suppression observed the exception.
+    assert conn._pending_drain.done()