Bound pool.close() drain wall-clock with an aggregate deadline

``pool.close()`` previously called ``self._drain_idle()`` with no
``deadline=`` kwarg. A queue of pathologically-slow ``close()``
peers (leader flip mid-shutdown, half-closed sockets whose
``wait_closed`` never returns, FIN/ACK loss on a downed peer)
consumed ``max_size × close_timeout × cap_multiplier`` of loop
time before close returned — ~20 s at defaults, ~200 s with
``max_size=100``. Sibling acquire-path callers already passed
``deadline=...``; the close path was the structural outlier.

Compute an aggregate close-budget deadline up-front and thread
it through both ``_drain_idle`` and ``_drain_remaining_after_cancel``
(which gains the matching kwarg + between-iteration gate). The
cancel-recovery arm now honours the same envelope so an outer
cancel cannot re-amplify the gap via the recovery sweep. The
second-caller arm in ``close()`` mirrors the shape so re-entry
under a stuck queue is bounded too.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: antoineleclair

src/dqliteclient/pool.py

@@ -1642,9 +1642,14 @@ async def _drain_idle(self, *, deadline: float | None = None) -> None:
             # the next dead-conn detect, or ``pool.close()``.
             if not deadline_exit:
                 with contextlib.suppress(asyncio.CancelledError):
-                    await asyncio.shield(self._drain_remaining_after_cancel())
-
-    async def _drain_remaining_after_cancel(self) -> None:
+                    # Forward the same deadline so the
+                    # cancel-recovery sweep honours the aggregate
+                    # close-budget; without this the recovery arm
+                    # re-amplifies the gap the main-loop deadline
+                    # was meant to bound.
+                    await asyncio.shield(self._drain_remaining_after_cancel(deadline=deadline))
+
+    async def _drain_remaining_after_cancel(self, *, deadline: float | None = None) -> None:
         """Best-effort sweep for connections still in the queue after
         the main drain loop exited (typically via outer cancel).
 
@@ -1654,8 +1659,30 @@ async def _drain_remaining_after_cancel(self) -> None:
         does not abort the cleanup of the rest. Re-entry from a second
         ``close()`` caller is safe — the queue empties on the first
         sweep and subsequent calls find ``self._pool.empty() is True``.
+
+        Optional ``deadline`` (monotonic ``loop.time()`` value): when
+        the running loop's clock advances past it, the sweep stops
+        attempting further connections. The check is between
+        iterations so an already-started per-connection close still
+        gets its per-iteration cap. Mirrors the discipline at
+        ``_drain_idle``. Callers from ``pool.close()`` pass an
+        aggregate close-budget deadline so the SIGTERM contract is
+        honoured end-to-end (cancel-recovery arm cannot extend close
+        past the operator's budget).
         """
+        loop = asyncio.get_running_loop()
         while not self._pool.empty():
+            if deadline is not None and loop.time() >= deadline:
+                # Aggregate deadline exceeded between iterations.
+                # The remaining queue stays for the next caller
+                # (next ``pool.close()`` re-entry, or the
+                # post-close finalizer). Mirrors ``_drain_idle``'s
+                # between-iteration deadline gate.
+                logger.debug(
+                    "pool: _drain_remaining_after_cancel deadline "
+                    "exhausted between iterations; abandoning"
+                )
+                return
             try:
                 conn = self._pool.get_nowait()
             except asyncio.QueueEmpty:  # pragma: no cover
@@ -2898,7 +2925,15 @@ async def close(self) -> None:
             # being silently broken.
             if not self._drain_complete:
                 with contextlib.suppress(asyncio.CancelledError):
-                    await asyncio.shield(self._drain_remaining_after_cancel())
+                    # Mirror the first-caller's close-budget shape so
+                    # the second-caller's sweep cannot extend past
+                    # the same aggregate envelope.
+                    second_close_deadline = asyncio.get_running_loop().time() + (
+                        self._close_timeout * self._max_size * _DRAIN_PER_CONN_CAP_MULTIPLIER
+                    )
+                    await asyncio.shield(
+                        self._drain_remaining_after_cancel(deadline=second_close_deadline)
+                    )
             return
         # Publish the drain-done event BEFORE flipping the closed flag
         # so any second caller observing ``_closed=True`` is guaranteed
@@ -2928,7 +2963,24 @@ async def close(self) -> None:
                 self._pool.qsize(),
                 max(self._size - self._pool.qsize(), 0),
             )
-            await self._drain_idle()
+            # Aggregate close-budget deadline. Bounds the drain at the
+            # documented worst-case shape (``close_timeout × max_size
+            # × _DRAIN_PER_CONN_CAP_MULTIPLIER``) — same envelope the
+            # operator would compose manually around ``pool.close()``
+            # via ``asyncio.timeout``. Without this gate, a queue of
+            # stuck-close peers (leader-flip mid-shutdown, half-closed
+            # sockets whose ``wait_closed`` never returns) keeps the
+            # close path waiting the full
+            # ``N × per_iter_cap`` even though the operator's
+            # ``close_timeout`` was sized per-peer, not per-batch. The
+            # cancel-recovery arm in ``_drain_idle`` forwards the same
+            # deadline so an outer cancel still truncates predictably
+            # rather than re-amplifying via
+            # ``_drain_remaining_after_cancel``.
+            close_deadline = asyncio.get_running_loop().time() + (
+                self._close_timeout * self._max_size * _DRAIN_PER_CONN_CAP_MULTIPLIER
+            )
+            await self._drain_idle(deadline=close_deadline)
             # Drain completed normally. Set BEFORE the finally so a
             # cancel landing between the drain return and the flag
             # assignment leaves ``_drain_complete=False`` — siblings

tests/test_pool_close_drain_honours_aggregate_deadline.py

@@ -0,0 +1,119 @@
+"""Pin: ``ConnectionPool.close()`` derives an aggregate close-budget
+deadline and threads it through both ``_drain_idle`` and
+``_drain_remaining_after_cancel`` so the SIGTERM contract bounds
+end-to-end wall-clock regardless of stuck-close stampedes.
+
+The prior shape called ``await self._drain_idle()`` (no
+``deadline=`` kwarg). With the wire-default ``close_timeout=0.5``
+and ``max_size=10``, a queue of pathologically-slow ``close()``
+peers consumed ~20 s before the close path returned; with
+``max_size=100``, ~200 s. Sibling acquire-path call sites
+already passed ``deadline=...``; the close path was the
+structural outlier.
+
+``_drain_remaining_after_cancel`` (the cancel-recovery sweep)
+gains the same ``deadline`` kwarg so an outer cancel mid-drain
+cannot re-amplify the gap via the recovery arm.
+"""
+
+from __future__ import annotations
+
+import ast
+import inspect
+import textwrap
+
+from dqliteclient.pool import ConnectionPool
+
+
+def _close_source() -> str:
+    return textwrap.dedent(inspect.getsource(ConnectionPool.close))
+
+
+def _drain_remaining_after_cancel_signature() -> inspect.Signature:
+    return inspect.signature(ConnectionPool._drain_remaining_after_cancel)
+
+
+def test_pool_close_passes_deadline_to_drain_idle() -> None:
+    """Structural pin: ``pool.close()`` must call
+    ``self._drain_idle(deadline=...)`` (not the bare
+    ``self._drain_idle()`` shape that omits the aggregate
+    close-budget gate).
+    """
+    src = _close_source()
+    tree = ast.parse(src)
+    found_with_deadline = False
+    found_bare = False
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Call):
+            continue
+        func = node.func
+        if (
+            isinstance(func, ast.Attribute)
+            and func.attr == "_drain_idle"
+            and isinstance(func.value, ast.Name)
+            and func.value.id == "self"
+        ):
+            if any(kw.arg == "deadline" for kw in node.keywords):
+                found_with_deadline = True
+            else:
+                found_bare = True
+    assert found_with_deadline, (
+        "pool.close() must pass ``deadline=<aggregate_close_budget>`` "
+        "to ``self._drain_idle``. The aggregate-budget gate is what "
+        "bounds end-to-end close wall-clock under a stuck-close "
+        "stampede; without it, ``close()`` waits ``N × per_iter_cap`` "
+        "even when ``close_timeout`` was sized per-peer."
+    )
+    assert not found_bare, (
+        "pool.close() should NOT contain a bare ``self._drain_idle()`` "
+        "call. Every drain invocation from the close path must thread "
+        "the aggregate deadline through."
+    )
+
+
+def test_drain_remaining_after_cancel_accepts_deadline_kwarg() -> None:
+    """The sibling cancel-recovery sweep must accept the same
+    ``deadline`` plumbing so the close-budget gate is honoured
+    end-to-end (outer cancel can't re-amplify via the recovery arm).
+    """
+    sig = _drain_remaining_after_cancel_signature()
+    assert "deadline" in sig.parameters, (
+        "_drain_remaining_after_cancel must accept a ``deadline`` "
+        "kwarg so the cancel-recovery sweep honours the same "
+        "aggregate close-budget the main drain loop enforces. "
+        "Without it, an outer cancel re-amplifies the gap the "
+        "main-loop deadline was meant to bound."
+    )
+
+
+def test_pool_close_forwards_deadline_into_drain_remaining_after_cancel() -> None:
+    """Pin: the second-caller arm in ``pool.close()`` calls
+    ``self._drain_remaining_after_cancel(deadline=...)`` (not the
+    bare form).
+    """
+    src = _close_source()
+    tree = ast.parse(src)
+    found_with_deadline = False
+    found_bare = False
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Call):
+            continue
+        func = node.func
+        if (
+            isinstance(func, ast.Attribute)
+            and func.attr == "_drain_remaining_after_cancel"
+            and isinstance(func.value, ast.Name)
+            and func.value.id == "self"
+        ):
+            if any(kw.arg == "deadline" for kw in node.keywords):
+                found_with_deadline = True
+            else:
+                found_bare = True
+    assert found_with_deadline, (
+        "pool.close()'s second-caller arm must pass a ``deadline=`` "
+        "to ``self._drain_remaining_after_cancel`` so the recovery "
+        "sweep honours the close-budget envelope."
+    )
+    assert not found_bare, (
+        "pool.close() should not contain a bare ``self._drain_remaining_after_cancel()`` call."
+    )