3 files changed, 113 insertions, 1 deletions

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;