Add naive async profiling functionality and proof of concept program

1 files changed, 15 insertions, 2 deletions

diff --git a/mini-scalene.py b/mini-scalene.py
index 2fa5b59..01a8de2 100644
--- a/mini-scalene.py
+++ b/mini-scalene.py
@@ -4,6 +4,8 @@ import traceback
 import runpy
 import atexit
 import signal
+import asyncio
+import inspect
 from typing import cast
 from types import FrameType
 from collections import defaultdict
@@ -55,6 +57,7 @@ class mini_scalene:
         frames = [this_frame]
         frames += [sys._current_frames().get(t.ident, None)
                    for t in threading.enumerate()]
+        frames += mini_scalene.get_async_frames()
         # Process all the frames to remove ones we aren't going to track.
         new_frames = []
         for frame in frames:
@@ -91,6 +94,12 @@ class mini_scalene:
         return filename + '\t' + func_name + '\t' + str(line_no)
 
     @staticmethod
+    def get_async_frames():
+        if mini_scalene.is_event_loop_running():
+            return [task.get_coro().cr_frame for task in asyncio.all_tasks()]
+        return []
+
+    @staticmethod
     def should_trace(filename):
         # We're assuming Guix System. That makes it easy.
         if '/gnu/store' in filename:
@@ -103,10 +112,14 @@ class mini_scalene:
             return False
         return True
 
+    @staticmethod
+    def is_event_loop_running() -> bool:
+        return asyncio.get_event_loop_policy()._local._loop is not None
+
+
 
 def main():
-    assert len(
-        sys.argv) >= 2, "(Usage): python3 mini-scalene.py file.py {args ...}"
+    assert len(sys.argv) >= 2, "(Usage): python3 mini-scalene.py file.py {args ...}"
 
     script = sys.argv[1]
     mini_scalene().start()