fix: drain idle connections when a broken connection is detected

When one connection to the dqlite server is found to be broken (e.g.,
after a leader change or server restart), other idle connections in
the pool are likely stale too — they point to the same server.

Previously, only the single broken connection was discarded. Other
idle connections would be handed out and fail one-by-one, each
requiring a full leader discovery + reconnect cycle. In a pool of
size N, this caused N-1 unnecessary failures after a leader change.

Add _drain_idle() which closes all idle connections in the queue.
Call it whenever a broken connection is detected — both at checkout
(dead connection replacement) and during use (connection invalidated
by execute/fetch error handlers). Subsequent acquire() calls create
fresh connections via leader discovery to the new leader.

Closes #093 (Phase 1: bulk invalidation)

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

Author: antoineleclair

src/dqliteclient/pool.py

@@ -73,6 +73,22 @@ async def _create_connection(self) -> DqliteConnection:
         self._size += 1
         return conn
 
+    async def _drain_idle(self) -> None:
+        """Close all idle connections in the pool.
+
+        Called when a connection is found to be broken (e.g., after a
+        leader change or server restart), since other idle connections
+        are likely stale too.
+        """
+        while not self._pool.empty():
+            try:
+                conn = self._pool.get_nowait()
+                with contextlib.suppress(Exception):
+                    await conn.close()
+                self._size -= 1
+            except asyncio.QueueEmpty:
+                break
+
     @asynccontextmanager
     async def acquire(self) -> AsyncIterator[DqliteConnection]:
         """Acquire a connection from the pool."""
@@ -113,8 +129,10 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
                 # 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 connection is dead, discard and create a fresh one with leader discovery.
+        # Also drain other idle connections — they likely point to the same dead server.
         if not conn.is_connected:
+            await self._drain_idle()
             async with self._lock:
                 self._size -= 1
                 conn = await self._create_connection()
@@ -134,6 +152,8 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
                     self._size -= 1
             else:
                 # Connection is broken (invalidated by execute/fetch error handlers).
+                # Drain other idle connections — they likely point to the same dead server.
+                await self._drain_idle()
                 with contextlib.suppress(BaseException):
                     await conn.close()
                 self._size -= 1

tests/test_pool.py

@@ -342,6 +342,48 @@ async def tracking_create():
 
         await pool.close()
 
+    async def test_broken_connection_drains_idle_pool(self) -> None:
+        """When a connection breaks, all idle connections should be drained."""
+        pool = ConnectionPool(["localhost:9001"], min_size=3, max_size=5)
+
+        conns = []
+        for _ in range(3):
+            mock_conn = MagicMock()
+            mock_conn.is_connected = True
+            mock_conn.connect = AsyncMock()
+            mock_conn.close = AsyncMock()
+            conns.append(mock_conn)
+
+        new_conn = MagicMock()
+        new_conn.is_connected = True
+        new_conn.connect = AsyncMock()
+        new_conn.close = AsyncMock()
+
+        all_conns = [*conns, new_conn]
+        conn_iter = iter(all_conns)
+        with patch.object(pool._cluster, "connect", side_effect=lambda **kw: next(conn_iter)):
+            await pool.initialize()
+
+            assert pool._size == 3
+            assert pool._pool.qsize() == 3
+
+            # Acquire a connection — the first one from the queue
+            with pytest.raises(DqliteConnectionError):
+                async with pool.acquire() as acquired:
+                    # Mark as broken (simulates leader change / server error)
+                    acquired.is_connected = False
+                    raise DqliteConnectionError("connection lost")
+
+        # The broken connection was discarded. The 2 idle connections should
+        # also have been drained (they likely point to the same dead server).
+        idle_closed = sum(1 for c in conns[1:] if c.close.called)
+        assert idle_closed == 2, (
+            f"Expected 2 idle connections to be drained, but only {idle_closed} were closed"
+        )
+        assert pool._pool.qsize() == 0
+
+        await pool.close()
+
     async def test_dead_connection_triggers_leader_rediscovery(self) -> None:
         """A dead connection should be replaced via leader discovery, not reconnected."""
         pool = ConnectionPool(["localhost:9001"], max_size=1)