gh-127124: Pass optional state to context watcher callback

Doc/c-api/contextvars.rst

@@ -101,21 +101,25 @@ 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)
+.. c:function:: int PyContext_AddWatcher(PyContext_WatchCallback *callback, PyObject *cbarg)
 
-   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.
+   Registers *callback* as a context object watcher for the current
+   interpreter, and *cbarg* (which may be NULL; if not, this function creates a
+   new reference) as the object to pass as the callback's first parameter.  On
+   success, returns a non-negative ID which may be passed to
+   :c:func:`PyContext_ClearWatcher` to unregister the callback and remove the
+   added reference to *cbarg*.  Sets an exception and returns ``-1`` on error
+   (e.g., no more watcher IDs available).
 
    .. 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.)
+   Clears the watcher identified by *watcher_id* previously returned from
+   :c:func:`PyContext_AddWatcher` for the current interpreter, and removes the
+   reference created for the *cbarg* object that was registered with the
+   callback.  Returns ``0`` on success, or sets an exception and returns ``-1``
+   on error (e.g., if the given *watcher_id* was never registered).
 
    .. versionadded:: 3.14
 
@@ -130,10 +134,12 @@ Context object management functions:
 
    .. versionadded:: 3.14
 
-.. c:type:: int (*PyContext_WatchCallback)(PyContextEvent event, PyObject *obj)
+.. c:type:: int PyContext_WatchCallback(PyObject *cbarg, PyContextEvent event, PyObject *obj)
 
-   Context object watcher callback function.  The object passed to the callback
-   is event-specific; see :c:type:`PyContextEvent` for details.
+   Context object watcher callback function.  *cbarg* is the same object
+   registered in the call to :c:func:`PyContext_AddWatcher`, as a borrowed
+   reference if non-NULL.  The *obj* object is event-specific; see
+   :c:type:`PyContextEvent` for details.
 
    If the callback returns with an exception set, it must return ``-1``; this
    exception will be printed as an unraisable exception using

Include/cpython/context.h

@@ -28,37 +28,13 @@ PyAPI_FUNC(int) PyContext_Enter(PyObject *);
 PyAPI_FUNC(int) PyContext_Exit(PyObject *);
 
 typedef enum {
-    /*
-     * The current context has switched to a different context.  The object
-     * passed to the watch callback is the now-current contextvars.Context
-     * object, or None if no context is current.
-     */
     Py_CONTEXT_SWITCHED = 1,
 } PyContextEvent;
 
-/*
- * Context object watcher callback function.  The object passed to the callback
- * is event-specific; see PyContextEvent for details.
- *
- * if the callback returns with an exception set, it must return -1. Otherwise
- * it should return 0
- */
-typedef int (*PyContext_WatchCallback)(PyContextEvent, PyObject *);
-
-/*
- * 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.
- */
+typedef int PyContext_WatchCallback(
+    PyObject *cbarg, PyContextEvent event, PyObject *obj);
+PyAPI_FUNC(int) PyContext_AddWatcher(
+    PyContext_WatchCallback *callback, PyObject *cbarg);
 PyAPI_FUNC(int) PyContext_ClearWatcher(int watcher_id);
 
 /* Create a new context variable.

Include/internal/pycore_interp.h

@@ -242,7 +242,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];
+    struct {
+        PyContext_WatchCallback *callback;
+        PyObject *arg;
+    } 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;

Lib/test/test_capi/test_watchers.py

@@ -1,3 +1,4 @@
+import contextlib
 import unittest
 import contextvars
 
@@ -589,60 +590,43 @@ def test_allocate_too_many_watchers(self):
 
 class TestContextObjectWatchers(unittest.TestCase):
     @contextmanager
-    def context_watcher(self, which_watcher):
-        wid = _testcapi.add_context_watcher(which_watcher)
+    def context_watcher(self, cfg=None):
+        if cfg is None:
+            cfg = {}
+        cfg.setdefault('log', [])
+        wid = _testcapi.add_context_watcher(cfg)
         try:
-            switches = _testcapi.get_context_switches(which_watcher)
-        except ValueError:
-            switches = None
-        try:
-            yield switches
+            yield cfg.get('log', None)
         finally:
             _testcapi.clear_context_watcher(wid)
 
-    def assert_event_counts(self, want_0, want_1):
-        self.assertEqual(len(_testcapi.get_context_switches(0)), want_0)
-        self.assertEqual(len(_testcapi.get_context_switches(1)), want_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)
-
-        # 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)
-        self.assert_event_counts(0, 0)
-
-        # verify counts are as expected when first watcher is registered
-        with self.context_watcher(0):
-            self.assert_event_counts(0, 0)
-            ctx.run(self.assert_event_counts, 1, 0)
-            self.assert_event_counts(2, 0)
-
-            # again with second watcher registered
-            with self.context_watcher(1):
-                self.assert_event_counts(2, 0)
-                ctx.run(self.assert_event_counts, 3, 1)
-                self.assert_event_counts(4, 2)
-
-        # verify counts are reset and don't change after both watchers are cleared
-        ctx.run(self.assert_event_counts, 0, 0)
-        self.assert_event_counts(0, 0)
+        with self.context_watcher() as switches_0:
+            self.assertEqual(len(switches_0), 0)
+            ctx.run(lambda: self.assertEqual(len(switches_0), 1))
+            self.assertEqual(len(switches_0), 2)
+            with self.context_watcher() as switches_1:
+                self.assertEqual((len(switches_0), len(switches_1)), (2, 0))
+                ctx.run(lambda: self.assertEqual(
+                    (len(switches_0), len(switches_1)), (3, 1)))
+                self.assertEqual((len(switches_0), len(switches_1)), (4, 2))
 
     def test_callback_error(self):
         ctx_outer = contextvars.copy_context()
         ctx_inner = contextvars.copy_context()
         unraisables = []
 
         def _in_outer():
-            with self.context_watcher(2):
+            with self.context_watcher(cfg={'err': RuntimeError('boom!')}):
                 with catch_unraisable_exception() as cm:
                     ctx_inner.run(lambda: unraisables.append(cm.unraisable))
                     unraisables.append(cm.unraisable)
 
         try:
             ctx_outer.run(_in_outer)
+            self.assertEqual([x is not None for x in unraisables],
+                             [True, True])
             self.assertEqual([x.err_msg for x in unraisables],
                              ["Exception ignored in Py_CONTEXT_SWITCHED "
                               f"watcher callback for {ctx!r}"
@@ -656,21 +640,24 @@ def _in_outer():
     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
+        with self.assertRaisesRegex(ValueError, f"Invalid context watcher ID {_testcapi.CONTEXT_MAX_WATCHERS}"):
+            _testcapi.clear_context_watcher(_testcapi.CONTEXT_MAX_WATCHERS)
 
     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()
+        with contextlib.ExitStack() as stack:
+            for i in range(_testcapi.CONTEXT_MAX_WATCHERS):
+                stack.enter_context(self.context_watcher())
+            with self.assertRaisesRegex(RuntimeError, r"no more context watcher IDs available"):
+                stack.enter_context(self.context_watcher())
 
     def test_exit_base_context(self):
         ctx = contextvars.Context()
         _testcapi.clear_context_stack()
-        with self.context_watcher(0) as switches:
+        with self.context_watcher() as switches:
             ctx.run(lambda: None)
         self.assertEqual(switches, [ctx, None])
 

Misc/NEWS.d/next/C_API/2024-11-22-18-38-33.gh-issue-127124.6k6Qj7.rst

@@ -0,0 +1,2 @@
+Added optional callback state to :c:func:`PyContext_AddWatcher` and
+:c:type:`PyContext_WatchCallback`.