gh-159312: Resolve CPU attribution from collected stacks

Author: pablogsal

Lib/test/test_external_inspection.py

@@ -36,6 +36,7 @@
 THREAD_STATUS_ON_CPU = 1 << 1
 THREAD_STATUS_UNKNOWN = 1 << 2
 THREAD_STATUS_HAS_EXCEPTION = 1 << 4
+THREAD_STATUS_MAIN_THREAD = 1 << 5
 
 # Maximum number of retry attempts for operations that may fail transiently
 MAX_TRIES = 10
@@ -2627,7 +2628,7 @@ def suspended():
                 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)
+                self.assertEqual(suspended.status, THREAD_STATUS_MAIN_THREAD)
                 break
 
     def _run_thread_status_test(self, mode, check_condition):

Modules/_remote_debugging/_remote_debugging.h

@@ -179,6 +179,9 @@ typedef enum _WIN32_THREADSTATE {
 #define THREAD_STATUS_HAS_EXCEPTION       (1 << 4)
 #define THREAD_STATUS_MAIN_THREAD         (1 << 5)
 
+// Internal marker, removed before returning the sample.
+#define THREAD_STATUS_UNBOUND             (1 << 6)
+
 /* Exception cause macro */
 #define set_exception_cause(unwinder, exc_type, message)                              \
     do {                                                                              \
@@ -686,7 +689,7 @@ extern int get_thread_status(RemoteUnwinderObject *unwinder, uint64_t tid, uint6
 extern PyObject* unwind_stack_for_thread(
     RemoteUnwinderObject *unwinder,
     uintptr_t *current_tstate,
-    uintptr_t gil_holder_tstate,
+    int *has_unbound,
     uintptr_t gc_frame,
     uintptr_t main_thread_tstate,
     const RemoteReadPrefetch *prefetch

Modules/_remote_debugging/module.c

@@ -650,6 +650,66 @@ read_interp_state_and_maybe_thread_frame(
         interp_state_buffer);
 }
 
+static int
+filter_unbound_cpu_samples(PyObject *result, int skip)
+{
+    PyObject *bound = PySet_New(NULL);
+    if (bound == NULL) {
+        return -1;
+    }
+    // An unbound stack may still be doing native work after a callback.
+    // Exclude it only when a bound stack shares its native thread ID.
+    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 = 0; j < PyList_GET_SIZE(threads); j++) {
+            PyObject *thread = PyList_GET_ITEM(threads, j);
+            long status = PyLong_AsLong(PyStructSequence_GET_ITEM(thread, 1));
+            if (!(status & THREAD_STATUS_UNBOUND) &&
+                PyList_GET_SIZE(PyStructSequence_GET_ITEM(thread, 2)) != 0 &&
+                PySet_Add(bound, PyStructSequence_GET_ITEM(thread, 0)) < 0) {
+                goto error;
+            }
+        }
+    }
+    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);
+            long status = PyLong_AsLong(PyStructSequence_GET_ITEM(thread, 1));
+            if (!(status & THREAD_STATUS_UNBOUND)) {
+                continue;
+            }
+            status &= ~THREAD_STATUS_UNBOUND;
+            int found = PySet_Contains(bound, PyStructSequence_GET_ITEM(thread, 0));
+            if (found < 0) {
+                goto error;
+            }
+            if (found) {
+                status &= ~THREAD_STATUS_ON_CPU;
+            }
+            if (skip && !(status & THREAD_STATUS_ON_CPU)) {
+                if (PySequence_DelItem(threads, j) < 0) {
+                    goto error;
+                }
+            }
+            else {
+                PyObject *py_status = PyLong_FromLong(status);
+                if (py_status == NULL) {
+                    goto error;
+                }
+                Py_DECREF(PyStructSequence_GET_ITEM(thread, 1));
+                PyStructSequence_SET_ITEM(thread, 1, py_status);
+            }
+        }
+    }
+    Py_DECREF(bound);
+    return 0;
+
+error:
+    Py_DECREF(bound);
+    return -1;
+}
+
 /*[clinic input]
 @critical_section
 _remote_debugging.RemoteUnwinder.get_stack_trace
@@ -715,6 +775,8 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
         return NULL;
     }
 
+    int has_unbound = 0;
+
     // Iterate over all interpreters
     uintptr_t current_interpreter = self->interpreter_addr;
     while (current_interpreter != 0) {
@@ -803,7 +865,7 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
         while (current_tstate != 0) {
             uintptr_t prev_tstate = current_tstate;
             PyObject* frame_info = unwind_stack_for_thread(self, &current_tstate,
-                                                           gil_holder_tstate,
+                                                           &has_unbound,
                                                            gc_frame,
                                                            main_thread_tstate,
                                                            &prefetch);
@@ -890,6 +952,11 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
         }
     }
 
+    if (has_unbound && filter_unbound_cpu_samples(result,
+            self->mode == PROFILING_MODE_CPU && self->skip_non_matching_threads) < 0) {
+        Py_CLEAR(result);
+    }
+
 exit:
     // Invalidate cache entries for threads not seen in this sample.
     // Only do this every 1024 iterations to avoid performance overhead.

Modules/_remote_debugging/threads.c

@@ -293,62 +293,6 @@ 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,
@@ -388,7 +332,7 @@ PyObject*
 unwind_stack_for_thread(
     RemoteUnwinderObject *unwinder,
     uintptr_t *current_tstate,
-    uintptr_t gil_holder_tstate,
+    int *has_unbound,
     uintptr_t gc_frame,
     uintptr_t main_thread_tstate,
     const RemoteReadPrefetch *prefetch
@@ -557,11 +501,7 @@ unwind_stack_for_thread(
 
     if (cpu_status == THREAD_STATE_UNKNOWN) {
         status_flags |= THREAD_STATUS_UNKNOWN;
-    } 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.
+    } else if (cpu_status == THREAD_STATE_RUNNING) {
         status_flags |= THREAD_STATUS_ON_CPU;
     }
 
@@ -671,6 +611,13 @@ unwind_stack_for_thread(
         goto error;
     }
 
+    if ((unwinder->mode == PROFILING_MODE_CPU || unwinder->mode == PROFILING_MODE_ALL) &&
+        !GET_MEMBER(_thread_status, ts, unwinder->debug_offsets.thread_state.status).bound_gilstate &&
+        PyList_GET_SIZE(frame_info) != 0) {
+        status_flags |= THREAD_STATUS_UNBOUND;
+        *has_unbound = 1;
+    }
+
     // Always use status_flags
     PyObject *py_status = PyLong_FromLong(status_flags);
     if (py_status == NULL) {