diff --git a/Lib/profiling/sampling/constants.py b/Lib/profiling/sampling/constants.py index 58a57700fbdd4a..a364d0b8fde1e0 100644 --- a/Lib/profiling/sampling/constants.py +++ b/Lib/profiling/sampling/constants.py @@ -37,6 +37,7 @@ THREAD_STATUS_UNKNOWN, THREAD_STATUS_GIL_REQUESTED, THREAD_STATUS_HAS_EXCEPTION, + THREAD_STATUS_MAIN_THREAD, ) except ImportError: # Fallback for tests or when module is not available @@ -45,3 +46,4 @@ THREAD_STATUS_UNKNOWN = (1 << 2) THREAD_STATUS_GIL_REQUESTED = (1 << 3) THREAD_STATUS_HAS_EXCEPTION = (1 << 4) + THREAD_STATUS_MAIN_THREAD = (1 << 5) diff --git a/Lib/profiling/sampling/gecko_collector.py b/Lib/profiling/sampling/gecko_collector.py index 28ef9b69bf7968..8986194268b3ce 100644 --- a/Lib/profiling/sampling/gecko_collector.py +++ b/Lib/profiling/sampling/gecko_collector.py @@ -9,7 +9,7 @@ from .collector import Collector, filter_internal_frames from .opcode_utils import get_opcode_info, format_opcode try: - from _remote_debugging import THREAD_STATUS_HAS_GIL, THREAD_STATUS_ON_CPU, THREAD_STATUS_UNKNOWN, THREAD_STATUS_GIL_REQUESTED, THREAD_STATUS_HAS_EXCEPTION + from _remote_debugging import THREAD_STATUS_HAS_GIL, THREAD_STATUS_ON_CPU, THREAD_STATUS_UNKNOWN, THREAD_STATUS_GIL_REQUESTED, THREAD_STATUS_HAS_EXCEPTION, THREAD_STATUS_MAIN_THREAD except ImportError: # Fallback if module not available (shouldn't happen in normal use) THREAD_STATUS_HAS_GIL = (1 << 0) @@ -17,6 +17,7 @@ THREAD_STATUS_UNKNOWN = (1 << 2) THREAD_STATUS_GIL_REQUESTED = (1 << 3) THREAD_STATUS_HAS_EXCEPTION = (1 << 4) + THREAD_STATUS_MAIN_THREAD = (1 << 5) # Categories matching Firefox Profiler expectations @@ -174,15 +175,16 @@ def collect(self, stack_frames, timestamps_us=None): for thread_info in interpreter_info.threads: frames = filter_internal_frames(thread_info.frame_info) tid = thread_info.thread_id + status_flags = thread_info.status + is_main_thread = bool(status_flags & THREAD_STATUS_MAIN_THREAD) # Initialize thread if needed if tid not in self.threads: - self.threads[tid] = self._create_thread(tid) + self.threads[tid] = self._create_thread(tid, is_main_thread) thread_data = self.threads[tid] # Decode status flags - status_flags = thread_info.status has_gil = bool(status_flags & THREAD_STATUS_HAS_GIL) on_cpu = bool(status_flags & THREAD_STATUS_ON_CPU) gil_requested = bool(status_flags & THREAD_STATUS_GIL_REQUESTED) @@ -288,18 +290,12 @@ def collect(self, stack_frames, timestamps_us=None): self.sample_count += len(times) - def _create_thread(self, tid): + def _create_thread(self, tid, is_main_thread): """Create a new thread structure with processed profile format.""" - # Determine if this is the main thread - try: - is_main = tid == threading.main_thread().ident - except (RuntimeError, AttributeError): - is_main = False - thread = { "name": f"Thread-{tid}", - "isMainThread": is_main, + "isMainThread": is_main_thread, "processStartupTime": 0, "processShutdownTime": None, "registerTime": 0, diff --git a/Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py b/Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py index 033a533fe5444e..29f83c843561cd 100644 --- a/Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py +++ b/Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py @@ -18,9 +18,11 @@ THREAD_STATUS_UNKNOWN, THREAD_STATUS_GIL_REQUESTED, THREAD_STATUS_HAS_EXCEPTION, + THREAD_STATUS_MAIN_THREAD, ) from profiling.sampling.binary_collector import BinaryCollector from profiling.sampling.binary_reader import BinaryReader + from profiling.sampling.gecko_collector import GeckoCollector ZSTD_AVAILABLE = _remote_debugging.zstd_available() except ImportError: @@ -318,6 +320,7 @@ def test_status_flags_preserved(self): THREAD_STATUS_UNKNOWN, THREAD_STATUS_GIL_REQUESTED, THREAD_STATUS_HAS_EXCEPTION, + THREAD_STATUS_MAIN_THREAD, THREAD_STATUS_HAS_GIL | THREAD_STATUS_ON_CPU, THREAD_STATUS_HAS_GIL | THREAD_STATUS_HAS_EXCEPTION, THREAD_STATUS_HAS_GIL @@ -342,6 +345,35 @@ def test_status_flags_preserved(self): self.assertEqual(count, len(statuses)) self.assert_samples_equal(samples, collector) + def test_binary_replay_preserves_main_thread_for_gecko(self): + """Binary replay preserves main thread identity for GeckoCollector.""" + samples = [ + [ + make_interpreter( + 0, + [ + make_thread( + 1, + [make_frame("main.py", 10, "main")], + THREAD_STATUS_MAIN_THREAD, + ), + make_thread(2, [make_frame("worker.py", 20, "worker")]), + ], + ) + ] + ] + filename = self.create_binary_file(samples) + collector = GeckoCollector(1000) + + with BinaryReader(filename) as reader: + count = reader.replay_samples(collector) + + self.assertEqual(count, 2) + profile = collector._build_profile() + threads = {thread["tid"]: thread for thread in profile["threads"]} + self.assertTrue(threads[1]["isMainThread"]) + self.assertFalse(threads[2]["isMainThread"]) + def test_multiple_threads_per_sample(self): """Multiple threads in one sample roundtrip exactly.""" threads = [ diff --git a/Lib/test/test_profiling/test_sampling_profiler/test_collectors.py b/Lib/test/test_profiling/test_sampling_profiler/test_collectors.py index 8e6afa91e89daf..06c9e51e0c9c55 100644 --- a/Lib/test/test_profiling/test_sampling_profiler/test_collectors.py +++ b/Lib/test/test_profiling/test_sampling_profiler/test_collectors.py @@ -28,6 +28,7 @@ THREAD_STATUS_HAS_GIL, THREAD_STATUS_ON_CPU, THREAD_STATUS_GIL_REQUESTED, + THREAD_STATUS_MAIN_THREAD, ) except ImportError: raise unittest.SkipTest( @@ -524,6 +525,7 @@ def test_gecko_collector_basic(self): MockThreadInfo( 1, [MockFrameInfo("file.py", 10, "func1"), MockFrameInfo("file.py", 20, "func2")], + status=THREAD_STATUS_MAIN_THREAD, ) ], ) @@ -556,6 +558,7 @@ def test_gecko_collector_basic(self): threads = profile_data["threads"] self.assertEqual(len(threads), 1) thread_data = threads[0] + self.assertTrue(thread_data["isMainThread"]) # Verify thread structure self.assertIn("samples", thread_data) diff --git a/Modules/_remote_debugging/_remote_debugging.h b/Modules/_remote_debugging/_remote_debugging.h index 7bcb2f483234ec..570f6b23b75849 100644 --- a/Modules/_remote_debugging/_remote_debugging.h +++ b/Modules/_remote_debugging/_remote_debugging.h @@ -172,6 +172,7 @@ typedef enum _WIN32_THREADSTATE { #define THREAD_STATUS_UNKNOWN (1 << 2) #define THREAD_STATUS_GIL_REQUESTED (1 << 3) #define THREAD_STATUS_HAS_EXCEPTION (1 << 4) +#define THREAD_STATUS_MAIN_THREAD (1 << 5) /* Exception cause macro */ #define set_exception_cause(unwinder, exc_type, message) \ @@ -575,7 +576,8 @@ extern PyObject* unwind_stack_for_thread( RemoteUnwinderObject *unwinder, uintptr_t *current_tstate, uintptr_t gil_holder_tstate, - uintptr_t gc_frame + uintptr_t gc_frame, + uintptr_t main_thread_tstate ); /* Thread stopping functions (for blocking mode) */ diff --git a/Modules/_remote_debugging/module.c b/Modules/_remote_debugging/module.c index 040bd3db377315..4f294b80ba0739 100644 --- a/Modules/_remote_debugging/module.c +++ b/Modules/_remote_debugging/module.c @@ -583,11 +583,16 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self current_tstate = self->tstate_addr; } + // Acquire main thread state information + uintptr_t main_thread_tstate = GET_MEMBER(uintptr_t, interp_state_buffer, + self->debug_offsets.interpreter_state.threads_main); + while (current_tstate != 0) { uintptr_t prev_tstate = current_tstate; PyObject* frame_info = unwind_stack_for_thread(self, ¤t_tstate, gil_holder_tstate, - gc_frame); + gc_frame, + main_thread_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 || @@ -1207,6 +1212,9 @@ _remote_debugging_exec(PyObject *m) if (PyModule_AddIntConstant(m, "THREAD_STATUS_HAS_EXCEPTION", THREAD_STATUS_HAS_EXCEPTION) < 0) { return -1; } + if (PyModule_AddIntConstant(m, "THREAD_STATUS_MAIN_THREAD", THREAD_STATUS_MAIN_THREAD) < 0) { + return -1; + } if (RemoteDebugging_InitState(st) < 0) { return -1; diff --git a/Modules/_remote_debugging/threads.c b/Modules/_remote_debugging/threads.c index 3100b83c8f4899..527957c6fef067 100644 --- a/Modules/_remote_debugging/threads.c +++ b/Modules/_remote_debugging/threads.c @@ -291,7 +291,8 @@ unwind_stack_for_thread( RemoteUnwinderObject *unwinder, uintptr_t *current_tstate, uintptr_t gil_holder_tstate, - uintptr_t gc_frame + uintptr_t gc_frame, + uintptr_t main_thread_tstate ) { PyObject *frame_info = NULL; PyObject *thread_id = NULL; @@ -395,6 +396,10 @@ unwind_stack_for_thread( status_flags |= THREAD_STATUS_ON_CPU; } + if (*current_tstate == main_thread_tstate) { + status_flags |= THREAD_STATUS_MAIN_THREAD; + } + // Check if we should skip this thread based on mode int should_skip = 0; if (unwinder->skip_non_matching_threads) {