GH-91054: Add code object watchers API (GH-99859)

* Add API to allow extensions to set callback function on creation and destruction of PyCodeObject

Co-authored-by: Ye11ow-Flash <janshah@cs.stonybrook.edu>

Doc/c-api/code.rst

@@ -115,3 +115,51 @@ bound into a function.
    the free variables. On error, ``NULL`` is returned and an exception is raised.
 
    .. versionadded:: 3.11
+
+.. c:function:: int PyCode_AddWatcher(PyCode_WatchCallback callback)
+
+   Register *callback* as a code object watcher for the current interpreter.
+   Return an ID which may be passed to :c:func:`PyCode_ClearWatcher`.
+   In case of error (e.g. no more watcher IDs available),
+   return ``-1`` and set an exception.
+
+   .. versionadded:: 3.12
+
+.. c:function:: int PyCode_ClearWatcher(int watcher_id)
+
+   Clear watcher identified by *watcher_id* previously returned from
+   :c:func:`PyCode_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.12
+
+.. c:type:: PyCodeEvent
+
+   Enumeration of possible code object watcher events:
+   - ``PY_CODE_EVENT_CREATE``
+   - ``PY_CODE_EVENT_DESTROY``
+
+   .. versionadded:: 3.12
+
+.. c:type:: int (*PyCode_WatchCallback)(PyCodeEvent event, PyCodeObject* co)
+
+   Type of a code object watcher callback function.
+
+   If *event* is ``PY_CODE_EVENT_CREATE``, then the callback is invoked
+   after `co` has been fully initialized. Otherwise, the callback is invoked
+   before the destruction of *co* takes place, so the prior state of *co*
+   can be inspected.
+
+   Users of this API should not rely on internal runtime implementation
+   details. Such details may include, but are not limited to, the exact
+   order and timing of creation and destruction of code objects. While
+   changes in these details may result in differences observable by watchers
+   (including whether a callback is invoked or not), it does not change
+   the semantics of the Python code being executed.
+
+   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_WriteUnraisable`. Otherwise it should return ``0``.
+
+   .. versionadded:: 3.12

Doc/whatsnew/3.12.rst

@@ -773,6 +773,10 @@ New Features
   callbacks to receive notification on changes to a type.
   (Contributed by Carl Meyer in :gh:`91051`.)
 
+* Added :c:func:`PyCode_AddWatcher` and :c:func:`PyCode_ClearWatcher`
+  APIs to register callbacks to receive notification on creation and
+  destruction of code objects.
+  (Contributed by Itamar Ostricher in :gh:`91054`.)
 
 * Add :c:func:`PyFrame_GetVar` and :c:func:`PyFrame_GetVarString` functions to
   get a frame variable by its name.

Include/cpython/code.h

@@ -181,6 +181,41 @@ PyAPI_FUNC(int) PyCode_Addr2Line(PyCodeObject *, int);
 
 PyAPI_FUNC(int) PyCode_Addr2Location(PyCodeObject *, int, int *, int *, int *, int *);
 
+typedef enum PyCodeEvent {
+  PY_CODE_EVENT_CREATE,
+  PY_CODE_EVENT_DESTROY
+} PyCodeEvent;
+
+
+/*
+ * A callback that is invoked for different events in a code object's lifecycle.
+ *
+ * The callback is invoked with a borrowed reference to co, after it is
+ * created and before it is destroyed.
+ *
+ * If the callback returns with an exception set, it must return -1. Otherwise
+ * it should return 0.
+ */
+typedef int (*PyCode_WatchCallback)(
+  PyCodeEvent event,
+  PyCodeObject* co);
+
+/*
+ * Register a per-interpreter callback that will be invoked for code object
+ * lifecycle events.
+ *
+ * Returns a handle that may be passed to PyCode_ClearWatcher on success,
+ * or -1 and sets an error if no more handles are available.
+ */
+PyAPI_FUNC(int) PyCode_AddWatcher(PyCode_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) PyCode_ClearWatcher(int watcher_id);
+
 /* for internal use only */
 struct _opaque {
     int computed_line;

Include/internal/pycore_code.h

@@ -4,6 +4,8 @@
 extern "C" {
 #endif
 
+#define CODE_MAX_WATCHERS 8
+
 /* PEP 659
  * Specialization and quickening structs and helper functions
  */

Include/internal/pycore_interp.h

@@ -191,6 +191,9 @@ struct _is {
 
     PyObject *audit_hooks;
     PyType_WatchCallback type_watchers[TYPE_MAX_WATCHERS];
+    PyCode_WatchCallback code_watchers[CODE_MAX_WATCHERS];
+    // One bit is set for each non-NULL entry in code_watchers
+    uint8_t active_code_watchers;
 
     struct _Py_unicode_state unicode;
     struct _Py_float_state float_state;

Lib/test/test_capi/test_watchers.py

@@ -336,6 +336,74 @@ class TestTypeWatchers(unittest.TestCase):
                 self.add_watcher()
 
 
+class TestCodeObjectWatchers(unittest.TestCase):
+    @contextmanager
+    def code_watcher(self, which_watcher):
+        wid = _testcapi.add_code_watcher(which_watcher)
+        try:
+            yield wid
+        finally:
+            _testcapi.clear_code_watcher(wid)
+
+    def assert_event_counts(self, exp_created_0, exp_destroyed_0,
+                            exp_created_1, exp_destroyed_1):
+        self.assertEqual(
+            exp_created_0, _testcapi.get_code_watcher_num_created_events(0))
+        self.assertEqual(
+            exp_destroyed_0, _testcapi.get_code_watcher_num_destroyed_events(0))
+        self.assertEqual(
+            exp_created_1, _testcapi.get_code_watcher_num_created_events(1))
+        self.assertEqual(
+            exp_destroyed_1, _testcapi.get_code_watcher_num_destroyed_events(1))
+
+    def test_code_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 code object is
+        # created and destroyed with no watchers registered
+        co1 = _testcapi.code_newempty("test_watchers", "dummy1", 0)
+        self.assert_event_counts(0, 0, 0, 0)
+        del co1
+        self.assert_event_counts(0, 0, 0, 0)
+
+        # verify counts are as expected when first watcher is registered
+        with self.code_watcher(0):
+            self.assert_event_counts(0, 0, 0, 0)
+            co2 = _testcapi.code_newempty("test_watchers", "dummy2", 0)
+            self.assert_event_counts(1, 0, 0, 0)
+            del co2
+            self.assert_event_counts(1, 1, 0, 0)
+
+            # again with second watcher registered
+            with self.code_watcher(1):
+                self.assert_event_counts(1, 1, 0, 0)
+                co3 = _testcapi.code_newempty("test_watchers", "dummy3", 0)
+                self.assert_event_counts(2, 1, 1, 0)
+                del co3
+                self.assert_event_counts(2, 2, 1, 1)
+
+        # verify counts remain as they were after both watchers are cleared
+        co4 = _testcapi.code_newempty("test_watchers", "dummy4", 0)
+        self.assert_event_counts(2, 2, 1, 1)
+        del co4
+        self.assert_event_counts(2, 2, 1, 1)
+
+    def test_clear_out_of_range_watcher_id(self):
+        with self.assertRaisesRegex(ValueError, r"Invalid code watcher ID -1"):
+            _testcapi.clear_code_watcher(-1)
+        with self.assertRaisesRegex(ValueError, r"Invalid code watcher ID 8"):
+            _testcapi.clear_code_watcher(8)  # CODE_MAX_WATCHERS = 8
+
+    def test_clear_unassigned_watcher_id(self):
+        with self.assertRaisesRegex(ValueError, r"No code watcher set for ID 1"):
+            _testcapi.clear_code_watcher(1)
+
+    def test_allocate_too_many_watchers(self):
+        with self.assertRaisesRegex(RuntimeError, r"no more code watcher IDs available"):
+            _testcapi.allocate_too_many_code_watchers()
+
+
 class TestFuncWatchers(unittest.TestCase):
     @contextmanager
     def add_watcher(self, func):

Misc/ACKS

@@ -1320,6 +1320,7 @@ Michele Orrù
 Tomáš Orsava
 Oleg Oshmyan
 Denis Osipov
+Itamar Ostricher
 Denis S. Otkidach
 Peter Otten
 Michael Otteneder
@@ -1627,6 +1628,7 @@ Silas Sewell
 Ian Seyer
 Dmitry Shachnev
 Anish Shah
+Jaineel Shah
 Daniel Shahaf
 Hui Shang
 Geoff Shannon

Misc/NEWS.d/next/Core and Builtins/2022-11-27-13-50-13.gh-issue-91054.oox_kW.rst

@@ -0,0 +1,3 @@
+Add :c:func:`PyCode_AddWatcher` and :c:func:`PyCode_ClearWatcher` APIs to
+register callbacks to receive notification on creation and destruction of
+code objects.

Modules/_testcapi/watchers.c

@@ -2,6 +2,7 @@
 
 #define Py_BUILD_CORE
 #include "pycore_function.h"  // FUNC_MAX_WATCHERS
+#include "pycore_code.h"  // CODE_MAX_WATCHERS
 
 // Test dict watching
 static PyObject *g_dict_watch_events;
@@ -277,6 +278,126 @@ unwatch_type(PyObject *self, PyObject *args)
     Py_RETURN_NONE;
 }
 
+
+// Test code object watching
+
+#define NUM_CODE_WATCHERS 2
+static int num_code_object_created_events[NUM_CODE_WATCHERS] = {0, 0};
+static int num_code_object_destroyed_events[NUM_CODE_WATCHERS] = {0, 0};
+
+static int
+handle_code_object_event(int which_watcher, PyCodeEvent event, PyCodeObject *co) {
+    if (event == PY_CODE_EVENT_CREATE) {
+        num_code_object_created_events[which_watcher]++;
+    }
+    else if (event == PY_CODE_EVENT_DESTROY) {
+        num_code_object_destroyed_events[which_watcher]++;
+    }
+    else {
+        return -1;
+    }
+    return 0;
+}
+
+static int
+first_code_object_callback(PyCodeEvent event, PyCodeObject *co)
+{
+    return handle_code_object_event(0, event, co);
+}
+
+static int
+second_code_object_callback(PyCodeEvent event, PyCodeObject *co)
+{
+    return handle_code_object_event(1, event, co);
+}
+
+static int
+noop_code_event_handler(PyCodeEvent event, PyCodeObject *co)
+{
+    return 0;
+}
+
+static PyObject *
+add_code_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 = PyCode_AddWatcher(first_code_object_callback);
+    }
+    else if (which_l == 1) {
+        watcher_id = PyCode_AddWatcher(second_code_object_callback);
+    }
+    else {
+        return NULL;
+    }
+    if (watcher_id < 0) {
+        return NULL;
+    }
+    return PyLong_FromLong(watcher_id);
+}
+
+static PyObject *
+clear_code_watcher(PyObject *self, PyObject *watcher_id)
+{
+    assert(PyLong_Check(watcher_id));
+    long watcher_id_l = PyLong_AsLong(watcher_id);
+    if (PyCode_ClearWatcher(watcher_id_l) < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+get_code_watcher_num_created_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_CODE_WATCHERS);
+    return PyLong_FromLong(num_code_object_created_events[watcher_id_l]);
+}
+
+static PyObject *
+get_code_watcher_num_destroyed_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_CODE_WATCHERS);
+    return PyLong_FromLong(num_code_object_destroyed_events[watcher_id_l]);
+}
+
+static PyObject *
+allocate_too_many_code_watchers(PyObject *self, PyObject *args)
+{
+    int watcher_ids[CODE_MAX_WATCHERS + 1];
+    int num_watchers = 0;
+    for (unsigned long i = 0; i < sizeof(watcher_ids) / sizeof(int); i++) {
+        int watcher_id = PyCode_AddWatcher(noop_code_event_handler);
+        if (watcher_id == -1) {
+            break;
+        }
+        watcher_ids[i] = watcher_id;
+        num_watchers++;
+    }
+    PyObject *type, *value, *traceback;
+    PyErr_Fetch(&type, &value, &traceback);
+    for (int i = 0; i < num_watchers; i++) {
+        if (PyCode_ClearWatcher(watcher_ids[i]) < 0) {
+            PyErr_WriteUnraisable(Py_None);
+            break;
+        }
+    }
+    if (type) {
+        PyErr_Restore(type, value, traceback);
+        return NULL;
+    }
+    else if (PyErr_Occurred()) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
 // Test function watchers
 
 #define NUM_FUNC_WATCHERS 2
@@ -509,6 +630,16 @@ static PyMethodDef test_methods[] = {
     {"unwatch_type",             unwatch_type,            METH_VARARGS, NULL},
     {"get_type_modified_events", get_type_modified_events, METH_NOARGS, NULL},
 
+    // Code object watchers.
+    {"add_code_watcher",         add_code_watcher,        METH_O,       NULL},
+    {"clear_code_watcher",       clear_code_watcher,      METH_O,       NULL},
+    {"get_code_watcher_num_created_events",
+     get_code_watcher_num_created_events,                 METH_O,       NULL},
+    {"get_code_watcher_num_destroyed_events",
+     get_code_watcher_num_destroyed_events,               METH_O,       NULL},
+    {"allocate_too_many_code_watchers",
+     (PyCFunction) allocate_too_many_code_watchers,       METH_NOARGS,  NULL},
+
     // Function watchers.
     {"add_func_watcher",         add_func_watcher,        METH_O,       NULL},
     {"clear_func_watcher",       clear_func_watcher,      METH_O,       NULL},

Objects/codeobject.c

@@ -12,6 +12,66 @@
 #include "clinic/codeobject.c.h"
 
 
+static void
+notify_code_watchers(PyCodeEvent event, PyCodeObject *co)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    if (interp->active_code_watchers) {
+        assert(interp->_initialized);
+        for (int i = 0; i < CODE_MAX_WATCHERS; i++) {
+            PyCode_WatchCallback cb = interp->code_watchers[i];
+            if ((cb != NULL) && (cb(event, co) < 0)) {
+                PyErr_WriteUnraisable((PyObject *) co);
+            }
+        }
+    }
+}
+
+int
+PyCode_AddWatcher(PyCode_WatchCallback callback)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    assert(interp->_initialized);
+
+    for (int i = 0; i < CODE_MAX_WATCHERS; i++) {
+        if (!interp->code_watchers[i]) {
+            interp->code_watchers[i] = callback;
+            interp->active_code_watchers |= (1 << i);
+            return i;
+        }
+    }
+
+    PyErr_SetString(PyExc_RuntimeError, "no more code watcher IDs available");
+    return -1;
+}
+
+static inline int
+validate_watcher_id(PyInterpreterState *interp, int watcher_id)
+{
+    if (watcher_id < 0 || watcher_id >= CODE_MAX_WATCHERS) {
+        PyErr_Format(PyExc_ValueError, "Invalid code watcher ID %d", watcher_id);
+        return -1;
+    }
+    if (!interp->code_watchers[watcher_id]) {
+        PyErr_Format(PyExc_ValueError, "No code watcher set for ID %d", watcher_id);
+        return -1;
+    }
+    return 0;
+}
+
+int
+PyCode_ClearWatcher(int watcher_id)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    assert(interp->_initialized);
+    if (validate_watcher_id(interp, watcher_id) < 0) {
+        return -1;
+    }
+    interp->code_watchers[watcher_id] = NULL;
+    interp->active_code_watchers &= ~(1 << watcher_id);
+    return 0;
+}
+
 /******************
  * generic helpers
  ******************/
@@ -355,6 +415,7 @@ init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
     }
     co->_co_firsttraceable = entry_point;
     _PyCode_Quicken(co);
+    notify_code_watchers(PY_CODE_EVENT_CREATE, co);
 }
 
 static int
@@ -1615,6 +1676,8 @@ code_new_impl(PyTypeObject *type, int argcount, int posonlyargcount,
 static void
 code_dealloc(PyCodeObject *co)
 {
+    notify_code_watchers(PY_CODE_EVENT_DESTROY, co);
+
     if (co->co_extra != NULL) {
         PyInterpreterState *interp = _PyInterpreterState_GET();
         _PyCodeObjectExtra *co_extra = co->co_extra;

Python/pystate.c

@@ -466,6 +466,11 @@ interpreter_clear(PyInterpreterState *interp, PyThreadState *tstate)
     }
     interp->active_func_watchers = 0;
 
+    for (int i=0; i < CODE_MAX_WATCHERS; i++) {
+        interp->code_watchers[i] = NULL;
+    }
+    interp->active_code_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.