COUTIME

Author: pablogsal

Include/internal/pycore_global_objects_fini_generated.h

@@ -387,6 +387,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(coro)
         STRUCT_FOR_ID(count)
         STRUCT_FOR_ID(covariant)
+        STRUCT_FOR_ID(cpu_time)
         STRUCT_FOR_ID(ctx)
         STRUCT_FOR_ID(cwd)
         STRUCT_FOR_ID(d_parameter_type)

Include/internal/pycore_global_strings.h

@@ -1,5 +1,17 @@
 from abc import ABC, abstractmethod
 
+# Enums are slow
+THREAD_STATE_RUNNING = 0
+THREAD_STATE_IDLE = 1
+THREAD_STATE_GIL_WAIT = 2
+THREAD_STATE_UNKNOWN = 3
+
+STATUS = {
+    THREAD_STATE_RUNNING: "running",
+    THREAD_STATE_IDLE: "idle",
+    THREAD_STATE_GIL_WAIT: "gil_wait",
+    THREAD_STATE_UNKNOWN: "unknown",
+}
 
 class Collector(ABC):
     @abstractmethod
@@ -12,10 +24,8 @@ def export(self, filename):
 
     def _iter_all_frames(self, stack_frames, skip_idle=False):
         """Iterate over all frame stacks from all interpreters and threads."""
-        print("Skipping idle threads" if skip_idle else "Including idle threads")
         for interpreter_info in stack_frames:
             for thread_info in interpreter_info.threads:
-                print(thread_info.status)
                 frames = thread_info.frame_info
                 if frames:
                     yield frames

Include/internal/pycore_runtime_init_generated.h

@@ -114,18 +114,18 @@ def _run_with_sync(original_cmd):
 
 
 class SampleProfiler:
-    def __init__(self, pid, sample_interval_usec, all_threads):
+    def __init__(self, pid, sample_interval_usec, all_threads, *, cpu_time=False):
         self.pid = pid
         self.sample_interval_usec = sample_interval_usec
         self.all_threads = all_threads
         if _FREE_THREADED_BUILD:
             self.unwinder = _remote_debugging.RemoteUnwinder(
-                self.pid, all_threads=self.all_threads
+                self.pid, all_threads=self.all_threads, cpu_time=cpu_time
             )
         else:
             only_active_threads = bool(self.all_threads)
             self.unwinder = _remote_debugging.RemoteUnwinder(
-                self.pid, only_active_thread=only_active_threads
+                self.pid, only_active_thread=only_active_threads, cpu_time=cpu_time
             )
         # Track sample intervals and total sample count
         self.sample_intervals = deque(maxlen=100)
@@ -586,7 +586,7 @@ def sample(
     skip_idle=False,
 ):
     profiler = SampleProfiler(
-        pid, sample_interval_usec, all_threads=all_threads
+        pid, sample_interval_usec, all_threads=all_threads, cpu_time=skip_idle
     )
     profiler.realtime_stats = realtime_stats
 

Include/internal/pycore_unicodeobject_generated.h

@@ -1744,10 +1744,16 @@ def busy():
                                 except Exception:
                                     pass
                         break
-
-                unwinder = RemoteUnwinder(p.pid, all_threads=True)
-                time.sleep(0.2)  # Give a bit of time to let threads settle
-                traces = unwinder.get_stack_trace()
+                    
+                try:
+                    unwinder = RemoteUnwinder(p.pid, all_threads=True, cpu_time=True)
+                    time.sleep(0.2)  # Give a bit of time to let threads settle
+                    traces = unwinder.get_stack_trace()
+                except PermissionError:
+                    self.skipTest(
+                        "Insufficient permissions to read the stack trace"
+                    )
+    
 
                 # Find threads and their statuses
                 statuses = {}

Lib/profiling/sampling/collector.py

@@ -41,6 +41,24 @@
 #define MAX_NATIVE_THREADS 4096
 #endif
 
+#ifdef MS_WINDOWS
+#include <windows.h>
+#include <winternl.h>
+// Define the NTSTATUS values we need
+#define STATUS_SUCCESS                   ((NTSTATUS)0x00000000L)
+#define STATUS_INFO_LENGTH_MISMATCH      ((NTSTATUS)0xC0000004L)
+typedef enum _WIN32_THREADSTATE {
+    WIN32_THREADSTATE_INITIALIZED = 0,   // Recognized by the kernel
+    WIN32_THREADSTATE_READY       = 1,   // Prepared to run on the next available processor
+    WIN32_THREADSTATE_RUNNING     = 2,   // Currently executing
+    WIN32_THREADSTATE_STANDBY     = 3,   // About to run, only one thread may be in this state at a time
+    WIN32_THREADSTATE_TERMINATED  = 4,   // Finished executing
+    WIN32_THREADSTATE_WAITING     = 5,   // Not ready for the processor, when ready, it will be rescheduled
+    WIN32_THREADSTATE_TRANSITION  = 6,   // Waiting for resources other than the processor
+    WIN32_THREADSTATE_UNKNOWN     = 7    // Thread state is unknown
+} WIN32_THREADSTATE;
+#endif
+
 /* ============================================================================
  * TYPE DEFINITIONS AND STRUCTURES
  * ============================================================================ */
@@ -219,9 +237,10 @@ typedef struct {
     PyTypeObject *AwaitedInfo_Type;
 } RemoteDebuggingState;
 
-enum ThreadState {
+enum _ThreadState {
     THREAD_STATE_RUNNING,
     THREAD_STATE_IDLE,
+    THREAD_STATE_GIL_WAIT,
     THREAD_STATE_UNKNOWN
 };
 
@@ -248,6 +267,10 @@ typedef struct {
 #ifdef __APPLE__
     uint64_t thread_id_offset;
 #endif
+#ifdef MS_WINDOWS
+    PVOID win_process_buffer;
+    ULONG win_process_buffer_size;
+#endif
 } RemoteUnwinderObject;
 
 #define RemoteUnwinder_CAST(op) ((RemoteUnwinderObject *)(op))
@@ -2472,7 +2495,7 @@ process_frame_chain(
 }
 
 static int
-get_thread_status(RemoteUnwinderObject *unwinder, uint64_t tid) {
+get_thread_status(RemoteUnwinderObject *unwinder, uint64_t tid, uint64_t pthread_id) {
 #ifdef __APPLE__
    if (unwinder->thread_id_offset == 0) {
         // Apple is a bit of a circus and doesn't give us a straightforward way to get thread information
@@ -2493,7 +2516,7 @@ get_thread_status(RemoteUnwinderObject *unwinder, uint64_t tid) {
         }
         uint64_t min_offset = UINT64_MAX;
         for (int i = 0; i < n; i++) {
-            uint64_t offset = tids[i] - tid;
+            uint64_t offset = tids[i] - pthread_id;
             if (offset < min_offset) {
                 min_offset = offset;
             }
@@ -2502,14 +2525,92 @@ get_thread_status(RemoteUnwinderObject *unwinder, uint64_t tid) {
         PyMem_Free(tids);
     }
     struct proc_threadinfo ti;
-    uint64_t tid_with_offset = tid + unwinder->thread_id_offset;
+    uint64_t tid_with_offset = pthread_id + unwinder->thread_id_offset;
     if (proc_pidinfo(unwinder->handle.pid, PROC_PIDTHREADINFO, tid_with_offset, &ti, sizeof(ti)) != sizeof(ti)) {
         return THREAD_STATE_UNKNOWN;
     }
     if (ti.pth_run_state == TH_STATE_RUNNING) {
         return THREAD_STATE_RUNNING;
     }
     return THREAD_STATE_IDLE;
+#elif defined(__linux__)
+    char stat_path[256];
+    char buffer[2048] = "";
+    
+    snprintf(stat_path, sizeof(stat_path), "/proc/%d/task/%lu/stat", unwinder->handle.pid, tid);
+    
+    int fd = open(stat_path, O_RDONLY);
+    if (fd == -1) {
+        return THREAD_STATE_UNKNOWN;
+    }
+    
+    if (read(fd, buffer, 2047) == 0) {
+        close(fd);
+        return THREAD_STATE_UNKNOWN;
+    }
+    close(fd);
+    
+    char *p = strchr(buffer, ')');
+    if (!p) {
+        return THREAD_STATE_UNKNOWN;
+    }
+    
+    p += 2;  // Skip ") "
+    if (*p == ' ') {
+        p++;
+    }
+    
+    switch (*p) {
+        case 'R':  // Running
+            return THREAD_STATE_RUNNING;
+        case 'S':  // Interruptible sleep
+        case 'D':  // Uninterruptible sleep
+        case 'T':  // Stopped
+        case 'Z':  // Zombie
+        case 'I':  // Idle kernel thread
+            return THREAD_STATE_IDLE;
+        default:
+            return THREAD_STATE_UNKNOWN;
+    }
+#elif defined(MS_WINDOWS)
+    ULONG n;
+    NTSTATUS status = NtQuerySystemInformation(
+        SystemProcessInformation,
+        unwinder->win_process_buffer,
+        unwinder->win_process_buffer_size,
+        &n
+    );
+    if (status == STATUS_INFO_LENGTH_MISMATCH) {
+        // Buffer was too small so we reallocate a larger one and try again.
+        unwinder->win_process_buffer_size = n;
+        PVOID new_buffer = PyMem_Realloc(unwinder->win_process_buffer, n);
+        if (!new_buffer) {
+            return -1;
+        }
+        unwinder->win_process_buffer = new_buffer;
+        return get_thread_status(unwinder, tid, pthread_id);
+    }
+    if (status != STATUS_SUCCESS) {
+        return -1;
+    }
+
+    SYSTEM_PROCESS_INFORMATION *pi = (SYSTEM_PROCESS_INFORMATION *)unwinder->win_process_buffer;
+    while ((ULONG)(ULONG_PTR)pi->UniqueProcessId != unwinder->handle.pid) {
+        if (pi->NextEntryOffset == 0) {
+            // We didn't find the process
+            return -1;
+        }
+        pi = (SYSTEM_PROCESS_INFORMATION *)(((BYTE *)pi) + pi->NextEntryOffset);
+    }
+
+    SYSTEM_THREAD_INFORMATION *ti = (SYSTEM_THREAD_INFORMATION *)((char *)pi + sizeof(SYSTEM_PROCESS_INFORMATION));
+    for (Py_ssize_t i = 0; i < pi->NumberOfThreads; i++, ti++) {
+        if (ti->ClientId.UniqueThread == (HANDLE)tid) {
+            return ti->ThreadState != WIN32_THREADSTATE_RUNNING ? THREAD_STATE_IDLE : THREAD_STATE_RUNNING;
+        }
+    }
+
+    return -1;
 #else
     return THREAD_STATE_UNKNOWN;
 #endif
@@ -2518,7 +2619,8 @@ get_thread_status(RemoteUnwinderObject *unwinder, uint64_t tid) {
 static PyObject*
 unwind_stack_for_thread(
     RemoteUnwinderObject *unwinder,
-    uintptr_t *current_tstate
+    uintptr_t *current_tstate,
+    uintptr_t gil_holder_tstate
 ) {
     PyObject *frame_info = NULL;
     PyObject *thread_id = NULL;
@@ -2546,14 +2648,27 @@ unwind_stack_for_thread(
         goto error;
     }
 
+    long tid = GET_MEMBER(long, ts, unwinder->debug_offsets.thread_state.native_thread_id);
+    int status = THREAD_STATE_UNKNOWN;
+    if (unwinder->cpu_time == 1) {
+        long pthread_id = GET_MEMBER(long, ts, unwinder->debug_offsets.thread_state.thread_id);
+        status = get_thread_status(unwinder, tid, pthread_id);
+        if (status == -1) {
+            PyErr_Print();
+            PyErr_SetString(PyExc_RuntimeError, "Failed to get thread status");
+            goto error;
+        }
+    } else {
+        status = (*current_tstate == gil_holder_tstate) ? THREAD_STATE_RUNNING : THREAD_STATE_GIL_WAIT;
+    }
+
     if (process_frame_chain(unwinder, frame_addr, &chunks, frame_info) < 0) {
         set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to process frame chain");
         goto error;
     }
 
     *current_tstate = GET_MEMBER(uintptr_t, ts, unwinder->debug_offsets.thread_state.next);
 
-    long tid = GET_MEMBER(long, ts, unwinder->debug_offsets.thread_state.native_thread_id);
     thread_id = PyLong_FromLongLong(tid);
     if (thread_id == NULL) {
         set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to create thread ID");
@@ -2567,13 +2682,6 @@ unwind_stack_for_thread(
         goto error;
     }
 
-    long pthread_id = GET_MEMBER(long, ts, unwinder->debug_offsets.thread_state.thread_id);
-    int status = get_thread_status(unwinder, pthread_id);
-    if (status == -1) {
-        PyErr_Print();
-        PyErr_SetString(PyExc_RuntimeError, "Failed to get thread status");
-        goto error;
-    }
     PyObject *py_status = PyLong_FromLong(status);
     if (py_status == NULL) {
         set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to create thread status");
@@ -2740,6 +2848,11 @@ _remote_debugging_RemoteUnwinder___init___impl(RemoteUnwinderObject *self,
     self->thread_id_offset = 0;
 #endif
 
+#ifdef MS_WINDOWS
+    self->win_process_buffer = NULL;
+    self->win_process_buffer_size = 0;
+#endif
+
     return 0;
 }
 
@@ -2845,21 +2958,25 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
             goto exit;
         }
 
+        // Get the GIL holder for this interpreter (needed for GIL_WAIT logic)
+        uintptr_t gil_holder_tstate = 0;
+        int gil_locked = GET_MEMBER(int, interp_state_buffer,
+            self->debug_offsets.interpreter_state.gil_runtime_state_locked);
+        if (gil_locked) {
+            gil_holder_tstate = (uintptr_t)GET_MEMBER(PyThreadState*, interp_state_buffer,
+                self->debug_offsets.interpreter_state.gil_runtime_state_holder);
+        }
+
         uintptr_t current_tstate;
         if (self->only_active_thread) {
             // Find the GIL holder for THIS interpreter
-            int gil_locked = GET_MEMBER(int, interp_state_buffer,
-                self->debug_offsets.interpreter_state.gil_runtime_state_locked);
-
             if (!gil_locked) {
                 // This interpreter's GIL is not locked, skip it
                 Py_DECREF(interpreter_threads);
                 goto next_interpreter;
             }
 
-            // Get the GIL holder for this interpreter
-            current_tstate = (uintptr_t)GET_MEMBER(PyThreadState*, interp_state_buffer,
-                self->debug_offsets.interpreter_state.gil_runtime_state_holder);
+            current_tstate = gil_holder_tstate;
         } else if (self->tstate_addr == 0) {
             // Get all threads for this interpreter
             current_tstate = GET_MEMBER(uintptr_t, interp_state_buffer,
@@ -2870,7 +2987,7 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
         }
 
         while (current_tstate != 0) {
-            PyObject* frame_info = unwind_stack_for_thread(self, &current_tstate);
+            PyObject* frame_info = unwind_stack_for_thread(self, &current_tstate, gil_holder_tstate);
             if (!frame_info) {
                 Py_DECREF(interpreter_threads);
                 set_exception_cause(self, PyExc_RuntimeError, "Failed to unwind stack for thread");
@@ -3122,6 +3239,12 @@ RemoteUnwinder_dealloc(PyObject *op)
     if (self->code_object_cache) {
         _Py_hashtable_destroy(self->code_object_cache);
     }
+#ifdef MS_WINDOWS
+    if(self->win_process_buffer != NULL) {
+        PyMem_Free(self->win_process_buffer);
+    }
+#endif
+
 #ifdef Py_GIL_DISABLED
     if (self->tlbc_cache) {
         _Py_hashtable_destroy(self->tlbc_cache);