profiling(gecko): retrieve main thread id directly from interpreter state

Signed-off-by: Sofia Donato Ferreira <flowlnlnln@gmail.com>

Author: flowln

Lib/profiling/sampling/gecko_collector.py

@@ -171,11 +171,7 @@ def collect(self, stack_frames, timestamps_us=None):
 
         # Process threads
         for interpreter_info in stack_frames:
-            # Since 'threads' is in order from newest to oldest,
-            # we know the first thread must be the main thread.
-            main_tid = None
-            if len(interpreter_info.threads) != 0:
-                main_tid = interpreter_info.threads[-1].thread_id
+            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

Lib/test/test_profiling/test_sampling_profiler/test_binary_format.py

@@ -47,9 +47,9 @@ def make_thread(thread_id, frames, status=0):
     return ThreadInfo((thread_id, status, frames))
 
 
-def make_interpreter(interp_id, threads):
+def make_interpreter(interp_id, threads, main_thread_id=0):
     """Create an InterpreterInfo struct sequence."""
-    return InterpreterInfo((interp_id, threads))
+    return InterpreterInfo((interp_id, main_thread_id, threads))
 
 
 def extract_lineno(location):

Modules/_remote_debugging/_remote_debugging.h

@@ -575,7 +575,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,
+    uint64_t *current_thread_id
 );
 
 /* 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, 1, thread_list);
+    PyStructSequence_SetItem(interp_info, 2, 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, 1);
+        PyObject *threads = PyStructSequence_GET_ITEM(interp_info, 2);
 
         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,6 +106,7 @@ 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}
 };
@@ -114,7 +115,7 @@ PyStructSequence_Desc InterpreterInfo_desc = {
     "_remote_debugging.InterpreterInfo",
     "Information about an interpreter",
     InterpreterInfo_fields,
-    2
+    3
 };
 
 // AwaitedInfo structseq type
@@ -583,11 +584,19 @@ _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);
+
+        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);
+                                                           gc_frame,
+                                                           &prev_thread_id);
             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 ||
@@ -613,6 +622,10 @@ _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);
@@ -648,7 +661,8 @@ _remote_debugging_RemoteUnwinder_get_stack_trace_impl(RemoteUnwinderObject *self
         }
 
         PyStructSequence_SetItem(interpreter_info, 0, interp_id);  // steals reference
-        PyStructSequence_SetItem(interpreter_info, 1, interpreter_threads);  // steals reference
+        PyStructSequence_SetItem(interpreter_info, 1, main_thread_id);  // steals reference
+        PyStructSequence_SetItem(interpreter_info, 2, interpreter_threads);  // steals reference
 
         // Add this interpreter to the result list
         if (PyList_Append(result, interpreter_info) == -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,
+    uint64_t *current_tid
 ) {
     PyObject *frame_info = NULL;
     PyObject *thread_id = NULL;
@@ -309,6 +310,7 @@ 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);