gh-116750: Add clear_tool_id function to unregister events and callbacks (#124568)

8 files changed, 165 insertions, 8 deletions

diff --git a/Doc/library/sys.monitoring.rst b/Doc/library/sys.monitoring.rst
index ac8bcce..f7140af 100644
--- a/Doc/library/sys.monitoring.rst
+++ b/Doc/library/sys.monitoring.rst
@@ -50,16 +50,14 @@ Registering and using tools
    *tool_id* must be in the range 0 to 5 inclusive.
    Raises a :exc:`ValueError` if *tool_id* is in use.
 
-.. function:: free_tool_id(tool_id: int, /) -> None
+.. function:: clear_tool_id(tool_id: int, /) -> None
 
-   Should be called once a tool no longer requires *tool_id*.
+   Unregister all events and callback functions associated with *tool_id*.
 
-.. note::
+.. function:: free_tool_id(tool_id: int, /) -> None
 
-   :func:`free_tool_id` will not disable global or local events associated
-   with *tool_id*, nor will it unregister any callback functions. This
-   function is only intended to be used to notify the VM that the
-   particular *tool_id* is no longer in use.
+   Should be called once a tool no longer requires *tool_id*.
+   Will call :func:`clear_tool_id` before releasing *tool_id*.
 
 .. function:: get_tool(tool_id: int, /) -> str | None
 
diff --git a/Include/cpython/code.h b/Include/cpython/code.h
index 58d93fc..0362269 100644
--- a/Include/cpython/code.h
+++ b/Include/cpython/code.h
@@ -8,6 +8,8 @@
 extern "C" {
 #endif
 
+/* Total tool ids available */
+#define  _PY_MONITORING_TOOL_IDS 8
 /* Count of all local monitoring events */
 #define  _PY_MONITORING_LOCAL_EVENTS 10
 /* Count of all "real" monitoring events (not derived from other events) */
@@ -57,6 +59,8 @@ typedef struct {
     _Py_LocalMonitors active_monitors;
     /* The tools that are to be notified for events for the matching code unit */
     uint8_t *tools;
+    /* The version of tools when they instrument the code */
+    uintptr_t tool_versions[_PY_MONITORING_TOOL_IDS];
     /* Information to support line events */
     _PyCoLineInstrumentationData *lines;
     /* The tools that are to be notified for line events for the matching code unit */
diff --git a/Include/internal/pycore_interp.h b/Include/internal/pycore_interp.h
index ade69be..d7e5840 100644
--- a/Include/internal/pycore_interp.h
+++ b/Include/internal/pycore_interp.h
@@ -272,6 +272,7 @@ struct _is {
     Py_ssize_t sys_tracing_threads; /* Count of threads with c_tracefunc set */
     PyObject *monitoring_callables[PY_MONITORING_TOOL_IDS][_PY_MONITORING_EVENTS];
     PyObject *monitoring_tool_names[PY_MONITORING_TOOL_IDS];
+    uintptr_t monitoring_tool_versions[PY_MONITORING_TOOL_IDS];
 
     struct _Py_interp_cached_objects cached_objects;
     struct _Py_interp_static_objects static_objects;
diff --git a/Lib/test/test_monitoring.py b/Lib/test/test_monitoring.py
index 351f106..2a32668 100644
--- a/Lib/test/test_monitoring.py
+++ b/Lib/test/test_monitoring.py
@@ -46,10 +46,14 @@ def nth_line(func, offset):
 
 class MonitoringBasicTest(unittest.TestCase):
 
+    def tearDown(self):
+        sys.monitoring.free_tool_id(TEST_TOOL)
+
     def test_has_objects(self):
         m = sys.monitoring
         m.events
         m.use_tool_id
+        m.clear_tool_id
         m.free_tool_id
         m.get_tool
         m.get_events
@@ -77,6 +81,43 @@ class MonitoringBasicTest(unittest.TestCase):
         with self.assertRaises(ValueError):
             sys.monitoring.set_events(TEST_TOOL, sys.monitoring.events.CALL)
 
+    def test_clear(self):
+        events = []
+        sys.monitoring.use_tool_id(TEST_TOOL, "MonitoringTest.Tool")
+        sys.monitoring.register_callback(TEST_TOOL, E.PY_START, lambda *args: events.append(args))
+        sys.monitoring.register_callback(TEST_TOOL, E.LINE, lambda *args: events.append(args))
+        def f():
+            a = 1
+        sys.monitoring.set_local_events(TEST_TOOL, f.__code__, E.LINE)
+        sys.monitoring.set_events(TEST_TOOL, E.PY_START)
+
+        f()
+        sys.monitoring.clear_tool_id(TEST_TOOL)
+        f()
+
+        # the first f() should trigger a PY_START and a LINE event
+        # the second f() after clear_tool_id should not trigger any event
+        # the callback function should be cleared as well
+        self.assertEqual(len(events), 2)
+        callback = sys.monitoring.register_callback(TEST_TOOL, E.LINE, None)
+        self.assertIs(callback, None)
+
+        sys.monitoring.free_tool_id(TEST_TOOL)
+
+        events = []
+        sys.monitoring.use_tool_id(TEST_TOOL, "MonitoringTest.Tool")
+        sys.monitoring.register_callback(TEST_TOOL, E.LINE, lambda *args: events.append(args))
+        sys.monitoring.set_local_events(TEST_TOOL, f.__code__, E.LINE)
+        f()
+        sys.monitoring.free_tool_id(TEST_TOOL)
+        sys.monitoring.use_tool_id(TEST_TOOL, "MonitoringTest.Tool")
+        f()
+        # the first f() should trigger a LINE event, and even if we use the
+        # tool id immediately after freeing it, the second f() should not
+        # trigger any event
+        self.assertEqual(len(events), 1)
+        sys.monitoring.free_tool_id(TEST_TOOL)
+
 
 class MonitoringTestBase:
 
diff --git a/Misc/NEWS.d/next/Library/2024-09-26-00-35-24.gh-issue-116750.X1aMHI.rst b/Misc/NEWS.d/next/Library/2024-09-26-00-35-24.gh-issue-116750.X1aMHI.rst
new file mode 100644
index 0000000..cf9dacf
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2024-09-26-00-35-24.gh-issue-116750.X1aMHI.rst
@@ -0,0 +1 @@
+Provide :func:`sys.monitoring.clear_tool_id` to unregister all events and callbacks set by the tool.
diff --git a/Python/clinic/instrumentation.c.h b/Python/clinic/instrumentation.c.h
index 8dae747..9b3373b 100644
--- a/Python/clinic/instrumentation.c.h
+++ b/Python/clinic/instrumentation.c.h
@@ -36,6 +36,33 @@ exit:
     return return_value;
 }
 
+PyDoc_STRVAR(monitoring_clear_tool_id__doc__,
+"clear_tool_id($module, tool_id, /)\n"
+"--\n"
+"\n");
+
+#define MONITORING_CLEAR_TOOL_ID_METHODDEF    \
+    {"clear_tool_id", (PyCFunction)monitoring_clear_tool_id, METH_O, monitoring_clear_tool_id__doc__},
+
+static PyObject *
+monitoring_clear_tool_id_impl(PyObject *module, int tool_id);
+
+static PyObject *
+monitoring_clear_tool_id(PyObject *module, PyObject *arg)
+{
+    PyObject *return_value = NULL;
+    int tool_id;
+
+    tool_id = PyLong_AsInt(arg);
+    if (tool_id == -1 && PyErr_Occurred()) {
+        goto exit;
+    }
+    return_value = monitoring_clear_tool_id_impl(module, tool_id);
+
+exit:
+    return return_value;
+}
+
 PyDoc_STRVAR(monitoring_free_tool_id__doc__,
 "free_tool_id($module, tool_id, /)\n"
 "--\n"
@@ -304,4 +331,4 @@ monitoring__all_events(PyObject *module, PyObject *Py_UNUSED(ignored))
 {
     return monitoring__all_events_impl(module);
 }
-/*[clinic end generated code: output=14ffc0884a6de50a input=a9049054013a1b77]*/
+/*[clinic end generated code: output=8f81876c6aba9be8 input=a9049054013a1b77]*/
diff --git a/Python/instrumentation.c b/Python/instrumentation.c
index 5e51a9c..8fd7c08 100644
--- a/Python/instrumentation.c
+++ b/Python/instrumentation.c
@@ -1660,6 +1660,16 @@ update_instrumentation_data(PyCodeObject *code, PyInterpreterState *interp)
     if (allocate_instrumentation_data(code)) {
         return -1;
     }
+    // If the local monitors are out of date, clear them up
+    _Py_LocalMonitors *local_monitors = &code->_co_monitoring->local_monitors;
+    for (int i = 0; i < PY_MONITORING_TOOL_IDS; i++) {
+        if (code->_co_monitoring->tool_versions[i] != interp->monitoring_tool_versions[i]) {
+            for (int j = 0; j < _PY_MONITORING_LOCAL_EVENTS; j++) {
+                local_monitors->tools[j] &= ~(1 << i);
+            }
+        }
+    }
+
     _Py_LocalMonitors all_events = local_union(
         interp->monitors,
         code->_co_monitoring->local_monitors);
@@ -2004,6 +2014,8 @@ _PyMonitoring_SetLocalEvents(PyCodeObject *code, int tool_id, _PyMonitoringEvent
         goto done;
     }
 
+    code->_co_monitoring->tool_versions[tool_id] = interp->monitoring_tool_versions[tool_id];
+
     _Py_LocalMonitors *local = &code->_co_monitoring->local_monitors;
     uint32_t existing_events = get_local_events(local, tool_id);
     if (existing_events == events) {
@@ -2036,6 +2048,43 @@ _PyMonitoring_GetLocalEvents(PyCodeObject *code, int tool_id, _PyMonitoringEvent
     return 0;
 }
 
+int _PyMonitoring_ClearToolId(int tool_id)
+{
+    assert(0 <= tool_id && tool_id < PY_MONITORING_TOOL_IDS);
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+
+    for (int i = 0; i < _PY_MONITORING_EVENTS; i++) {
+        PyObject *func = _PyMonitoring_RegisterCallback(tool_id, i, NULL);
+        if (func != NULL) {
+            Py_DECREF(func);
+        }
+    }
+
+    if (_PyMonitoring_SetEvents(tool_id, 0) < 0) {
+        return -1;
+    }
+
+    _PyEval_StopTheWorld(interp);
+    uint32_t version = global_version(interp) + MONITORING_VERSION_INCREMENT;
+    if (version == 0) {
+        PyErr_Format(PyExc_OverflowError, "events set too many times");
+        _PyEval_StartTheWorld(interp);
+        return -1;
+    }
+
+    // monitoring_tool_versions[tool_id] is set to latest global version here to
+    //   1. invalidate local events on all existing code objects
+    //   2. be ready for the next call to set local events
+    interp->monitoring_tool_versions[tool_id] = version;
+
+    // Set the new global version so all the code objects can refresh the
+    // instrumentation.
+    set_global_version(_PyThreadState_GET(), version);
+    int res = instrument_all_executing_code_objects(interp);
+    _PyEval_StartTheWorld(interp);
+    return res;
+}
+
 /*[clinic input]
 module monitoring
 [clinic start generated code]*/
@@ -2084,6 +2133,33 @@ monitoring_use_tool_id_impl(PyObject *module, int tool_id, PyObject *name)
 }
 
 /*[clinic input]
+monitoring.clear_tool_id
+
+    tool_id: int
+    /
+
+[clinic start generated code]*/
+
+static PyObject *
+monitoring_clear_tool_id_impl(PyObject *module, int tool_id)
+/*[clinic end generated code: output=04defc23470b1be7 input=af643d6648a66163]*/
+{
+    if (check_valid_tool(tool_id))  {
+        return NULL;
+    }
+
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+
+    if (interp->monitoring_tool_names[tool_id] != NULL) {
+        if (_PyMonitoring_ClearToolId(tool_id) < 0) {
+            return NULL;
+        }
+    }
+
+    Py_RETURN_NONE;
+}
+
+/*[clinic input]
 monitoring.free_tool_id
 
     tool_id: int
@@ -2099,6 +2175,13 @@ monitoring_free_tool_id_impl(PyObject *module, int tool_id)
         return NULL;
     }
     PyInterpreterState *interp = _PyInterpreterState_GET();
+
+    if (interp->monitoring_tool_names[tool_id] != NULL) {
+        if (_PyMonitoring_ClearToolId(tool_id) < 0) {
+            return NULL;
+        }
+    }
+
     Py_CLEAR(interp->monitoring_tool_names[tool_id]);
     Py_RETURN_NONE;
 }
@@ -2376,6 +2459,7 @@ monitoring__all_events_impl(PyObject *module)
 
 static PyMethodDef methods[] = {
     MONITORING_USE_TOOL_ID_METHODDEF
+    MONITORING_CLEAR_TOOL_ID_METHODDEF
     MONITORING_FREE_TOOL_ID_METHODDEF
     MONITORING_GET_TOOL_METHODDEF
     MONITORING_REGISTER_CALLBACK_METHODDEF
diff --git a/Python/pystate.c b/Python/pystate.c
index 9d11e2d..45e79ad 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -654,6 +654,7 @@ init_interpreter(PyInterpreterState *interp,
             interp->monitoring_callables[t][e] = NULL;
 
         }
+        interp->monitoring_tool_versions[t] = 0;
     }
     interp->sys_profile_initialized = false;
     interp->sys_trace_initialized = false;