add DISABLE support for PY_UNWIND

Author: P403n1x87

Lib/test/test_monitoring.py

@@ -198,11 +198,11 @@ def test_c_return_count(self):
 
 EXCEPT_EVENTS = [
     (E.RAISE, "raise"),
-    (E.PY_UNWIND, "unwind"),
     (E.EXCEPTION_HANDLED, "exception_handled"),
 ]
 
 SIMPLE_EVENTS = INSTRUMENTED_EVENTS + EXCEPT_EVENTS + [
+    (E.PY_UNWIND, "unwind"),  # local event (not in EXCEPT_EVENTS: DISABLE is legal)
     (E.C_RAISE, "c_raise"),
     (E.C_RETURN, "c_return"),
 ]
@@ -1484,6 +1484,72 @@ def test_set_non_local_event(self):
         with self.assertRaises(ValueError):
             sys.monitoring.set_local_events(TEST_TOOL, just_call.__code__, E.RAISE)
 
+    def test_local_py_unwind(self):
+        """PY_UNWIND fires as a local event only for the instrumented code object."""
+
+        def foo():
+            raise RuntimeError("test")
+
+        def bar():
+            raise RuntimeError("test")
+
+        events = []
+
+        def callback(code, offset, exc):
+            events.append(code.co_name)
+
+        try:
+            sys.monitoring.register_callback(TEST_TOOL, E.PY_UNWIND, callback)
+            sys.monitoring.set_local_events(TEST_TOOL, foo.__code__, E.PY_UNWIND)
+
+            try:
+                foo()
+            except RuntimeError:
+                pass
+
+            try:
+                bar()  # should NOT trigger the callback
+            except RuntimeError:
+                pass
+
+            self.assertEqual(events, ['foo'])
+        finally:
+            sys.monitoring.set_local_events(TEST_TOOL, foo.__code__, 0)
+            sys.monitoring.register_callback(TEST_TOOL, E.PY_UNWIND, None)
+
+    def test_local_py_unwind_disable(self):
+        """Returning DISABLE from a PY_UNWIND callback disables it for that code object."""
+
+        call_count = 0
+
+        def foo():
+            raise RuntimeError("test")
+
+        def callback(code, offset, exc):
+            nonlocal call_count
+            call_count += 1
+            return sys.monitoring.DISABLE
+
+        try:
+            sys.monitoring.register_callback(TEST_TOOL, E.PY_UNWIND, callback)
+            sys.monitoring.set_local_events(TEST_TOOL, foo.__code__, E.PY_UNWIND)
+
+            try:
+                foo()
+            except RuntimeError:
+                pass
+            self.assertEqual(call_count, 1)  # fired once
+
+            try:
+                foo()
+            except RuntimeError:
+                pass
+            self.assertEqual(call_count, 1)  # not fired again — disabled by DISABLE return
+
+        finally:
+            sys.monitoring.set_local_events(TEST_TOOL, foo.__code__, 0)
+            sys.monitoring.register_callback(TEST_TOOL, E.PY_UNWIND, None)
+
 def line_from_offset(code, offset):
     for start, end, line in code.co_lines():
         if start <= offset < end:

Python/instrumentation.c

@@ -1139,6 +1139,25 @@ static const char *const event_names [] = {
     [PY_MONITORING_EVENT_STOP_ITERATION] = "STOP_ITERATION",
 };
 
+/* Disable a local-but-not-instrumented event (e.g. PY_UNWIND) for a single
+ * tool on this code object.  Must be called with the world stopped or the
+ * code lock held. */
+static void
+remove_local_tool(PyCodeObject *code, PyInterpreterState *interp,
+                  int event, int tool)
+{
+    ASSERT_WORLD_STOPPED_OR_LOCKED(code);
+    assert(PY_MONITORING_IS_LOCAL_EVENT(event));
+    assert(!PY_MONITORING_IS_INSTRUMENTED_EVENT(event));
+    assert(code->_co_monitoring);
+    code->_co_monitoring->local_monitors.tools[event] &= ~(1 << tool);
+    /* Recompute active_monitors for this event as the union of global and
+     * (now updated) local monitors. */
+    code->_co_monitoring->active_monitors.tools[event] =
+        interp->monitors.tools[event] |
+        code->_co_monitoring->local_monitors.tools[event];
+}
+
 static int
 call_instrumentation_vector(
     _Py_CODEUNIT *instr, PyThreadState *tstate, int event,
@@ -1183,7 +1202,19 @@ call_instrumentation_vector(
         }
         else {
             /* DISABLE */
-            if (!PY_MONITORING_IS_INSTRUMENTED_EVENT(event)) {
+            if (PY_MONITORING_IS_INSTRUMENTED_EVENT(event)) {
+                _PyEval_StopTheWorld(interp);
+                remove_tools(code, offset, event, 1 << tool);
+                _PyEval_StartTheWorld(interp);
+            }
+            else if (PY_MONITORING_IS_LOCAL_EVENT(event)) {
+                /* Local but not tied to a bytecode instruction: disable for
+                 * this code object entirely. */
+                _PyEval_StopTheWorld(interp);
+                remove_local_tool(code, interp, event, tool);
+                _PyEval_StartTheWorld(interp);
+            }
+            else {
                 PyErr_Format(PyExc_ValueError,
                               "Cannot disable %s events. Callback removed.",
                              event_names[event]);
@@ -1192,12 +1223,6 @@ call_instrumentation_vector(
                 err = -1;
                 break;
             }
-            else {
-                PyInterpreterState *interp = tstate->interp;
-                _PyEval_StopTheWorld(interp);
-                remove_tools(code, offset, event, 1 << tool);
-                _PyEval_StartTheWorld(interp);
-            }
         }
     }
     Py_DECREF(arg2_obj);