Use live get_task state in pool acquire post-wait demux

The capacity-wait branch's post-wait demux used ``if get_task in done:``,
but ``done`` is the snapshot taken before the post-wait
``await closed_task`` yield. A sibling ``_release`` ``put_nowait``
during that yield resolves ``get_task``; the stale snapshot still says
False, so the else-arm cancels (no-op on a done task) and ``await
get_task`` returns the conn — silently discarded by ``continue``. The
reservation slot leaks because ``_release`` only fires for connections
that flow back through the user's context manager.

Replace the snapshot membership test with a live-state predicate
(``done() and not cancelled() and exception() is None``) and add a
second live-state recheck after the cancel-and-await in the else-arm.
Factor the existing put_nowait-or-close-and-release fallback (already
present in the ``except BaseException`` arm) into a private helper
``_put_back_or_release_late_winner`` so both call sites share one code
path.

Author: antoineleclair

src/dqliteclient/pool.py

@@ -441,6 +441,48 @@ async def _create_connection(self) -> DqliteConnection:
             max_attempts=self._max_attempts,
         )
 
+    async def _put_back_or_release_late_winner(self, conn: DqliteConnection) -> None:
+        """Put a connection back on the queue, or close + release the
+        reservation if the queue is full.
+
+        Used in two places in ``acquire()``:
+
+        1. The ``except BaseException`` arm — outer cancel raced with
+           a successful ``get_task`` (a sibling ``_release``
+           ``put_nowait`` ran between our snapshot and the cancel).
+        2. The post-wait demux's else-arm — timeout snapshot raced a
+           winning ``get_task`` during the post-wait
+           ``await closed_task``.
+
+        Without this routing, the conn is referenced only by the
+        soon-to-be-GC'd ``get_task`` and silently disappears. Its
+        reservation slot is never released because ``_release`` only
+        fires for connections that flow back through the user's
+        context manager — so the pool permanently loses one slot of
+        capacity per occurrence.
+        """
+        try:
+            self._pool.put_nowait(conn)
+        except asyncio.QueueFull:
+            # Invariant violation: reservations should track queue
+            # capacity exactly, so a full queue on return is
+            # "impossible." If it happens anyway, silently dropping
+            # the reference would leak a live reader task and a
+            # socket. Close explicitly and adjust the reservation
+            # count so the pool shrinks cleanly instead of leaking.
+            # Suppression of close's own errors is narrow — OSError on
+            # an already-dead writer is expected; anything else
+            # propagates.
+            with contextlib.suppress(OSError):
+                await conn.close()
+            # Route through the helper so the counter stays
+            # lock-protected and sibling acquirers parked on
+            # ``closed_event.wait()`` get woken via
+            # ``_signal_state_change``. Shield so a nested cancel
+            # cannot leave ``_size`` inconsistent.
+            with contextlib.suppress(asyncio.CancelledError):
+                await asyncio.shield(self._release_reservation())
+
     async def _release_reservation(self) -> None:
         """Decrement ``_size`` under the lock, waking waiters.
 
@@ -704,30 +746,7 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
                     # valid; return it to the queue so the next
                     # acquirer can use it instead of closing and
                     # releasing (which would shrink _size).
-                    conn_won = get_task.result()
-                    try:
-                        self._pool.put_nowait(conn_won)
-                    except asyncio.QueueFull:
-                        # Invariant violation: reservations should track
-                        # queue capacity exactly, so a full queue on
-                        # return is "impossible." If it happens anyway,
-                        # silently dropping the reference would leak a
-                        # live reader task and a socket. Close
-                        # explicitly and adjust the reservation count
-                        # so the pool shrinks cleanly instead of
-                        # leaking. Suppression of close's own errors is
-                        # narrow — OSError on an already-dead writer is
-                        # expected; anything else propagates.
-                        with contextlib.suppress(OSError):
-                            await conn_won.close()
-                        # Route through the helper so the counter
-                        # stays lock-protected and sibling acquirers
-                        # parked on ``closed_event.wait()`` get
-                        # woken via ``_signal_state_change``. Shield
-                        # so a nested cancel cannot leave ``_size``
-                        # inconsistent.
-                        with contextlib.suppress(asyncio.CancelledError):
-                            await asyncio.shield(self._release_reservation())
+                    await self._put_back_or_release_late_winner(get_task.result())
                 elif get_task is not None and not get_task.done():
                     get_task.cancel()
                     with contextlib.suppress(asyncio.CancelledError):
@@ -755,14 +774,29 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
                 # must still propagate.
                 with contextlib.suppress(asyncio.CancelledError):
                     await closed_task
-            if get_task in done:
+            if get_task.done() and not get_task.cancelled() and get_task.exception() is None:
+                # Live-state check: ``done`` is the snapshot taken
+                # before the post-wait ``await closed_task`` yield
+                # above, during which a sibling ``_release`` can
+                # ``put_nowait`` and resolve ``get_task``. Trusting
+                # ``get_task in done`` would silently route the
+                # winning conn into the cancel-and-discard arm,
+                # leaking one slot of capacity per occurrence.
                 conn = get_task.result()
             else:
                 # Either close fired or the poll timer fired; either way,
                 # cancel the queue wait cleanly and let the loop re-check.
                 get_task.cancel()
                 with contextlib.suppress(asyncio.CancelledError):
                     await get_task
+                if get_task.done() and not get_task.cancelled() and get_task.exception() is None:
+                    # Cancel raced a successful get during the await
+                    # above (a sibling ``_release`` put_nowait between
+                    # our cancel call and the cancel actually
+                    # delivering). Route the conn back via the same
+                    # put-back-or-release path the outer-cancel arm
+                    # uses, so the slot is not leaked.
+                    await self._put_back_or_release_late_winner(get_task.result())
                 continue
 
         # If connection is dead, discard and create a fresh one with leader discovery.

tests/test_pool_acquire_timeout_late_get_task.py

@@ -0,0 +1,201 @@
+"""Pin: ``acquire()``'s capacity-wait timeout demux must not silently
+discard a connection that resolved ``get_task`` during the post-wait
+``await closed_task`` yield.
+
+When ``asyncio.wait`` returns the timeout (``done == set()``) and the
+post-wait code cancels and awaits ``closed_task``, that await yields to
+the scheduler. A sibling ``_pool.put_nowait(conn)`` running at that
+yield resolves the still-pending ``get_task``. The original code's
+demux test ``if get_task in done`` uses the *snapshot* taken before the
+yield, so it incorrectly takes the else-arm: cancels (no-op on a done
+task), ``await get_task`` returns the connection, ``continue`` discards
+it. The reservation slot is never released because ``_release`` only
+fires for connections that flow back through the user's context
+manager. The pool permanently loses one slot of capacity per
+occurrence.
+
+The fix replaces the snapshot-membership check with a live state check
+``get_task.done() and not get_task.cancelled() and get_task.exception()
+is None`` and routes a winning conn through the same put-back-or-release
+helper used by the existing ``except BaseException`` arm.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from dqliteclient.cluster import ClusterClient
+from dqliteclient.pool import ConnectionPool
+
+
+class _FakeConn:
+    def __init__(self, name: str = "fake") -> None:
+        self.name = name
+        self._address = "localhost:9001"
+        self._in_transaction = False
+        self._tx_owner = None
+        self._pool_released = False
+        self._protocol = MagicMock()
+        self._protocol._writer = MagicMock()
+        self._protocol._writer.transport = MagicMock()
+        self._protocol._writer.transport.is_closing = lambda: False
+        self._protocol._reader = MagicMock()
+        self._protocol._reader.at_eof = lambda: False
+        self.close_called = False
+
+    @property
+    def is_connected(self) -> bool:
+        return self._protocol is not None
+
+    async def close(self) -> None:
+        self.close_called = True
+        self._protocol = None  # type: ignore[assignment]
+
+
+def _make_pool() -> ConnectionPool:
+    async def _connect(**_: Any) -> _FakeConn:
+        return _FakeConn()
+
+    cluster = MagicMock(spec=ClusterClient)
+    cluster.connect = _connect
+    return ConnectionPool(
+        addresses=["localhost:9001"],
+        min_size=0,
+        max_size=1,
+        timeout=0.1,
+        cluster=cluster,
+    )
+
+
+@pytest.mark.asyncio
+async def test_acquire_timeout_race_does_not_discard_late_winning_get_task() -> None:
+    """``asyncio.wait`` times out (``done == set()``) but a sibling
+    ``put_nowait`` resolves ``get_task`` during the post-wait
+    ``await closed_task`` yield. The conn must be put back on the
+    queue (or routed to the user), not silently discarded by the
+    stale-snapshot demux.
+
+    Setup: simulate "pool already at max_size" by directly setting
+    ``_size = 1`` (no real held connection). When ``acquire()`` enters
+    the capacity-wait branch, our patched ``asyncio.wait`` drops a
+    phantom connection into the queue synchronously before returning
+    a timeout-shaped result (done=empty, both tasks still pending).
+    The next loop iteration delivered to ``await closed_task`` then
+    runs ``get_task.__step`` before our coroutine resumes, so
+    ``get_task`` consumes the phantom and becomes done. The buggy
+    post-wait demux at ``pool.py`` ``if get_task in done`` sees the
+    stale empty snapshot, takes the else-arm, and discards the
+    phantom on ``continue``. The fix re-checks ``get_task.done()``
+    live and routes the conn through put-back-or-release.
+    """
+    pool = _make_pool()
+
+    # Pretend max_size is reached so ``acquire()`` can't reserve and
+    # enters the capacity-wait branch on its first iteration. Avoids
+    # the asynccontextmanager-finalization noise from holding a real
+    # acquire's slot across a bare ``__aenter__()``.
+    pool._size = 1
+
+    phantom = _FakeConn(name="phantom")
+    original_put_nowait = pool._pool.put_nowait
+
+    import dqliteclient.pool as pool_mod
+
+    real_wait = asyncio.wait
+    call_count = 0
+
+    async def fake_wait(
+        tasks: Any, *, timeout: Any = None, return_when: Any = None
+    ) -> tuple[set[Any], set[Any]]:
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            # Drop the phantom into the queue while ``get_task`` has
+            # not yet had its first ``__step`` run (we are inside
+            # ``await asyncio.wait`` synchronously — the loop hasn't
+            # iterated since ``create_task``). When the post-wait code
+            # subsequently yields on ``await closed_task``, the loop
+            # runs ``get_task.__step`` first and ``get_task`` consumes
+            # ``phantom``, becoming done with a real result, before
+            # our coroutine resumes. The post-wait demux's
+            # ``if get_task in done`` then sees the stale empty
+            # snapshot and routes ``get_task`` into the
+            # cancel-and-discard arm.
+            original_put_nowait(phantom)  # type: ignore[arg-type]
+            # Return timeout: done=empty, both tasks still pending
+            # from the snapshot's perspective.
+            return set(), set(tasks)
+        # Subsequent calls: defer to the real wait so the deadline
+        # actually consumes time and the loop terminates promptly.
+        return await real_wait(tasks, timeout=timeout, return_when=return_when)
+
+    pool_mod.asyncio.wait = fake_wait  # type: ignore[attr-defined]
+    received: object | None = None
+    try:
+        # With the fix: the live-state recheck after ``await
+        # closed_task`` finds get_task done with phantom and routes it
+        # to the user (no timeout). Without the fix: stale snapshot
+        # demux drops phantom on the floor; subsequent iterations time
+        # out with an empty queue.
+        async with pool.acquire() as conn:
+            received = conn
+    finally:
+        pool_mod.asyncio.wait = real_wait  # type: ignore[attr-defined]
+
+    # The phantom that ``put_nowait`` deposited during the
+    # capacity-wait race must reach the user (or round-trip back to
+    # the queue), never be silently discarded by the stale ``done``
+    # snapshot demux.
+    assert received is phantom, (
+        f"acquire returned {received!r}, not the phantom that was put "
+        "into the queue during the capacity-wait race — the post-wait "
+        "demux's stale 'done' snapshot dropped phantom on the floor"
+    )
+
+    # Reset _size to the value used to simulate at-capacity so close()
+    # does not hit the underflow guard. ``_release`` already
+    # decremented _size by routing through ``_release_reservation`` on
+    # __aexit__.
+    pool._size = 0
+    await pool.close()
+
+
+@pytest.mark.asyncio
+async def test_put_back_or_release_late_winner_queuefull_falls_back_to_close() -> None:
+    """If the queue is full when the late-winner helper tries to put,
+    it must close the conn and release the reservation rather than
+    silently leak it.
+
+    The QueueFull branch represents an "impossible" reservation-vs-
+    capacity violation; the helper must handle it without dropping
+    the conn on the floor or skipping the ``_size`` decrement that
+    wakes sibling acquirers.
+    """
+    pool = _make_pool()
+
+    # Pre-fill the bounded queue (max_size=1) so the helper's
+    # put_nowait immediately raises QueueFull.
+    pool._size = 1
+    occupant = _FakeConn(name="occupant")
+    pool._pool.put_nowait(occupant)  # type: ignore[arg-type]
+    assert pool._pool.full()
+
+    late_winner = _FakeConn(name="late_winner")
+    await pool._put_back_or_release_late_winner(late_winner)  # type: ignore[arg-type]
+
+    # The late_winner must have been close()'d (because put_nowait
+    # raised QueueFull, the helper falls back to close + release).
+    assert late_winner.close_called is True
+
+    # The reservation must have been released (size -= 1) so a
+    # sibling acquirer can replace the slot.
+    assert pool._size == 0
+
+    # Cleanup: drain the occupant from the queue.
+    queued = pool._pool.get_nowait()
+    assert queued is occupant
+    await pool.close()