backport 3.12 code watchers

Summary: Backport of python/cpython#99859 (plus some more recent changes.)

Reviewed By: alexmalyshev

Differential Revision: D47201113

fbshipit-source-id: 9390ca289c043bc56ebfdd2ade63305347190a99

Author: Carl Meyer

Cinder/module/known-core-python-exported-symbols

@@ -382,6 +382,7 @@ PyCMethod_New
 PyCMethod_Type
 _Py_c_neg
 PyCode_Addr2Line
+PyCode_AddWatcher
 _PyCodeCache_RefType
 PyCodec_BackslashReplaceErrors
 PyCodec_Decode
@@ -398,6 +399,7 @@ PyCodec_IncrementalEncoder
 _PyCodecInfo_GetIncrementalDecoder
 _PyCodecInfo_GetIncrementalEncoder
 PyCodec_KnownEncoding
+PyCode_ClearWatcher
 _PyCodec_Lookup
 PyCodec_LookupError
 _PyCodec_LookupTextEncoding

Include/cpython/code.h

@@ -136,6 +136,46 @@ PyCode_NewEmpty(const char *filename, const char *funcname, int firstlineno);
    use PyFrame_GetLineNumber() instead. */
 PyAPI_FUNC(int) PyCode_Addr2Line(PyCodeObject *, int);
 
+#define PY_FOREACH_CODE_EVENT(V) \
+    V(CREATE)                 \
+    V(DESTROY)
+
+typedef enum {
+    #define PY_DEF_EVENT(op) PY_CODE_EVENT_##op,
+    PY_FOREACH_CODE_EVENT(PY_DEF_EVENT)
+    #undef PY_DEF_EVENT
+} 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 sets an exception, 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
+
 typedef struct {
     PyObject *ptr;  /* Cached pointer (borrowed reference) */
     uint64_t globals_ver;  /* ma_version of global dict */

Include/internal/pycore_interp.h

@@ -7,6 +7,7 @@
 
 #include "pycore_atomic.h"        // _Py_atomic_address
 #include "pycore_ast_state.h"     // struct ast_state
+#include "pycore_code.h"          // CODE_MAX_WATCHERS
 #include "pycore_gil.h"           // struct _gil_runtime_state
 #include "pycore_gc.h"            // struct _gc_runtime_state
 #include "pycore_warnings.h"      // struct _warnings_runtime_state
@@ -297,6 +298,9 @@ struct _is {
     struct atexit_state atexit;
 
     PyObject *audit_hooks;
+    PyCode_WatchCallback code_watchers[CODE_MAX_WATCHERS];
+    // One bit is set for each non-NULL entry in code_watchers
+    uint8_t active_code_watchers;
 
     /* Small integers are preallocated in this array so that they
        can be shared.

Lib/test/test_capi.py

@@ -1221,6 +1221,93 @@ def test_clear_unassigned_watcher_id(self):
             self.clear_watcher(1)
 
 
+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 are reset and don't change after both watchers are cleared
+        co4 = _testcapi.code_newempty("test_watchers", "dummy4", 0)
+        self.assert_event_counts(0, 0, 0, 0)
+        del co4
+        self.assert_event_counts(0, 0, 0, 0)
+
+    def test_error(self):
+        with self.code_watcher(2):
+            with catch_unraisable_exception() as cm:
+                co = _testcapi.code_newempty("test_watchers", "dummy0", 0)
+
+                self.assertEqual(
+                    cm.unraisable.object,
+                    f"PY_CODE_EVENT_CREATE watcher callback for {co!r}"
+                )
+                self.assertEqual(str(cm.unraisable.exc_value), "boom!")
+
+    def test_dealloc_error(self):
+        co = _testcapi.code_newempty("test_watchers", "dummy0", 0)
+        with self.code_watcher(2):
+            with catch_unraisable_exception() as cm:
+                del co
+
+                self.assertEqual(str(cm.unraisable.exc_value), "boom!")
+
+    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):

Modules/_testcapimodule.c

@@ -20,6 +20,7 @@
 #include "Python.h"
 #include "datetime.h"
 #include "marshal.h"
+#include "pycore_code.h"
 #include "structmember.h"         // PyMemberDef
 #include <float.h>
 #include <signal.h>
@@ -6252,6 +6253,145 @@ get_dict_watcher_events(PyObject *self, PyObject *Py_UNUSED(args))
     return Py_NewRef(g_dict_watch_events);
 }
 
+// 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 int
+error_code_event_handler(PyCodeEvent event, PyCodeObject *co)
+{
+    PyErr_SetString(PyExc_RuntimeError, "boom!");
+    return -1;
+}
+
+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);
+        num_code_object_created_events[0] = 0;
+        num_code_object_destroyed_events[0] = 0;
+    }
+    else if (which_l == 1) {
+        watcher_id = PyCode_AddWatcher(second_code_object_callback);
+        num_code_object_created_events[1] = 0;
+        num_code_object_destroyed_events[1] = 0;
+    }
+    else if (which_l == 2) {
+        watcher_id = PyCode_AddWatcher(error_code_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_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;
+    }
+    // reset static events counters
+    if (watcher_id_l >= 0 && watcher_id_l < NUM_CODE_WATCHERS) {
+        num_code_object_created_events[watcher_id_l] = 0;
+        num_code_object_destroyed_events[watcher_id_l] = 0;
+    }
+    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 *err_type, *err_val, *err_tb;
+    PyErr_Fetch(&err_type, &err_val, &err_tb);
+    for (int i = 0; i < num_watchers; i++) {
+        if (PyCode_ClearWatcher(watcher_ids[i]) < 0) {
+            PyErr_WriteUnraisable(Py_None);
+            break;
+        }
+    }
+    if (err_val) {
+        PyErr_Restore(err_type, err_val, err_tb);
+        return NULL;
+    }
+    else if (PyErr_Occurred()) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
 // Test function watchers
 
 #define NUM_TEST_FUNC_WATCHERS 2
@@ -6764,6 +6904,14 @@ static PyMethodDef TestMethods[] = {
     {"watch_dict", watch_dict, METH_VARARGS},
     {"unwatch_dict", unwatch_dict, METH_VARARGS},
     {"get_dict_watcher_events", get_dict_watcher_events, METH_NOARGS},
+    {"add_code_watcher", add_code_watcher, METH_O},
+    {"clear_code_watcher", clear_code_watcher, METH_O},
+    {"get_code_watcher_num_created_events",
+     get_code_watcher_num_created_events, METH_O},
+    {"get_code_watcher_num_destroyed_events",
+     get_code_watcher_num_destroyed_events, METH_O},
+    {"allocate_too_many_code_watchers",
+     (PyCFunction) allocate_too_many_code_watchers, METH_NOARGS},
     {"add_func_watcher", add_func_watcher, METH_O},
     {"clear_func_watcher", clear_func_watcher, METH_O},
     {"set_func_defaults_via_capi", set_func_defaults_via_capi, METH_VARARGS},