pythongh-93503: Add APIs to set profiling and tracing functions in all threads in the C-API

Author: pablogsal

Doc/c-api/init.rst

@@ -1774,6 +1774,15 @@ Python-level trace functions in previous versions.
 
    The caller must hold the :term:`GIL`.
 
+.. c:function:: void PyEval_SetProfileAllThreads(Py_tracefunc func, PyObject *obj)
+
+   Like :c:func:`PyEval_SetProfile` but sets the profile function in all running threads
+   instead of the setting it only on the current thread.
+
+   The caller must hold the :term:`GIL`.
+
+.. versionadded:: 3.12
+
 
 .. c:function:: void PyEval_SetTrace(Py_tracefunc func, PyObject *obj)
 
@@ -1788,6 +1797,15 @@ Python-level trace functions in previous versions.
 
    The caller must hold the :term:`GIL`.
 
+.. c:function:: void PyEval_SetTraceAllThreads(Py_tracefunc func, PyObject *obj)
+
+   Like :c:func:`PyEval_SetTrace` but sets the tracing function in all running threads
+   instead of the setting it only on the current thread.
+
+   The caller must hold the :term:`GIL`.
+
+.. versionadded:: 3.12
+
 
 .. _advanced-debugging:
 

Doc/data/refcounts.dat

@@ -796,10 +796,18 @@ PyEval_SetProfile:void:::
 PyEval_SetProfile:Py_tracefunc:func::
 PyEval_SetProfile:PyObject*:obj:+1:
 
+PyEval_SetProfileAllThreads:void:::
+PyEval_SetProfileAllThreads:Py_tracefunc:func::
+PyEval_SetProfileAllThreads:PyObject*:obj:+1:
+
 PyEval_SetTrace:void:::
 PyEval_SetTrace:Py_tracefunc:func::
 PyEval_SetTrace:PyObject*:obj:+1:
 
+PyEval_SetTraceAllThreads:void:::
+PyEval_SetTraceAllThreads:Py_tracefunc:func::
+PyEval_SetTraceAllThreads:PyObject*:obj:+1:
+
 PyEval_EvalCode:PyObject*::+1:
 PyEval_EvalCode:PyObject*:co:0:
 PyEval_EvalCode:PyObject*:globals:0:

Doc/library/threading.rst

@@ -138,14 +138,19 @@ This module defines the following functions:
    .. versionadded:: 3.4
 
 
-.. function:: settrace(func)
+.. function:: settrace(func, *, running_threads=False)
 
    .. index:: single: trace function
 
    Set a trace function for all threads started from the :mod:`threading` module.
    The *func* will be passed to  :func:`sys.settrace` for each thread, before its
    :meth:`~Thread.run` method is called.
 
+   If *running_threads* is set, the tracing function will be installed in all running
+   threads.
+
+   .. versionchanged:: 3.12
+      The *running_threads* parameter was added.
 
 .. function:: gettrace()
 
@@ -158,14 +163,19 @@ This module defines the following functions:
    .. versionadded:: 3.10
 
 
-.. function:: setprofile(func)
+.. function:: setprofile(func, *, running_threads=False)
 
    .. index:: single: profile function
 
    Set a profile function for all threads started from the :mod:`threading` module.
    The *func* will be passed to  :func:`sys.setprofile` for each thread, before its
    :meth:`~Thread.run` method is called.
 
+   If *running_threads* is set, the profile function will be installed in all running
+   threads.
+
+   .. versionchanged:: 3.12
+      The *running_threads* parameter was added.
 
 .. function:: getprofile()
 

Include/cpython/ceval.h

@@ -3,8 +3,10 @@
 #endif
 
 PyAPI_FUNC(void) PyEval_SetProfile(Py_tracefunc, PyObject *);
+PyAPI_FUNC(void) PyEval_SetProfileAllThreads(Py_tracefunc, PyObject *);
 PyAPI_DATA(int) _PyEval_SetProfile(PyThreadState *tstate, Py_tracefunc func, PyObject *arg);
 PyAPI_FUNC(void) PyEval_SetTrace(Py_tracefunc, PyObject *);
+PyAPI_FUNC(void) PyEval_SetTraceAllThreads(Py_tracefunc, PyObject *);
 PyAPI_FUNC(int) _PyEval_SetTrace(PyThreadState *tstate, Py_tracefunc func, PyObject *arg);
 
 /* Helper to look up a builtin object */

Lib/test/test_threading.py

@@ -866,6 +866,34 @@ def noop_trace(frame, event, arg):
         finally:
             threading.settrace(old_trace)
 
+    def test_gettrace_all_threads(self):
+        def fn(*args): pass
+        old_trace = threading.gettrace()
+        first_check = threading.Event()
+        second_check = threading.Event()
+
+        trace_funcs = []
+        def checker():
+            trace_funcs.append(sys.gettrace())
+            first_check.set()
+            second_check.wait()
+            trace_funcs.append(sys.gettrace())
+
+        try:
+            t = threading.Thread(target=checker)
+            t.start()
+            first_check.wait()
+            threading.settrace(fn, running_threads=True)
+            second_check.set()
+            t.join()
+            self.assertEqual(trace_funcs, [None, fn])
+            self.assertEqual(threading.gettrace(), fn)
+            self.assertEqual(sys.gettrace(), fn)
+        finally:
+            threading.settrace(old_trace, running_threads=True)
+        self.assertEqual(threading.gettrace(), old_trace)
+        self.assertEqual(sys.gettrace(), old_trace)
+
     def test_getprofile(self):
         def fn(*args): pass
         old_profile = threading.getprofile()
@@ -875,6 +903,35 @@ def fn(*args): pass
         finally:
             threading.setprofile(old_profile)
 
+    def test_getprofile_all_threads(self):
+        def fn(*args): pass
+        old_profile = threading.getprofile()
+        first_check = threading.Event()
+        second_check = threading.Event()
+
+        profile_funcs = []
+        def checker():
+            profile_funcs.append(sys.getprofile())
+            first_check.set()
+            second_check.wait()
+            profile_funcs.append(sys.getprofile())
+
+        try:
+            t = threading.Thread(target=checker)
+            t.start()
+            first_check.wait()
+            threading.setprofile(fn, running_threads=True)
+            second_check.set()
+            t.join()
+            self.assertEqual(profile_funcs, [None, fn])
+            self.assertEqual(threading.getprofile(), fn)
+            self.assertEqual(sys.getprofile(), fn)
+        finally:
+            threading.setprofile(old_profile, running_threads=True)
+
+        self.assertEqual(threading.getprofile(), old_profile)
+        self.assertEqual(sys.getprofile(), old_profile)
+
     @cpython_only
     def test_shutdown_locks(self):
         for daemon in (False, True):

Lib/threading.py

@@ -55,30 +55,40 @@
 _profile_hook = None
 _trace_hook = None
 
-def setprofile(func):
+def setprofile(func, *, running_threads=False):
     """Set a profile function for all threads started from the threading module.
 
     The func will be passed to sys.setprofile() for each thread, before its
     run() method is called.
 
+    If running_threads is set, the profile function will be installed in all running
+    threads.
     """
     global _profile_hook
     _profile_hook = func
 
+    if running_threads:
+        _sys._setprofileallthreads(func)
+
 def getprofile():
     """Get the profiler function as set by threading.setprofile()."""
     return _profile_hook
 
-def settrace(func):
+def settrace(func, *, running_threads=False):
     """Set a trace function for all threads started from the threading module.
 
     The func will be passed to sys.settrace() for each thread, before its run()
     method is called.
 
+    If running_threads is set, the trace function will be installed in all running
+    threads.
     """
     global _trace_hook
     _trace_hook = func
 
+    if running_threads:
+        _sys._settraceallthreads(func)
+
 def gettrace():
     """Get the trace function as set by threading.settrace()."""
     return _trace_hook

Misc/NEWS.d/next/C API/2022-06-06-16-04-14.gh-issue-93503.MHJTu8.rst

@@ -0,0 +1,7 @@
+Add two new public functions to the public C-API,
+:c:func:`PyEval_SetProfileAllThreads and
+:c:func:`PyEval_SetTraceAllThreads`, that allow to set tracking and
+profiling functions in all running threads in addition to the calling one.
+Also, add a new *running_threads* parameter to :func:`threading.setprofile`
+and :func:`threading.settrace` that allows to do the same from Python. Patch
+by Pablo Galindo

Python/ceval.c

@@ -7033,6 +7033,20 @@ PyEval_SetProfile(Py_tracefunc func, PyObject *arg)
     }
 }
 
+void
+PyEval_SetProfileAllThreads(Py_tracefunc func, PyObject *arg)
+{
+    PyThreadState *this_tstate = _PyThreadState_GET();
+    PyInterpreterState* interp = this_tstate->interp;
+    PyThreadState* ts = PyInterpreterState_ThreadHead(interp);
+    while (ts) {
+        if (_PyEval_SetProfile(ts, func, arg) < 0) {
+            _PyErr_WriteUnraisableMsg("in PyEval_SetProfileAllThreads", NULL);
+        }
+        ts = PyThreadState_Next(ts);
+    }
+}
+
 int
 _PyEval_SetTrace(PyThreadState *tstate, Py_tracefunc func, PyObject *arg)
 {
@@ -7086,6 +7100,19 @@ PyEval_SetTrace(Py_tracefunc func, PyObject *arg)
     }
 }
 
+void
+PyEval_SetTraceAllThreads(Py_tracefunc func, PyObject *arg)
+{
+    PyThreadState *this_tstate = _PyThreadState_GET();
+    PyInterpreterState* interp = this_tstate->interp;
+    PyThreadState* ts = PyInterpreterState_ThreadHead(interp);
+    while (ts) {
+        if (_PyEval_SetTrace(ts, func, arg) < 0) {
+            _PyErr_WriteUnraisableMsg("in PyEval_SetTraceAllThreads", NULL);
+        }
+        ts = PyThreadState_Next(ts);
+    }
+}
 
 int
 _PyEval_SetCoroutineOriginTrackingDepth(int depth)