It might work but the code is bad

Author: savannahostrowski

Lib/profiling/sampling/collector.py

@@ -1,5 +1,8 @@
 from abc import ABC, abstractmethod
 
+from _remote_debugging import FrameInfo
+
+
 # Enums are slow
 THREAD_STATE_RUNNING = 0
 THREAD_STATE_IDLE = 1
@@ -31,3 +34,85 @@ def _iter_all_frames(self, stack_frames, skip_idle=False):
                 frames = thread_info.frame_info
                 if frames:
                     yield frames, thread_info.thread_id
+
+    def _iter_async_frames(self, awaited_info_list):
+        """
+        Iterate over all async frame stacks from awaited info.
+        Yields one stack per task in LEAF→ROOT order: [task body, task marker, parent body, ..., Program Root]
+        """
+        # Index all tasks by ID
+        all_tasks = {}
+        for awaited_info in awaited_info_list:
+            for task_info in awaited_info.awaited_by:
+                all_tasks[task_info.task_id] = (task_info, awaited_info.thread_id)
+
+        cache = {}  # Memoize parent chains
+
+        def build_parent_chain(task_id, parent_id):
+            """Build ancestor chain: [await-site frames, grandparent chain..., Program Root]"""
+            if parent_id in cache:
+                return cache[parent_id]
+
+            if parent_id not in all_tasks:
+                return []
+
+            parent_info, _ = all_tasks[parent_id]
+
+            # Find the await-site frames for this parent relationship
+            await_frames = []
+            for coro_info in all_tasks[task_id][0].awaited_by:
+                if coro_info.task_name == parent_id:
+                    await_frames = list(coro_info.call_stack or [])
+                    break
+
+            # Recursively build grandparent chain, or terminate with Program Root
+            if (parent_info.awaited_by and parent_info.awaited_by[0].task_name and
+                parent_info.awaited_by[0].task_name in all_tasks):
+                grandparent_id = parent_info.awaited_by[0].task_name
+                chain = await_frames + build_parent_chain(parent_id, grandparent_id)
+            else:
+                # Parent is root or grandparent not tracked
+                root_frame = FrameInfo(("<thread>", 0, "Program Root"))
+                chain = await_frames + [root_frame]
+
+            cache[parent_id] = chain
+            return chain
+
+        # Find all parent task IDs (tasks that have children)
+        parent_task_ids = {
+            coro.task_name
+            for task_info, _ in all_tasks.values()
+            for coro in (task_info.awaited_by or [])
+            if coro.task_name
+        }
+
+        # Yield one stack per leaf task (tasks that are not parents)
+        for task_id, (task_info, thread_id) in all_tasks.items():
+            # Skip parent tasks - they'll be included in their children's stacks
+            if task_id in parent_task_ids:
+                continue
+            # Collect task's coroutine frames
+            body_frames = [
+                frame
+                for coro in (task_info.coroutine_stack or [])
+                for frame in (coro.call_stack or [])
+            ]
+
+            # Add synthetic task marker
+            task_name = task_info.task_name or f"Task-{task_id}"
+            synthetic = FrameInfo(("<task>", 0, f"running {task_name}"))
+
+            # Build complete stack with parent chain
+            if task_info.awaited_by and task_info.awaited_by[0].task_name:
+                parent_id = task_info.awaited_by[0].task_name
+                if parent_id in all_tasks:
+                    parent_chain = build_parent_chain(task_id, parent_id)
+                    yield body_frames + [synthetic] + parent_chain, thread_id, 0
+                else:
+                    # Parent not tracked, treat as root task
+                    root = FrameInfo(("<thread>", 0, "Program Root"))
+                    yield body_frames + [synthetic, root], thread_id, 0
+            else:
+                # Root task (no parents or empty awaited_by)
+                root = FrameInfo(("<thread>", 0, "Program Root"))
+                yield body_frames + [synthetic, root], thread_id, 0

Lib/profiling/sampling/pstats_collector.py

@@ -41,8 +41,14 @@ def _process_frames(self, frames):
             self.callers[callee][caller] += 1
 
     def collect(self, stack_frames):
-        for frames, thread_id in self._iter_all_frames(stack_frames, skip_idle=self.skip_idle):
-            self._process_frames(frames)
+        if stack_frames and hasattr(stack_frames[0], "awaited_by"):
+            # Async frame processing
+            for frames, thread_id, _depth in self._iter_async_frames(stack_frames):
+                self._process_frames(frames)
+        else:
+            # Regular frame processing
+            for frames, thread_id in self._iter_all_frames(stack_frames, skip_idle=self.skip_idle):
+                self._process_frames(frames)
 
     def export(self, filename):
         self.create_stats()

Lib/profiling/sampling/sample.py

@@ -154,7 +154,7 @@ def __init__(self, pid, sample_interval_usec, all_threads, *, mode=PROFILING_MOD
         self.total_samples = 0
         self.realtime_stats = False
 
-    def sample(self, collector, duration_sec=10):
+    def sample(self, collector, async_aware, duration_sec=10):
         sample_interval_sec = self.sample_interval_usec / 1_000_000
         running_time = 0
         num_samples = 0
@@ -168,7 +168,10 @@ def sample(self, collector, duration_sec=10):
             current_time = time.perf_counter()
             if next_time < current_time:
                 try:
-                    stack_frames = self.unwinder.get_stack_trace()
+                    if async_aware:
+                        stack_frames = self.unwinder.get_all_awaited_by()
+                    else:
+                        stack_frames = self.unwinder.get_stack_trace()
                     collector.collect(stack_frames)
                 except ProcessLookupError:
                     duration_sec = current_time - start_time
@@ -613,6 +616,7 @@ def sample(
     output_format="pstats",
     realtime_stats=False,
     mode=PROFILING_MODE_WALL,
+    async_aware=False,
 ):
     profiler = SampleProfiler(
         pid, sample_interval_usec, all_threads=all_threads, mode=mode
@@ -638,7 +642,7 @@ def sample(
         case _:
             raise ValueError(f"Invalid output format: {output_format}")
 
-    profiler.sample(collector, duration_sec)
+    profiler.sample(collector, async_aware, duration_sec)
 
     if output_format == "pstats" and not filename:
         stats = pstats.SampledStats(collector).strip_dirs()
@@ -706,6 +710,7 @@ def wait_for_process_and_sample(pid, sort_value, args):
         output_format=args.format,
         realtime_stats=args.realtime_stats,
         mode=mode,
+        async_aware=args.async_aware,
     )
 
 
@@ -760,6 +765,13 @@ def main():
         help="Print real-time sampling statistics (Hz, mean, min, max, stdev) during profiling",
     )
 
+    sampling_group.add_argument(
+        "--async-aware",
+        action="store_true",
+        default=False,
+        help="Enable async-aware sampling (experimental)",
+    )
+
     # Mode options
     mode_group = parser.add_argument_group("Mode options")
     mode_group.add_argument(
@@ -915,6 +927,7 @@ def main():
             output_format=args.format,
             realtime_stats=args.realtime_stats,
             mode=mode,
+            async_aware=args.async_aware,
         )
     elif args.module or args.args:
         if args.module:

Lib/profiling/sampling/stack_collector.py

@@ -15,10 +15,16 @@ def __init__(self, *, skip_idle=False):
         self.skip_idle = skip_idle
 
     def collect(self, stack_frames, skip_idle=False):
-        for frames, thread_id in self._iter_all_frames(stack_frames, skip_idle=skip_idle):
-            if not frames:
-                continue
-            self.process_frames(frames, thread_id)
+        if stack_frames and hasattr(stack_frames[0], "awaited_by"):
+            for frames, thread_id, _depth in self._iter_async_frames(stack_frames):
+                if not frames:
+                    continue
+                self.process_frames(frames, thread_id)
+        else:
+            for frames, thread_id in self._iter_all_frames(stack_frames, skip_idle=skip_idle):
+                if not frames:
+                    continue
+                self.process_frames(frames, thread_id)
 
     def process_frames(self, frames, thread_id):
         pass