fix(async): propagate cancellation and interrupt only the cancelled query

AsyncConnection.execute() caught asyncio.CancelledError, interrupted the
connection and returned None. As a result:

- a cancelled task looked like a normal return, so loops kept running and
  asyncio.wait_for()/asyncio.timeout() returned None instead of raising
  TimeoutError;
- cancelling a task that was still queued interrupted whatever query was
  running on its connection, i.e. another task's query;
- the connection counter was released when the awaiting task exited, while
  the worker could still be running the query.

Now the call is submitted to the executor directly and CancelledError is
re-raised. Calls on one connection run one at a time (the engine serializes
them anyway), and the connection is interrupted only while the cancelled call
is the one running on it. A cancelled call that has not started is skipped.
The interrupt is repeated until the call returns, because an interrupt that
arrives while the query is still compiling is cleared when execution starts.
The counter is released from the executor future's done callback.

execute() with a PreparedStatement whose connection is not part of the pool
no longer raises UnboundLocalError.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: 2 people

src_py/async_connection.py

@@ -3,6 +3,7 @@
 import asyncio
 import threading
 import warnings
+import weakref
 from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, Any
 
@@ -51,6 +52,14 @@ def __init__(
         self.connections = [Connection(database) for _ in range(max_concurrent_queries)]
         self.connections_counter = [0 for _ in range(max_concurrent_queries)]
         self.lock = threading.Lock()
+        # Calls on one connection run one at a time (the engine serializes them
+        # anyway), so the call holding a connection's lock is the one an
+        # interrupt would hit. Queries run directly on a connection obtained
+        # from acquire_connection() are not tracked.
+        self._connection_locks: weakref.WeakKeyDictionary[
+            Connection, threading.Lock
+        ] = weakref.WeakKeyDictionary()
+        self._running: dict[Connection, _Call] = {}
 
         for conn in self.connections:
             conn.init_connection()
@@ -114,10 +123,9 @@ def release_connection(self, conn: Connection) -> None:
 
 
         """
-        for i, existing_conn in enumerate(self.connections):
-            if existing_conn == conn:
-                self.__decrement_connection_counter(i)
-                break
+        conn_index = self.__index_of(conn)
+        if conn_index is not None:
+            self.__decrement_connection_counter(conn_index)
 
     def set_query_timeout(self, timeout_in_ms: int) -> None:
         """
@@ -153,28 +161,22 @@ async def execute(
         QueryResult
             Query result.
 
+        Raises
+        ------
+        asyncio.CancelledError
+            If the awaiting task is cancelled. A running query is interrupted; a
+            query still waiting for a free connection is not started.
+
         """
-        loop = asyncio.get_running_loop()
-        # If the query is a prepared statement, use the connection associated with it
         if isinstance(query, PreparedStatement):
             conn = query._connection
-            for i, existing_conn in enumerate(self.connections):
-                if existing_conn == conn:
-                    conn_index = i
-                    with self.lock:
-                        self.connections_counter[conn_index] += 1
-                    break
+            conn_index = self.__index_of(conn)
+            if conn_index is not None:
+                with self.lock:
+                    self.connections_counter[conn_index] += 1
         else:
             conn, conn_index = self.__get_connection_with_least_queries()
-
-        try:
-            return await loop.run_in_executor(
-                self.executor, conn.execute, query, parameters
-            )
-        except asyncio.CancelledError:
-            conn.interrupt()
-        finally:
-            self.__decrement_connection_counter(conn_index)
+        return await self.__submit(conn, conn_index, conn.execute, query, parameters)
 
     async def _prepare(
         self, query: str, parameters: dict[str, Any] | None = None
@@ -183,16 +185,64 @@ async def _prepare(
         The only parameters supported during prepare are dataframes.
         Any remaining parameters will be ignored and should be passed to execute().
         """  # noqa: D401
-        loop = asyncio.get_running_loop()
         conn, conn_index = self.__get_connection_with_least_queries()
+        return await self.__submit(conn, conn_index, conn._prepare, query, parameters)
 
+    def __index_of(self, conn: Connection) -> int | None:
+        for i, existing_conn in enumerate(self.connections):
+            if existing_conn is conn:
+                return i
+        return None
+
+    async def __submit(
+        self, conn: Connection, conn_index: int | None, fn: Any, *args: Any
+    ) -> Any:
+        call = _Call()
         try:
-            preparedStatement = await loop.run_in_executor(
-                self.executor, conn.prepare, query, parameters
+            future = self.executor.submit(self.__run, conn, call, fn, *args)
+        except RuntimeError:
+            if conn_index is not None:
+                self.__decrement_connection_counter(conn_index)
+            raise
+        if conn_index is not None:
+            # Release the slot when the work has actually finished (or was
+            # cancelled before it started), not when the awaiting task exits.
+            future.add_done_callback(
+                lambda _: self.__decrement_connection_counter(conn_index)
             )
-            return preparedStatement
-        finally:
-            self.__decrement_connection_counter(conn_index)
+        try:
+            return await asyncio.wrap_future(future)
+        except asyncio.CancelledError:
+            self.__cancel(conn, call)
+            raise
+
+    def __run(self, conn: Connection, call: _Call, fn: Any, *args: Any) -> Any:
+        with self.lock:
+            connection_lock = self._connection_locks.setdefault(conn, threading.Lock())
+        with connection_lock:
+            with self.lock:
+                if call.cancelled:
+                    return None
+                self._running[conn] = call
+            try:
+                return fn(*args)
+            finally:
+                with self.lock:
+                    del self._running[conn]
+
+    def __cancel(self, conn: Connection, call: _Call) -> None:
+        # Interrupting acts on whatever the connection is running, so only
+        # interrupt while this call is the one running. A call still waiting
+        # for its connection is skipped when its turn comes.
+        with self.lock:
+            call.cancelled = True
+            running = self._running.get(conn) is call
+            if running:
+                conn.interrupt()
+        if running:
+            # An interrupt that arrives while the query is still compiling is
+            # cleared when execution starts, so repeat it until the call returns.
+            asyncio.get_running_loop().call_later(0.05, self.__cancel, conn, call)
 
     async def prepare(
         self, query: str, parameters: dict[str, Any] | None = None
@@ -232,3 +282,8 @@ def close(self) -> None:
             conn.close()
 
         self.executor.shutdown(wait=True)
+
+
+class _Call:
+    def __init__(self) -> None:
+        self.cancelled = False

test/test_async_connection.py

@@ -1,5 +1,6 @@
 import asyncio
 import time
+from typing import Any
 
 import ladybug as lb
 import pyarrow as pa
@@ -159,12 +160,102 @@ def test_acquire_connection(async_connection_readonly):
         assert i == 0
 
 
+# range() is bound once and the two UNWINDs multiply it, so these run for long
+# enough to cancel mid-flight without materializing large lists.
+LONG_QUERY = "WITH range(1, 30000) AS r UNWIND r AS x UNWIND r AS y RETURN sum(x * y)"
+
+
+async def _wait_until_idle(
+    async_connection: lb.AsyncConnection, timeout: float = 5.0
+) -> None:
+    deadline = time.monotonic() + timeout
+    while any(async_connection.connections_counter):
+        assert time.monotonic() < deadline, async_connection.connections_counter
+        await asyncio.sleep(0.01)
+
+
 @pytest.mark.asyncio
-async def test_async_connection_interrupt(async_connection_readonly) -> None:
-    query = "UNWIND RANGE(1,1000000) AS x UNWIND RANGE(1, 1000000) AS y RETURN COUNT(x + y);"
-    async_connection_readonly.set_query_timeout(100 * 1000)
-    task = asyncio.create_task(async_connection_readonly.execute(query))
-    time.sleep(5)
+async def test_async_cancel_running_query_raises(
+    async_connection_readonly: lb.AsyncConnection,
+) -> None:
+    task = asyncio.create_task(async_connection_readonly.execute(LONG_QUERY))
+    await asyncio.sleep(0.3)
     task.cancel()
     with pytest.raises(asyncio.CancelledError):
         await task
+    assert task.cancelled()
+    # The query was interrupted, so its connection frees up long before the
+    # query could have finished on its own.
+    await _wait_until_idle(async_connection_readonly)
+
+
+@pytest.mark.asyncio
+async def test_async_wait_for_timeout_raises(
+    async_connection_readonly: lb.AsyncConnection,
+) -> None:
+    with pytest.raises(asyncio.TimeoutError):
+        await asyncio.wait_for(async_connection_readonly.execute(LONG_QUERY), 0.3)
+    await _wait_until_idle(async_connection_readonly)
+
+
+@pytest.mark.asyncio
+async def test_async_cancel_queued_query_spares_running_query(
+    async_connection_readonly: lb.AsyncConnection,
+) -> None:
+    async_connection = lb.AsyncConnection(
+        async_connection_readonly.database,
+        max_concurrent_queries=1,
+        max_threads_per_query=4,
+    )
+    try:
+        running = asyncio.create_task(async_connection.execute(LONG_QUERY))
+        queued = asyncio.create_task(async_connection.execute("RETURN 1;"))
+        await asyncio.sleep(0.1)
+        queued.cancel()
+        await asyncio.sleep(0.1)
+        running_alive = not running.done()
+        running.cancel()
+        outcomes = await asyncio.gather(queued, running, return_exceptions=True)
+        assert all(isinstance(o, asyncio.CancelledError) for o in outcomes), outcomes
+        assert running_alive
+        await _wait_until_idle(async_connection)
+    finally:
+        async_connection.close()
+
+
+@pytest.mark.asyncio
+async def test_async_cancel_during_compile_still_interrupts(
+    async_connection_readonly: lb.AsyncConnection,
+) -> None:
+    # An interrupt sent before execution starts is cleared by the engine, so a
+    # cancel that lands in that window must be repeated.
+    conn = async_connection_readonly.connections[0]
+    execute = conn.execute
+
+    def slow_start(*args: Any, **kwargs: Any) -> Any:
+        time.sleep(0.2)
+        return execute(*args, **kwargs)
+
+    conn.execute = slow_start
+    task = asyncio.create_task(async_connection_readonly.execute(LONG_QUERY))
+    await asyncio.sleep(0.1)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    started = time.monotonic()
+    await _wait_until_idle(async_connection_readonly)
+    assert time.monotonic() - started < 2
+
+
+@pytest.mark.asyncio
+async def test_async_execute_prepared_statement_from_other_connection(
+    async_connection_readonly: lb.AsyncConnection,
+) -> None:
+    conn = lb.Connection(async_connection_readonly.database)
+    prepared_statement = conn._prepare("RETURN $x;")
+    result = await async_connection_readonly.execute(prepared_statement, {"x": 7})
+    assert result.get_next() == [7]
+    result.close()
+    for i in async_connection_readonly.connections_counter:
+        assert i == 0
+    conn.close()