path: root/nemesis/nemesis.py

#!/usr/bin/env python3
'''
       _/      _/                                                _/
      _/_/    _/    _/_/    _/_/_/  _/_/      _/_/      _/_/_/        _/_/_/
     _/  _/  _/  _/_/_/_/  _/    _/    _/  _/_/_/_/  _/_/      _/  _/_/
    _/    _/_/  _/        _/    _/    _/  _/            _/_/  _/      _/_/
   _/      _/    _/_/_/  _/    _/    _/    _/_/_/  _/_/_/    _/  _/_/_/


Copyright 2025 bdunahu

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

Commentary:
Code:
'''
from typing import Any, Self
from causal_event_loop import CausalEventLoop
from collections import defaultdict
from html_gen import plot_results
import argparse
import asyncio
import os
import inspect
import random
import signal
import sys
import threading
import time
import traceback
import types
from pathlib import Path

CO_COROUTINE = inspect.CO_COROUTINE

class Experiment:

    # event loops participating in this this experiment
    _loops: list[Any] = []

    def __init__(self: Self, loops: list[Any]) -> None:
        self._loops = loops

    def get_loops(self: Self) -> list[Any]:
        return [l for l in self._loops if l.is_running()]

class Nemesis(object):

    # the name of the target program
    prog: str

    # the (ideal) interval between samples
    signal_interval: float
    # the timestamp which the last sample was taken
    last_sample: float

    # the current experiment being run
    experiment_data: Experiment
    # the coroutine the current experiment is speeding up
    experiment_coro: str
    # the speedup of the current experiment
    experiment_spdp: float
    # the total time this experiment has been running
    experiment_time: float

    # results from previous experiments. Keys represent names of coroutines.
    results = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: [])))

    # the file to write results to
    filename: str

    # The base duration of each performance experiment
    e_duration: int

    # A mapping of event loops to the previous running coroutine.
    prev_coro = defaultdict(lambda: None)

    # Path of files to profile.
    path_include: list(Path)
    # Path of files to exclude
    path_exclude: list(Path)

    @staticmethod
    def __init__(e_duration: int,
                 path_include: None,
                 path_exclude: None,
                 filename: str,
                 prog: str,
                 signal_interval: float=0.01) -> None:
        Nemesis.signal_interval = signal_interval
        Nemesis.e_duration = e_duration
        Nemesis.path_include = path_include
        Nemesis.filename = filename
        Nemesis.prog = prog

    @staticmethod
    def start() -> None:
        Nemesis.last_sample = time.perf_counter()
        signal.signal(signal.SIGALRM,
                      Nemesis._signal_handler)
        signal.setitimer(signal.ITIMER_REAL,
                         Nemesis.signal_interval,
                         Nemesis.signal_interval)

    @staticmethod
    def print_results():
        for coro_name, x_values in Nemesis.results.items():
            print(f'Results for {coro_name:}')
            for speedup, experiments in x_values.items():
                print(f'  {speedup * 100}% speedup:')
                for experiment in experiments:
                    num_callbacks = len(experiment)
                    if num_callbacks > 0:
                        total_wait = sum([cb[1] for cb in experiment])
                        latency = total_wait / num_callbacks
                        print(f'    latency: {latency}')
                        print(f'    callbacks processed: {num_callbacks}')
            print(f'')

    @staticmethod
    def stop() -> None:
        signal.setitimer(signal.ITIMER_REAL, 0)
        plot_results(Nemesis.results, Nemesis.filename, Nemesis.prog)
        print(f"Wrote {Nemesis.filename}")

    @staticmethod
    def _start_experiment(coro: str, speedup: float) -> None:
        Nemesis.prev_coro = defaultdict(lambda: None)
        Nemesis.experiment_coro = coro
        Nemesis.experiment_spdp = speedup
        Nemesis.experiment_time = 0

        loops = Nemesis._get_event_loops()
        for loop in loops:
            if not isinstance(loop, CausalEventLoop):
                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)

    @staticmethod
    def _stop_experiment() -> None:
        if Nemesis.experiment_data is not None:
            loops = Nemesis.experiment_data.get_loops()

            latency = []
            virtual_run_time = []
            for loop in loops:
                results = Nemesis.results[Nemesis.experiment_coro][Nemesis.experiment_spdp][loop._thread_id]
                results.append((loop.get_completed_coros(), loop.get_run_time()))

            print(f'Ran {Nemesis.experiment_coro} at {Nemesis.experiment_spdp} speed')
            del Nemesis.experiment_data

    @staticmethod
    def _signal_handler(sig: int, frame: types.FrameType) -> None:
        curr_sample = time.perf_counter()
        passed_time = curr_sample - Nemesis.last_sample
        Nemesis.last_sample = curr_sample

        if getattr(Nemesis, 'experiment_data', None):
            loops = Nemesis.experiment_data.get_loops()
            exp_coro = Nemesis.experiment_coro
            for loop in loops:
                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:
                        loop.ping_exit_coro()
                    elif coro == exp_coro:
                        loop.ping_enter_coro()
                    Nemesis.prev_coro[loop] = coro

                loop.collect_ready_events()

            Nemesis.experiment_time += passed_time
            if (Nemesis.e_duration <= Nemesis.experiment_time):
                Nemesis._stop_experiment()

        else:
            frames = []
            loops = Nemesis._get_event_loops()
            for loop in loops:
                frame = Nemesis._get_current_frame(loop)
                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())

    @staticmethod
    def _parse_handle(handle):
        cb = handle._callback
        if isinstance(getattr(cb, '__self__', None), asyncio.tasks.Task):
            task = cb.__self__
            coro = task.get_coro()
            return [task.get_name(), Nemesis._get_coro_name(coro)]
        else:
            return [str(type(handle).__name__), cb.__name__]

    @staticmethod
    def _get_current_frame(loop: Any) -> types.FrameType:
        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:
            return frame
        return None

    @staticmethod
    def _is_child_of_async(frame: types.FrameType) -> bool:
        '''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() -> list[Any]:
        '''Returns each thread's event loop, if it exists.'''
        loops = []
        for t in threading.enumerate():
            frame = sys._current_frames().get(t.ident)
            if frame:
                loop = Nemesis._walk_back_until_loop(frame)
                if loop and loop not in loops:
                    loops.append(loop)
        return loops

    @staticmethod
    def _walk_back_until_loop(frame: types.FrameType) -> Any:
        '''Walks back the callstack until we are in a method named '_run_once'.
        If this is ever true, we assume we are in an Asyncio event loop method,
        and check to see if the 'self' variable is indeed and instance of
        AbstractEventLoop. Return this variable if true.'''
        while frame:
            if frame.f_code.co_name == '_run_once' and 'self' in frame.f_locals:
                loop = frame.f_locals['self']
                if isinstance(loop, asyncio.AbstractEventLoop):
                    return loop
            else:
                frame = frame.f_back
        return None

    @staticmethod
    def _should_trace(filename: str) -> bool:
        '''Returns FALSE if filename is uninteresting to the user.
        Don't depend on this. It kind of sucks.'''
        if not filename:
            return False
        if '/gnu/store' in filename:
            return False
        if '/usr/local/lib/python' in filename:
            return False
        if 'site-packages' in filename:
            return False
        if 'propcache' in filename:
            return False
        if '.pyx' in filename:
            return False
        if filename[0] == '<':
            return False
        if 'nemesis' in filename:
            return False
        return True

    def _select_speedup() -> float:
        '''
        Returns a random speedup between 0% to 100%, in multiples of 5%.
        Because a baseline is needed to calculate effect on program
        performance, selects a speedup of 0 with 50% probability.
        '''
        r1 = random.randint(-19, 20)
        return max(0, r1) / 20

    @staticmethod
    def _get_coro_name(coro):
        '''
        Stolen from _format_coroutine in cpython/Lib/asyncio/coroutines.py
        '''
        # Coroutines compiled with Cython sometimes don't have
        # proper __qualname__ or __name__.  While that is a bug
        # in Cython, asyncio shouldn't crash with an AttributeError
        # in its __repr__ functions.
        if hasattr(coro, '__qualname__') and coro.__qualname__:
            coro_name = coro.__qualname__
        elif hasattr(coro, '__name__') and coro.__name__:
            coro_name = coro.__name__
        else:
            # Stop masking Cython bugs, expose them in a friendly way.
            coro_name = f'<{type(coro).__name__} without __name__>'
        return f'{coro_name}()'

the_globals = {
    '__name__': '__main__',
    '__doc__': None,
    '__package__': None,
    '__loader__': globals()['__loader__'],
    '__spec__': None,
    '__annotations__': {},
    '__builtins__': globals()['__builtins__'],
    '__file__': None,
    '__cached__': None,
}

def validate_dir(path_str: str) -> Path:
    p = Path(path_str).expanduser().resolve()
    if not p.exists():
        raise ValueError(f"Profile path does not exist: {p}")
    if not p.is_dir():
        raise ValueError(f"Can't profile a non-dir: {p}")
    return p

if __name__ == "__main__":
    # parses CLI arguments and facilitates profiler runtime.
    parser = argparse.ArgumentParser(
        usage='%(prog)s [args] -- prog'
    )

    parser.add_argument('-i', '--interval',
                        help='The minimum amount of time inbetween \
                        samples in seconds.',
                        metavar='',
                        type=float,
                        default=0.01)
    parser.add_argument('-e', '--experiment-duration',
                        help='The performance experiment duration. Defaults to 3 seconds.',
                        metavar='',
                        type=float,
                        default=3)
    parser.add_argument('-f', '--filename',
                        help='The filename to write results to.',
                        metavar='',
                        type=str,
                        default="results.html")
    # parser.add_argument('--include-paths',
    #                     help='Specify the path(s) containing files to profile. If a file is in this path, it is a candidate for optimization.',
    #                     nargs="+",
    #                     type=validate_dir,
    #                     required=True)
    # parser.add_argument('--exclude-paths',
    #                     help='Specify the path(s) containing files to exclude profile. Takes priority over --include-paths.',
    #                     nargs="*",
    #                     type=validate_dir,
    #                     required=False)
    parser.add_argument('prog',
                        type=str,
                        nargs='*',
                        help='Path to the python script and its arguments.')
    args = parser.parse_args()

    sys.argv = args.prog
    try:
        with open(args.prog[0], 'r', encoding='utf-8') as fp:
            code = compile(fp.read(), args.prog[0], "exec")
            Nemesis(args.experiment_duration,
                    # args.include_paths,
                    # args.exclude_paths,
                    None, None,
                    args.filename,
                    args.prog[0],
                    args.interval).start()
            exec(code, the_globals)
            Nemesis.stop()
    except Exception:
        traceback.print_exc()