1+ import collections
2+ import os
3+ from .stack_collectors import StackTraceCollector
4+
5+
6+ class AsyncStackCollector (StackTraceCollector ):
7+ def __init__ (self , unwinder ):
8+ super ().__init__ ()
9+ self .unwinder = unwinder
10+ self .async_stacks = []
11+
12+ def collect (self , stack_frames ):
13+ # First collect regular stack frames
14+ super ().collect (stack_frames )
15+
16+ # Now collect async task information
17+ try :
18+ awaited_info = self .unwinder .get_all_awaited_by ()
19+ if awaited_info :
20+ self ._process_async_stacks (awaited_info )
21+ except Exception :
22+ # If async collection fails, continue with regular profiling
23+ pass
24+
25+ def _process_async_stacks (self , awaited_info ):
26+ """Process the async task information and reconstruct full stacks."""
27+ # Build a map of task_id to TaskInfo for easy lookup
28+ task_map = {}
29+
30+ for thread_info in awaited_info :
31+ if thread_info .thread_id == 0 :
32+ continue # Skip empty thread info
33+
34+ for task_info in thread_info .awaited_by :
35+ task_map [task_info .task_id ] = task_info
36+
37+ # For each task, reconstruct the full async stack
38+ for thread_info in awaited_info :
39+ if thread_info .thread_id == 0 :
40+ continue
41+
42+ for task_info in thread_info .awaited_by :
43+ full_stacks = self ._reconstruct_task_stacks (task_info , task_map )
44+ self .async_stacks .extend (full_stacks )
45+
46+ def _reconstruct_task_stacks (self , task_info , task_map , visited = None ):
47+ """Recursively reconstruct full async stacks for a task."""
48+ if visited is None :
49+ visited = set ()
50+
51+ if task_info .task_id in visited :
52+ return [] # Avoid infinite recursion
53+
54+ visited .add (task_info .task_id )
55+ full_stacks = []
56+
57+ # Get the current task's stack from its coroutine_stack
58+ for coro_info in task_info .coroutine_stack :
59+ current_stack = []
60+ for frame in coro_info .call_stack :
61+ # Convert FrameInfo to tuple format (filename, lineno, funcname)
62+ current_stack .append ((frame .filename , frame .lineno , frame .funcname ))
63+
64+ # If this task is awaited by others, we need to connect to awaiter stacks
65+ if task_info .awaited_by :
66+ for awaiter_coro in task_info .awaited_by :
67+ # The task_name in CoroInfo is actually the task_id of the awaiting task
68+ awaiter_task_id = awaiter_coro .task_name
69+
70+ if awaiter_task_id in task_map :
71+ awaiter_task = task_map [awaiter_task_id ]
72+ # Recursively get the awaiter's full stacks
73+ awaiter_stacks = self ._reconstruct_task_stacks (
74+ awaiter_task , task_map , visited
75+ )
76+ if awaiter_stacks :
77+ # For each awaiter stack, create a combined stack
78+ for awaiter_stack in awaiter_stacks :
79+ # Build: current_stack + [current_task_name] + awaiter_stack
80+ full_stack = current_stack [:]
81+ # Insert task name as transition between this task and its awaiter
82+ if task_info .task_name and not task_info .task_name .startswith ('Task-' ):
83+ full_stack .append (("" , 0 , f"[Task:{ task_info .task_name } ))
84+ full_stack .extend (awaiter_stack )
85+ full_stacks .append (full_stack )
86+ else :
87+ # Awaiter has no further awaiters, build final stack
88+ full_stack = current_stack [:]
89+ # Add current task name
90+ if task_info .task_name and not task_info .task_name .startswith ('Task-' ):
91+ full_stack .append (("" , 0 , f"[Task:{ task_info .task_name } ))
92+ # Add awaiter's coroutine stack
93+ for awaiter_task_coro in awaiter_task .coroutine_stack :
94+ for frame in awaiter_task_coro .call_stack :
95+ full_stack .append ((frame .filename , frame .lineno , frame .funcname ))
96+ # Add awaiter task name at the end if it's a leaf
97+ if awaiter_task .task_name and not awaiter_task .task_name .startswith ('Task-' ):
98+ full_stack .append (("" , 0 , f"[Task:{ awaiter_task .task_name } ))
99+ full_stacks .append (full_stack )
100+ else :
101+ # Can't find awaiter task in map, use the coroutine info directly
102+ full_stack = current_stack [:]
103+ # Add current task name
104+ if task_info .task_name and not task_info .task_name .startswith ('Task-' ):
105+ full_stack .append (("" , 0 , f"[Task:{ task_info .task_name } ))
106+ # Add awaiter coroutine's stack
107+ for frame in awaiter_coro .call_stack :
108+ full_stack .append ((frame .filename , frame .lineno , frame .funcname ))
109+ full_stacks .append (full_stack )
110+ else :
111+ # No awaiters, this is a leaf task - just add task name at the end
112+ full_stack = current_stack [:]
113+ if task_info .task_name and not task_info .task_name .startswith ('Task-' ):
114+ full_stack .append (("" , 0 , f"[Task:{ task_info .task_name } ))
115+ full_stacks .append (full_stack )
116+
117+ visited .remove (task_info .task_id )
118+ return full_stacks
119+
120+ def export (self , filename ):
121+ """Export both regular and async stacks in collapsed format."""
122+ stack_counter = collections .Counter ()
123+
124+ # Add regular stacks
125+ for call_tree in self .call_trees :
126+ stack_str = ";" .join (
127+ f"{ os .path .basename (f [0 ])} { f [2 ]} { f [1 ]} for f in call_tree
128+ )
129+ stack_counter [stack_str ] += 1
130+
131+ # Add async stacks with [async] prefix
132+ for async_stack in self .async_stacks :
133+ # Reverse to get root->leaf order
134+ async_stack_reversed = list (reversed (async_stack ))
135+ stack_parts = ["[async]" ]
136+ for f in async_stack_reversed :
137+ if f [0 ] == "" and f [1 ] == 0 and f [2 ].startswith ("[Task:" ):
138+ # This is a task name pseudo-frame
139+ stack_parts .append (f [2 ])
140+ else :
141+ # Regular frame
142+ stack_parts .append (f"{ os .path .basename (f [0 ])} { f [2 ]} { f [1 ]} )
143+ stack_str = ";" .join (stack_parts )
144+ stack_counter [stack_str ] += 1
145+
146+ with open (filename , "w" ) as f :
147+ for stack , count in stack_counter .items ():
148+ f .write (f"{ stack } { count } \n " )
149+
150+ total_regular = len (self .call_trees )
151+ total_async = len (self .async_stacks )
152+ print (f"Collapsed stack output written to { filename } )
153+ print (f" Regular stacks: { total_regular } )
154+ print (f" Async stacks: { total_async } )
0 commit comments