Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions Lib/profiling/sampling/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
18 changes: 7 additions & 11 deletions Lib/profiling/sampling/gecko_collector.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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 = [
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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,
)
],
)
Expand Down Expand Up @@ -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)
Expand Down
4 changes: 3 additions & 1 deletion Modules/_remote_debugging/_remote_debugging.h
Original file line number Diff line number Diff line change
Expand Up @@ -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) \
Expand Down Expand Up @@ -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) */
Expand Down
10 changes: 9 additions & 1 deletion Modules/_remote_debugging/module.c
Original file line number Diff line number Diff line change
Expand Up @@ -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 ||
Expand Down Expand Up @@ -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;
Expand Down
7 changes: 6 additions & 1 deletion Modules/_remote_debugging/threads.c
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down
Loading