Use a more typical strategy for profiling current tasks

1 files changed, 55 insertions, 19 deletions

diff --git a/aergia/aergia.py b/aergia/aergia.py
index 77be0fc..d3c525d 100755
--- a/aergia/aergia.py
+++ b/aergia/aergia.py
@@ -97,10 +97,13 @@ class Aergia(object):
     # the (ideal) interval between samples
 
     signal_interval = 0.0
+    # if we should profile currently running tasks
+    do_profile_current = False
 
     @staticmethod
-    def __init__(signal_interval):
+    def __init__(signal_interval, do_profile_current):
         Aergia.signal_interval = signal_interval
+        Aergia.do_profile_current = do_profile_current
 
     @staticmethod
     def start():
@@ -154,15 +157,16 @@ class Aergia(object):
     @staticmethod
     def _idle_signal_handler(sig, frame):
         '''Obtains and records which lines are currently being waited on.'''
-        keys = Aergia._compute_frames_to_record()
+        keys = Aergia._compute_frames_to_record(frame)
         for key in keys:
             Aergia.samples[Aergia._frame_to_tuple(key)] += 1
             Aergia.total_samples += 1
 
     @staticmethod
-    def _compute_frames_to_record():
+    def _compute_frames_to_record(frame):
         '''Collects all stack frames which are currently being awaited on
-        during a given timestamp, and
+        during a given timestamp, as well as those which are currently
+        executing.
 
         Note that we do NOT need to walk back up the call-stack to find
         which of the user's lines caused the await call. There is NEVER
@@ -171,7 +175,11 @@ class Aergia(object):
         Luckily, the event loop and asyncio.all_tasks keeps track of
         what is running for us.'''
         loops = Aergia._get_event_loops()
+        # idle tasks
         frames = Aergia._get_frames_from_loops(loops)
+        # current running frames
+        if Aergia.do_profile_current:
+            frames += Aergia._get_frames_from_threads(frame)
         return frames
 
     @staticmethod
@@ -211,6 +219,30 @@ class Aergia(object):
         ]
 
     @staticmethod
+    def _get_frames_from_threads(frame):
+        frames = [frame]
+        frames += [sys._current_frames().get(t.ident, None)
+                  for t in threading.enumerate()]
+        # process frames to remove those we do not track
+        new_frames = []
+        for f in frames:
+            if f is None:
+                continue
+            fname = frame.f_code.co_filename
+            while not Aergia._should_trace(fname):
+                # walk the stack backwards until we hit a frame that is one
+                # we should trace.
+                if frame:
+                    frame = frame.f_back
+                else:
+                    break
+                if frame:
+                    fname = frame.f_code.co_filename
+            if frame:
+                new_frames.append(frame)
+        return new_frames
+
+    @staticmethod
     def _frame_to_tuple(frame):
         '''Given a frame, constructs a sample key for tallying lines.'''
         co = frame.f_code
@@ -238,7 +270,10 @@ class Aergia(object):
         '''Given an asyncio event loop, returns the list of idle task frames.
         A task is considered 'idle' if it is not currently executing.'''
         idle = []
+        current = asyncio.current_task(loop)
         for task in asyncio.all_tasks(loop):
+            if task == current:
+                continue
             coro = task.get_coro()
             if coro:
                 f = Aergia._get_deepest_traceable_frame(coro)
@@ -278,10 +313,10 @@ class Aergia(object):
         # third case:  we'd like to return the current value of deepest_frame.
 
         previous_frame = None
-        print('-' * 60)
-        print('--- Frame Information for ---')
-        print(deepest_frame)
-        print(f"{'Filename':<30} {'Function':<30} {'Line':<10}")
+        # print('-' * 60)
+        # print('--- Frame Information for ---')
+        # print(deepest_frame)
+        # print(f"{'Filename':<30} {'Function':<30} {'Line':<10}")
         while previous_frame is not deepest_frame:
             previous_frame = deepest_frame
             refs = gc.get_referrers(deepest_frame)
@@ -289,13 +324,13 @@ class Aergia(object):
                 if inspect.isframe(r) and \
                    Aergia._should_trace(r.f_code.co_filename):
                     deepest_frame = r
-                    co = r.f_code
-                    func_name = co.co_name
-                    filename = co.co_filename
-                    filename = (filename if len(filename) <= 30 else filename[-30:])
-                    line_no = r.f_lineno
-                    print(f"{filename:<30} {func_name:<30} {line_no:<10}")
-        print(f'finished with {deepest_frame}')
+                    # co = r.f_code
+                    # func_name = co.co_name
+                    # filename = co.co_filename
+                    # filename = (filename if len(filename) <= 30 else filename[-30:])
+                    # line_no = r.f_lineno
+        #             print(f"{filename:<30} {func_name:<30} {line_no:<10}")
+        # print(f'finished with {deepest_frame}')
         return deepest_frame
 
     @staticmethod
@@ -346,9 +381,10 @@ if __name__ == "__main__":
                         default=0.01)
     parser.add_argument('-d', '--debug',
                         help='Turn on debug information for the event loop.',
-                        metavar='',
-                        type=bool,
-                        default=False)
+                        action='store_true')
+    parser.add_argument('-a', '--async-only',
+                        help='Do not profile currently running tasks.',
+                        action='store_true')
     parser.add_argument('script', help='A python script to run.')
     parser.add_argument('s_args', nargs=argparse.REMAINDER,
                         help='python script args')
@@ -360,7 +396,7 @@ if __name__ == "__main__":
     try:
         with open(args.script, 'r', encoding='utf-8') as fp:
             code = compile(fp.read(), args.script, "exec")
-            Aergia(args.interval).start()
+            Aergia(args.interval, not args.async_only).start()
             exec(code, the_globals)
             Aergia.print_samples()
             Aergia.stop()