gh-159312: Prefer bound stacks for CPU sample attribution

Author: pablogsal

Lib/test/test_external_inspection.py

@@ -2596,6 +2596,109 @@ def test_unsupported_platform_error(self):
 
 @requires_remote_subprocess_debugging()
 class TestDetectionOfThreadStatus(RemoteInspectionTestBase):
+    def _check_cpu_status_with_multiple_thread_states(
+        self, *, native=False, subinterpreter=False,
+    ):
+        import_module("_testinternalcapi")
+        import_module("_interpreters")
+        if native:
+            import_module("_hashlib")
+            work = "_hashlib.pbkdf2_hmac('sha256', b'password', b'salt', 10_000_000)"
+        else:
+            work = "pass"
+        code = textwrap.dedent(f"""
+            import socket
+            {"import _hashlib" if native else ""}
+            active_sock = socket.socket(fileno={{fd}})
+
+            def active():
+                active_sock.sendall(b"ready")
+                while True:
+                    {work}
+
+            active()
+        """)
+        script = f"""
+            import _interpreters, _testinternalcapi
+
+            def suspended():
+                if {subinterpreter}:
+                    interpid = _interpreters.create()
+                else:
+                    interpid = _interpreters.get_current()[0]
+                # Create another state on this OS thread.
+                code = {code!r}.format(fd=sock.fileno())
+                _testinternalcapi.exec_interpreter(interpid, code)
+
+            suspended()
+        """
+        with self._target_process(script) as (p, sock, _):
+            _wait_for_signal(sock, b"ready")
+            for mode in (PROFILING_MODE_CPU, PROFILING_MODE_ALL):
+                with self.subTest(mode=mode):
+                    unwinder = RemoteUnwinder(
+                        p.pid, all_threads=True, mode=mode,
+                        skip_non_matching_threads=(mode == PROFILING_MODE_CPU),
+                    )
+                    for _ in busy_retry(SHORT_TIMEOUT):
+                        try:
+                            trace = unwinder.get_stack_trace()
+                        except TRANSIENT_ERRORS:
+                            continue
+                        active = self._find_thread_with_frame(
+                            trace, lambda f: f.funcname == "active")
+                        if active is None or not active.status & THREAD_STATUS_ON_CPU:
+                            continue
+                        if native and active.status & THREAD_STATUS_HAS_GIL:
+                            continue
+                        suspended = self._find_thread_with_frame(
+                            trace, lambda f: f.funcname == "suspended")
+                        if mode == PROFILING_MODE_CPU:
+                            self.assertIsNone(suspended)
+                        else:
+                            self.assertIsNotNone(suspended)
+                            self.assertEqual(suspended.thread_id, active.thread_id)
+                            self.assertFalse(suspended.status & THREAD_STATUS_ON_CPU)
+                        break
+
+    @skip_if_not_supported
+    def test_cpu_status_with_multiple_thread_states(self):
+        for native in (False, True):
+            for subinterpreter in (False, True):
+                with self.subTest(native=native, subinterpreter=subinterpreter):
+                    self._check_cpu_status_with_multiple_thread_states(
+                        native=native, subinterpreter=subinterpreter)
+
+    @skip_if_not_supported
+    def test_cpu_status_after_tstate_callback(self):
+        import_module("_testinternalcapi")
+        for retain in (False, True):
+            with self.subTest(retain=retain):
+                script = f"""
+                    import _testinternalcapi
+
+                    def ready():
+                        sock.sendall(b"ready")
+
+                    def native_work():
+                        _testinternalcapi.native_work_after_tstate_callback(
+                            ready, {SHORT_TIMEOUT * 3}, {retain})
+
+                    native_work()
+                """
+                with self._target_process(script) as (p, sock, _):
+                    _wait_for_signal(sock, b"ready")
+                    for mode in (PROFILING_MODE_CPU, PROFILING_MODE_ALL):
+                        with self.subTest(mode=mode):
+                            unwinder = RemoteUnwinder(p.pid, all_threads=True, mode=mode)
+                            for _ in busy_retry(SHORT_TIMEOUT):
+                                trace = _get_stack_trace_with_retry(unwinder)
+                                thread = self._find_thread_with_frame(
+                                    trace, lambda f: f.funcname == "native_work")
+                                if thread and thread.status & THREAD_STATUS_ON_CPU:
+                                    self.assertFalse(thread.status & THREAD_STATUS_HAS_GIL)
+                                    break
+
     def _run_thread_status_test(self, mode, check_condition):
         """
         Common pattern for thread status detection tests.

Misc/NEWS.d/next/Library/2026-10-13-18-00-00.gh-issue-159312.Lk3w9a.rst

@@ -0,0 +1,2 @@
+Fix :mod:`profiling.sampling` attributing CPU samples to suspended thread states
+when multiple thread states share the same native thread.

Modules/_remote_debugging/_remote_debugging.h

@@ -689,7 +689,8 @@ extern PyObject* unwind_stack_for_thread(
     uintptr_t gil_holder_tstate,
     uintptr_t gc_frame,
     uintptr_t main_thread_tstate,
-    const RemoteReadPrefetch *prefetch
+    const RemoteReadPrefetch *prefetch,
+    int *has_bound_stack
 );
 
 /* Thread stopping functions (for blocking mode) */

Modules/_remote_debugging/module.c

@@ -650,6 +650,45 @@ read_interp_state_and_maybe_thread_frame(
         interp_state_buffer);
 }
 
+static int
+filter_cpu_thread_states(PyObject *result, PyObject *bound_threads, int skip)
+{
+    // Prefer a bound stack when states share an OS thread. There may be no
+    // bound state after a native callback deletes its temporary thread state.
+    for (Py_ssize_t i = 0; i < PyList_GET_SIZE(result); i++) {
+        PyObject *threads = PyStructSequence_GET_ITEM(PyList_GET_ITEM(result, i), 1);
+        for (Py_ssize_t j = PyList_GET_SIZE(threads); j-- > 0;) {
+            PyObject *thread = PyList_GET_ITEM(threads, j);
+            PyObject *bound = PyDict_GetItemWithError(
+                bound_threads, PyStructSequence_GET_ITEM(thread, 0));
+            if (bound == NULL) {
+                if (PyErr_Occurred()) {
+                    return -1;
+                }
+                continue;
+            }
+            long status = PyLong_AsLong(PyStructSequence_GET_ITEM(thread, 1));
+            if (bound == thread || !(status & THREAD_STATUS_ON_CPU)) {
+                continue;
+            }
+            if (skip) {
+                if (PySequence_DelItem(threads, j) < 0) {
+                    return -1;
+                }
+            }
+            else {
+                PyObject *new_status = PyLong_FromLong(status & ~THREAD_STATUS_ON_CPU);
+                if (new_status == NULL) {
+                    return -1;
+                }
+                Py_DECREF(PyStructSequence_GET_ITEM(thread, 1));
+                PyStructSequence_SET_ITEM(thread, 1, new_status);
+            }
+        }
+    }
+    return 0;
+}
+
 /*[clinic input]
 @critical_section
 _remote_debugging.RemoteUnwinder.get_stack_trace
@@ -714,6 +753,15 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
         set_exception_cause(self, PyExc_MemoryError, "Failed to create stack trace result list");
         return NULL;
     }
+    PyObject *bound_threads = NULL;
+    if (self->tstate_addr == 0 &&
+        (self->mode == PROFILING_MODE_CPU || self->mode == PROFILING_MODE_ALL)) {
+        bound_threads = PyDict_New();
+        if (bound_threads == NULL) {
+            Py_DECREF(result);
+            return NULL;
+        }
+    }
 
     // Iterate over all interpreters
     uintptr_t current_interpreter = self->interpreter_addr;
@@ -802,11 +850,13 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
 
         while (current_tstate != 0) {
             uintptr_t prev_tstate = current_tstate;
+            int has_bound_stack;
             PyObject* frame_info = unwind_stack_for_thread(self, &current_tstate,
                                                            gil_holder_tstate,
                                                            gc_frame,
                                                            main_thread_tstate,
-                                                           &prefetch);
+                                                           &prefetch,
+                                                           &has_bound_stack);
             if (!frame_info) {
                 // Check if this was an intentional skip due to mode-based filtering
                 if ((self->mode == PROFILING_MODE_CPU || self->mode == PROFILING_MODE_GIL ||
@@ -832,10 +882,13 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
                 goto exit;
             }
 
-            if (PyList_Append(interpreter_threads, frame_info) == -1) {
+            if (PyList_Append(interpreter_threads, frame_info) == -1 ||
+                (bound_threads && has_bound_stack &&
+                 PyDict_SetItem(bound_threads,
+                                PyStructSequence_GET_ITEM(frame_info, 0), frame_info) < 0)) {
                 Py_DECREF(frame_info);
                 Py_DECREF(interpreter_threads);
-                set_exception_cause(self, PyExc_RuntimeError, "Failed to append thread frame info");
+                set_exception_cause(self, PyExc_RuntimeError, "Failed to collect thread frame info");
                 Py_CLEAR(result);
                 goto exit;
             }
@@ -890,7 +943,14 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
         }
     }
 
+    if (bound_threads != NULL &&
+        filter_cpu_thread_states(result, bound_threads,
+            self->mode == PROFILING_MODE_CPU && self->skip_non_matching_threads) < 0) {
+        Py_CLEAR(result);
+    }
+
 exit:
+    Py_XDECREF(bound_threads);
     // Invalidate cache entries for threads not seen in this sample.
     // Only do this every 1024 iterations to avoid performance overhead.
     if (self->cache_frames && result) {

Modules/_remote_debugging/threads.c

@@ -335,7 +335,8 @@ unwind_stack_for_thread(
     uintptr_t gil_holder_tstate,
     uintptr_t gc_frame,
     uintptr_t main_thread_tstate,
-    const RemoteReadPrefetch *prefetch
+    const RemoteReadPrefetch *prefetch,
+    int *has_bound_stack
 ) {
     PyObject *frame_info = NULL;
     PyObject *thread_id = NULL;
@@ -618,6 +619,10 @@ unwind_stack_for_thread(
         goto error;
     }
 
+    *has_bound_stack = GET_MEMBER(
+        _thread_status, ts, unwinder->debug_offsets.thread_state.status).bound_gilstate
+        && PyList_GET_SIZE(frame_info) != 0;
+
     // py_status contains status flags (bitfield)
     PyStructSequence_SetItem(result, 0, thread_id);
     PyStructSequence_SetItem(result, 1, py_status);  // Steals reference

Modules/_testinternalcapi.c

@@ -2372,6 +2372,49 @@ destroy_interpreter(PyObject *self, PyObject *args, PyObject *kwargs)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+native_work_after_tstate_callback(PyObject *self, PyObject *args)
+{
+    PyObject *callback;
+    double seconds;
+    int retain;
+    if (!PyArg_ParseTuple(args, "Odp", &callback, &seconds, &retain)) {
+        return NULL;
+    }
+    PyTime_t start, now;
+    if (PyTime_Monotonic(&start) < 0) {
+        return NULL;
+    }
+    PyThreadState *temporary = PyThreadState_New(PyInterpreterState_Get());
+    if (temporary == NULL) {
+        return PyErr_NoMemory();
+    }
+    PyThreadState *saved = PyThreadState_Swap(temporary);
+    PyObject *result = PyObject_CallNoArgs(callback);
+    PyObject *exc = result ? NULL : PyErr_GetRaisedException();
+    if (retain) {
+        PyThreadState_Swap(NULL);
+    }
+    else {
+        PyThreadState_Clear(temporary);
+        PyThreadState_DeleteCurrent();
+    }
+    // Continue native work without reattaching the original thread state.
+    while (result != NULL && PyTime_MonotonicRaw(&now) == 0 &&
+           PyTime_AsSecondsDouble(now - start) < seconds) {
+    }
+    if (retain) {
+        PyThreadState_Swap(temporary);
+        PyThreadState_Clear(temporary);
+        PyThreadState_DeleteCurrent();
+    }
+    PyThreadState_Swap(saved);
+    if (exc != NULL) {
+        PyErr_SetRaisedException(exc);
+    }
+    return result;
+}
+
 // This exists mostly for testing the _interpreters module, as an
 // alternative to _interpreters.destroy()
 static PyObject *
@@ -3409,6 +3452,8 @@ static PyMethodDef module_functions[] = {
      METH_VARARGS | METH_KEYWORDS},
     {"destroy_interpreter", _PyCFunction_CAST(destroy_interpreter),
      METH_VARARGS | METH_KEYWORDS},
+    {"native_work_after_tstate_callback", native_work_after_tstate_callback,
+     METH_VARARGS, NULL},
     {"exec_interpreter", _PyCFunction_CAST(exec_interpreter),
      METH_VARARGS | METH_KEYWORDS},
     {"run_in_subinterp_with_config",