From d87482bc4ee9458d6ba16140e7bc008637dbbb16 Mon Sep 17 00:00:00 2001 From: Jason Fried Date: Mon, 23 Sep 2024 20:40:17 -0700 Subject: [PATCH] gh-119333: Add C api to have contextvar enter/exit callbacks (#119335) Co-authored-by: Erlend E. Aasland --- Doc/c-api/contextvars.rst | 46 ++++++ Include/cpython/context.h | 32 ++++ Include/internal/pycore_context.h | 1 + Include/internal/pycore_interp.h | 2 + Lib/test/test_capi/test_watchers.py | 83 ++++++++++ ...-05-21-18-28-44.gh-issue-119333.OTsYVX.rst | 2 + Modules/_testcapi/watchers.c | 152 ++++++++++++++++++ Python/context.c | 76 +++++++++ Python/pystate.c | 5 + Tools/c-analyzer/cpython/ignored.tsv | 3 + 10 files changed, 402 insertions(+) create mode 100644 Misc/NEWS.d/next/C API/2024-05-21-18-28-44.gh-issue-119333.OTsYVX.rst diff --git a/Doc/c-api/contextvars.rst b/Doc/c-api/contextvars.rst index fe7b8f93f2c..0de135b232a 100644 --- a/Doc/c-api/contextvars.rst +++ b/Doc/c-api/contextvars.rst @@ -101,6 +101,52 @@ Context object management functions: current context for the current thread. Returns ``0`` on success, and ``-1`` on error. +.. c:function:: int PyContext_AddWatcher(PyContext_WatchCallback callback) + + Register *callback* as a context object watcher for the current interpreter. + Return an ID which may be passed to :c:func:`PyContext_ClearWatcher`. + In case of error (e.g. no more watcher IDs available), + return ``-1`` and set an exception. + + .. versionadded:: 3.14 + +.. c:function:: int PyContext_ClearWatcher(int watcher_id) + + Clear watcher identified by *watcher_id* previously returned from + :c:func:`PyContext_AddWatcher` for the current interpreter. + Return ``0`` on success, or ``-1`` and set an exception on error + (e.g. if the given *watcher_id* was never registered.) + + .. versionadded:: 3.14 + +.. c:type:: PyContextEvent + + Enumeration of possible context object watcher events: + - ``Py_CONTEXT_EVENT_ENTER`` + - ``Py_CONTEXT_EVENT_EXIT`` + + .. versionadded:: 3.14 + +.. c:type:: int (*PyContext_WatchCallback)(PyContextEvent event, PyContext* ctx) + + Type of a context object watcher callback function. + If *event* is ``Py_CONTEXT_EVENT_ENTER``, then the callback is invoked + after *ctx* has been set as the current context for the current thread. + Otherwise, the callback is invoked before the deactivation of *ctx* as the current context + and the restoration of the previous contex object for the current thread. + + If the callback returns with an exception set, it must return ``-1``; this + exception will be printed as an unraisable exception using + :c:func:`PyErr_FormatUnraisable`. Otherwise it should return ``0``. + + There may already be a pending exception set on entry to the callback. In + this case, the callback should return ``0`` with the same exception still + set. This means the callback may not call any other API that can set an + exception unless it saves and clears the exception state first, and restores + it before returning. + + .. versionadded:: 3.14 + Context variable functions: diff --git a/Include/cpython/context.h b/Include/cpython/context.h index a3249fc29b0..a509f4eaba3 100644 --- a/Include/cpython/context.h +++ b/Include/cpython/context.h @@ -27,6 +27,38 @@ PyAPI_FUNC(PyObject *) PyContext_CopyCurrent(void); PyAPI_FUNC(int) PyContext_Enter(PyObject *); PyAPI_FUNC(int) PyContext_Exit(PyObject *); +typedef enum { + Py_CONTEXT_EVENT_ENTER, + Py_CONTEXT_EVENT_EXIT, +} PyContextEvent; + +/* + * A Callback to clue in non-python contexts impls about a + * change in the active python context. + * + * The callback is invoked with the event and a reference to = + * the context after its entered and before its exited. + * + * if the callback returns with an exception set, it must return -1. Otherwise + * it should return 0 + */ +typedef int (*PyContext_WatchCallback)(PyContextEvent, PyContext *); + +/* + * Register a per-interpreter callback that will be invoked for context object + * enter/exit events. + * + * Returns a handle that may be passed to PyContext_ClearWatcher on success, + * or -1 and sets and error if no more handles are available. + */ +PyAPI_FUNC(int) PyContext_AddWatcher(PyContext_WatchCallback callback); + +/* + * Clear the watcher associated with the watcher_id handle. + * + * Returns 0 on success or -1 if no watcher exists for the provided id. + */ +PyAPI_FUNC(int) PyContext_ClearWatcher(int watcher_id); /* Create a new context variable. diff --git a/Include/internal/pycore_context.h b/Include/internal/pycore_context.h index 2ecb40ba584..c2b98d15da6 100644 --- a/Include/internal/pycore_context.h +++ b/Include/internal/pycore_context.h @@ -7,6 +7,7 @@ #include "pycore_hamt.h" // PyHamtObject +#define CONTEXT_MAX_WATCHERS 8 extern PyTypeObject _PyContextTokenMissing_Type; diff --git a/Include/internal/pycore_interp.h b/Include/internal/pycore_interp.h index a1c1dd0c957..36366429e8d 100644 --- a/Include/internal/pycore_interp.h +++ b/Include/internal/pycore_interp.h @@ -240,8 +240,10 @@ struct _is { PyObject *audit_hooks; PyType_WatchCallback type_watchers[TYPE_MAX_WATCHERS]; PyCode_WatchCallback code_watchers[CODE_MAX_WATCHERS]; + PyContext_WatchCallback context_watchers[CONTEXT_MAX_WATCHERS]; // One bit is set for each non-NULL entry in code_watchers uint8_t active_code_watchers; + uint8_t active_context_watchers; struct _py_object_state object_state; struct _Py_unicode_state unicode; diff --git a/Lib/test/test_capi/test_watchers.py b/Lib/test/test_capi/test_watchers.py index 709b5e1c4b7..f21d2627c60 100644 --- a/Lib/test/test_capi/test_watchers.py +++ b/Lib/test/test_capi/test_watchers.py @@ -1,4 +1,5 @@ import unittest +import contextvars from contextlib import contextmanager, ExitStack from test.support import ( @@ -571,5 +572,87 @@ class TestFuncWatchers(unittest.TestCase): _testcapi.allocate_too_many_func_watchers() +class TestContextObjectWatchers(unittest.TestCase): + @contextmanager + def context_watcher(self, which_watcher): + wid = _testcapi.add_context_watcher(which_watcher) + try: + yield wid + finally: + _testcapi.clear_context_watcher(wid) + + def assert_event_counts(self, exp_enter_0, exp_exit_0, + exp_enter_1, exp_exit_1): + self.assertEqual( + exp_enter_0, _testcapi.get_context_watcher_num_enter_events(0)) + self.assertEqual( + exp_exit_0, _testcapi.get_context_watcher_num_exit_events(0)) + self.assertEqual( + exp_enter_1, _testcapi.get_context_watcher_num_enter_events(1)) + self.assertEqual( + exp_exit_1, _testcapi.get_context_watcher_num_exit_events(1)) + + def test_context_object_events_dispatched(self): + # verify that all counts are zero before any watchers are registered + self.assert_event_counts(0, 0, 0, 0) + + # verify that all counts remain zero when a context object is + # entered and exited with no watchers registered + ctx = contextvars.copy_context() + ctx.run(self.assert_event_counts, 0, 0, 0, 0) + self.assert_event_counts(0, 0, 0, 0) + + # verify counts are as expected when first watcher is registered + with self.context_watcher(0): + self.assert_event_counts(0, 0, 0, 0) + ctx.run(self.assert_event_counts, 1, 0, 0, 0) + self.assert_event_counts(1, 1, 0, 0) + + # again with second watcher registered + with self.context_watcher(1): + self.assert_event_counts(1, 1, 0, 0) + ctx.run(self.assert_event_counts, 2, 1, 1, 0) + self.assert_event_counts(2, 2, 1, 1) + + # verify counts are reset and don't change after both watchers are cleared + ctx.run(self.assert_event_counts, 0, 0, 0, 0) + self.assert_event_counts(0, 0, 0, 0) + + def test_enter_error(self): + with self.context_watcher(2): + with catch_unraisable_exception() as cm: + ctx = contextvars.copy_context() + ctx.run(int, 0) + self.assertEqual( + cm.unraisable.err_msg, + "Exception ignored in " + f"Py_CONTEXT_EVENT_EXIT watcher callback for {ctx!r}" + ) + self.assertEqual(str(cm.unraisable.exc_value), "boom!") + + def test_exit_error(self): + ctx = contextvars.copy_context() + def _in_context(stack): + stack.enter_context(self.context_watcher(2)) + + with catch_unraisable_exception() as cm: + with ExitStack() as stack: + ctx.run(_in_context, stack) + self.assertEqual(str(cm.unraisable.exc_value), "boom!") + + def test_clear_out_of_range_watcher_id(self): + with self.assertRaisesRegex(ValueError, r"Invalid context watcher ID -1"): + _testcapi.clear_context_watcher(-1) + with self.assertRaisesRegex(ValueError, r"Invalid context watcher ID 8"): + _testcapi.clear_context_watcher(8) # CONTEXT_MAX_WATCHERS = 8 + + def test_clear_unassigned_watcher_id(self): + with self.assertRaisesRegex(ValueError, r"No context watcher set for ID 1"): + _testcapi.clear_context_watcher(1) + + def test_allocate_too_many_watchers(self): + with self.assertRaisesRegex(RuntimeError, r"no more context watcher IDs available"): + _testcapi.allocate_too_many_context_watchers() + if __name__ == "__main__": unittest.main() diff --git a/Misc/NEWS.d/next/C API/2024-05-21-18-28-44.gh-issue-119333.OTsYVX.rst b/Misc/NEWS.d/next/C API/2024-05-21-18-28-44.gh-issue-119333.OTsYVX.rst new file mode 100644 index 00000000000..6fb6013c4d4 --- /dev/null +++ b/Misc/NEWS.d/next/C API/2024-05-21-18-28-44.gh-issue-119333.OTsYVX.rst @@ -0,0 +1,2 @@ +Add :c:func:`PyContext_AddWatcher` and :c:func:`PyContext_ClearWatcher` APIs to +register callbacks to receive notification on enter and exit of context objects. diff --git a/Modules/_testcapi/watchers.c b/Modules/_testcapi/watchers.c index 1eb0db2c2e6..689863d098a 100644 --- a/Modules/_testcapi/watchers.c +++ b/Modules/_testcapi/watchers.c @@ -8,6 +8,7 @@ #define Py_BUILD_CORE #include "pycore_function.h" // FUNC_MAX_WATCHERS #include "pycore_code.h" // CODE_MAX_WATCHERS +#include "pycore_context.h" // CONTEXT_MAX_WATCHERS /*[clinic input] module _testcapi @@ -622,6 +623,147 @@ allocate_too_many_func_watchers(PyObject *self, PyObject *args) Py_RETURN_NONE; } +// Test contexct object watchers +#define NUM_CONTEXT_WATCHERS 2 +static int context_watcher_ids[NUM_CONTEXT_WATCHERS] = {-1, -1}; +static int num_context_object_enter_events[NUM_CONTEXT_WATCHERS] = {0, 0}; +static int num_context_object_exit_events[NUM_CONTEXT_WATCHERS] = {0, 0}; + +static int +handle_context_watcher_event(int which_watcher, PyContextEvent event, PyContext *ctx) { + if (event == Py_CONTEXT_EVENT_ENTER) { + num_context_object_enter_events[which_watcher]++; + } + else if (event == Py_CONTEXT_EVENT_EXIT) { + num_context_object_exit_events[which_watcher]++; + } + else { + return -1; + } + return 0; +} + +static int +first_context_watcher_callback(PyContextEvent event, PyContext *ctx) { + return handle_context_watcher_event(0, event, ctx); +} + +static int +second_context_watcher_callback(PyContextEvent event, PyContext *ctx) { + return handle_context_watcher_event(1, event, ctx); +} + +static int +noop_context_event_handler(PyContextEvent event, PyContext *ctx) { + return 0; +} + +static int +error_context_event_handler(PyContextEvent event, PyContext *ctx) { + PyErr_SetString(PyExc_RuntimeError, "boom!"); + return -1; +} + +static PyObject * +add_context_watcher(PyObject *self, PyObject *which_watcher) +{ + int watcher_id; + assert(PyLong_Check(which_watcher)); + long which_l = PyLong_AsLong(which_watcher); + if (which_l == 0) { + watcher_id = PyContext_AddWatcher(first_context_watcher_callback); + context_watcher_ids[0] = watcher_id; + num_context_object_enter_events[0] = 0; + num_context_object_exit_events[0] = 0; + } + else if (which_l == 1) { + watcher_id = PyContext_AddWatcher(second_context_watcher_callback); + context_watcher_ids[1] = watcher_id; + num_context_object_enter_events[1] = 0; + num_context_object_exit_events[1] = 0; + } + else if (which_l == 2) { + watcher_id = PyContext_AddWatcher(error_context_event_handler); + } + else { + PyErr_Format(PyExc_ValueError, "invalid watcher %d", which_l); + return NULL; + } + if (watcher_id < 0) { + return NULL; + } + return PyLong_FromLong(watcher_id); +} + +static PyObject * +clear_context_watcher(PyObject *self, PyObject *watcher_id) +{ + assert(PyLong_Check(watcher_id)); + long watcher_id_l = PyLong_AsLong(watcher_id); + if (PyContext_ClearWatcher(watcher_id_l) < 0) { + return NULL; + } + // reset static events counters + if (watcher_id_l >= 0) { + for (int i = 0; i < NUM_CONTEXT_WATCHERS; i++) { + if (watcher_id_l == context_watcher_ids[i]) { + context_watcher_ids[i] = -1; + num_context_object_enter_events[i] = 0; + num_context_object_exit_events[i] = 0; + } + } + } + Py_RETURN_NONE; +} + +static PyObject * +get_context_watcher_num_enter_events(PyObject *self, PyObject *watcher_id) +{ + assert(PyLong_Check(watcher_id)); + long watcher_id_l = PyLong_AsLong(watcher_id); + assert(watcher_id_l >= 0 && watcher_id_l < NUM_CONTEXT_WATCHERS); + return PyLong_FromLong(num_context_object_enter_events[watcher_id_l]); +} + +static PyObject * +get_context_watcher_num_exit_events(PyObject *self, PyObject *watcher_id) +{ + assert(PyLong_Check(watcher_id)); + long watcher_id_l = PyLong_AsLong(watcher_id); + assert(watcher_id_l >= 0 && watcher_id_l < NUM_CONTEXT_WATCHERS); + return PyLong_FromLong(num_context_object_exit_events[watcher_id_l]); +} + +static PyObject * +allocate_too_many_context_watchers(PyObject *self, PyObject *args) +{ + int watcher_ids[CONTEXT_MAX_WATCHERS + 1]; + int num_watchers = 0; + for (unsigned long i = 0; i < sizeof(watcher_ids) / sizeof(int); i++) { + int watcher_id = PyContext_AddWatcher(noop_context_event_handler); + if (watcher_id == -1) { + break; + } + watcher_ids[i] = watcher_id; + num_watchers++; + } + PyObject *exc = PyErr_GetRaisedException(); + for (int i = 0; i < num_watchers; i++) { + if (PyContext_ClearWatcher(watcher_ids[i]) < 0) { + PyErr_WriteUnraisable(Py_None); + break; + } + } + if (exc) { + PyErr_SetRaisedException(exc); + return NULL; + } + else if (PyErr_Occurred()) { + return NULL; + } + Py_RETURN_NONE; +} + /*[clinic input] _testcapi.set_func_defaults_via_capi func: object @@ -689,6 +831,16 @@ static PyMethodDef test_methods[] = { _TESTCAPI_SET_FUNC_KWDEFAULTS_VIA_CAPI_METHODDEF {"allocate_too_many_func_watchers", allocate_too_many_func_watchers, METH_NOARGS, NULL}, + + // Code object watchers. + {"add_context_watcher", add_context_watcher, METH_O, NULL}, + {"clear_context_watcher", clear_context_watcher, METH_O, NULL}, + {"get_context_watcher_num_enter_events", + get_context_watcher_num_enter_events, METH_O, NULL}, + {"get_context_watcher_num_exit_events", + get_context_watcher_num_exit_events, METH_O, NULL}, + {"allocate_too_many_context_watchers", + (PyCFunction) allocate_too_many_context_watchers, METH_NOARGS, NULL}, {NULL}, }; diff --git a/Python/context.c b/Python/context.c index 5cafde4dab9..e52efbb6516 100644 --- a/Python/context.c +++ b/Python/context.c @@ -99,6 +99,80 @@ PyContext_CopyCurrent(void) return (PyObject *)context_new_from_vars(ctx->ctx_vars); } +static const char * +context_event_name(PyContextEvent event) { + switch (event) { + case Py_CONTEXT_EVENT_ENTER: + return "Py_CONTEXT_EVENT_ENTER"; + case Py_CONTEXT_EVENT_EXIT: + return "Py_CONTEXT_EVENT_EXIT"; + default: + return "?"; + } + Py_UNREACHABLE(); +} + +static void notify_context_watchers(PyContextEvent event, PyContext *ctx) +{ + assert(Py_REFCNT(ctx) > 0); + PyInterpreterState *interp = _PyInterpreterState_GET(); + assert(interp->_initialized); + uint8_t bits = interp->active_context_watchers; + int i = 0; + while (bits) { + assert(i < CONTEXT_MAX_WATCHERS); + if (bits & 1) { + PyContext_WatchCallback cb = interp->context_watchers[i]; + assert(cb != NULL); + if (cb(event, ctx) < 0) { + PyErr_FormatUnraisable( + "Exception ignored in %s watcher callback for %R", + context_event_name(event), ctx); + } + } + i++; + bits >>= 1; + } +} + + +int +PyContext_AddWatcher(PyContext_WatchCallback callback) +{ + PyInterpreterState *interp = _PyInterpreterState_GET(); + assert(interp->_initialized); + + for (int i = 0; i < CONTEXT_MAX_WATCHERS; i++) { + if (!interp->context_watchers[i]) { + interp->context_watchers[i] = callback; + interp->active_context_watchers |= (1 << i); + return i; + } + } + + PyErr_SetString(PyExc_RuntimeError, "no more context watcher IDs available"); + return -1; +} + + +int +PyContext_ClearWatcher(int watcher_id) +{ + PyInterpreterState *interp = _PyInterpreterState_GET(); + assert(interp->_initialized); + if (watcher_id < 0 || watcher_id >= CONTEXT_MAX_WATCHERS) { + PyErr_Format(PyExc_ValueError, "Invalid context watcher ID %d", watcher_id); + return -1; + } + if (!interp->context_watchers[watcher_id]) { + PyErr_Format(PyExc_ValueError, "No context watcher set for ID %d", watcher_id); + return -1; + } + interp->context_watchers[watcher_id] = NULL; + interp->active_context_watchers &= ~(1 << watcher_id); + return 0; +} + static int _PyContext_Enter(PyThreadState *ts, PyObject *octx) @@ -118,6 +192,7 @@ _PyContext_Enter(PyThreadState *ts, PyObject *octx) ts->context = Py_NewRef(ctx); ts->context_ver++; + notify_context_watchers(Py_CONTEXT_EVENT_ENTER, ctx); return 0; } @@ -151,6 +226,7 @@ _PyContext_Exit(PyThreadState *ts, PyObject *octx) return -1; } + notify_context_watchers(Py_CONTEXT_EVENT_EXIT, ctx); Py_SETREF(ts->context, (PyObject *)ctx->ctx_prev); ts->context_ver++; diff --git a/Python/pystate.c b/Python/pystate.c index 54caf373e91..6bf7ebeb75f 100644 --- a/Python/pystate.c +++ b/Python/pystate.c @@ -906,6 +906,11 @@ interpreter_clear(PyInterpreterState *interp, PyThreadState *tstate) interp->code_watchers[i] = NULL; } interp->active_code_watchers = 0; + + for (int i=0; i < CONTEXT_MAX_WATCHERS; i++) { + interp->context_watchers[i] = NULL; + } + interp->active_context_watchers = 0; // XXX Once we have one allocator per interpreter (i.e. // per-interpreter GC) we must ensure that all of the interpreter's // objects have been cleaned up at the point. diff --git a/Tools/c-analyzer/cpython/ignored.tsv b/Tools/c-analyzer/cpython/ignored.tsv index b002ada1a8d..f4dc807198a 100644 --- a/Tools/c-analyzer/cpython/ignored.tsv +++ b/Tools/c-analyzer/cpython/ignored.tsv @@ -453,6 +453,9 @@ Modules/_testcapi/watchers.c - num_code_object_destroyed_events - Modules/_testcapi/watchers.c - pyfunc_watchers - Modules/_testcapi/watchers.c - func_watcher_ids - Modules/_testcapi/watchers.c - func_watcher_callbacks - +Modules/_testcapi/watchers.c - context_watcher_ids - +Modules/_testcapi/watchers.c - num_context_object_enter_events - +Modules/_testcapi/watchers.c - num_context_object_exit_events - Modules/_testcapimodule.c - BasicStaticTypes - Modules/_testcapimodule.c - num_basic_static_types_used - Modules/_testcapimodule.c - ContainerNoGC_members -