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