COUTIME

Author: pablogsal

Lib/profiling/sampling/collector.py

@@ -10,10 +10,12 @@ def collect(self, stack_frames):
     def export(self, filename):
         """Export collected data to a file."""
 
-    def _iter_all_frames(self, stack_frames):
+    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

Lib/profiling/sampling/pstats_collector.py

@@ -5,7 +5,7 @@
 
 
 class PstatsCollector(Collector):
-    def __init__(self, sample_interval_usec):
+    def __init__(self, sample_interval_usec, *, skip_idle=False):
         self.result = collections.defaultdict(
             lambda: dict(total_rec_calls=0, direct_calls=0, cumulative_calls=0)
         )
@@ -14,6 +14,7 @@ def __init__(self, sample_interval_usec):
         self.callers = collections.defaultdict(
             lambda: collections.defaultdict(int)
         )
+        self.skip_idle = skip_idle
 
     def _process_frames(self, frames):
         """Process a single thread's frame stack."""
@@ -40,7 +41,7 @@ def _process_frames(self, frames):
             self.callers[callee][caller] += 1
 
     def collect(self, stack_frames):
-        for frames in self._iter_all_frames(stack_frames):
+        for frames in self._iter_all_frames(stack_frames, skip_idle=self.skip_idle):
             self._process_frames(frames)
 
     def export(self, filename):

Lib/profiling/sampling/sample.py

@@ -583,6 +583,7 @@ def sample(
     show_summary=True,
     output_format="pstats",
     realtime_stats=False,
+    skip_idle=False,
 ):
     profiler = SampleProfiler(
         pid, sample_interval_usec, all_threads=all_threads
@@ -592,9 +593,9 @@ def sample(
     collector = None
     match output_format:
         case "pstats":
-            collector = PstatsCollector(sample_interval_usec)
+            collector = PstatsCollector(sample_interval_usec, skip_idle=skip_idle)
         case "collapsed":
-            collector = CollapsedStackCollector()
+            collector = CollapsedStackCollector(skip_idle=skip_idle)
             filename = filename or f"collapsed.{pid}.txt"
         case _:
             raise ValueError(f"Invalid output format: {output_format}")
@@ -644,6 +645,7 @@ 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
 
     sample(
         pid,
@@ -656,6 +658,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,
     )
 
 
@@ -710,6 +713,15 @@ def main():
         help="Print real-time sampling statistics (Hz, mean, min, max, stdev) during profiling",
     )
 
+    # Mode options
+    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)",
+    )
+
     # Output format selection
     output_group = parser.add_argument_group("Output options")
     output_format = output_group.add_mutually_exclusive_group()
@@ -826,6 +838,9 @@ 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
+
     if args.pid:
         sample(
             args.pid,
@@ -838,6 +853,7 @@ def main():
             show_summary=not args.no_summary,
             output_format=args.format,
             realtime_stats=args.realtime_stats,
+            skip_idle=skip_idle,
         )
     elif args.module or args.args:
         if args.module:

Lib/profiling/sampling/stack_collector.py

@@ -5,9 +5,10 @@
 
 
 class StackTraceCollector(Collector):
-    def __init__(self):
+    def __init__(self, *, skip_idle=False):
         self.call_trees = []
         self.function_samples = collections.defaultdict(int)
+        self.skip_idle = skip_idle
 
     def _process_frames(self, frames):
         """Process a single thread's frame stack."""
@@ -23,7 +24,7 @@ def _process_frames(self, frames):
             self.function_samples[frame] += 1
 
     def collect(self, stack_frames):
-        for frames in self._iter_all_frames(stack_frames):
+        for frames in self._iter_all_frames(stack_frames, skip_idle=self.skip_idle):
             self._process_frames(frames)
 
 

Lib/test/test_external_inspection.py

@@ -1670,6 +1670,103 @@ def test_unsupported_platform_error(self):
             str(cm.exception)
         )
 
+class TestDetectionOfThreadStatus(unittest.TestCase):
+    @unittest.skipIf(
+        sys.platform not in ("linux", "darwin", "win32"),
+        "Test only runs on unsupported platforms (not Linux, macOS, or Windows)",
+    )
+    @unittest.skipIf(sys.platform == "android", "Android raises Linux-specific exception")
+    def test_thread_status_detection(self):
+        port = find_unused_port()
+        script = textwrap.dedent(
+            f"""\
+            import time, sys, socket, threading
+            import os
+
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.connect(('localhost', {port}))
+
+            def sleeper():
+                tid = threading.get_native_id()
+                sock.sendall(f'ready:sleeper:{{tid}}\\n'.encode())
+                time.sleep(10000)
+
+            def busy():
+                tid = threading.get_native_id()
+                sock.sendall(f'ready:busy:{{tid}}\\n'.encode())
+                x = 0
+                while True:
+                    x = x + 1
+                time.sleep(0.5)
+
+            t1 = threading.Thread(target=sleeper)
+            t2 = threading.Thread(target=busy)
+            t1.start()
+            t2.start()
+            sock.sendall(b'ready:main\\n')
+            t1.join()
+            t2.join()
+            sock.close()
+            """
+        )
+        with os_helper.temp_dir() as work_dir:
+            script_dir = os.path.join(work_dir, "script_pkg")
+            os.mkdir(script_dir)
+            server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            server_socket.bind(("localhost", port))
+            server_socket.settimeout(SHORT_TIMEOUT)
+            server_socket.listen(1)
+
+            script_name = _make_test_script(script_dir, "thread_status_script", script)
+            client_socket = None
+            try:
+                p = subprocess.Popen([sys.executable, script_name])
+                client_socket, _ = server_socket.accept()
+                server_socket.close()
+                response = b""
+                sleeper_tid = None
+                busy_tid = None
+                while True:
+                    chunk = client_socket.recv(1024)
+                    response += chunk
+                    if b"ready:main" in response and b"ready:sleeper" in response and b"ready:busy" in response:
+                        # Parse TIDs from the response
+                        for line in response.split(b"\n"):
+                            if line.startswith(b"ready:sleeper:"):
+                                try:
+                                    sleeper_tid = int(line.split(b":")[-1])
+                                except Exception:
+                                    pass
+                            elif line.startswith(b"ready:busy:"):
+                                try:
+                                    busy_tid = int(line.split(b":")[-1])
+                                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()
+
+                # Find threads and their statuses
+                statuses = {}
+                for interpreter_info in traces:
+                    for thread_info in interpreter_info.threads:
+                        statuses[thread_info.thread_id] = thread_info.status
+
+                self.assertIsNotNone(sleeper_tid, "Sleeper thread id not received")
+                self.assertIsNotNone(busy_tid, "Busy thread id not received")
+                self.assertIn(sleeper_tid, statuses, "Sleeper tid not found in sampled threads")
+                self.assertIn(busy_tid, statuses, "Busy tid not found in sampled threads")
+                self.assertEqual(statuses[sleeper_tid], 1, "Sleeper thread should be idle (1)")
+                self.assertEqual(statuses[busy_tid], 0, "Busy thread should be running (0)")
+
+            finally:
+                if client_socket is not None:
+                    client_socket.close()
+                p.terminate()
+                p.wait(timeout=SHORT_TIMEOUT)
 
 if __name__ == "__main__":
     unittest.main()