fix: track in-use connections in pool and close them on shutdown

Pool now tracks checked-out connections in an _in_use set. close()
closes both idle and in-use connections. Connections returned after
close() are closed immediately instead of being put back in the queue.

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

Author: antoineleclair

src/dqliteclient/pool.py

@@ -39,6 +39,7 @@ def __init__(
 
         self._cluster = ClusterClient.from_addresses(addresses, timeout=timeout)
         self._pool: asyncio.Queue[DqliteConnection] = asyncio.Queue(maxsize=max_size)
+        self._in_use: set[DqliteConnection] = set()
         self._size = 0
         self._lock = asyncio.Lock()
         self._closed = False
@@ -84,6 +85,7 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
                     f"(max_size={self._max_size}, timeout={self._timeout}s)"
                 ) from None
 
+        self._in_use.add(conn)
         try:
             # Verify connection is still good
             if not conn.is_connected:
@@ -96,16 +98,22 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
 
             with contextlib.suppress(BaseException):
                 await conn.close()
+            self._in_use.discard(conn)
             self._size -= 1
             raise
         else:
+            self._in_use.discard(conn)
             # Return to pool
-            try:
-                self._pool.put_nowait(conn)
-            except asyncio.QueueFull:
-                # Pool full, close connection
+            if self._closed:
                 await conn.close()
                 self._size -= 1
+            else:
+                try:
+                    self._pool.put_nowait(conn)
+                except asyncio.QueueFull:
+                    # Pool full, close connection
+                    await conn.close()
+                    self._size -= 1
 
     async def execute(self, sql: str, params: Sequence[Any] | None = None) -> tuple[int, int]:
         """Execute a SQL statement using a pooled connection."""
@@ -118,16 +126,25 @@ async def fetch(self, sql: str, params: Sequence[Any] | None = None) -> list[dic
             return await conn.fetch(sql, params)
 
     async def close(self) -> None:
-        """Close all connections in the pool."""
+        """Close all connections (both idle and in-use)."""
         self._closed = True
 
+        # Close idle connections
         while not self._pool.empty():
             try:
                 conn = self._pool.get_nowait()
                 await conn.close()
             except asyncio.QueueEmpty:
                 break
 
+        # Close in-use connections
+        for conn in list(self._in_use):
+            import contextlib
+
+            with contextlib.suppress(Exception):
+                await conn.close()
+        self._in_use.clear()
+
         self._size = 0
 
     async def __aenter__(self) -> "ConnectionPool":

tests/test_pool.py

@@ -85,6 +85,35 @@ async def test_acquire_timeout_when_pool_exhausted(self) -> None:
                     pass
 
 
+    async def test_close_handles_checked_out_connections(self) -> None:
+        """close() should close in-flight connections, not just idle ones."""
+        pool = ConnectionPool(["localhost:9001"], max_size=2)
+
+        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)):
+            await pool.initialize()  # Creates mock_conn1
+
+        # Acquire a connection (checks it out)
+        ctx = pool.acquire()
+        conn = await ctx.__aenter__()
+
+        # Close the pool while connection is checked out
+        await pool.close()
+
+        # The checked-out connection should have been closed
+        mock_conn1.close.assert_called()
+
+
 class TestConnectionPoolIntegration:
     """Integration tests requiring mocked connections."""