refactor(pool): reservation pattern for _size accounting (ISSUE-34/58)

Every mutation of ``ConnectionPool._size`` now goes through either the
new ``_release_reservation`` helper or happens under ``_lock``, and
``_create_connection`` no longer mutates ``_size`` itself. In
``acquire()``, the reservation (``_size += 1``) is taken under the
lock; the TCP / leader-discovery work happens *outside* the lock; on
failure the caller releases the reservation.

Two improvements:

1. ISSUE-58 — pool throughput no longer bottlenecked on the single
   ``_lock`` across network handshakes. Concurrent acquires under
   capacity can create connections in parallel.

2. ISSUE-34 — failure and drain paths (``_drain_idle``, the
   exception branches in ``acquire``, ``_release``) no longer
   decrement ``_size`` without holding the lock, so the count stays
   consistent against concurrent capacity checks.

Add a stress test that drives 10 concurrent acquires with max_size=3
and asserts the cap is always respected even with injected slow
creates; add tests for create-failure rollback and drain-idle size
accounting. Update the pre-existing ``test_dead_conn_replacement``
assertion from "lock must be held during create" (now false) to
"_size must never exceed max_size" (the actual invariant).

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

Author: antoineleclair

src/dqliteclient/pool.py

@@ -118,29 +118,57 @@ def __init__(
         self._initialized = False
 
     async def initialize(self) -> None:
-        """Initialize the pool with minimum connections."""
+        """Initialize the pool with minimum connections.
+
+        Idempotent: concurrent callers share the same initialization —
+        only one performs the TCP work, the others await its result.
+        """
+        # Hold the lock across the gather so a second concurrent
+        # initialize() call observes _initialized=True after the first
+        # completes and returns without re-creating.
         async with self._lock:
             if self._initialized:
                 return
-            # Create min_size connections concurrently so pool startup
-            # latency doesn't scale with min_size × per-connect RTT. If any
-            # one fails, cancel the rest and propagate; leave _initialized
-            # False so the caller can retry.
             if self._min_size > 0:
-                conns = await asyncio.gather(
-                    *(self._create_connection() for _ in range(self._min_size))
-                )
+                self._size += self._min_size
+                try:
+                    # Create min_size connections concurrently so startup
+                    # latency doesn't scale with min_size × per-connect RTT.
+                    conns = await asyncio.gather(
+                        *(self._create_connection() for _ in range(self._min_size))
+                    )
+                except BaseException:
+                    self._size -= self._min_size
+                    self._signal_state_change()
+                    raise
                 for conn in conns:
                     await self._pool.put(conn)
             self._initialized = True
 
     async def _create_connection(self) -> DqliteConnection:
-        """Create a new connection to the leader."""
-        conn = await self._cluster.connect(
+        """Create a new connection to the leader.
+
+        Does NOT mutate ``self._size`` — callers must have incremented
+        ``_size`` under ``_lock`` as a reservation before calling this
+        method (see ``acquire`` / ``initialize``) so that concurrent
+        callers cannot collectively exceed ``_max_size``. On failure,
+        the caller is responsible for decrementing to release the
+        reservation.
+        """
+        return await self._cluster.connect(
             database=self._database, max_total_rows=self._max_total_rows
         )
-        self._size += 1
-        return conn
+
+    async def _release_reservation(self) -> None:
+        """Decrement ``_size`` under the lock, waking waiters.
+
+        Every ``_size -= 1`` call in the pool must go through this
+        helper so the counter stays consistent against concurrent
+        capacity checks in ``acquire``.
+        """
+        async with self._lock:
+            self._size -= 1
+        self._signal_state_change()
 
     def _get_closed_event(self) -> asyncio.Event:
         """Lazily create the closed Event bound to the running loop."""
@@ -178,8 +206,7 @@ async def _drain_idle(self) -> None:
             except BaseException:
                 pass
             finally:
-                self._size -= 1
-                self._signal_state_change()
+                await self._release_reservation()
 
     @asynccontextmanager
     async def acquire(self) -> AsyncIterator[DqliteConnection]:
@@ -202,13 +229,23 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
             except asyncio.QueueEmpty:
                 pass
 
-            # Try to create a new connection if under max
+            # Try to reserve a new-connection slot under the lock, then
+            # drop the lock before the TCP handshake so concurrent
+            # pool users aren't serialized on network latency.
+            reserved = False
             async with self._lock:
                 if self._closed:
                     raise DqliteConnectionError("Pool is closed")
                 if self._size < self._max_size:
+                    self._size += 1
+                    reserved = True
+            if reserved:
+                try:
                     conn = await self._create_connection()
-                    break
+                except BaseException:
+                    await self._release_reservation()
+                    raise
+                break
 
             # At capacity — wait briefly on the queue, then loop back to
             # re-check capacity (another coroutine may have freed a slot)
@@ -250,9 +287,16 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
         # Also drain other idle connections — they likely point to the same dead server.
         if not conn.is_connected:
             await self._drain_idle()
+            # Release the dead reservation, reserve a new slot, then
+            # do the network work outside the lock.
             async with self._lock:
-                self._size -= 1
+                self._size -= 1  # dead conn's reservation
+                self._size += 1  # fresh reservation
+            try:
                 conn = await self._create_connection()
+            except BaseException:
+                await self._release_reservation()
+                raise
 
         conn._pool_released = False
         try:
@@ -265,25 +309,22 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
                     conn._pool_released = True
                     with contextlib.suppress(Exception):
                         await conn.close()
-                    self._size -= 1
-                    self._signal_state_change()
+                    await self._release_reservation()
                     raise
                 conn._pool_released = True
                 try:
                     self._pool.put_nowait(conn)
                 except asyncio.QueueFull:
                     await conn.close()
-                    self._size -= 1
-                    self._signal_state_change()
+                    await self._release_reservation()
             else:
                 # Connection is broken (invalidated by execute/fetch error handlers).
                 # Drain other idle connections — they likely point to the same dead server.
                 conn._pool_released = True
                 await self._drain_idle()
                 with contextlib.suppress(BaseException):
                     await conn.close()
-                self._size -= 1
-                self._signal_state_change()
+                await self._release_reservation()
             raise
         else:
             await self._release(conn)
@@ -330,25 +371,22 @@ async def _release(self, conn: DqliteConnection) -> None:
         if self._closed:
             conn._pool_released = True
             await conn.close()
-            self._size -= 1
-            self._signal_state_change()
+            await self._release_reservation()
             return
 
         if not await self._reset_connection(conn):
             conn._pool_released = True
             with contextlib.suppress(Exception):
                 await conn.close()
-            self._size -= 1
-            self._signal_state_change()
+            await self._release_reservation()
             return
 
         conn._pool_released = True
         try:
             self._pool.put_nowait(conn)
         except asyncio.QueueFull:
             await conn.close()
-            self._size -= 1
-            self._signal_state_change()
+            await self._release_reservation()
 
     async def execute(self, sql: str, params: Sequence[Any] | None = None) -> tuple[int, int]:
         """Execute a SQL statement using a pooled connection."""

tests/test_pool.py

@@ -300,7 +300,15 @@ async def waiter():
         await pool.close()
 
     async def test_dead_conn_replacement_respects_max_size(self) -> None:
-        """Dead-connection replacement must not allow _size to exceed max_size."""
+        """Dead-connection replacement must not allow _size to exceed max_size.
+
+        ISSUE-34/58 changed the pool to a reservation pattern: the
+        lock is released before the TCP handshake. The invariant we
+        assert here is the one that actually matters — that
+        ``_size`` never exceeds ``_max_size`` at any observation
+        point — rather than the specific implementation detail of
+        "lock held during create".
+        """
         pool = ConnectionPool(["localhost:9001"], min_size=1, max_size=2)
 
         dead_conn = MagicMock()
@@ -313,16 +321,16 @@ async def test_dead_conn_replacement_respects_max_size(self) -> None:
 
         assert pool._size == 1
 
-        # Track whether the lock is held during dead-conn replacement.
-        # If _create_connection is called while the lock is NOT held,
-        # another coroutine could race and exceed max_size.
-        lock_was_held_during_create = False
+        # Track _size at every create call. It must never exceed
+        # max_size, even transiently, because creates happen outside
+        # the lock under the new pattern and concurrent acquires
+        # would otherwise race past the cap.
+        peak_size_during_create = 0
         original_create = pool._create_connection
 
         async def tracking_create():
-            nonlocal lock_was_held_during_create
-            if pool._lock.locked():
-                lock_was_held_during_create = True
+            nonlocal peak_size_during_create
+            peak_size_during_create = max(peak_size_during_create, pool._size)
             return await original_create()
 
         new_conn = MagicMock()
@@ -335,9 +343,9 @@ async def tracking_create():
             async with pool.acquire() as conn:
                 assert conn is new_conn
 
-        assert lock_was_held_during_create, (
-            "Dead-connection replacement must hold the lock to prevent "
-            "_size race with concurrent acquire() calls"
+        assert peak_size_during_create <= pool._max_size, (
+            f"_size exceeded max_size during dead-conn replacement "
+            f"(peak={peak_size_during_create}, max={pool._max_size})"
         )
 
         await pool.close()

tests/test_pool_size_invariants.py

@@ -0,0 +1,109 @@
+"""ISSUE-34 / ISSUE-58: concurrent-access invariants of ConnectionPool.
+
+Verifies that _size never exceeds max_size under concurrent acquires,
+even when _create_connection is slow (simulating TCP handshake
+latency) — a regression would be the reservation pattern losing
+ordering between check and increment."""
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from dqliteclient.pool import ConnectionPool
+
+
+async def _make_conn() -> MagicMock:
+    mock = MagicMock()
+    mock.is_connected = True
+    mock.connect = AsyncMock()
+    mock.close = AsyncMock()
+    mock._in_transaction = False
+    mock._pool_released = False
+    mock._tx_owner = None
+    mock.execute = AsyncMock()
+    return mock
+
+
+class TestPoolSizeInvariants:
+    async def test_concurrent_acquires_respect_max_size(self) -> None:
+        """10 concurrent acquires with max_size=3 must create at most 3.
+
+        With the pre-fix implementation that held the lock across the
+        network call, this would pass — but slowly. The new
+        reservation pattern drops the lock for the create; this test
+        ensures it still respects the cap.
+        """
+        pool = ConnectionPool(["localhost:9001"], min_size=0, max_size=3, timeout=5.0)
+
+        create_count = 0
+        peak_size = 0
+
+        async def slow_create(**kwargs):
+            nonlocal create_count, peak_size
+            create_count += 1
+            peak_size = max(peak_size, pool._size)
+            # Simulate TCP handshake latency; allow other tasks to run.
+            for _ in range(5):
+                await asyncio.sleep(0)
+            return await _make_conn()
+
+        with patch.object(pool._cluster, "connect", side_effect=slow_create):
+
+            async def borrow() -> None:
+                async with pool.acquire() as c:
+                    assert c is not None
+
+            await asyncio.gather(*(borrow() for _ in range(10)))
+
+        assert create_count == 3, f"expected 3 creates (max_size), got {create_count}"
+        assert peak_size <= pool._max_size, (
+            f"_size exceeded max_size: peak={peak_size}, max={pool._max_size}"
+        )
+        await pool.close()
+
+    async def test_create_failure_rolls_back_reservation(self) -> None:
+        """If _create_connection raises, the reservation must be released
+        so a subsequent acquire can succeed."""
+        pool = ConnectionPool(["localhost:9001"], min_size=0, max_size=1, timeout=1.0)
+
+        call_count = 0
+
+        async def flaky_create(**kwargs):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise OSError("boom")
+            return await _make_conn()
+
+        with patch.object(pool._cluster, "connect", side_effect=flaky_create):
+
+            async def borrow() -> None:
+                async with pool.acquire() as c:
+                    assert c is not None
+
+            # First call fails — reservation must be released.
+            import pytest
+
+            with pytest.raises(OSError):
+                await borrow()
+            assert pool._size == 0, f"reservation leaked after create failure; _size={pool._size}"
+
+            # Second call should succeed.
+            await borrow()
+
+        await pool.close()
+
+    async def test_drained_idle_releases_size(self) -> None:
+        """_drain_idle must decrement _size by exactly the number of drained connections."""
+        pool = ConnectionPool(["localhost:9001"], min_size=3, max_size=5, timeout=1.0)
+
+        async def ok_create(**kwargs):
+            return await _make_conn()
+
+        with patch.object(pool._cluster, "connect", side_effect=ok_create):
+            await pool.initialize()
+            assert pool._size == 3
+
+            await pool._drain_idle()
+            assert pool._size == 0, f"drain_idle left _size={pool._size}"
+
+        await pool.close()