profiling(gecko): use status flag for main thread instead of InterpreterInfo field

Replace the main_thread_id field in InterpreterInfo with a
THREAD_STATUS_MAIN_THREAD status flag set directly in the thread
unwinding code. This simplifies the struct from 3 to 2 fields and
removes the need to compare thread IDs at the Python level.

Author: pablogsal

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
@@ -171,19 +172,19 @@ def collect(self, stack_frames, timestamps_us=None):
 
         # Process threads
         for interpreter_info in stack_frames:
-            main_tid = interpreter_info.main_thread_id
             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, main_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)
@@ -289,14 +290,12 @@ def collect(self, stack_frames, timestamps_us=None):
 
         self.sample_count += len(times)
 
-    def _create_thread(self, tid, main_tid):
+    def _create_thread(self, tid, is_main_thread):
         """Create a new thread structure with processed profile format."""
 
-        is_main = tid == main_tid
-
         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/mocks.py

@@ -50,13 +50,12 @@ def __repr__(self):
 class MockInterpreterInfo:
     """Mock InterpreterInfo for testing since the real one isn't accessible."""
 
-    def __init__(self, interpreter_id, threads, main_thread_id=None):
+    def __init__(self, interpreter_id, threads):
         self.interpreter_id = interpreter_id
-        self.main_thread_id = main_thread_id
         self.threads = threads
 
     def __repr__(self):
-        return f"MockInterpreterInfo(interpreter_id={self.interpreter_id}, main_thread_id={self.main_thread_id}, threads={self.threads})"
+        return f"MockInterpreterInfo(interpreter_id={self.interpreter_id}, threads={self.threads})"
 
 
 class MockCoroInfo:

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:
@@ -47,9 +49,9 @@ def make_thread(thread_id, frames, status=0):
     return ThreadInfo((thread_id, status, frames))
 
 
-def make_interpreter(interp_id, threads, main_thread_id=0):
+def make_interpreter(interp_id, threads):
     """Create an InterpreterInfo struct sequence."""
-    return InterpreterInfo((interp_id, main_thread_id, threads))
+    return InterpreterInfo((interp_id, threads))
 
 
 def extract_lineno(location):
@@ -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,9 +525,9 @@ def test_gecko_collector_basic(self):
                     MockThreadInfo(
                         1,
                         [MockFrameInfo("file.py", 10, "func1"), MockFrameInfo("file.py", 20, "func2")],
+                        status=THREAD_STATUS_MAIN_THREAD,
                     )
                 ],
-                main_thread_id=1,
             )
         ]
         collector.collect(test_frames)

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)                              \
@@ -576,7 +577,7 @@ extern PyObject* unwind_stack_for_thread(
     uintptr_t *current_tstate,
     uintptr_t gil_holder_tstate,
     uintptr_t gc_frame,
-    uint64_t *current_thread_id
+    uintptr_t main_thread_tstate
 );
 
 /* Thread stopping functions (for blocking mode) */

Modules/_remote_debugging/binary_io_reader.c

@@ -828,7 +828,7 @@ build_sample_list(RemoteDebuggingState *state, BinaryReader *reader,
         goto error;
     }
     PyStructSequence_SetItem(interp_info, 0, iid);
-    PyStructSequence_SetItem(interp_info, 2, thread_list);
+    PyStructSequence_SetItem(interp_info, 1, thread_list);
     thread_list = NULL;
 
     sample_list = PyList_New(1);

Modules/_remote_debugging/binary_io_writer.c

@@ -1008,7 +1008,7 @@ binary_writer_write_sample(BinaryWriter *writer, PyObject *stack_frames, uint64_
         PyObject *interp_info = PyList_GET_ITEM(stack_frames, i);
 
         PyObject *interp_id_obj = PyStructSequence_GET_ITEM(interp_info, 0);
-        PyObject *threads = PyStructSequence_GET_ITEM(interp_info, 2);
+        PyObject *threads = PyStructSequence_GET_ITEM(interp_info, 1);
 
         unsigned long interp_id_long = PyLong_AsUnsignedLong(interp_id_obj);
         if (interp_id_long == (unsigned long)-1 && PyErr_Occurred()) {

Modules/_remote_debugging/module.c

@@ -106,7 +106,6 @@ PyStructSequence_Desc ThreadInfo_desc = {
 // InterpreterInfo structseq type
 static PyStructSequence_Field InterpreterInfo_fields[] = {
     {"interpreter_id", "Interpreter ID"},
-    {"main_thread_id", "Main thread ID"},
     {"threads", "List of threads in this interpreter"},
     {NULL}
 };
@@ -115,7 +114,7 @@ PyStructSequence_Desc InterpreterInfo_desc = {
     "_remote_debugging.InterpreterInfo",
     "Information about an interpreter",
     InterpreterInfo_fields,
-    3
+    2
 };
 
 // AwaitedInfo structseq type
@@ -588,15 +587,12 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
         uintptr_t main_thread_tstate = GET_MEMBER(uintptr_t, interp_state_buffer,
                 self->debug_offsets.interpreter_state.threads_main);
 
-        PyObject *main_thread_id = NULL;
-
-        uint64_t prev_thread_id = 0;
         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,
-                                                           &prev_thread_id);
+                                                           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 ||
@@ -622,10 +618,6 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
                 goto exit;
             }
 
-            if (prev_tstate == main_thread_tstate) {
-                main_thread_id = PyLong_FromUnsignedLongLong(prev_thread_id);
-            }
-
             if (PyList_Append(interpreter_threads, frame_info) == -1) {
                 Py_DECREF(frame_info);
                 Py_DECREF(interpreter_threads);
@@ -661,8 +653,7 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
         }
 
         PyStructSequence_SetItem(interpreter_info, 0, interp_id);  // steals reference
-        PyStructSequence_SetItem(interpreter_info, 1, main_thread_id);  // steals reference
-        PyStructSequence_SetItem(interpreter_info, 2, interpreter_threads);  // steals reference
+        PyStructSequence_SetItem(interpreter_info, 1, interpreter_threads);  // steals reference
 
         // Add this interpreter to the result list
         if (PyList_Append(result, interpreter_info) == -1) {
@@ -1221,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

@@ -292,7 +292,7 @@ unwind_stack_for_thread(
     uintptr_t *current_tstate,
     uintptr_t gil_holder_tstate,
     uintptr_t gc_frame,
-    uint64_t *current_tid
+    uintptr_t main_thread_tstate
 ) {
     PyObject *frame_info = NULL;
     PyObject *thread_id = NULL;
@@ -310,7 +310,6 @@ unwind_stack_for_thread(
     STATS_ADD(unwinder, memory_bytes_read, unwinder->debug_offsets.thread_state.size);
 
     long tid = GET_MEMBER(long, ts, unwinder->debug_offsets.thread_state.native_thread_id);
-    *current_tid = tid;
 
     // Read GC collecting state from the interpreter (before any skip checks)
     uintptr_t interp_addr = GET_MEMBER(uintptr_t, ts, unwinder->debug_offsets.thread_state.interp);
@@ -397,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) {