path: root/nemesis/causal_event_loop.py

'''
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:

   This event loop facilitates testing virtual speedups for specific coroutines by providing functions
   to insert delays into I/O operations and synchronous callback logic.

   A `_time_dilation' less than 1.0 will cause the event loop to slow down I/O and scheduled
   callbacks proportionately. For this, some inspiration taken from aiodebug's 'time_dilated_loop'.

   `_time_dilation' also slows down the execution of synchronous logic at the same proportion,
   by modifying the `_run_once' to insert a calculated delay at the end of each callback.

   The amount of delay to insert is dependent on how long the callback ran, and how long it spent running
   inside the target coroutine. It relies on an external profiler (Nemesis) to signal when and how long
   the target coroutine is running through the `ping_enter_coro' and `ping_exit_coro' methods.

   Finally, this event loop provides an `_update_ready' method. This allows the attached profiler to
   have access to which tasks are currently being blocked at any given point in time.

Code:
'''

import asyncio
import collections
import heapq
import selectors
import time
from asyncio.log import logger
from asyncio import Task, events
from asyncio.base_events import _format_handle

_MIN_SCHEDULED_TIMER_HANDLES = 100
_MIN_CANCELLED_TIMER_HANDLES_FRACTION = 0.5
MAXIMUM_SELECT_TIMEOUT = 24 * 3600

class CausalEventLoop(asyncio.SelectorEventLoop):

    _time_fd_registered = dict()
    _time_entered_coro = None
    _accumulated_time = 0
    _num_processed_tasks = 0
    _total_time_in_coro = 0

    _time_dilation = 1.0
    _processing = collections.deque()
    _last_objective_time = None
    _last_subjective_time = None

    def __init__(self) -> None:
        super().__init__()

        _select = self._selector.select
        def select(timeout: float):
            return _select(None if timeout == None else timeout / self._time_dilation)
        self._selector.select = select

    def set_dilation(self, dilation):
        self._time_dilation = dilation

        # reset experiment counters
        self._time_entered_coro = None
        self._num_processed_tasks = 0
        self._total_time_in_coro = 0

    def get_time_in_coro(self):
        return self.total_time_in_coro

    def get_num_processed_tasks(self):
        return self._num_processed_tasks

    def _get_time_in_coro(self):
        t = self._accumulated_time
        prev = self._time_entered_coro
        if prev:
            t += super().time() - prev
        return t

    def _update_ready(self, sampling=False):
        '''
        Polls the IO selector, schedules resulting callbacks, and schedules
        'call_later' callbacks.

        This logic was separated out of `run_once` so that the list of `ready`
        tasks may be updated more frequently than once per iteration.

        If SAMPLING is true, the timeout passed to the selector will always be 0.
        '''
        sched_count = len(self._scheduled)
        if (sched_count > _MIN_SCHEDULED_TIMER_HANDLES and
            self._timer_cancelled_count / sched_count >
            _MIN_CANCELLED_TIMER_HANDLES_FRACTION):
            # Remove delayed calls that were cancelled if their number
            # is too high
            new_scheduled = []
            for handle in self._scheduled:
                if handle._cancelled:
                    handle._scheduled = False
                else:
                    new_scheduled.append(handle)

            heapq.heapify(new_scheduled)
            self._scheduled = new_scheduled
            self._timer_cancelled_count = 0
        else:
            # Remove delayed calls that were cancelled from head of queue.
            while self._scheduled and self._scheduled[0]._cancelled:
                self._timer_cancelled_count -= 1
                handle = heapq.heappop(self._scheduled)
                handle._scheduled = False

        timeout = None
        if sampling or self._ready or self._stopping:
            timeout = 0
        elif self._scheduled:
            # Compute the desired timeout.
            timeout = self._scheduled[0]._when - self.time()
            if timeout > MAXIMUM_SELECT_TIMEOUT:
                timeout = MAXIMUM_SELECT_TIMEOUT
            elif timeout < 0:
                timeout = 0

        event_list = self._selector.select(timeout)
        self._process_events(event_list)
        # Needed to break cycles when an exception occurs.
        event_list = None

        # Handle 'later' callbacks that are ready.
        end_time = self.time() + self._clock_resolution
        while self._scheduled:
            handle = self._scheduled[0]
            if handle._when >= end_time:
                break
            handle = heapq.heappop(self._scheduled)
            handle._scheduled = False
            self._ready.append(handle)

    def _run_once(self):
        """
        Run one full iteration of the event loop.

        This calls all currently ready callbacks.
        """
        self._update_ready()
        self._processing.extend(self._ready)
        # This is the only place where callbacks are actually *called*.
        # All other places just add them to ready.
        # Note: We run all currently scheduled callbacks, but not any
        # callbacks scheduled by callbacks run this time around --
        # they will be run the next time (after another I/O poll).
        # Use an idiom that is thread-safe without using locks.
        ntodo = len(self._processing)
        for i in range(ntodo):
            handle = self._processing.popleft()
            try:
                self._ready.remove(handle)
            except ValueError:
                pass
            if handle._cancelled:
                continue
            try:
                self._current_handle = handle
                t0 = super().time()
                handle._run()
                dt = super().time() - t0
                if self._debug and dt >= self.slow_callback_duration:
                    logger.warning('Executing %s took %.3f seconds',
                                   _format_handle(handle), dt * 1/self._time_dilation)

                time_inside_coro = self._get_time_in_coro()

                # update stat counters
                self._total_time_in_coro += time_inside_coro
                self._num_processed_tasks += 1

                # calculate the amount of time to slow this callback down by. We only want
                # to add time proportionate to the amount of time spent OUTSIDE of the
                # coroutine of interest.
                time_outside_coro = dt - time_inside_coro
                delay = time_outside_coro * (1 / self._time_dilation - 1)
                t0 = super().time()
                # do it this way so the python interpreter still receives signals.
                while super().time() - t0 < delay:
                    time.sleep(0.001)

                self._time_entered_coro = super().time() if self._time_entered_coro else None
                self._accumulated_time = 0
            finally:
                self._current_handle = None
            handle = None  # Needed to break cycles when an exception occurs.

    def _process_events(self, event_list):
        end_time = super().time()
        for key, mask in event_list:
            fileobj, (reader, writer) = key.fileobj, key.data
            if mask & selectors.EVENT_READ and reader is not None:
                if reader._cancelled:
                    self._remove_reader(fileobj)
                elif (start_time := self._time_fd_registered[fileobj]) is not None:
                    self._time_fd_registered[fileobj] = None
                    dt = end_time - start_time
                    when = dt * (1 / self._time_dilation - 1) + super().time()
                    self.call_at(when, reader._callback, *reader._args, context=reader._context)
            if mask & selectors.EVENT_WRITE and writer is not None:
                if writer._cancelled:
                    self._remove_writer(fileobj)
                elif (start_time := self._time_fd_registered[fileobj]) is not None:
                    self._time_fd_registered[fileobj] = None
                    dt = end_time - start_time
                    when = dt * (1 / self._time_dilation - 1) + super().time()
                    self.call_at(when, writer._callback, *writer._args, context=reader._context)

    def _add_reader(self, fd, callback, *args):
        self._time_fd_registered[fd] = super().time()
        return super()._add_reader(fd, callback, *args)

    def _remove_reader(self, fd):
        self._time_fd_registered[fd] = None
        return super()._remove_reader(fd)

    def _add_writer(self, fd, callback, *args):
        self._time_fd_registered[fd] = super().time()
        return super()._add_writer(fd, callback, *args)

    def _remove_writer(self, fd):
        self._time_fd_registered[fd] = None
        return super()._remove_writer(fd)

    def ping_enter_coro(self):
        self._time_entered_coro = super().time()

    def ping_exit_coro(self):
        assert isinstance(self._time_entered_coro, float), f"Tried to exit coro before recorded entry!"
        self._accumulated_time += super().time() - self._time_entered_coro
        self._time_entered_coro = None

    def time(self):
        obj = super().time()
        if self._last_objective_time is None:
            self._last_objective_time = self._last_subjective_time = obj
            return obj
        else:
            sub = self._last_subjective_time + (obj - self._last_objective_time) * self._time_dilation
            self._last_subjective_time = sub
            self._last_objective_time = obj
            return sub


class CausalEventLoopPolicy(asyncio.DefaultEventLoopPolicy):
    def new_event_loop(self):
        return CausalEventLoop()


asyncio.set_event_loop_policy(CausalEventLoopPolicy())