diff --git a/src_py/async_connection.py b/src_py/async_connection.py index d53268b..c1a32aa 100644 --- a/src_py/async_connection.py +++ b/src_py/async_connection.py @@ -225,6 +225,8 @@ def close(self) -> None: """ Close all connections and shutdown the thread pool. + Queries still running are interrupted and fail with "Interrupted.". + Note: Call to this method is optional. The connections and thread pool will be closed automatically when the instance is garbage collected. """ diff --git a/src_py/connection.py b/src_py/connection.py index b528195..c881922 100644 --- a/src_py/connection.py +++ b/src_py/connection.py @@ -1,6 +1,7 @@ from __future__ import annotations import contextlib +import functools import inspect import json import re @@ -8,7 +9,7 @@ import uuid import warnings import weakref -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypeVar, cast from weakref import WeakSet from ._backend import get_capi_module, get_pybind_module @@ -308,6 +309,82 @@ def _pybind_param_signature(parameters: dict[str, Any]) -> tuple: ) +_Method = TypeVar("_Method", bound="Callable[..., Any]") + +# Per thread: the connections whose calls are running, innermost last, and a +# marker for each Python UDF entered on top of them. +_active = threading.local() +_UDF = object() + + +def _active_stack() -> list[object]: + stack: list[object] | None = getattr(_active, "stack", None) + if stack is None: + stack = _active.stack = [] + return stack + + +def _called_from_udf(conn: Connection) -> bool: + """Whether this thread is in a UDF entered after its last call on conn.""" + for entry in reversed(_active_stack()): + if entry is conn: + return False + if entry is _UDF: + return True + return False + + +def _tracks_call(method: _Method) -> _Method: + """Count the call as in flight so that close() can wait for it.""" + + @functools.wraps(method) + def wrapper(self: Connection, *args: Any, **kwargs: Any) -> Any: + with self._calls: + if self._closing: + msg = "Connection is closed." + raise RuntimeError(msg) + if self._calls_in_flight and _called_from_udf(self): + # UDFs run inside the query in flight, which holds the + # connection, so the call could wait for itself forever. + msg = "A UDF cannot use a connection while it runs a query." + raise RuntimeError(msg) + self._calls_in_flight += 1 + stack = _active_stack() + stack.append(self) + try: + return method(self, *args, **kwargs) + finally: + stack.pop() + with self._calls: + self._calls_in_flight -= 1 + self._calls.notify_all() + + return cast("_Method", wrapper) + + +def _stacklevel_outside_module() -> int: + """Return the warnings.warn() stacklevel of the first caller outside this module.""" + frame, level = inspect.currentframe(), 0 + while frame is not None and frame.f_code.co_filename == __file__: + frame, level = frame.f_back, level + 1 + return level + + +def _marks_udf_call(udf: Callable[..., Any]) -> Callable[..., Any]: + """Mark the thread as running a UDF so that calls into a busy connection raise.""" + + @functools.wraps(udf) + def wrapper(*args: Any, **kwargs: Any) -> Any: + stack = _active_stack() + stack.append(_UDF) + try: + return udf(*args, **kwargs) + finally: + stack.pop() + + return wrapper + + class Connection: """Connection to a database.""" @@ -329,6 +406,9 @@ def __init__(self, database: Database, num_threads: int = 0): self.database = database self.num_threads = num_threads self.is_closed = False + self._closing = False + self._calls = threading.Condition() + self._calls_in_flight = 0 self._prefer_pybind = False self._query_timeout_ms = 0 self._query_results: WeakSet[QueryResult] = WeakSet() @@ -379,6 +459,7 @@ def _using_pybind_backend(self) -> bool: self.database._use_pybind_backend and get_pybind_module() is not None ) + @_tracks_call def set_max_threads_for_exec(self, num_threads: int) -> None: """ Set the maximum number of threads for executing queries. @@ -404,9 +485,39 @@ def close(self) -> None: Note: Call to this method is optional. The connection will be closed automatically when the object goes out of scope. + + A query running on this connection in another thread is interrupted, + and close() returns once it has left the engine. Called from a UDF + while the connection runs a query, close() raises instead. """ - if self.is_closed: - return + if any(x is self for x in _active_stack()): + msg = "close() cannot be called from a call on the same connection." + raise RuntimeError(msg) + + with self._calls: + # Idempotent under concurrent close(): check-and-set inside the + # lock so two threads cannot both run the teardown tail. + if self.is_closed or self._closing: + return + if self._calls_in_flight and _called_from_udf(self): + # UDFs run inside the query in flight, so waiting for it would + # never end. + msg = "close() cannot be called from a UDF while the connection runs a query." + raise RuntimeError(msg) + self._closing = True + try: + while self._calls_in_flight: + # An interrupt that arrives while the query is still compiling is + # cleared when execution starts, so repeat it until the call returns. + with contextlib.suppress(RuntimeError): + self.interrupt() + self._calls.wait(timeout=0.05) + except BaseException: + # Don't brick the connection if the wait is aborted + # (KeyboardInterrupt, ...): teardown never ran, so allow reuse. + self._closing = False + self._calls.notify_all() + raise for query_result in list(self._query_results): query_result.close() @@ -814,11 +925,9 @@ def _maybe_raise_scan_unsupported_object(self, query: str) -> None: return var_name = match.group(1) - frame = inspect.currentframe() - if frame is None or frame.f_back is None: - return - - caller = frame.f_back.f_back + caller = inspect.currentframe() + while caller is not None and caller.f_code.co_filename == __file__: + caller = caller.f_back if caller is None: return @@ -837,6 +946,7 @@ def _maybe_raise_scan_unsupported_object(self, query: str) -> None: ) raise RuntimeError(msg) + @_tracks_call def execute( self, query: str | PreparedStatement, @@ -932,7 +1042,7 @@ def execute( warnings.warn( drop_result.getErrorMessage(), RuntimeWarning, - stacklevel=2, + stacklevel=_stacklevel_outside_module(), ) finally: self._capi_scan_tables.discard(table_name) @@ -950,6 +1060,7 @@ def execute( all_query_results.append(next_query_result) return all_query_results + @_tracks_call def query_as_arrow(self, query: str, chunk_size: int) -> ArrowQueryResult: """ Execute a query with the native Arrow collector path. @@ -971,6 +1082,7 @@ def query_as_arrow(self, query: str, chunk_size: int) -> ArrowQueryResult: self._register_query_result(current_query_result) return current_query_result + @_tracks_call def _prepare( self, query: str, @@ -1069,6 +1181,7 @@ def _get_rel_table_names(self) -> list[dict[str, Any]]: results.append({"name": name, "src": src_node, "dst": dst_node}) return results + @_tracks_call def set_query_timeout(self, timeout_in_ms: int) -> None: """ Set the query timeout value in ms for executing queries. @@ -1089,8 +1202,11 @@ def interrupt(self) -> None: If there is no currently executing query, this function does nothing. """ - self._connection.interrupt() + for native in (self._connection, self._py_connection): + if native is not None: + native.interrupt() + @_tracks_call def create_function( self, name: str, @@ -1124,12 +1240,16 @@ def create_function( catch_exceptions: Optional[bool] if true, when an exception is thrown from python, the function output will be null Otherwise, the exception will be rethrown + + The UDF must not call a connection that is running a query: such calls raise + RuntimeError, because the query running the UDF may be on that connection. """ if params_type is None: params_type = [] parsed_params_type = [x if type(x) is str else x.value for x in params_type] if type(return_type) is not str: return_type = return_type.value + udf = _marks_udf_call(udf) try: self._connection.create_function( @@ -1154,6 +1274,7 @@ def create_function( catch_exceptions=catch_exceptions, ) + @_tracks_call def remove_function(self, name: str) -> None: """ Remove a User Defined Function (UDF). @@ -1172,6 +1293,7 @@ def remove_function(self, name: str) -> None: self._prefer_pybind = True py_connection.remove_function(name) + @_tracks_call def create_arrow_table( self, table_name: str, @@ -1211,6 +1333,7 @@ def create_arrow_table( raise RuntimeError(query_result_internal.getErrorMessage()) return QueryResult(self, query_result_internal) + @_tracks_call def drop_arrow_table(self, table_name: str) -> QueryResult: """ Drop an Arrow memory-backed table. @@ -1239,6 +1362,7 @@ def drop_arrow_table(self, table_name: str) -> QueryResult: raise RuntimeError(query_result_internal.getErrorMessage()) return QueryResult(self, query_result_internal) + @_tracks_call def create_arrow_rel_table( self, table_name: str, diff --git a/test/capi_xfails.py b/test/capi_xfails.py index c928e6f..61e5cef 100644 --- a/test/capi_xfails.py +++ b/test/capi_xfails.py @@ -75,6 +75,10 @@ "test/test_udf.py::test_udf_null", "test/test_udf.py::test_udf_except", "test/test_udf.py::test_udf_remove", + "test/test_connection_close.py::test_close_with_query_in_flight[close_from_a_udf]", + "test/test_udf_reentrancy.py::test_udf_calling_a_connection[idle_other_connection]", + "test/test_udf_reentrancy.py::test_udf_calling_a_connection[same_connection_on_a_worker]", + "test/test_udf_reentrancy.py::test_udf_calling_a_connection[same_connection_while_binding]", # C API Arrow export for JSON types is not yet implemented. get_as_df() # tests that need the pybind NPArrayWrapper path are xfailed. "test/test_json.py::test_get_as_df_json_scalar", diff --git a/test/test_connection_close.py b/test/test_connection_close.py new file mode 100644 index 0000000..4c62685 --- /dev/null +++ b/test/test_connection_close.py @@ -0,0 +1,182 @@ +from __future__ import annotations + +import subprocess +import sys +import textwrap +from typing import TYPE_CHECKING + +import ladybug as lb +import pytest + +if TYPE_CHECKING: + from pathlib import Path + +# range() is bound once and the two UNWINDs multiply it, so this runs for tens of +# seconds without materializing large lists. +LONG_QUERY = "WITH range(1, 60000) AS r UNWIND r AS x UNWIND r AS y RETURN sum(x * y)" + +# Closing a connection while a query runs used to hang or crash the process, so +# each scenario runs in its own interpreter. +PRELUDE = """ +import asyncio, sys, threading, time +import ladybug as lb + +db = lb.Database(sys.argv[1]) +setup = lb.Connection(db) +setup.execute("CREATE NODE TABLE T(id INT64, PRIMARY KEY(id))") +setup.close() +LONG_QUERY = {long_query!r} +""" + +SCENARIOS = { + "close_from_inside_a_call": """ +conn = lb.Connection(db) + +def close_inside(query, parameters): + conn.close() + +# Stands in for anything that closes the connection from inside one of its own +# calls on the same thread, such as a finalizer that runs mid-call. +conn._rewrite_local_scan_object = close_inside +started = time.monotonic() +try: + conn.execute("RETURN 1") + outcome = "finished" +except RuntimeError as e: + outcome = str(e) +print("CLOSE", time.monotonic() - started) +print("OUTCOME", outcome) +""", + "close_from_a_udf": """ +conn = lb.Connection(db) + +def closer(x: int) -> int: + if x == 3: + conn.close() + return x + +# UDFs run on engine threads, so this is not caught as a call on the same thread. +conn.create_function("closer", closer) +started = time.monotonic() +try: + conn.execute("UNWIND range(1, 5) AS i RETURN closer(i)").get_all() + outcome = "finished" +except RuntimeError as e: + outcome = str(e) +print("CLOSE", time.monotonic() - started) +print("OUTCOME", outcome) +""", + "query_without_parameters": """ +conn = lb.Connection(db) +outcome = [] + +def run(): + try: + conn.execute(LONG_QUERY) + outcome.append("finished") + except RuntimeError as e: + outcome.append(str(e)) + +worker = threading.Thread(target=run) +worker.start() +time.sleep(0.3) +started = time.monotonic() +conn.close() +print("CLOSE", time.monotonic() - started) +worker.join() +print("OUTCOME", outcome[0]) +""", + "query_with_parameters": """ +conn = lb.Connection(db) +outcome = [] + +def run(): + try: + conn.execute(LONG_QUERY.replace("60000", "$n"), {"n": 60000}) + outcome.append("finished") + except RuntimeError as e: + outcome.append(str(e)) + +worker = threading.Thread(target=run) +worker.start() +time.sleep(0.3) +started = time.monotonic() +conn.close() +print("CLOSE", time.monotonic() - started) +worker.join() +print("OUTCOME", outcome[0]) +""", + "write": """ +conn = lb.Connection(db) +outcome = [] + +def run(): + try: + conn.execute("UNWIND range(1, 500000) AS i CREATE (:T {id: i})") + outcome.append("finished") + except RuntimeError as e: + outcome.append(str(e)) + +worker = threading.Thread(target=run) +worker.start() +time.sleep(0.3) +started = time.monotonic() +conn.close() +print("CLOSE", time.monotonic() - started) +worker.join() +print("OUTCOME", outcome[0]) +""", + "async_connection": """ +async def main(): + async_conn = lb.AsyncConnection(db, max_concurrent_queries=2) + task = asyncio.create_task(async_conn.execute(LONG_QUERY)) + await asyncio.sleep(0.3) + started = time.monotonic() + async_conn.close() + print("CLOSE", time.monotonic() - started) + try: + await task + print("OUTCOME finished") + except RuntimeError as e: + print("OUTCOME", e) + +asyncio.run(main()) +""", +} + + +@pytest.mark.parametrize("scenario", sorted(SCENARIOS)) +def test_close_with_query_in_flight(tmp_path: Path, scenario: str) -> None: + script = ( + textwrap.dedent(PRELUDE.format(long_query=LONG_QUERY)) + SCENARIOS[scenario] + ) + try: + proc = subprocess.run( + [sys.executable, "-c", script, str(tmp_path / "db")], + capture_output=True, + text=True, + timeout=120, + ) + except subprocess.TimeoutExpired: + pytest.fail(f"{scenario}: close() hung") + assert proc.returncode == 0, proc.stderr[-2000:] + assert "OUTCOME" in proc.stdout + if scenario == "close_from_inside_a_call": + assert "same connection" in proc.stdout, proc.stdout + if scenario == "close_from_a_udf": + assert "from a UDF" in proc.stdout, proc.stdout + if scenario != "write": + # A running read is interrupted rather than waited for. A write is + # interrupted too, but undoing the rows it already inserted takes seconds, + # so its close time is not bounded here. + close_seconds = float(proc.stdout.split("CLOSE")[1].split()[0]) + assert close_seconds < 5, close_seconds + + +def test_execute_after_close_raises(tmp_path: Path) -> None: + db = lb.Database(tmp_path / "db") + conn = lb.Connection(db) + conn.close() + with pytest.raises(RuntimeError, match="Connection is closed"): + conn.execute("RETURN 1") + db.close() diff --git a/test/test_udf_reentrancy.py b/test/test_udf_reentrancy.py new file mode 100644 index 0000000..e7669fd --- /dev/null +++ b/test/test_udf_reentrancy.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import subprocess +import sys +import textwrap +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from pathlib import Path + +# A UDF that calls back into the connection running its query used to wait for +# that query forever, so each scenario runs in its own interpreter. +PRELUDE = """ +import sys +import ladybug as lb + +db = lb.Database(sys.argv[1]) +conn = lb.Connection(db) +other = lb.Connection(db) +target = {target} + +def reenter(x: int) -> int: + return target.execute("RETURN 1").get_next()[0] + x + +conn.create_function("reenter", reenter) +try: + print("OUTCOME", conn.execute({query!r}).get_next()[0]) +except RuntimeError as e: + print("OUTCOME", e) +""" + +SCENARIOS = { + # Arguments from a column: the UDF runs on an engine worker thread. + "same_connection_on_a_worker": ( + "conn", + "UNWIND range(1, 3) AS i RETURN sum(reenter(i))", + "A UDF cannot use a connection while it runs a query.", + ), + # A literal argument: the UDF is folded while binding, on the calling thread. + "same_connection_while_binding": ( + "conn", + "RETURN reenter(1)", + "A UDF cannot use a connection while it runs a query.", + ), + "idle_other_connection": ( + "other", + "UNWIND range(1, 3) AS i RETURN sum(reenter(i))", + "OUTCOME 9", + ), +} + + +@pytest.mark.parametrize("scenario", sorted(SCENARIOS)) +def test_udf_calling_a_connection(tmp_path: Path, scenario: str) -> None: + target, query, expected = SCENARIOS[scenario] + script = textwrap.dedent(PRELUDE.format(target=target, query=query)) + try: + proc = subprocess.run( + [sys.executable, "-c", script, str(tmp_path / "db")], + capture_output=True, + text=True, + timeout=120, + ) + except subprocess.TimeoutExpired: + pytest.fail(f"{scenario}: hung") + assert proc.returncode == 0, proc.stderr[-2000:] + assert expected in proc.stdout, proc.stdout