path: root/nemesis/nemesis.py

'''
Copyright:

   Copyright © 2025 bdunahu <bdunahu@operationnull.com>

   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 nemesis.causal_event_loop import CausalEventLoop
from collections import defaultdict
from nemesis.html_gen import plot_results
from pathlib import Path
from typing import Any, Self
import asyncio
import inspect
import os
import random
import signal
import sys
import threading
import time
import types

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: list[Path],
                 path_exclude: list[Path],
                 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.path_exclude = path_exclude
        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 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)'. Received: {type(loop)}")
            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._frame_to_name(Nemesis._get_current_frame(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.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(Nemesis._frame_to_name(frame))
            if frames:
                next_frame = random.choice(frames)
                Nemesis._start_experiment(next_frame, 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(fname: str) -> bool:
        '''Returns FALSE if filename is uninteresting to the user.'''
        # never profile the profiler.
        if 'nemesis' in fname:
            return False

        # written this way to avoid hitting max recusion depth
        f = Path(fname).resolve()
        f_par = set(f.parents)
        f_par.add(f)

        if not any(p in f_par for p in Nemesis.path_include):
            return False

        if any(p in f_par for p in Nemesis.path_exclude):
            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}()'

    @staticmethod
    def _frame_to_name(frame):
        """
        """
        fname = frame.f_code.co_filename
        func = frame.f_code.co_name
        lineno = frame.f_code.co_firstlineno
        return f"{fname}:{func}.{lineno}"