From fc9883669ca3bea94094312a311e8ced8f5cd1e3 Mon Sep 17 00:00:00 2001 From: Darek Date: Mon, 5 Oct 2026 23:08:22 +0200 Subject: [PATCH 1/2] fix(connection): make close() safe while a query is running Closing a connection while another thread was executing a query on it destroyed the native connection under the running query. Depending on the query this hung the process, aborted it with heap corruption, or segfaulted (a write). AsyncConnection.close() hit the same path, since it closes its connections before shutting down its executor. Connection now counts calls in flight in every method that reaches the native connection. close() refuses new calls, interrupts a running query (repeating the interrupt until the call returns, because an interrupt that arrives while the query is still compiling is cleared when execution starts), and waits for in-flight calls to leave the engine before releasing the native connection. Calls made after close() raise "Connection is closed." The guard lives in the Python layer, so it covers both backends. close() called from inside one of the connection's own calls on the same thread raises instead of waiting for itself. A Python UDF that calls into a connection while it runs a query (close(), execute() or any other call) raises too: UDFs run inside the query in flight, so such calls used to wait for themselves forever, on main as well. A per-thread stack of active calls and UDFs tells these apart from the connection's own nested calls, and covers UDFs folded on the calling thread while binding. UDFs are shared by all connections of a database and the driver cannot tell which query runs one, so a UDF may not call any connection that has a call in flight; create_function() documents this. interrupt() also reaches the pybind connection used as a fallback by the C-API backend, and is a no-op on a closed connection. The scan-object lookup for LOAD FROM no longer assumes a fixed call depth and walks up to the first frame outside this module. Co-Authored-By: Claude Opus 5.5 --- src_py/async_connection.py | 2 + src_py/connection.py | 132 ++++++++++++++++++++++-- test/capi_xfails.py | 4 + test/test_connection_close.py | 182 ++++++++++++++++++++++++++++++++++ test/test_udf_reentrancy.py | 69 +++++++++++++ 5 files changed, 381 insertions(+), 8 deletions(-) create mode 100644 test/test_connection_close.py create mode 100644 test/test_udf_reentrancy.py 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..b1ded0c 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,10 +485,32 @@ 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 self in _active_stack(): + msg = "close() cannot be called from a call on the same connection." + raise RuntimeError(msg) + + with self._calls: + 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 + 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) + for query_result in list(self._query_results): query_result.close() self._query_results.clear() @@ -814,11 +917,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 +938,7 @@ def _maybe_raise_scan_unsupported_object(self, query: str) -> None: ) raise RuntimeError(msg) + @_tracks_call def execute( self, query: str | PreparedStatement, @@ -932,7 +1034,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 +1052,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 +1074,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 +1173,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 +1194,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 +1232,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 +1266,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 +1285,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 +1325,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 +1354,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 From c17381b570591f27ab6dbe2ac22cb6bec9228782 Mon Sep 17 00:00:00 2001 From: Arun Sharma Date: Mon, 12 Oct 2026 13:11:13 -0700 Subject: [PATCH 2/2] fix(connection): make close() idempotent and abort-safe - check-and-set is_closed/_closing inside _calls lock so concurrent close() cannot double-teardown natives - reset _closing on aborted wait so connection stays usable - identity check for re-entrant close() on same connection --- src_py/connection.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/src_py/connection.py b/src_py/connection.py index b1ded0c..c881922 100644 --- a/src_py/connection.py +++ b/src_py/connection.py @@ -490,26 +490,34 @@ def close(self) -> None: 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 self in _active_stack(): + 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 - 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) + 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()