gh-159312: Confine CPU attribution checks to unbound states

Author: pablogsal

Lib/test/test_external_inspection.py

@@ -2596,108 +2596,39 @@ 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,
-    ):
+    @skip_if_not_supported
+    def test_cpu_status_with_multiple_thread_states(self):
         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}})
-
+        import_module("_hashlib")
+        code = textwrap.dedent("""
+            import _hashlib
             def active():
-                active_sock.sendall(b"ready")
                 while True:
-                    {work}
-
+                    _hashlib.pbkdf2_hmac('sha256', b'password', b'salt', 10_000_000)
             active()
         """)
         script = f"""
-            import _interpreters, _testinternalcapi
-
+            import _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)
-
+                sock.sendall(b"ready")
+                _testinternalcapi.exec_interpreter(0, {code!r})
             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
+            unwinder = RemoteUnwinder(p.pid, all_threads=True, mode=PROFILING_MODE_ALL)
+            for _ in busy_retry(SHORT_TIMEOUT):
+                trace = _get_stack_trace_with_retry(unwinder)
+                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 active.status & THREAD_STATUS_HAS_GIL:
+                    continue
+                suspended = self._find_thread_with_frame(trace, lambda f: f.funcname == "suspended")
+                self.assertIsNotNone(suspended)
+                self.assertEqual(suspended.thread_id, active.thread_id)
+                self.assertFalse(suspended.status & THREAD_STATUS_ON_CPU)
+                break
 
     def _run_thread_status_test(self, mode, check_condition):
         """

Modules/_remote_debugging/_remote_debugging.h

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

Modules/_remote_debugging/module.c

@@ -650,45 +650,6 @@ 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
@@ -753,15 +714,6 @@ _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;
@@ -850,13 +802,11 @@ _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,
-                                                           &has_bound_stack);
+                                                           &prefetch);
             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 ||
@@ -882,13 +832,10 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
                 goto exit;
             }
 
-            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)) {
+            if (PyList_Append(interpreter_threads, frame_info) == -1) {
                 Py_DECREF(frame_info);
                 Py_DECREF(interpreter_threads);
-                set_exception_cause(self, PyExc_RuntimeError, "Failed to collect thread frame info");
+                set_exception_cause(self, PyExc_RuntimeError, "Failed to append thread frame info");
                 Py_CLEAR(result);
                 goto exit;
             }
@@ -943,14 +890,7 @@ _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

@@ -293,6 +293,62 @@ typedef struct {
     unsigned int :24;
 } _thread_status;
 
+static int
+has_bound_stack(RemoteUnwinderObject *unwinder, long tid)
+{
+    uintptr_t interp, tstate;
+    int remaining = 8192;
+    if (read_ptr(unwinder, unwinder->runtime_start_address +
+                 unwinder->debug_offsets.runtime_state.interpreters_head, &interp) < 0) {
+        goto unavailable;
+    }
+    while (interp && remaining-- > 0) {
+        if (read_ptr(unwinder, interp +
+                     unwinder->debug_offsets.interpreter_state.threads_head, &tstate) < 0) {
+            goto unavailable;
+        }
+        while (tstate && remaining-- > 0) {
+            char ts[SIZEOF_THREAD_STATE];
+            if (_Py_RemoteDebug_ReadRemoteMemory(&unwinder->handle, tstate,
+                    unwinder->debug_offsets.thread_state.size, ts) < 0) {
+                goto unavailable;
+            }
+            if (GET_MEMBER(long, ts, unwinder->debug_offsets.thread_state.native_thread_id) == tid &&
+                GET_MEMBER(_thread_status, ts, unwinder->debug_offsets.thread_state.status).bound_gilstate) {
+                uintptr_t frame = GET_MEMBER(uintptr_t, ts, unwinder->debug_offsets.thread_state.current_frame);
+                uintptr_t base = GET_MEMBER(uintptr_t, ts, unwinder->debug_offsets.thread_state.base_frame);
+                while (frame && frame != base && remaining-- > 0) {
+                    char data[SIZEOF_INTERP_FRAME];
+                    if (_Py_RemoteDebug_ReadRemoteMemory(&unwinder->handle, frame,
+                                                        sizeof(data), data) < 0) {
+                        goto unavailable;
+                    }
+                    int valid = is_frame_valid(unwinder, (uintptr_t)data,
+                        GET_MEMBER_NO_TAG(uintptr_t, data, unwinder->debug_offsets.interpreter_frame.executable));
+                    if (valid < 0) {
+                        goto unavailable;
+                    }
+                    if (valid == 1) {
+                        return 1;
+                    }
+                    frame = GET_MEMBER(uintptr_t, data, unwinder->debug_offsets.interpreter_frame.previous);
+                }
+            }
+            tstate = GET_MEMBER(uintptr_t, ts, unwinder->debug_offsets.thread_state.next);
+        }
+        if (read_ptr(unwinder, interp +
+                     unwinder->debug_offsets.interpreter_state.next, &interp) < 0) {
+            goto unavailable;
+        }
+    }
+    return 0;
+
+unavailable:
+    // Keep the sample when remote state changes during the lookup.
+    PyErr_Clear();
+    return 0;
+}
+
 static int
 read_thread_state_and_maybe_frame(
     RemoteUnwinderObject *unwinder,
@@ -335,8 +391,7 @@ unwind_stack_for_thread(
     uintptr_t gil_holder_tstate,
     uintptr_t gc_frame,
     uintptr_t main_thread_tstate,
-    const RemoteReadPrefetch *prefetch,
-    int *has_bound_stack
+    const RemoteReadPrefetch *prefetch
 ) {
     PyObject *frame_info = NULL;
     PyObject *thread_id = NULL;
@@ -502,7 +557,11 @@ unwind_stack_for_thread(
 
     if (cpu_status == THREAD_STATE_UNKNOWN) {
         status_flags |= THREAD_STATUS_UNKNOWN;
-    } else if (cpu_status == THREAD_STATE_RUNNING) {
+    } else if (cpu_status == THREAD_STATE_RUNNING &&
+               (GET_MEMBER(_thread_status, ts, unwinder->debug_offsets.thread_state.status).bound_gilstate ||
+                !has_bound_stack(unwinder, tid))) {
+        // An unbound state may still be doing native work after a callback.
+        // Exclude it only if another bound stack shares the same OS thread.
         status_flags |= THREAD_STATUS_ON_CPU;
     }
 
@@ -619,10 +678,6 @@ 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