GH-93503: Add thread-specific APIs to set profiling and tracing functions in the C-API (#93504)

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

* Use a separate API

* Fix NEWS entry

* Add locks around the loop

* Document ignoring exceptions

* Use the new APIs in the sys module

* Update docs

10 files changed, 271 insertions, 4 deletions

diff --git a/Doc/c-api/init.rst b/Doc/c-api/init.rst
index 038498f..2a9cf0e 100644
--- a/Doc/c-api/init.rst
+++ b/Doc/c-api/init.rst
@@ -1774,6 +1774,18 @@ 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
+   belonging to the current interpreter instead of the setting it only on the current thread.
+
+   The caller must hold the :term:`GIL`.
+
+   As :c:func:`PyEval_SetProfile`, this function ignores any exceptions raised while
+   setting the profile functions in all threads.
+
+.. versionadded:: 3.12
+
 
 .. c:function:: void PyEval_SetTrace(Py_tracefunc func, PyObject *obj)
 
@@ -1788,6 +1800,18 @@ 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
+   belonging to the current interpreter instead of the setting it only on the current thread.
+
+   The caller must hold the :term:`GIL`.
+
+   As :c:func:`PyEval_SetTrace`, this function ignores any exceptions raised while
+   setting the trace functions in all threads.
+
+.. versionadded:: 3.12
+
 
 .. _advanced-debugging:
 
diff --git a/Doc/data/refcounts.dat b/Doc/data/refcounts.dat
index 1694cad..51ccacf 100644
--- a/Doc/data/refcounts.dat
+++ b/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:
diff --git a/Doc/library/threading.rst b/Doc/library/threading.rst
index b22d487..b352125 100644
--- a/Doc/library/threading.rst
+++ b/Doc/library/threading.rst
@@ -158,6 +158,15 @@ This module defines the following functions:
    The *func* will be passed to  :func:`sys.settrace` for each thread, before its
    :meth:`~Thread.run` method is called.
 
+.. function:: settrace_all_threads(func)
+
+   Set a trace function for all threads started from the :mod:`threading` module
+   and all Python threads that are currently executing.
+
+   The *func* will be passed to  :func:`sys.settrace` for each thread, before its
+   :meth:`~Thread.run` method is called.
+
+   .. versionadded:: 3.12
 
 .. function:: gettrace()
 
@@ -178,6 +187,15 @@ This module defines the following functions:
    The *func* will be passed to  :func:`sys.setprofile` for each thread, before its
    :meth:`~Thread.run` method is called.
 
+.. function:: setprofile_all_threads(func)
+
+   Set a profile function for all threads started from the :mod:`threading` module
+   and all Python threads that are currently executing.
+
+   The *func* will be passed to  :func:`sys.setprofile` for each thread, before its
+   :meth:`~Thread.run` method is called.
+
+   .. versionadded:: 3.12
 
 .. function:: getprofile()
 
diff --git a/Include/cpython/ceval.h b/Include/cpython/ceval.h
index 9d4eeaf..74665c9 100644
--- a/Include/cpython/ceval.h
+++ b/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 */
diff --git a/Lib/test/test_threading.py b/Lib/test/test_threading.py
index dcd2769..c664996 100644
--- a/Lib/test/test_threading.py
+++ b/Lib/test/test_threading.py
@@ -853,6 +853,7 @@ class ThreadTests(BaseTestCase):
                 callback()
         finally:
             sys.settrace(old_trace)
+            threading.settrace(old_trace)
 
     def test_gettrace(self):
         def noop_trace(frame, event, arg):
@@ -866,6 +867,35 @@ class ThreadTests(BaseTestCase):
         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_all_threads(fn)
+            second_check.set()
+            t.join()
+            self.assertEqual(trace_funcs, [None, fn])
+            self.assertEqual(threading.gettrace(), fn)
+            self.assertEqual(sys.gettrace(), fn)
+        finally:
+            threading.settrace_all_threads(old_trace)
+
+        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 +905,35 @@ class ThreadTests(BaseTestCase):
         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_all_threads(fn)
+            second_check.set()
+            t.join()
+            self.assertEqual(profile_funcs, [None, fn])
+            self.assertEqual(threading.getprofile(), fn)
+            self.assertEqual(sys.getprofile(), fn)
+        finally:
+            threading.setprofile_all_threads(old_profile)
+
+        self.assertEqual(threading.getprofile(), old_profile)
+        self.assertEqual(sys.getprofile(), old_profile)
+
     @cpython_only
     def test_shutdown_locks(self):
         for daemon in (False, True):
diff --git a/Lib/threading.py b/Lib/threading.py
index e32ad14..f28597c 100644
--- a/Lib/threading.py
+++ b/Lib/threading.py
@@ -28,7 +28,8 @@ __all__ = ['get_ident', 'active_count', 'Condition', 'current_thread',
            'Event', 'Lock', 'RLock', 'Semaphore', 'BoundedSemaphore', 'Thread',
            'Barrier', 'BrokenBarrierError', 'Timer', 'ThreadError',
            'setprofile', 'settrace', 'local', 'stack_size',
-           'excepthook', 'ExceptHookArgs', 'gettrace', 'getprofile']
+           'excepthook', 'ExceptHookArgs', 'gettrace', 'getprofile',
+           'setprofile_all_threads','settrace_all_threads']
 
 # Rename some stuff so "from threading import *" is safe
 _start_new_thread = _thread.start_new_thread
@@ -60,11 +61,20 @@ def setprofile(func):
 
     The func will be passed to sys.setprofile() for each thread, before its
     run() method is called.
-
     """
     global _profile_hook
     _profile_hook = func
 
+def setprofile_all_threads(func):
+    """Set a profile function for all threads started from the threading module
+    and all Python threads that are currently executing.
+
+    The func will be passed to sys.setprofile() for each thread, before its
+    run() method is called.
+    """
+    setprofile(func)
+    _sys._setprofileallthreads(func)
+
 def getprofile():
     """Get the profiler function as set by threading.setprofile()."""
     return _profile_hook
@@ -74,11 +84,20 @@ def settrace(func):
 
     The func will be passed to sys.settrace() for each thread, before its run()
     method is called.
-
     """
     global _trace_hook
     _trace_hook = func
 
+def settrace_all_threads(func):
+    """Set a trace function for all threads started from the threading module
+    and all Python threads that are currently executing.
+
+    The func will be passed to sys.settrace() for each thread, before its run()
+    method is called.
+    """
+    settrace(func)
+    _sys._settraceallthreads(func)
+
 def gettrace():
     """Get the trace function as set by threading.settrace()."""
     return _trace_hook
diff --git a/Misc/NEWS.d/next/C API/2022-06-06-16-04-14.gh-issue-93503.MHJTu8.rst b/Misc/NEWS.d/next/C API/2022-06-06-16-04-14.gh-issue-93503.MHJTu8.rst
new file mode 100644
index 0000000..6df9f95
--- /dev/null
+++ b/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
diff --git a/Python/ceval.c b/Python/ceval.c
index 1ab104c..ac77ab8 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -96,6 +96,10 @@
 #define _Py_atomic_load_relaxed_int32(ATOMIC_VAL) _Py_atomic_load_relaxed(ATOMIC_VAL)
 #endif
 
+#define HEAD_LOCK(runtime) \
+    PyThread_acquire_lock((runtime)->interpreters.mutex, WAIT_LOCK)
+#define HEAD_UNLOCK(runtime) \
+    PyThread_release_lock((runtime)->interpreters.mutex)
 
 /* Forward declarations */
 static PyObject *trace_call_function(
@@ -6455,6 +6459,27 @@ 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;
+
+    _PyRuntimeState *runtime = &_PyRuntime;
+    HEAD_LOCK(runtime);
+    PyThreadState* ts = PyInterpreterState_ThreadHead(interp);
+    HEAD_UNLOCK(runtime);
+
+    while (ts) {
+        if (_PyEval_SetProfile(ts, func, arg) < 0) {
+            _PyErr_WriteUnraisableMsg("in PyEval_SetProfileAllThreads", NULL);
+        }
+        HEAD_LOCK(runtime);
+        ts = PyThreadState_Next(ts);
+        HEAD_UNLOCK(runtime);
+    }
+}
+
 int
 _PyEval_SetTrace(PyThreadState *tstate, Py_tracefunc func, PyObject *arg)
 {
@@ -6508,6 +6533,26 @@ 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;
+
+    _PyRuntimeState *runtime = &_PyRuntime;
+    HEAD_LOCK(runtime);
+    PyThreadState* ts = PyInterpreterState_ThreadHead(interp);
+    HEAD_UNLOCK(runtime);
+
+    while (ts) {
+        if (_PyEval_SetTrace(ts, func, arg) < 0) {
+            _PyErr_WriteUnraisableMsg("in PyEval_SetTraceAllThreads", NULL);
+        }
+        HEAD_LOCK(runtime);
+        ts = PyThreadState_Next(ts);
+        HEAD_UNLOCK(runtime);
+    }
+}
 
 int
 _PyEval_SetCoroutineOriginTrackingDepth(int depth)
diff --git a/Python/clinic/sysmodule.c.h b/Python/clinic/sysmodule.c.h
index beaf21c..0f96366 100644
--- a/Python/clinic/sysmodule.c.h
+++ b/Python/clinic/sysmodule.c.h
@@ -292,6 +292,18 @@ exit:
     return return_value;
 }
 
+PyDoc_STRVAR(sys__settraceallthreads__doc__,
+"_settraceallthreads($module, arg, /)\n"
+"--\n"
+"\n"
+"Set the global debug tracing function in all running threads belonging to the current interpreter.\n"
+"\n"
+"It will be called on each function call. See the debugger chapter\n"
+"in the library manual.");
+
+#define SYS__SETTRACEALLTHREADS_METHODDEF    \
+    {"_settraceallthreads", (PyCFunction)sys__settraceallthreads, METH_O, sys__settraceallthreads__doc__},
+
 PyDoc_STRVAR(sys_gettrace__doc__,
 "gettrace($module, /)\n"
 "--\n"
@@ -312,6 +324,18 @@ sys_gettrace(PyObject *module, PyObject *Py_UNUSED(ignored))
     return sys_gettrace_impl(module);
 }
 
+PyDoc_STRVAR(sys__setprofileallthreads__doc__,
+"_setprofileallthreads($module, arg, /)\n"
+"--\n"
+"\n"
+"Set the profiling function in all running threads belonging to the current interpreter.\n"
+"\n"
+"It will be called on each function call and return.  See the profiler chapter\n"
+"in the library manual.");
+
+#define SYS__SETPROFILEALLTHREADS_METHODDEF    \
+    {"_setprofileallthreads", (PyCFunction)sys__setprofileallthreads, METH_O, sys__setprofileallthreads__doc__},
+
 PyDoc_STRVAR(sys_getprofile__doc__,
 "getprofile($module, /)\n"
 "--\n"
@@ -1170,4 +1194,4 @@ sys_getandroidapilevel(PyObject *module, PyObject *Py_UNUSED(ignored))
 #ifndef SYS_GETANDROIDAPILEVEL_METHODDEF
     #define SYS_GETANDROIDAPILEVEL_METHODDEF
 #endif /* !defined(SYS_GETANDROIDAPILEVEL_METHODDEF) */
-/*[clinic end generated code: output=38446a4c76e2f3b6 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=322fb0409e376ad4 input=a9049054013a1b77]*/
diff --git a/Python/sysmodule.c b/Python/sysmodule.c
index b8009b2..c286438 100644
--- a/Python/sysmodule.c
+++ b/Python/sysmodule.c
@@ -1022,6 +1022,36 @@ function call.  See the debugger chapter in the library manual."
 );
 
 /*[clinic input]
+sys._settraceallthreads
+
+    arg: object
+    /
+
+Set the global debug tracing function in all running threads belonging to the current interpreter.
+
+It will be called on each function call. See the debugger chapter
+in the library manual.
+[clinic start generated code]*/
+
+static PyObject *
+sys__settraceallthreads(PyObject *module, PyObject *arg)
+/*[clinic end generated code: output=161cca30207bf3ca input=5906aa1485a50289]*/
+{
+    PyObject* argument = NULL;
+    Py_tracefunc func = NULL;
+
+    if (arg != Py_None) {
+        func = trace_trampoline;
+        argument = arg;
+    }
+
+
+    PyEval_SetTraceAllThreads(func, argument);
+
+    Py_RETURN_NONE;
+}
+
+/*[clinic input]
 sys.gettrace
 
 Return the global debug tracing function set with sys.settrace.
@@ -1067,6 +1097,35 @@ and return.  See the profiler chapter in the library manual."
 );
 
 /*[clinic input]
+sys._setprofileallthreads
+
+    arg: object
+    /
+
+Set the profiling function in all running threads belonging to the current interpreter.
+
+It will be called on each function call and return.  See the profiler chapter
+in the library manual.
+[clinic start generated code]*/
+
+static PyObject *
+sys__setprofileallthreads(PyObject *module, PyObject *arg)
+/*[clinic end generated code: output=2d61319e27b309fe input=d1a356d3f4f9060a]*/
+{
+    PyObject* argument = NULL;
+    Py_tracefunc func = NULL;
+
+    if (arg != Py_None) {
+        func = profile_trampoline;
+        argument = arg;
+    }
+
+    PyEval_SetProfileAllThreads(func, argument);
+
+    Py_RETURN_NONE;
+}
+
+/*[clinic input]
 sys.getprofile
 
 Return the profiling function set with sys.setprofile.
@@ -2035,9 +2094,11 @@ static PyMethodDef sys_methods[] = {
     SYS_GETSWITCHINTERVAL_METHODDEF
     SYS_SETDLOPENFLAGS_METHODDEF
     {"setprofile", sys_setprofile, METH_O, setprofile_doc},
+    SYS__SETPROFILEALLTHREADS_METHODDEF
     SYS_GETPROFILE_METHODDEF
     SYS_SETRECURSIONLIMIT_METHODDEF
     {"settrace", sys_settrace, METH_O, settrace_doc},
+    SYS__SETTRACEALLTHREADS_METHODDEF
     SYS_GETTRACE_METHODDEF
     SYS_CALL_TRACING_METHODDEF
     SYS__DEBUGMALLOCSTATS_METHODDEF