Skip to content
Merged
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
81 changes: 53 additions & 28 deletions Lib/profiling/sampling/gecko_collector.py
Original file line number Diff line number Diff line change
Expand Up @@ -250,6 +250,25 @@ def collect(self, stack_frames, timestamps_us=None):
self.interval = (times[-1] - self.last_sample_time) / self.sample_count
self.last_sample_time = times[-1]

# Process async tasks
if stack_frames and hasattr(stack_frames[0], "awaited_by"):
for frames, thread_id, _ in self._iter_async_frames(stack_frames):
frames = filter_internal_frames(frames)
if not frames:
continue

if thread_id not in self.threads:
self.threads[thread_id] = self._create_thread(
thread_id, False
)

self._record_stack_sample(
self.threads[thread_id], frames, thread_id, times, first_time
)

self.sample_count += len(times)
return

# Process threads
for interpreter_info in stack_frames:
for thread_info in interpreter_info.threads:
Expand Down Expand Up @@ -333,37 +352,43 @@ def collect(self, stack_frames, timestamps_us=None):
if not frames:
continue

# Process stack once to get stack_index
stack_index = self._process_stack(thread_data, frames)

# Add samples with timestamps
thread_spill = thread_data["_spill"]
for t in times:
thread_spill.append_sample(stack_index, t)

# Handle opcodes
if self.opcodes_enabled and frames:
leaf_frame = frames[0]
filename, location, funcname, opcode = leaf_frame
if isinstance(location, tuple):
lineno, _, col_offset, _ = location
else:
lineno = location
col_offset = -1

current_state = (opcode, lineno, col_offset, funcname, filename)

if tid not in self.opcode_state:
self.opcode_state[tid] = (*current_state, first_time)
elif self.opcode_state[tid][:5] != current_state:
prev_opcode, prev_lineno, prev_col, prev_funcname, prev_filename, prev_start = self.opcode_state[tid]
self._add_opcode_interval_marker(
tid, prev_opcode, prev_lineno, prev_col, prev_funcname, prev_start, first_time
)
self.opcode_state[tid] = (*current_state, first_time)
self._record_stack_sample(
thread_data, frames, tid, times, first_time
)

self.sample_count += len(times)

def _record_stack_sample(self, thread_data, frames, tid, times, first_time):
stack_index = self._process_stack(thread_data, frames)

thread_spill = thread_data["_spill"]
for t in times:
thread_spill.append_sample(stack_index, t)

if self.opcodes_enabled and frames:
leaf_frame = frames[0]
filename, location, funcname, opcode = leaf_frame
if isinstance(location, tuple):
lineno, _, col_offset, _ = location
else:
lineno = location
col_offset = -1

current_state = (opcode, lineno, col_offset, funcname, filename)

if tid not in self.opcode_state:
self.opcode_state[tid] = (*current_state, first_time)
elif self.opcode_state[tid][:5] != current_state:
(
prev_opcode, prev_lineno, prev_col, prev_funcname,
prev_filename, prev_start
) = self.opcode_state[tid]
self._add_opcode_interval_marker(
tid, prev_opcode, prev_lineno, prev_col, prev_funcname,
prev_start, first_time
)
self.opcode_state[tid] = (*current_state, first_time)

def _create_thread(self, tid, is_main_thread):
"""Create a new thread structure with processed profile format."""
if self.spill_dir is None:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,16 @@

from test.support import captured_stdout, captured_stderr

from .mocks import MockFrameInfo, MockThreadInfo, MockInterpreterInfo, LocationInfo, make_diff_collector_with_mock_baseline
from .mocks import (
MockAwaitedInfo,
MockCoroInfo,
MockFrameInfo,
MockInterpreterInfo,
MockTaskInfo,
MockThreadInfo,
LocationInfo,
make_diff_collector_with_mock_baseline,
)
from .helpers import close_and_unlink, jsonl_tables


Expand Down Expand Up @@ -657,6 +666,48 @@ def test_gecko_collector_basic(self):
self.assertGreater(stack_table["length"], 0)
self.assertGreater(len(stack_table["frame"]), 0)

def test_gecko_collector_async_aware(self):
collector = GeckoCollector(1000)

parent = MockTaskInfo(
task_id=1,
task_name="Parent",
coroutine_stack=[
MockCoroInfo(
task_name="Parent",
call_stack=[MockFrameInfo("parent.py", 10, "parent_fn")],
)
],
)
child = MockTaskInfo(
task_id=2,
task_name="Child",
coroutine_stack=[
MockCoroInfo(
task_name="Child",
call_stack=[MockFrameInfo("child.py", 20, "child_fn")],
)
],
awaited_by=[MockCoroInfo(task_name=1, call_stack=[])],
)

collector.collect(
[MockAwaitedInfo(thread_id=100, awaited_by=[parent, child])],
timestamps_us=[1000, 2000],
)
profile_data = export_gecko_profile(self, collector)

self.assertEqual(len(profile_data["threads"]), 1)
thread_data = profile_data["threads"][0]
self.assertEqual(thread_data["samples"]["length"], 2)

string_array = profile_data["shared"]["stringArray"]
self.assertIn("parent_fn", string_array)
self.assertIn("child_fn", string_array)
self.assertIn("Parent", string_array)
self.assertIn("Child", string_array)
self.assertEqual(thread_data["markers"]["length"], 0)

@unittest.skipIf(is_emscripten, "threads not available")
def test_gecko_collector_export(self):
"""Test Gecko profile export functionality."""
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
Fixed ``profiling.sampling --gecko`` with ``--async-aware`` by flattening
async task stacks before generating Gecko samples. ``--binary`` now rejects
``--async-aware`` until the binary format supports async task data.
Loading