fix: wake pool acquire() waiters promptly on close

pool.close() set _closed and drained idle connections but left tasks
already blocked on the queue waiting up to 500 ms (the polling
interval) to notice — effectively a silent close delay scaled by the
number of waiters.

Race the queue get against a _closed_event that close() now sets.
Re-check _closed at the top of the acquire loop (and under the create
branch's lock) so close takes effect immediately for in-flight waiters
whether the event races them or they loop back around.

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

Author: antoineleclair

src/dqliteclient/pool.py

@@ -84,6 +84,7 @@ def __init__(
         self._size = 0
         self._lock = asyncio.Lock()
         self._closed = False
+        self._closed_event: asyncio.Event | None = None
         self._initialized = False
 
     async def initialize(self) -> None:
@@ -102,6 +103,14 @@ async def _create_connection(self) -> DqliteConnection:
         self._size += 1
         return conn
 
+    def _get_closed_event(self) -> asyncio.Event:
+        """Lazily create the closed Event bound to the running loop."""
+        if self._closed_event is None:
+            self._closed_event = asyncio.Event()
+            if self._closed:
+                self._closed_event.set()
+        return self._closed_event
+
     async def _drain_idle(self) -> None:
         """Close all idle connections in the pool.
 
@@ -132,6 +141,9 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
         conn: DqliteConnection | None = None
 
         while conn is None:
+            if self._closed:
+                raise DqliteConnectionError("Pool is closed")
+
             # Try to get an idle connection from the queue
             try:
                 conn = self._pool.get_nowait()
@@ -141,6 +153,8 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
 
             # Try to create a new connection if under max
             async with self._lock:
+                if self._closed:
+                    raise DqliteConnectionError("Pool is closed")
                 if self._size < self._max_size:
                     conn = await self._create_connection()
                     break
@@ -153,10 +167,28 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
                     f"Timed out waiting for a connection from the pool "
                     f"(max_size={self._max_size}, timeout={self._timeout}s)"
                 )
+            # Race the queue against the closed event so close() wakes
+            # waiters promptly instead of leaving them on the polling loop.
+            closed_event = self._get_closed_event()
+            get_task: asyncio.Task[DqliteConnection] = asyncio.create_task(self._pool.get())
+            closed_task = asyncio.create_task(closed_event.wait())
             try:
-                conn = await asyncio.wait_for(self._pool.get(), timeout=min(remaining, 0.5))
-            except TimeoutError:
-                # Don't fail yet — loop back to re-check _size
+                done, _pending = await asyncio.wait(
+                    {get_task, closed_task},
+                    timeout=min(remaining, 0.5),
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+            finally:
+                if not closed_task.done():
+                    closed_task.cancel()
+            if get_task in done:
+                conn = get_task.result()
+            else:
+                # Either close fired or the poll timer fired; either way,
+                # cancel the queue wait cleanly and let the loop re-check.
+                get_task.cancel()
+                with contextlib.suppress(BaseException):
+                    await get_task
                 continue
 
         # If connection is dead, discard and create a fresh one with leader discovery.
@@ -277,6 +309,8 @@ async def close(self) -> None:
         in-flight tasks before calling close().
         """
         self._closed = True
+        if self._closed_event is not None:
+            self._closed_event.set()
         await self._drain_idle()
 
         # In-use connections are closed by acquire()'s cleanup when they

tests/test_pool.py

@@ -736,6 +736,64 @@ async def track_exec(*args, **kwargs):
         assert result is False
         assert not exec_called, "ROLLBACK must not be sent on a dead socket"
 
+    async def test_close_wakes_waiter_promptly(self) -> None:
+        """A task blocked in acquire() waiting for a connection must be
+        woken quickly when pool.close() is called — not sit on the queue
+        until its timeout expires.
+        """
+        import asyncio
+        import time
+
+        from dqliteclient.connection import DqliteConnection
+
+        pool = ConnectionPool(["localhost:9001"], max_size=1, timeout=5.0)
+
+        mock_conn = MagicMock(spec=DqliteConnection)
+        mock_conn.is_connected = True
+        mock_conn.close = AsyncMock()
+        mock_conn._in_transaction = False
+        mock_conn._in_use = False
+        mock_conn._bound_loop = None
+        mock_conn._pool_released = False
+        mock_conn._check_in_use = MagicMock()
+
+        with patch.object(pool._cluster, "connect", return_value=mock_conn):
+            await pool.initialize()
+
+            holder_acquired = asyncio.Event()
+            release_holder = asyncio.Event()
+
+            async def hold_connection() -> None:
+                async with pool.acquire():
+                    holder_acquired.set()
+                    await release_holder.wait()
+
+            holder = asyncio.create_task(hold_connection())
+            await holder_acquired.wait()
+
+            async def try_acquire() -> BaseException | None:
+                try:
+                    async with pool.acquire():
+                        return None
+                except BaseException as e:  # noqa: BLE001
+                    return e
+
+            waiter = asyncio.create_task(try_acquire())
+            await asyncio.sleep(0.05)
+
+            t0 = time.monotonic()
+            await pool.close()
+            err = await asyncio.wait_for(waiter, timeout=1.0)
+            elapsed = time.monotonic() - t0
+
+            release_holder.set()
+            await holder
+
+        assert isinstance(err, DqliteConnectionError)
+        assert elapsed < 0.3, (
+            f"acquire() should wake within ~100ms of pool.close(); took {elapsed:.3f}s"
+        )
+
     async def test_reset_connection_returns_false_on_cancelled_error(self) -> None:
         """_reset_connection must return False (not raise) when ROLLBACK is cancelled."""
         import asyncio