diff --git a/src_cpp/include/cached_import/py_cached_import.h b/src_cpp/include/cached_import/py_cached_import.h index 734550b..b85c74a 100644 --- a/src_cpp/include/cached_import/py_cached_import.h +++ b/src_cpp/include/cached_import/py_cached_import.h @@ -25,6 +25,8 @@ class PythonCachedImport { PandasCachedItem pandas; PolarsCachedItem polars; PyarrowCachedItem pyarrow; + SignalCachedItem signal; + ThreadingCachedItem threading; UUIDCachedItem uuid; private: diff --git a/src_cpp/include/cached_import/py_cached_modules.h b/src_cpp/include/cached_import/py_cached_modules.h index 0d18bde..dbf840e 100644 --- a/src_cpp/include/cached_import/py_cached_modules.h +++ b/src_cpp/include/cached_import/py_cached_modules.h @@ -154,6 +154,24 @@ class PyarrowCachedItem : public PythonCachedItem { LibCachedItem lib; }; +class SignalCachedItem : public PythonCachedItem { +public: + SignalCachedItem() + : PythonCachedItem("signal"), getsignal("getsignal", this), sigint("SIGINT", this), + default_int_handler("default_int_handler", this) {} + + PythonCachedItem getsignal; + PythonCachedItem sigint; + PythonCachedItem default_int_handler; +}; + +class ThreadingCachedItem : public PythonCachedItem { +public: + ThreadingCachedItem() : PythonCachedItem("threading"), main_thread("main_thread", this) {} + + PythonCachedItem main_thread; +}; + class UUIDCachedItem : public PythonCachedItem { public: UUIDCachedItem() : PythonCachedItem("uuid"), UUID("UUID", this) {} diff --git a/src_cpp/py_connection.cpp b/src_cpp/py_connection.cpp index de7c09d..d330063 100644 --- a/src_cpp/py_connection.cpp +++ b/src_cpp/py_connection.cpp @@ -1,9 +1,18 @@ #include "include/py_connection.h" #include +#include #include #include +#ifndef _WIN32 +#include +#endif + +#if PY_VERSION_HEX < 0x030D0000 +#define Py_IsFinalizing _Py_IsFinalizing +#endif + #include "cached_import/py_cached_import.h" #include "common/constants.h" #include "common/exception/not_implemented.h" @@ -183,6 +192,106 @@ void PyConnection::interrupt() { static std::unordered_map> transformPythonParameters( const py::dict& params, Connection* conn); +namespace { + +// Python runs its SIGINT handler only between bytecodes, so not while the main thread waits in +// the engine. During such a wait, SIGINT interrupts the connection's query and then calls +// Python's handler. This is only done on the main thread while Python's default handler +// (KeyboardInterrupt) is active, so custom handlers keep their behaviour. Only query(), +// execute() and query_as_arrow() install it. +// +// target is set before the handler is installed and cleared after the previous one is +// restored, so a SIGINT in between reaches Python's handler, as without this guard. A signal +// mask would not help: masks are per thread, and a process-directed SIGINT runs the handler on +// whichever thread does not block it. The statics are per process; this single-phase-init +// module cannot be imported into subinterpreters. +class SigintInterruptsQuery { +public: + explicit SigintInterruptsQuery([[maybe_unused]] main::ClientContext* context) { +#ifndef _WIN32 + if (target.load() != nullptr || !onMainThreadWithDefaultHandler()) { + return; + } + if (sigaction(SIGINT, nullptr, &previous) != 0) { + return; + } + struct sigaction action {}; + action.sa_sigaction = handle; + action.sa_flags = SA_SIGINFO | (previous.sa_flags & SA_ONSTACK); + sigemptyset(&action.sa_mask); + received = false; + target = context; + installed = sigaction(SIGINT, &action, nullptr) == 0; + if (!installed) { + target = nullptr; + } +#endif + } + + ~SigintInterruptsQuery() { +#ifndef _WIN32 + if (installed) { + sigaction(SIGINT, &previous, nullptr); + target = nullptr; + } +#endif + } + + // Call with the GIL held once the query has returned. + void raiseIfInterrupted() const { +#ifndef _WIN32 + if (installed && received && PyErr_CheckSignals() != 0) { + throw py::error_already_set(); + } +#endif + } + +private: +#ifndef _WIN32 + static bool onMainThreadWithDefaultHandler() { + // The check must never fail a query, e.g. once imports stop working at shutdown. + try { + auto mainThread = importCache->threading.main_thread()(); + if (mainThread.attr("ident").cast() != PyThread_get_thread_ident()) { + return false; + } + auto& signal = importCache->signal; + return signal.getsignal()(signal.sigint()).is(signal.default_int_handler()); + } catch (py::error_already_set& e) { + if (!Py_IsFinalizing()) { + e.discard_as_unraisable("checking the SIGINT handler before a query"); + } + return false; + } + } + + // Runs in signal context, so it may only use async-signal-safe operations: lock-free atomics + // and Python's own handler, which is written for signal context. ClientContext::interrupt() + // is a store to an std::atomic in the engine and must stay one. + static void handle(int signum, siginfo_t* info, void* ucontext) { + if (auto* context = target.load()) { + context->interrupt(); + } + received = true; + // Python's handler only records the signal; KeyboardInterrupt is raised later. + if (previous.sa_flags & SA_SIGINFO) { + previous.sa_sigaction(signum, info, ucontext); + } else if (previous.sa_handler != SIG_DFL && previous.sa_handler != SIG_IGN) { + previous.sa_handler(signum); + } + } + + static inline std::atomic target{nullptr}; + static inline std::atomic received{false}; + static_assert(std::atomic::is_always_lock_free); + static_assert(std::atomic::is_always_lock_free); + static inline struct sigaction previous {}; + bool installed = false; +#endif +}; + +} // namespace + std::unique_ptr PyConnection::execute(PyPreparedStatement* preparedStatement, const py::dict& params) { auto& stateRef = refState(); @@ -190,34 +299,47 @@ std::unique_ptr PyConnection::execute(PyPreparedStatement* prepar throw RuntimeException("Prepared statement is closed."); } auto parameters = transformPythonParameters(params, &stateRef.ref()); + // Keeps the ClientContext alive past the guard, even if another thread closes this + // connection while the GIL is released. + auto keepAlive = state; + SigintInterruptsQuery sigint(stateRef.ref().getClientContext()); py::gil_scoped_release release; auto queryResult = stateRef.ref().executeWithParams(&preparedStatement->state->ref(), std::move(parameters)); py::gil_scoped_acquire acquire; - return checkAndWrapQueryResult(queryResult, state); + sigint.raiseIfInterrupted(); + return checkAndWrapQueryResult(queryResult, keepAlive); } std::unique_ptr PyConnection::query(const std::string& statement) { auto& stateRef = refState(); + auto keepAlive = state; + SigintInterruptsQuery sigint(stateRef.ref().getClientContext()); py::gil_scoped_release release; auto queryResult = stateRef.ref().query(statement); py::gil_scoped_acquire acquire; - return checkAndWrapQueryResult(queryResult, state); + sigint.raiseIfInterrupted(); + return checkAndWrapQueryResult(queryResult, keepAlive); } std::unique_ptr PyConnection::queryAsArrow(const std::string& statement, int64_t chunkSize) { auto& stateRef = refState(); + auto keepAlive = state; + SigintInterruptsQuery sigint(stateRef.ref().getClientContext()); py::gil_scoped_release release; auto queryResult = stateRef.ref().queryAsArrow(statement, chunkSize); py::gil_scoped_acquire acquire; - return checkAndWrapQueryResult(queryResult, state); + sigint.raiseIfInterrupted(); + return checkAndWrapQueryResult(queryResult, keepAlive); } void PyConnection::setMaxNumThreadForExec(uint64_t numThreads) { refState().ref().setMaxNumThreadForExec(numThreads); } +// Not guarded by SigintInterruptsQuery: a Ctrl-C while the statement compiles takes effect when +// prepare() returns. PyPreparedStatement PyConnection::prepare(const std::string& query, const py::dict& parameters) { auto& stateRef = refState(); auto params = transformPythonParameters(parameters, &stateRef.ref()); diff --git a/test/capi_xfails.py b/test/capi_xfails.py index c928e6f..fe82ec7 100644 --- a/test/capi_xfails.py +++ b/test/capi_xfails.py @@ -69,6 +69,12 @@ # C API scan rewriting uses temporary Arrow-backed tables, which cannot be # created through a read-only connection. "test/test_scan_pyarrow.py::test_pyarrow_basic", + # Ctrl-C interrupts a running query only on the pybind backend. + "test/test_keyboard_interrupt.py::test_keyboard_interrupt_stops_the_query[execute_with_parameters]", + "test/test_keyboard_interrupt.py::test_keyboard_interrupt_stops_the_query[query]", + "test/test_keyboard_interrupt.py::test_keyboard_interrupt_stops_the_query[query_as_arrow]", + "test/test_keyboard_interrupt.py::test_keyboard_interrupt_after_close_from_another_thread", + "test/test_keyboard_interrupt.py::test_failed_handler_check_is_reported_not_raised", # UDF registration is still routed through pybind. "test/test_blob_parameter.py::test_bytes_param_udf", "test/test_udf.py::test_udf", diff --git a/test/test_keyboard_interrupt.py b/test/test_keyboard_interrupt.py new file mode 100644 index 0000000..a43a5f0 --- /dev/null +++ b/test/test_keyboard_interrupt.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +import subprocess +import sys +import textwrap +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from pathlib import Path + +# Each scenario sends SIGINT to itself one second into a query that runs for many +# seconds, and reports how the query ended. +PRELUDE = """ +import os, signal, sys, threading, time +import ladybug as lb + +db = lb.Database(sys.argv[1]) +conn = lb.Connection(db) +LONG_QUERY = "WITH range(1, $n) AS r UNWIND r AS x UNWIND r AS y RETURN sum(x * y)" +threading.Timer(1.0, os.kill, (os.getpid(), signal.SIGINT)).start() +""" + +SCENARIOS = { + "query": "conn.execute(LONG_QUERY.replace('$n', '40000'))", + "execute_with_parameters": "conn.execute(LONG_QUERY, {'n': 40000})", + "query_as_arrow": "conn.query_as_arrow(LONG_QUERY.replace('$n', '40000'), 1000)", +} + +RUN = """ +started = time.monotonic() +try: + {call} + print("OUTCOME finished") +except KeyboardInterrupt: + print("OUTCOME KeyboardInterrupt") +print("ELAPSED", time.monotonic() - started) +print("AFTER", conn.execute("RETURN 1").get_next()[0]) +""" + +CUSTOM_HANDLER = """ +received = [] +signal.signal(signal.SIGINT, lambda signum, frame: received.append(signum)) +result = conn.execute(LONG_QUERY.replace("$n", "20000")).get_next()[0] +time.sleep(0.2) +print("OUTCOME", "finished" if result else "empty", "handled" if received else "missed") +""" + + +def run_script(tmp_path: Path, body: str) -> str: + script = textwrap.dedent(PRELUDE) + textwrap.dedent(body) + try: + proc = subprocess.run( + [sys.executable, "-c", script, str(tmp_path / "db")], + capture_output=True, + text=True, + timeout=120, + ) + except subprocess.TimeoutExpired: + pytest.fail("the query was not interrupted") + assert proc.returncode == 0, proc.stderr[-2000:] + return proc.stdout + + +@pytest.mark.skipif(sys.platform == "win32", reason="SIGINT is POSIX only") +@pytest.mark.parametrize("scenario", sorted(SCENARIOS)) +def test_keyboard_interrupt_stops_the_query(tmp_path: Path, scenario: str) -> None: + stdout = run_script(tmp_path, RUN.format(call=SCENARIOS[scenario])) + assert "OUTCOME KeyboardInterrupt" in stdout, stdout + assert float(stdout.split("ELAPSED")[1].split()[0]) < 10, stdout + assert "AFTER 1" in stdout, stdout + + +@pytest.mark.skipif(sys.platform == "win32", reason="SIGINT is POSIX only") +def test_custom_sigint_handler_still_runs(tmp_path: Path) -> None: + stdout = run_script(tmp_path, CUSTOM_HANDLER) + # A custom handler is left alone: the query finishes, then the handler runs. + assert "OUTCOME finished handled" in stdout, stdout + + +CLOSE_DURING_QUERY = """ +threading.Timer(0.5, conn._connection.close).start() +started = time.monotonic() +try: + conn.execute(LONG_QUERY.replace('$n', '40000')) + print("OUTCOME finished") +except KeyboardInterrupt: + print("OUTCOME KeyboardInterrupt") +print("ELAPSED", time.monotonic() - started) +""" + + +@pytest.mark.skipif(sys.platform == "win32", reason="SIGINT is POSIX only") +def test_keyboard_interrupt_after_close_from_another_thread(tmp_path: Path) -> None: + # The running call keeps the connection alive, so closing it from another thread + # returns at once instead of waiting for the query with the GIL held. (Whether the + # context was freed under the running query would need ASan to observe.) + stdout = run_script(tmp_path, CLOSE_DURING_QUERY) + assert "OUTCOME KeyboardInterrupt" in stdout, stdout + assert float(stdout.split("ELAPSED")[1].split()[0]) < 10, stdout + + +@pytest.mark.skipif(sys.platform == "win32", reason="SIGINT is POSIX only") +def test_failed_handler_check_is_reported_not_raised(tmp_path: Path) -> None: + script = """ +import signal, sys +import ladybug as lb + +def broken(signum): + raise RuntimeError("getsignal is broken") + +signal.getsignal = broken +conn = lb.Connection(lb.Database(sys.argv[1])) +print("RESULT", conn.execute("RETURN 1").get_next()[0]) +""" + proc = subprocess.run( + [sys.executable, "-c", script, str(tmp_path / "db")], + capture_output=True, + text=True, + timeout=120, + ) + assert proc.returncode == 0, proc.stderr[-2000:] + assert "RESULT 1" in proc.stdout, proc.stdout + assert "getsignal is broken" in proc.stderr, proc.stderr