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

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()