Profile all function frames

1 files changed, 21 insertions, 11 deletions

diff --git a/nemesis/nemesis.py b/nemesis/nemesis.py
index 5901801..7eed743 100755
--- a/nemesis/nemesis.py
+++ b/nemesis/nemesis.py
@@ -157,7 +157,7 @@ class Nemesis(object):
             loops = Nemesis.experiment_data.get_loops()
             exp_coro = Nemesis.experiment_coro
             for loop in loops:
-                coro = Nemesis._get_current_coro(loop)
+                coro = Nemesis._get_current_frame(loop)
                 prev_coro = Nemesis.prev_coro[loop]
                 if not prev_coro == coro:
                     if prev_coro == exp_coro:
@@ -173,15 +173,15 @@ class Nemesis(object):
                 Nemesis._stop_experiment()
 
         else:
-            coros = []
+            frames = []
             loops = Nemesis._get_event_loops()
             for loop in loops:
-                coro = Nemesis._get_current_coro(loop)
-                if coro is not None:
-                    coros.append(coro)
-                if coros:
-                    Nemesis._start_experiment(random.choice(coros),
-                                              Nemesis._select_speedup())
+                frame = Nemesis._get_current_frame(loop)
+                if frame is not None:
+                    frames.append(frame)
+            if frames:
+                Nemesis._start_experiment(random.choice(frames),
+                                          Nemesis._select_speedup())
 
     @staticmethod
     def _parse_handle(handle):
@@ -194,21 +194,31 @@ class Nemesis(object):
             return [str(type(handle).__name__), cb.__name__]
 
     @staticmethod
-    def _get_current_coro(loop):
+    def _get_current_frame(loop):
         tid = loop._thread_id
         assert tid, f"{loop} is not running, yet we attempted to sample it!"
 
         frame = sys._current_frames().get(tid)
         fname = frame.f_code.co_filename
+        if not fname:
+            # 'eval/compile' gives no f_code.co_filename.  We have
+            # to look back into the outer frame in order to check
+            # the co_filename.
+            back = frame.f_back
+            fname = back.f_code.co_filename
         while not Nemesis._should_trace(fname):
+            # Walk the stack backwards until we hit a frame that
+            # IS one we should trace (if there is one).  i.e., if
+            # it's in the code being profiled, and it is just
+            # calling stuff deep in libraries.
             if frame:
                 frame = frame.f_back
             else:
                 break
             if frame:
                 fname = frame.f_code.co_filename
-        if frame and frame.f_generator:
-            return frame.f_generator.cr_code.co_name
+        if frame:
+            return frame.f_code.co_name
         return None
 
     @staticmethod