fix: prevent pool starvation by looping acquire with capacity re-check

Previously, acquire() made a one-shot decision: if the pool was at
max capacity, it committed to waiting on the queue for a returned
connection. If connections were dying (not returned) in error paths,
the _size counter decremented but no one notified the blocked waiter.
The waiter stayed stuck on the queue until timeout, even though the
pool had capacity to create a new connection.

Restructure acquire() as a loop with short sub-timeouts (0.5s). After
each sub-timeout, the loop re-checks whether _size < max_size and
creates a new connection if capacity is available. This ensures the
pool self-heals during cluster failover or connection failures.

Closes #081

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

Author: antoineleclair

src/dqliteclient/pool.py

@@ -79,28 +79,39 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
         if self._closed:
             raise DqliteConnectionError("Pool is closed")
 
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + self._timeout
         conn: DqliteConnection | None = None
 
-        # Try to get from pool
-        try:
-            conn = self._pool.get_nowait()
-        except asyncio.QueueEmpty:
-            # Create new if under max
+        while conn is None:
+            # Try to get an idle connection from the queue
+            try:
+                conn = self._pool.get_nowait()
+                break
+            except asyncio.QueueEmpty:
+                pass
+
+            # Try to create a new connection if under max
             async with self._lock:
                 if self._size < self._max_size:
                     conn = await self._create_connection()
+                    break
 
-        # Wait for one if at max
-        if conn is None:
+            # At capacity — wait briefly on the queue, then loop back to
+            # re-check capacity (another coroutine may have freed a slot)
+            remaining = deadline - loop.time()
+            if remaining <= 0:
+                raise DqliteConnectionError(
+                    f"Timed out waiting for a connection from the pool "
+                    f"(max_size={self._max_size}, timeout={self._timeout}s)"
+                )
             try:
                 conn = await asyncio.wait_for(
-                    self._pool.get(), timeout=self._timeout
+                    self._pool.get(), timeout=min(remaining, 0.5)
                 )
             except TimeoutError:
-                raise DqliteConnectionError(
-                    f"Timed out waiting for a connection from the pool "
-                    f"(max_size={self._max_size}, timeout={self._timeout}s)"
-                ) from None
+                # Don't fail yet — loop back to re-check _size
+                continue
 
         # If connection is dead, discard and create a fresh one with leader discovery
         if not conn.is_connected:

tests/test_pool.py

@@ -224,6 +224,76 @@ async def hold_connection():
         # _size must never go negative
         assert pool._size >= 0, f"_size went negative: {pool._size}"
 
+    async def test_acquire_recovers_after_connection_failure(self) -> None:
+        """A waiter blocked on the queue should recover when capacity frees up."""
+        import asyncio
+
+        pool = ConnectionPool(["localhost:9001"], min_size=0, max_size=1, timeout=2.0)
+
+        mock_conn1 = MagicMock()
+        mock_conn1.is_connected = True
+        mock_conn1.connect = AsyncMock()
+        mock_conn1.close = AsyncMock()
+
+        mock_conn2 = MagicMock()
+        mock_conn2.is_connected = True
+        mock_conn2.connect = AsyncMock()
+        mock_conn2.close = AsyncMock()
+
+        conns = iter([mock_conn1, mock_conn2])
+
+        with patch.object(pool._cluster, "connect", side_effect=lambda **kw: next(conns)):
+            # Fill the pool to max_size and keep the connection checked out
+            holder_acquired = asyncio.Event()
+            holder_release = asyncio.Event()
+
+            async def holder():
+                async with pool.acquire() as conn1:
+                    assert conn1 is mock_conn1
+                    holder_acquired.set()
+                    await holder_release.wait()
+                    raise RuntimeError("simulated failure")
+
+            holder_task = asyncio.create_task(holder())
+            await holder_acquired.wait()
+
+            # Now start the waiter — it MUST enter the queue.get() wait
+            # because _size == max_size and the queue is empty
+            waiter_done = asyncio.Event()
+            waiter_result: MagicMock | None = None
+
+            async def waiter():
+                nonlocal waiter_result
+                async with pool.acquire() as c:
+                    waiter_result = c
+                    waiter_done.set()
+
+            waiter_task = asyncio.create_task(waiter())
+            # Give the waiter time to pass get_nowait (empty) and the lock
+            # check (_size == max_size) and enter the queue.get() wait
+            await asyncio.sleep(0.1)
+
+            # Now release the holder — it fails, closing conn and decrementing _size
+            holder_release.set()
+            with contextlib.suppress(RuntimeError):
+                await holder_task
+            # At this point _size == 0. The waiter is blocked on queue.get().
+            # With the fix, the waiter should wake up and create a new connection.
+
+            try:
+                await asyncio.wait_for(waiter_done.wait(), timeout=1.5)
+            except TimeoutError:
+                waiter_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await waiter_task
+                pytest.fail(
+                    "Waiter timed out — pool starvation: waiter stuck on queue "
+                    "even though _size dropped below max_size"
+                )
+
+        assert waiter_result is mock_conn2
+        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."""
         pool = ConnectionPool(["localhost:9001"], min_size=1, max_size=2)