From 8d0f9060b780b6bc00f341bff1a57e9f19f7689a Mon Sep 17 00:00:00 2001 From: Darek Date: Mon, 5 Oct 2026 22:50:51 +0200 Subject: [PATCH] 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 --- src_py/async_connection.py | 109 +++++++++++++++++++++++++--------- test/test_async_connection.py | 101 +++++++++++++++++++++++++++++-- 2 files changed, 178 insertions(+), 32 deletions(-) diff --git a/src_py/async_connection.py b/src_py/async_connection.py index d53268b..6dbbfad 100644 --- a/src_py/async_connection.py +++ b/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 diff --git a/test/test_async_connection.py b/test/test_async_connection.py index 328a6a6..0e1d9dc 100644 --- a/test/test_async_connection.py +++ b/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()