gh-152434: Fix async-aware Gecko collection (GH-152442)

(cherry picked from commit 87ac0bc)

Co-authored-by: László Kiss Kollár <kiss.kollar.laszlo@gmail.com>

Author: lkollar

Lib/profiling/sampling/gecko_collector.py

@@ -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:
@@ -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:

Lib/test/test_profiling/test_sampling_profiler/test_collectors.py

@@ -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
 
 
@@ -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."""

Misc/NEWS.d/next/Library/2026-06-25-20-00-00.gh-issue-151496.kMnT3x.rst

@@ -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.