From ba1c72cedb56512f52c48ee947a2b11fa8a90c4d Mon Sep 17 00:00:00 2001
From: bd <bdunahu@operationnull.com>
Date: Mon, 8 Sep 2025 13:25:22 -0400
Subject: Perform speedups at coroutine granularity

---
 nemesis/nemesis.py | 64 +++++++++++++++++++++++++++++++++++++++---------------
 1 file changed, 46 insertions(+), 18 deletions(-)

(limited to 'nemesis/nemesis.py')

diff --git a/nemesis/nemesis.py b/nemesis/nemesis.py
index eaf41f3..69d0c2b 100755
--- a/nemesis/nemesis.py
+++ b/nemesis/nemesis.py
@@ -25,16 +25,17 @@ Commentary:
 Code:
 '''
 
+from asyncio.base_events import _format_handle
+from collections import defaultdict
+from experiment import Experiment
 import argparse
 import asyncio
+import os
 import signal
 import sys
-import traceback
 import time
+import traceback
 import types
-import os
-from experiment import Experiment
-from asyncio.base_events import _format_handle
 
 
 class Nemesis(object):
@@ -51,20 +52,21 @@ class Nemesis(object):
     e_duration = None
     # The number of seconds remaining in this performance experiment.
     r_duration = None
+    # A mapping of event loops to the previous running coroutine.
+    prev_coro = defaultdict(lambda: None)
 
     # temp
-    task = None
+    coro = None
     dilation = 1.0
 
     @staticmethod
-    def __init__(task, speedup, e_duration, w_time, signal_interval=0.01):
-        os.environ["PYTHONASYNCIODEBUG"] = "1"
+    def __init__(coro, speedup, e_duration, w_time, signal_interval=0.01):
         Nemesis.signal_interval = signal_interval
         Nemesis.e_duration = e_duration
-        Nemesis.r_duration = w_time
+        Nemesis.r_duration = 0
         # temporary
-        Nemesis.task = task
-        Nemesis.speedup = speedup
+        Nemesis.coro = coro
+        Nemesis.speedup = max(speedup, 1.0)
 
     @staticmethod
     def start():
@@ -86,12 +88,13 @@ class Nemesis(object):
     @staticmethod
     def _start_experiment():
         Nemesis.r_duration = Nemesis.e_duration
-        Nemesis.curr_experiment = Experiment(Nemesis.task, Nemesis.speedup)
+        Nemesis.prev_coro = defaultdict(lambda: None)
+        Nemesis.curr_experiment = Experiment(Nemesis.coro, Nemesis.speedup)
 
     @staticmethod
     def _stop_experiment():
         if Nemesis.curr_experiment is not None:
-            print(f'finished running {Nemesis.curr_experiment.get_task()} with speedup {Nemesis.curr_experiment.get_speedup()}')
+            print(f'finished running {Nemesis.curr_experiment.get_coro()} with speedup {Nemesis.curr_experiment.get_speedup()}')
             Nemesis.results.append(Nemesis.curr_experiment.get_results())
             del Nemesis.curr_experiment
 
@@ -102,8 +105,17 @@ class Nemesis(object):
         Nemesis.last_sample = curr_sample
         if Nemesis.curr_experiment:
             loops = Nemesis.curr_experiment.get_loops()
-            # print(loops)
+            exp_coro = Nemesis.curr_experiment.get_coro()
             for loop in loops:
+                coro = Nemesis._get_current_coro(loop)
+                prev_coro = Nemesis.prev_coro[loop]
+                if not prev_coro == coro:
+                    if prev_coro == exp_coro:
+                        loop.ping_exit_coro()
+                    elif coro == exp_coro:
+                        loop.ping_enter_coro()
+                    Nemesis.prev_coro[loop] = coro
+
                 loop._update_ready(True)
                 handles = Nemesis._get_waiting_handles(loop)
                 Nemesis.curr_experiment.add_handles(handles, loop, passed_time)
@@ -122,11 +134,27 @@ class Nemesis(object):
                 handles.append(fmt_handle)
         return handles
 
+    def _get_current_coro(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
+        while not Nemesis._should_trace(fname):
+            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
+        return None
+
     @staticmethod
     def _should_trace(filename):
         '''Returns FALSE if filename is uninteresting to the user.
-        Don't depend on this. It's good enough for testing.'''
-        # FIXME Assume GuixSD. Makes filtering easy
+        Don't depend on this. It kind of sucks.'''
         if not filename:
             return False
         if '/gnu/store' in filename:
@@ -172,10 +200,10 @@ if __name__ == "__main__":
                         type=float,
                         default=0.01)
     parser.add_argument('-s', '--speedup',
-                        help='The amount of virtual speedup.',
+                        help='The amount of virtual speedup. Cannot go below one. Default is 2.0.',
                         metavar='',
                         type=float,
-                        default=0.5)
+                        default=2.0)
     parser.add_argument('-c', '--task',
                         help='The task to virtually speedup.',
                         metavar='',
@@ -187,7 +215,7 @@ if __name__ == "__main__":
                         type=float,
                         default=4)
     parser.add_argument('-w', '--warmup-time',
-                        help='Amount of time to wait until the first performance experiment. Default is the minimum time of 100 milliseconds',
+                        help='Amount of time to wait until the first performance experiment. Default is 0 milliseconds',
                         metavar='',
                         type=float,
                         default=0.1)
-- 
cgit v1.2.3