Good shit

Author: pablogsal

Lib/profiling/sampling/collector.py

@@ -36,8 +36,16 @@ def _iter_all_frames(self, stack_frames, skip_idle=False):
                     yield frames, thread_info.thread_id
 
     def _iter_async_frames(self, awaited_info_list):
-        """Iterate over linear stacks for all leaf tasks (hot path optimized)."""
-        # Build adjacency graph (O(n))
+        # Phase 1: Index tasks and build parent relationships
+        task_map, child_to_parents, all_task_ids = self._build_task_graph(awaited_info_list)
+
+        # Phase 2: Find leaf tasks (tasks not awaited by anyone)
+        leaf_task_ids = self._find_leaf_tasks(child_to_parents, all_task_ids)
+
+        # Phase 3: Build linear stacks via BFS from each leaf to root
+        yield from self._build_linear_stacks(leaf_task_ids, task_map, child_to_parents)
+
+    def _build_task_graph(self, awaited_info_list):
         task_map = {}
         child_to_parents = {}
         all_task_ids = set()
@@ -48,70 +56,60 @@ def _iter_async_frames(self, awaited_info_list):
                 task_id = task_info.task_id
                 task_map[task_id] = (task_info, thread_id)
                 all_task_ids.add(task_id)
+
+                # Store parent task IDs (not frames - those are in task_info.coroutine_stack)
                 if task_info.awaited_by:
-                    # Store all parent coroutines, not just [0]
-                    child_to_parents[task_id] = task_info.awaited_by
+                    child_to_parents[task_id] = [p.task_name for p in task_info.awaited_by]
+
+        return task_map, child_to_parents, all_task_ids
 
-        # Identify leaf tasks (O(n))
-        # Collect all parent task IDs from all coroutines
+    def _find_leaf_tasks(self, child_to_parents, all_task_ids):
         all_parent_ids = set()
-        for parent_coros in child_to_parents.values():
-            for parent_coro in parent_coros:
-                all_parent_ids.add(parent_coro.task_name)
-        leaf_task_ids = all_task_ids - all_parent_ids
+        for parent_ids in child_to_parents.values():
+            all_parent_ids.update(parent_ids)
+        return all_task_ids - all_parent_ids
 
-        # Build linear stacks for each leaf (O(n × depth × num_paths))
-        # For tasks with multiple parents, we generate one stack per parent path
+    def _build_linear_stacks(self, leaf_task_ids, task_map, child_to_parents):
         for leaf_id in leaf_task_ids:
-            # Use BFS to explore all paths from leaf to root
-            # Queue items: (current_task_id, frames_accumulated)
-            queue = [(leaf_id, [])]
-            visited = set()
+            # BFS queue: (current_task_id, frames_so_far, path_for_cycle_detection)
+            queue = [(leaf_id, [], frozenset())]
 
             while queue:
-                current_id, frames = queue.pop(0)
+                current_id, frames, path = queue.pop(0)
 
-                # Avoid processing the same task twice in this path
-                if current_id in visited:
+                # Cycle detection
+                if current_id in path:
                     continue
-                visited.add(current_id)
 
+                # End of path (parent ID not in task_map)
                 if current_id not in task_map:
-                    # Reached end of path - yield if we have frames
                     if frames:
                         _, thread_id = task_map[leaf_id]
                         yield frames, thread_id, leaf_id
                     continue
 
+                # Process current task
                 task_info, tid = task_map[current_id]
-
-                # Add this task's frames
                 new_frames = list(frames)
+                new_path = path | {current_id}
+
+                # Add all frames from all coroutines in this task
                 if task_info.coroutine_stack:
-                    for frame in task_info.coroutine_stack[0].call_stack:
-                        new_frames.append(frame)
+                    for coro_info in task_info.coroutine_stack:
+                        for frame in coro_info.call_stack:
+                            new_frames.append(frame)
 
                 # Add task boundary marker
                 task_name = task_info.task_name or "Task-" + str(task_info.task_id)
                 new_frames.append(FrameInfo(("<task>", 0, task_name)))
 
-                # Get parent coroutines
-                parent_coros = child_to_parents.get(current_id)
-                if not parent_coros:
-                    # No parents - this is the root, yield the complete stack
-                    yield new_frames, tid, leaf_id
-                    continue
-
-                # For each parent coroutine, add its await frames and continue to parent task
-                for parent_coro in parent_coros:
-                    parent_task_id = parent_coro.task_name
-
-                    # Add the parent's await-site frames (where parent awaits this task)
-                    path_frames = list(new_frames)
-                    for frame in parent_coro.call_stack:
-                        path_frames.append(frame)
+                # Get parent task IDs
+                parent_ids = child_to_parents.get(current_id, [])
 
-                    # Continue BFS with parent task
-                    # Note: parent_coro.call_stack contains the frames from the parent task,
-                    # so we should NOT add parent task's coroutine_stack again
-                    queue.append((parent_task_id, path_frames))
+                if not parent_ids:
+                    # Root task - yield complete stack
+                    yield new_frames, tid, leaf_id
+                else:
+                    # Continue to each parent (creates multiple paths if >1 parent)
+                    for parent_id in parent_ids:
+                        queue.append((parent_id, new_frames, new_path))