gh-137400: Fix a thread-safety issues when profiling all threads

There were a few thread-safety issues when profiling or tracing all
threads via PyEval_SetProfileAllThreads or PyEval_SetTraceAllThreads:

* The loop over thread states could crash if a thread exits concurrently
  (in both the free threading and default build)
* The modification of `c_profilefunc` and `c_tracefunc` wasn't
  thread-safe on the free threading build.

Author: colesbury

Include/internal/pycore_ceval.h

@@ -22,8 +22,10 @@ struct _ceval_runtime_state;
 
 // Export for '_lsprof' shared extension
 PyAPI_FUNC(int) _PyEval_SetProfile(PyThreadState *tstate, Py_tracefunc func, PyObject *arg);
+extern int _PyEval_SetProfileAllThreads(PyInterpreterState *interp, Py_tracefunc func, PyObject *arg);
 
 extern int _PyEval_SetTrace(PyThreadState *tstate, Py_tracefunc func, PyObject *arg);
+extern int _PyEval_SetTraceAllThreads(PyInterpreterState *interp, Py_tracefunc func, PyObject *arg);
 
 extern int _PyEval_SetOpcodeTrace(PyFrameObject *f, bool enable);
 

Include/internal/pycore_interp_structs.h

@@ -98,7 +98,6 @@ struct _ceval_runtime_state {
     // For example, we use a preallocated array
     // for the list of pending calls.
     struct _pending_calls pending_mainthread;
-    PyMutex sys_trace_profile_mutex;
 };
 
 
@@ -950,8 +949,8 @@ struct _is {
     PyDict_WatchCallback builtins_dict_watcher;
 
     _Py_GlobalMonitors monitors;
-    bool sys_profile_initialized;
-    bool sys_trace_initialized;
+    _PyOnceFlag sys_profile_once_flag;
+    _PyOnceFlag sys_trace_once_flag;
     Py_ssize_t sys_profiling_threads; /* Count of threads with c_profilefunc set */
     Py_ssize_t sys_tracing_threads; /* Count of threads with c_tracefunc set */
     PyObject *monitoring_callables[PY_MONITORING_TOOL_IDS][_PY_MONITORING_EVENTS];

Lib/test/test_free_threading/test_monitoring.py

@@ -194,6 +194,32 @@ def during_threads(self):
         self.set = not self.set
 
 
+@threading_helper.requires_working_threading()
+class SetProfileAllThreadsMultiThreaded(InstrumentationMultiThreadedMixin, TestCase):
+    """Uses threading.setprofile_all_threads and repeatedly toggles instrumentation on and off"""
+
+    def setUp(self):
+        self.set = False
+        self.called = False
+
+    def after_test(self):
+        self.assertTrue(self.called)
+
+    def tearDown(self):
+        threading.setprofile_all_threads(None)
+
+    def trace_func(self, frame, event, arg):
+        self.called = True
+        return self.trace_func
+
+    def during_threads(self):
+        if self.set:
+            threading.setprofile_all_threads(self.trace_func)
+        else:
+            threading.setprofile_all_threads(None)
+        self.set = not self.set
+
+
 class TraceBuf:
     def __init__(self):
         self.traces = []

Python/ceval.c

@@ -2512,18 +2512,9 @@ PyEval_SetProfileAllThreads(Py_tracefunc func, PyObject *arg)
     PyThreadState *this_tstate = _PyThreadState_GET();
     PyInterpreterState* interp = this_tstate->interp;
 
-    _PyRuntimeState *runtime = &_PyRuntime;
-    HEAD_LOCK(runtime);
-    PyThreadState* ts = PyInterpreterState_ThreadHead(interp);
-    HEAD_UNLOCK(runtime);
-
-    while (ts) {
-        if (_PyEval_SetProfile(ts, func, arg) < 0) {
-            PyErr_FormatUnraisable("Exception ignored in PyEval_SetProfileAllThreads");
-        }
-        HEAD_LOCK(runtime);
-        ts = PyThreadState_Next(ts);
-        HEAD_UNLOCK(runtime);
+    if (_PyEval_SetProfileAllThreads(interp, func, arg) < 0) {
+        /* Log _PySys_Audit() error */
+        PyErr_FormatUnraisable("Exception ignored in PyEval_SetProfileAllThreads");
     }
 }
 
@@ -2543,18 +2534,9 @@ PyEval_SetTraceAllThreads(Py_tracefunc func, PyObject *arg)
     PyThreadState *this_tstate = _PyThreadState_GET();
     PyInterpreterState* interp = this_tstate->interp;
 
-    _PyRuntimeState *runtime = &_PyRuntime;
-    HEAD_LOCK(runtime);
-    PyThreadState* ts = PyInterpreterState_ThreadHead(interp);
-    HEAD_UNLOCK(runtime);
-
-    while (ts) {
-        if (_PyEval_SetTrace(ts, func, arg) < 0) {
-            PyErr_FormatUnraisable("Exception ignored in PyEval_SetTraceAllThreads");
-        }
-        HEAD_LOCK(runtime);
-        ts = PyThreadState_Next(ts);
-        HEAD_UNLOCK(runtime);
+    if (_PyEval_SetTraceAllThreads(interp, func, arg) < 0) {
+        /* Log _PySys_Audit() error */
+        PyErr_FormatUnraisable("Exception ignored in PyEval_SetTraceAllThreads");
     }
 }
 

Python/instrumentation.c

@@ -1040,6 +1040,8 @@ set_version_raw(uintptr_t *ptr, uint32_t version)
 static void
 set_global_version(PyThreadState *tstate, uint32_t version)
 {
+    ASSERT_WORLD_STOPPED();
+
     assert((version & _PY_EVAL_EVENTS_MASK) == 0);
     PyInterpreterState *interp = tstate->interp;
     set_version_raw(&interp->ceval.instrumentation_version, version);
@@ -1939,28 +1941,30 @@ _Py_Instrument(PyCodeObject *code, PyInterpreterState *interp)
 
 
 static int
-instrument_all_executing_code_objects(PyInterpreterState *interp) {
+instrument_all_executing_code_objects(PyInterpreterState *interp)
+{
     ASSERT_WORLD_STOPPED();
 
+    int err = 0;
     _PyRuntimeState *runtime = &_PyRuntime;
     HEAD_LOCK(runtime);
-    PyThreadState* ts = PyInterpreterState_ThreadHead(interp);
-    HEAD_UNLOCK(runtime);
+    PyThreadState *ts = PyInterpreterState_ThreadHead(interp);
     while (ts) {
         _PyInterpreterFrame *frame = ts->current_frame;
         while (frame) {
             if (frame->owner < FRAME_OWNED_BY_INTERPRETER) {
-                if (instrument_lock_held(_PyFrame_GetCode(frame), interp)) {
-                    return -1;
+                err = instrument_lock_held(_PyFrame_GetCode(frame), interp);
+                if (err) {
+                    goto done;
                 }
             }
             frame = frame->previous;
         }
-        HEAD_LOCK(runtime);
         ts = PyThreadState_Next(ts);
-        HEAD_UNLOCK(runtime);
     }
-    return 0;
+done:
+    HEAD_UNLOCK(runtime);
+    return err;
 }
 
 static void
@@ -2006,6 +2010,7 @@ check_tool(PyInterpreterState *interp, int tool_id)
 int
 _PyMonitoring_SetEvents(int tool_id, _PyMonitoringEventSet events)
 {
+    ASSERT_WORLD_STOPPED();
     assert(0 <= tool_id && tool_id < PY_MONITORING_TOOL_IDS);
     PyThreadState *tstate = _PyThreadState_GET();
     PyInterpreterState *interp = tstate->interp;
@@ -2014,33 +2019,28 @@ _PyMonitoring_SetEvents(int tool_id, _PyMonitoringEventSet events)
         return -1;
     }
 
-    int res;
-    _PyEval_StopTheWorld(interp);
     uint32_t existing_events = get_events(&interp->monitors, tool_id);
     if (existing_events == events) {
-        res = 0;
-        goto done;
+        return 0;
     }
     set_events(&interp->monitors, tool_id, events);
     uint32_t new_version = global_version(interp) + MONITORING_VERSION_INCREMENT;
     if (new_version == 0) {
         PyErr_Format(PyExc_OverflowError, "events set too many times");
-        res = -1;
-        goto done;
+        return -1;
     }
     set_global_version(tstate, new_version);
 #ifdef _Py_TIER2
     _Py_Executors_InvalidateAll(interp, 1);
 #endif
-    res = instrument_all_executing_code_objects(interp);
-done:
-    _PyEval_StartTheWorld(interp);
-    return res;
+    return instrument_all_executing_code_objects(interp);
 }
 
 int
 _PyMonitoring_SetLocalEvents(PyCodeObject *code, int tool_id, _PyMonitoringEventSet events)
 {
+    ASSERT_WORLD_STOPPED();
+
     assert(0 <= tool_id && tool_id < PY_MONITORING_TOOL_IDS);
     PyInterpreterState *interp = _PyInterpreterState_GET();
     assert(events < (1 << _PY_MONITORING_LOCAL_EVENTS));
@@ -2052,28 +2052,20 @@ _PyMonitoring_SetLocalEvents(PyCodeObject *code, int tool_id, _PyMonitoringEvent
         return -1;
     }
 
-    int res;
-    _PyEval_StopTheWorld(interp);
     if (allocate_instrumentation_data(code)) {
-        res = -1;
-        goto done;
+        return -1;
     }
 
     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) {
-        res = 0;
-        goto done;
+        return 0;
     }
     set_local_events(local, tool_id, events);
 
-    res = force_instrument_lock_held(code, interp);
-
-done:
-    _PyEval_StartTheWorld(interp);
-    return res;
+    return force_instrument_lock_held(code, interp);
 }
 
 int
@@ -2105,11 +2097,12 @@ int _PyMonitoring_ClearToolId(int tool_id)
         }
     }
 
+    _PyEval_StopTheWorld(interp);
     if (_PyMonitoring_SetEvents(tool_id, 0) < 0) {
+        _PyEval_StartTheWorld(interp);
         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");
@@ -2346,7 +2339,11 @@ monitoring_set_events_impl(PyObject *module, int tool_id, int event_set)
         event_set &= ~(1 << PY_MONITORING_EVENT_BRANCH);
         event_set |= (1 << PY_MONITORING_EVENT_BRANCH_RIGHT) | (1 << PY_MONITORING_EVENT_BRANCH_LEFT);
     }
-    if (_PyMonitoring_SetEvents(tool_id, event_set)) {
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    _PyEval_StopTheWorld(interp);
+    int err = _PyMonitoring_SetEvents(tool_id, event_set);
+    _PyEval_StartTheWorld(interp);
+    if (err) {
         return NULL;
     }
     Py_RETURN_NONE;
@@ -2427,7 +2424,11 @@ monitoring_set_local_events_impl(PyObject *module, int tool_id,
         return NULL;
     }
 
-    if (_PyMonitoring_SetLocalEvents((PyCodeObject*)code, tool_id, event_set)) {
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    _PyEval_StopTheWorld(interp);
+    int err = _PyMonitoring_SetLocalEvents((PyCodeObject*)code, tool_id, event_set);
+    _PyEval_StartTheWorld(interp);
+    if (err) {
         return NULL;
     }
     Py_RETURN_NONE;