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:

Code:
'''

import asyncio
import collections
from sortedcontainers import SortedList
import heapq
import selectors
import time
import traceback
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 TimeAwareMixin:

    # the timestamp this callback was registered
    register_time = None
    # the timestamp this callback completed i/o
    io_time = None
    # the timestamp this callback was called by the event loop
    process_start_time = None

    def __init__(self):
        self.register_time = time.monotonic()


orig_handle = asyncio.events.Handle


def create_subclass(base_class):
    class NewSubclass(base_class, TimeAwareMixin):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            TimeAwareMixin.__init__(self)

    return NewSubclass


# make all the subclasses inherit from TimeAwareHandle as well
for sc in [orig_handle] + orig_handle.__subclasses__():
    subclass = create_subclass(sc)
    setattr(asyncio.events, sc.__name__, subclass)

class CausalEventLoop(asyncio.SelectorEventLoop):

    # a value between 0 and 1. 0 means no optimization,
    # 1 means the target coroutine is optimized away entirely
    _speedup = 1.0
    # a list of callbacks which have recently completed
    _pause_buffer = []
    # a list of intervals in which the target coroutine has been active
    _coro_intervals = SortedList()
    # a list of completed callbacks, and their associated queue time
    _completed_coros = []
    # the last time we entered the target coro
    _time_entered_coro = None

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

    def set_speedup(self, speedup):
        self._speedup = speedup

        # reset experiment counters
        self._time_entered_coro = None
        self._coro_intervals.clear()
        self._completed_coros.clear()

    def get_completed_coros(self):
        return self._completed_coros

    def get_pause_time(self):
        if not self._coro_intervals:
            return 0

        start_interval = self._coro_intervals[0][0]
        end_interval = self.time() if self._time_entered_coro else self._coro_intervals[-1][1]
        interval = (start_interval, end_interval)
        return self._get_pause_time(interval)

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

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

    def update_ready(self, can_stall=True):
        '''
        Polls the IO selector, schedules resulting callbacks, and schedules
        'call_later' callbacks.

        This function can be called in the middle of an event loop iteration.

        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.
        '''
        curr_time = self.time()
        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
        # TODO this needs to be rewritten
        # We can't miss things placed in timeout either
        # if not can_stall 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

            time_interval = (handle.register_time, curr_time)
            time_to_buffer = curr_time + self._get_pause_time(time_interval)
            handle.io_time = time_to_buffer
            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()

        current_time = self.time()
        to_process = collections.deque([
            handle for handle in self._ready
            if handle.io_time < (current_time + self._clock_resolution)
        ])

        # 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(to_process)
        for i in range(ntodo):
            handle = to_process.popleft()
            self._ready.remove(handle)
            if handle._cancelled:
                continue
            try:
                self._current_handle = handle

                process_start_time = self.time()
                handle.process_start_time = process_start_time

                handle._run()

                process_end_time = self.time()
                dt = process_end_time - process_start_time
                if self._debug and dt >= self.slow_callback_duration:
                    logger.warning('Executing %s took %.3f seconds',
                                   _format_handle(handle), dt)

                time_interval = (handle.io_time, process_start_time)
                pause_time = self._get_pause_time(time_interval)
                adjusted_start_time = handle.process_start_time - \
                    pause_time
                wait_time = adjusted_start_time - handle.io_time
                assert wait_time >= -0.0001, f"wait time on {_format_handle(handle)} was found to be {wait_time:.4f}!"
                self._completed_coros.append((_format_handle(handle), wait_time))
            except Exception:
                traceback.print_exc()
            finally:
                self._current_handle = None
            handle = None  # Needed to break cycles when an exception occurs.

    def _process_events(self, event_list):
        curr_time = self.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)
                else:
                    time_interval = (reader.register_time, curr_time)
                    time_to_buffer = curr_time + \
                        self._get_pause_time(time_interval)
                    reader.io_time = time_to_buffer
                    self._add_callback(reader)
            if mask & selectors.EVENT_WRITE and writer is not None:
                if writer._cancelled:
                    self._remove_writer(fileobj)
                else:
                    time_interval = (writer.register_time, curr_time)
                    time_to_buffer = curr_time + \
                        self._get_pause_time(time_interval)
                    writer.io_time = time_to_buffer
                    self._add_callback(writer)

    def _call_soon(self, callback, args, context):
        handle = events.Handle(callback, args, self, context)
        if handle._source_traceback:
            del handle._source_traceback[-1]
        handle.io_time = self.time()
        self._ready.append(handle)
        return handle

    def call_soon_threadsafe(self, callback, *args, context=None):
        """Like call_soon(), but thread-safe."""
        self._check_closed()
        if self._debug:
            self._check_callback(callback, 'call_soon_threadsafe')
        handle = events._ThreadSafeHandle(callback, args, self, context)
        handle.io_time = self.time()
        self._ready.append(handle)
        if handle._source_traceback:
            del handle._source_traceback[-1]
        if handle._source_traceback:
            del handle._source_traceback[-1]
        self._write_to_self()
        return handle

    def _get_pause_time(self, cb_interval):
        time = 0
        start, end = cb_interval

        for coro_start, coro_end in self._coro_intervals:
            if start < coro_end and coro_start < end:
                time += self._get_overlap(start, end, coro_start, coro_end)
            # coro_intervals are sorted, so by this time all overlap has passed
            if end < coro_start:
                break

        curr_time = self.time()
        if self._time_entered_coro and \
           start < curr_time and self._time_entered_coro < end:
            time += self._get_overlap(start, end, self._time_entered_coro,
                                      curr_time)

        return time * self._speedup

    def _get_overlap(self, a_start, a_end, b_start, b_end):
        overlap_start = max(a_start, b_start)
        overlap_end = min(a_end, b_end)
        return overlap_end - overlap_start


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


asyncio.set_event_loop_policy(CausalEventLoopPolicy())