gh-135953: Properly obtain main thread identifier in Gecko Collector (#146045)

Author: flowln

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)

Lib/profiling/sampling/gecko_collector.py

@@ -9,14 +9,15 @@
 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)
     THREAD_STATUS_ON_CPU = (1 << 1)
     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,

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 = [

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)

Misc/NEWS.d/next/Tools-Demos/2026-03-22-00-00-00.gh-issue-135953.IptOwg.rst

@@ -0,0 +1,3 @@
+Properly identify the main thread in the Gecko profiler collector by
+using a status flag from the interpreter state instead of relying on
+:func:`threading.main_thread` in the collector process.

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) */

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, &current_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;

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) {