Numerous big fixes, target function only selects in event loop

1 files changed, 16 insertions, 5 deletions

diff --git a/nemesis/nemesis.py b/nemesis/nemesis.py
index 7eed743..03901e3 100755
--- a/nemesis/nemesis.py
+++ b/nemesis/nemesis.py
@@ -31,6 +31,7 @@ from html_gen import plot_results
 import argparse
 import asyncio
 import os
+import inspect
 import random
 import signal
 import sys
@@ -39,6 +40,8 @@ import time
 import traceback
 import types
 
+CO_COROUTINE = inspect.CO_COROUTINE
+
 class Experiment:
 
     # event loops participating in this this experiment
@@ -128,7 +131,7 @@ class Nemesis(object):
         loops = Nemesis._get_event_loops()
         for loop in loops:
             if not isinstance(loop, CausalEventLoop):
-                raise RuntimeError("Nemesis requires a custom event loop to insert slowdowns. It does not work on programs which change the event loop policy.")
+                raise RuntimeError(f"Nemesis requires a custom event loop to insert slowdowns. You must start the event loop with `asyncio.run(your_coro(), loop_factory=causal_loop_factory)'.")
             loop.set_speedup(speedup)
 
         Nemesis.experiment_data = Experiment(loops)
@@ -157,7 +160,7 @@ class Nemesis(object):
             loops = Nemesis.experiment_data.get_loops()
             exp_coro = Nemesis.experiment_coro
             for loop in loops:
-                coro = Nemesis._get_current_frame(loop)
+                coro = Nemesis._get_current_frame(loop).f_code.co_name
                 prev_coro = Nemesis.prev_coro[loop]
                 if not prev_coro == coro:
                     if prev_coro == exp_coro:
@@ -177,8 +180,8 @@ class Nemesis(object):
             loops = Nemesis._get_event_loops()
             for loop in loops:
                 frame = Nemesis._get_current_frame(loop)
-                if frame is not None:
-                    frames.append(frame)
+                if frame is not None and Nemesis._is_child_of_async(frame):
+                    frames.append(frame.f_code.co_name)
             if frames:
                 Nemesis._start_experiment(random.choice(frames),
                                           Nemesis._select_speedup())
@@ -218,10 +221,18 @@ class Nemesis(object):
             if frame:
                 fname = frame.f_code.co_filename
         if frame:
-            return frame.f_code.co_name
+            return frame
         return None
 
     @staticmethod
+    def _is_child_of_async(frame):
+        '''Returns TRUE if an async method is a parent of FRAME.'''
+        while frame:
+            if bool(frame.f_code.co_flags & CO_COROUTINE):
+                return True
+            frame = frame.f_back
+
+    @staticmethod
     def _get_event_loops():
         '''Returns each thread's event loop, if it exists.'''
         loops = []