Recover the queue slot when pool acquire is cancelled mid-wait

ConnectionPool.acquire() parks on asyncio.wait({get_task, closed_task})
when the pool is saturated. asyncio.wait does not cancel its argument
tasks on outer cancellation, so when the parent coroutine is cancelled
while suspended in that wait, `get_task` kept running. It could then
win a subsequent queue.put() race, orphan the connection inside its
own result, and silently shrink the pool's effective capacity — the
connection was never reachable again, but _size still counted it.

Restructure the wait block to run a cancellation-aware cleanup. On the
cancel path, stop closed_task, and either:

- if get_task already took a connection out of the queue, return it
  via put_nowait so the reservation backing it stays valid and the
  next acquirer can pick it up, or
- if get_task is still pending, cancel and drain it.

Re-raise to let cancellation propagate. The normal (non-cancel) path
is unchanged: cancel closed_task, take get_task's result if it won,
otherwise cancel the queue wait and re-check state.

Add two regression tests: a single cancel/release cycle asserting that
a subsequent acquire succeeds, and a 100-round fuzz loop asserting
_size stays consistent.

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

Author: antoineleclair

src/dqliteclient/pool.py

@@ -344,9 +344,31 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
                     timeout=remaining,
                     return_when=asyncio.FIRST_COMPLETED,
                 )
-            finally:
+            except BaseException:
+                # Outer cancellation of this coroutine while suspended in
+                # ``asyncio.wait``. ``asyncio.wait`` does not cancel its
+                # argument tasks, so both children are still alive — we
+                # MUST stop them before propagating, otherwise the
+                # abandoned get_task can win a later queue.put() and
+                # orphan a connection (silently shrinking pool capacity).
                 if not closed_task.done():
                     closed_task.cancel()
+                if get_task.done() and not get_task.cancelled() and get_task.exception() is None:
+                    # Outer cancel raced with a successful get. The
+                    # reservation that backed this connection is still
+                    # 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()
+                    with contextlib.suppress(asyncio.QueueFull):
+                        self._pool.put_nowait(conn_won)
+                elif not get_task.done():
+                    get_task.cancel()
+                    with contextlib.suppress(BaseException):
+                        await get_task
+                raise
+            if not closed_task.done():
+                closed_task.cancel()
             if get_task in done:
                 conn = get_task.result()
             else:

tests/test_pool.py

@@ -1303,3 +1303,117 @@ async def test_ordinary_exception_is_absorbed_and_logged(
         c2.close.assert_awaited_once()
         assert pool._size == 0
         assert any("boom" in rec.getMessage() for rec in caplog.records)
+
+
+class TestAcquireCancellationPreservesCapacity:
+    """An external cancellation of a coroutine parked in ``acquire()``'s
+    ``asyncio.wait(...)`` must not leak a connection or shrink effective
+    pool capacity. The abandoned ``get_task`` could win a subsequent
+    ``put()`` race; the fix must either cancel it or reclaim the
+    connection it took.
+    """
+
+    @pytest.fixture
+    def mock_connection(self) -> MagicMock:
+        conn = MagicMock()
+        conn.is_connected = True
+        conn._in_transaction = False
+        conn._pool_released = False
+        conn.connect = AsyncMock()
+        conn.close = AsyncMock()
+        conn._protocol = MagicMock()
+        conn._protocol._writer = MagicMock()
+        conn._protocol._writer.transport = MagicMock()
+        conn._protocol._writer.transport.is_closing = MagicMock(return_value=False)
+        conn._protocol._reader = MagicMock()
+        conn._protocol._reader.at_eof = MagicMock(return_value=False)
+        return conn
+
+    async def test_cancelled_waiter_does_not_wedge_pool(
+        self,
+        mock_connection: MagicMock,
+    ) -> None:
+        pool = ConnectionPool(["localhost:9001"], max_size=1, timeout=5.0)
+
+        with patch.object(pool._cluster, "connect", return_value=mock_connection):
+            await pool.initialize()
+
+            holder_released = asyncio.Event()
+            waiter_parked = asyncio.Event()
+
+            async def holder() -> None:
+                async with pool.acquire():
+                    waiter_parked.set()
+                    # Hold until the waiter has been cancelled AND we've
+                    # been signalled to exit.
+                    await holder_released.wait()
+
+            async def waiter() -> None:
+                async with pool.acquire():
+                    pass
+
+            holder_task = asyncio.create_task(holder())
+            await waiter_parked.wait()
+
+            waiter_task = asyncio.create_task(waiter())
+            # Give waiter() time to enter acquire() and park in asyncio.wait.
+            await asyncio.sleep(0.05)
+
+            waiter_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await waiter_task
+
+            # Release the in-use connection; the post-cancel put() must
+            # not be stolen by the abandoned get_task.
+            holder_released.set()
+            await holder_task
+
+            # Pool must still accept a new acquire within a short window —
+            # if the abandoned task captured the connection, this would
+            # time out.
+            async with asyncio.timeout(1.0):
+                async with pool.acquire() as conn:
+                    assert conn is mock_connection
+
+            assert pool._size == 1
+
+        await pool.close()
+
+    async def test_repeated_cancel_release_keeps_size_consistent(
+        self,
+        mock_connection: MagicMock,
+    ) -> None:
+        """Invariant: after any cancel/release cycle the ``_size`` counter
+        equals the number of connections actually reachable. Fuzz-test 100
+        rounds (smaller than the issue file's 1000 — kept light for CI)."""
+        pool = ConnectionPool(["localhost:9001"], max_size=1, timeout=5.0)
+
+        with patch.object(pool._cluster, "connect", return_value=mock_connection):
+            await pool.initialize()
+
+            for _ in range(100):
+                release_event = asyncio.Event()
+
+                async def holder(event: asyncio.Event = release_event) -> None:
+                    async with pool.acquire():
+                        await event.wait()
+
+                async def waiter() -> None:
+                    async with pool.acquire():
+                        pass
+
+                h = asyncio.create_task(holder())
+                await asyncio.sleep(0)
+                w = asyncio.create_task(waiter())
+                await asyncio.sleep(0.01)
+                w.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await w
+                release_event.set()
+                await h
+
+                # After every round the pool must have exactly one
+                # reachable reservation — queued or returned-by-holder.
+                assert pool._size == 1
+
+        await pool.close()