Implement GIL mode and filtering

Author: pablogsal

Lib/profiling/sampling/sample.py

@@ -15,6 +15,21 @@
 from .stack_collector import CollapsedStackCollector, FlamegraphCollector
 
 _FREE_THREADED_BUILD = sysconfig.get_config_var("Py_GIL_DISABLED") is not None
+
+# Profiling mode constants
+PROFILING_MODE_WALL = 0
+PROFILING_MODE_CPU = 1
+PROFILING_MODE_GIL = 2
+
+
+def _parse_mode(mode_string):
+    """Convert mode string to mode constant."""
+    mode_map = {
+        "wall": PROFILING_MODE_WALL,
+        "cpu": PROFILING_MODE_CPU,
+        "gil": PROFILING_MODE_GIL,
+    }
+    return mode_map[mode_string]
 _HELP_DESCRIPTION = """Sample a process's stack frames and generate profiling data.
 Supports the following target modes:
   - -p PID: Profile an existing process by PID
@@ -120,18 +135,18 @@ def _run_with_sync(original_cmd):
 
 
 class SampleProfiler:
-    def __init__(self, pid, sample_interval_usec, all_threads, *, cpu_time=False):
+    def __init__(self, pid, sample_interval_usec, all_threads, *, mode=PROFILING_MODE_WALL):
         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, cpu_time=cpu_time
+                self.pid, all_threads=self.all_threads, mode=mode
             )
         else:
             only_active_threads = bool(self.all_threads)
             self.unwinder = _remote_debugging.RemoteUnwinder(
-                self.pid, only_active_thread=only_active_threads, cpu_time=cpu_time
+                self.pid, only_active_thread=only_active_threads, mode=mode
             )
         # Track sample intervals and total sample count
         self.sample_intervals = deque(maxlen=100)
@@ -596,13 +611,16 @@ def sample(
     show_summary=True,
     output_format="pstats",
     realtime_stats=False,
-    skip_idle=False,
+    mode=PROFILING_MODE_WALL,
 ):
     profiler = SampleProfiler(
-        pid, sample_interval_usec, all_threads=all_threads, cpu_time=skip_idle
+        pid, sample_interval_usec, all_threads=all_threads, mode=mode
     )
     profiler.realtime_stats = realtime_stats
 
+    # Determine skip_idle for collector compatibility
+    skip_idle = mode != PROFILING_MODE_WALL
+
     collector = None
     match output_format:
         case "pstats":
@@ -661,7 +679,8 @@ def wait_for_process_and_sample(pid, sort_value, args):
     filename = args.outfile
     if not filename and args.format == "collapsed":
         filename = f"collapsed.{pid}.txt"
-    skip_idle = True if args.mode == "cpu" else False
+
+    mode = _parse_mode(args.mode)
 
     sample(
         pid,
@@ -674,7 +693,7 @@ def wait_for_process_and_sample(pid, sort_value, args):
         show_summary=not args.no_summary,
         output_format=args.format,
         realtime_stats=args.realtime_stats,
-        skip_idle=skip_idle,
+        mode=mode,
     )
 
 
@@ -733,9 +752,9 @@ def main():
     mode_group = parser.add_argument_group("Mode options")
     mode_group.add_argument(
         "--mode",
-        choices=["wall", "cpu"],
-        default="wall-time",
-        help="Sampling mode: wall-time (default, skip_idle=False) or cpu-time (skip_idle=True)",
+        choices=["wall", "cpu", "gil"],
+        default="wall",
+        help="Sampling mode: wall (all threads), cpu (only CPU-running threads), gil (only GIL-holding threads)",
     )
 
     # Output format selection
@@ -862,8 +881,7 @@ def main():
     elif target_count > 1:
         parser.error("only one target type can be specified: -p/--pid, -m/--module, or script")
 
-    # Set skip_idle based on mode
-    skip_idle = True if args.mode == "cpu" else False
+    mode = _parse_mode(args.mode)
 
     if args.pid:
         sample(
@@ -877,7 +895,7 @@ def main():
             show_summary=not args.no_summary,
             output_format=args.format,
             realtime_stats=args.realtime_stats,
-            skip_idle=skip_idle,
+            mode=mode,
         )
     elif args.module or args.args:
         if args.module:

Modules/_remote_debugging_module.c

@@ -244,6 +244,12 @@ enum _ThreadState {
     THREAD_STATE_UNKNOWN
 };
 
+enum _ProfilingMode {
+    PROFILING_MODE_WALL = 0,
+    PROFILING_MODE_CPU = 1,
+    PROFILING_MODE_GIL = 2
+};
+
 typedef struct {
     PyObject_HEAD
     proc_handle_t handle;
@@ -257,7 +263,7 @@ typedef struct {
     _Py_hashtable_t *code_object_cache;
     int debug;
     int only_active_thread;
-    int cpu_time;
+    int mode;  // Use enum _ProfilingMode values
     RemoteDebuggingState *cached_state;  // Cached module state
 #ifdef Py_GIL_DISABLED
     // TLBC cache invalidation tracking
@@ -2629,6 +2635,39 @@ unwind_stack_for_thread(
         goto error;
     }
 
+    long tid = GET_MEMBER(long, ts, unwinder->debug_offsets.thread_state.native_thread_id);
+
+    // Calculate thread status based on mode
+    int status = THREAD_STATE_UNKNOWN;
+    if (unwinder->mode == PROFILING_MODE_CPU) {
+        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 if (unwinder->mode == PROFILING_MODE_GIL) {
+        status = (*current_tstate == gil_holder_tstate) ? THREAD_STATE_RUNNING : THREAD_STATE_GIL_WAIT;
+    } else {
+        // PROFILING_MODE_WALL - all threads are considered running
+        status = THREAD_STATE_RUNNING;
+    }
+
+    // Check if we should skip this thread based on mode
+    int should_skip = 0;
+    if (unwinder->mode == PROFILING_MODE_CPU && status != THREAD_STATE_RUNNING) {
+        should_skip = 1;
+    } else if (unwinder->mode == PROFILING_MODE_GIL && status != THREAD_STATE_RUNNING) {
+        should_skip = 1;
+    }
+
+    if (should_skip) {
+        // Advance to next thread and return NULL to skip processing
+        *current_tstate = GET_MEMBER(uintptr_t, ts, unwinder->debug_offsets.thread_state.next);
+        return NULL;
+    }
+
     uintptr_t frame_addr = GET_MEMBER(uintptr_t, ts, unwinder->debug_offsets.thread_state.current_frame);
 
     frame_info = PyList_New(0);
@@ -2642,20 +2681,6 @@ 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;
@@ -2716,7 +2741,7 @@ _remote_debugging.RemoteUnwinder.__init__
     *
     all_threads: bool = False
     only_active_thread: bool = False
-    cpu_time: bool = False
+    mode: int = 0
     debug: bool = False
 
 Initialize a new RemoteUnwinder object for debugging a remote Python process.
@@ -2726,7 +2751,7 @@ Initialize a new RemoteUnwinder object for debugging a remote Python process.
     all_threads: If True, initialize state for all threads in the process.
                 If False, only initialize for the main thread.
     only_active_thread: If True, only sample the thread holding the GIL.
-    cpu_time: If True, enable CPU time tracking for unwinder operations.
+    mode: Profiling mode: 0=WALL (wall-time), 1=CPU (cpu-time), 2=GIL (gil-time).
                        Cannot be used together with all_threads=True.
     debug: If True, chain exceptions to explain the sequence of events that
            lead to the exception.
@@ -2745,8 +2770,8 @@ static int
 _remote_debugging_RemoteUnwinder___init___impl(RemoteUnwinderObject *self,
                                                int pid, int all_threads,
                                                int only_active_thread,
-                                               int cpu_time, int debug)
-/*[clinic end generated code: output=2598ce54f6335ac7 input=0cf2038cc304c165]*/
+                                               int mode, int debug)
+/*[clinic end generated code: output=784e9990115aa569 input=d082d792d2ba9924]*/
 {
     // Validate that all_threads and only_active_thread are not both True
     if (all_threads && only_active_thread) {
@@ -2765,7 +2790,7 @@ _remote_debugging_RemoteUnwinder___init___impl(RemoteUnwinderObject *self,
 
     self->debug = debug;
     self->only_active_thread = only_active_thread;
-    self->cpu_time = cpu_time;
+    self->mode = mode;
     self->cached_state = NULL;
     if (_Py_RemoteDebug_InitProcHandle(&self->handle, pid) < 0) {
         set_exception_cause(self, PyExc_RuntimeError, "Failed to initialize process handle");
@@ -2983,6 +3008,12 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
         while (current_tstate != 0) {
             PyObject* frame_info = unwind_stack_for_thread(self, &current_tstate, gil_holder_tstate);
             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) && !PyErr_Occurred()) {
+                    // Thread was skipped due to mode filtering, continue to next thread
+                    continue;
+                }
+                // This was an actual error
                 Py_DECREF(interpreter_threads);
                 set_exception_cause(self, PyExc_RuntimeError, "Failed to unwind stack for thread");
                 Py_CLEAR(result);