3 files changed, 51 insertions, 7 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()
diff --git a/tests/simult-busy.py b/tests/simult-busy.py
new file mode 100644
index 0000000..4cc4b2a
--- /dev/null
+++ b/tests/simult-busy.py
@@ -0,0 +1,25 @@
+import asyncio
+import time
+
+async def busy_task():
+    await asyncio.sleep(3.0)
+    return 1
+
+async def main():
+    tasks = [asyncio.create_task(busy_task()) for i in range(5)]
+
+    start_time = time.time()
+    try:
+        # this is to prevent waiting in 'select' all day,
+        # which makes the python intepreter not respond to
+        # mini-scalene / SCALENE
+        while True:
+            if time.time() - start_time > 3.5:
+                break
+            await asyncio.sleep(0)  # yield
+    except KeyboardInterrupt:
+        pass
+    finally:
+        print("Done.")
+
+asyncio.run(main())
diff --git a/tests/simult.py b/tests/simult.py
index 6228596..61a3792 100644
--- a/tests/simult.py
+++ b/tests/simult.py
@@ -1,14 +1,20 @@
 import asyncio
+import time
 
 
-async def count():
-    print("before")
-    await asyncio.sleep(1)
-    print("after")
+async def count(x):
+    i = 0
+    await asyncio.sleep(2)
+    while i < 100000:
+        z = x * x
+        z = z * z
+        z = z * z
+        i += 1
+    return z
 
 
 async def main():
-    await asyncio.gather(count(), count(), count())
+    await asyncio.gather(count(1), count(2), count(3))
     print("done")
 
 asyncio.run(main())