''' 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 _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(timeout / self._time_dilation) self._selector.select = select def _set_dilation(self, dilation): self._time_dilation = dilation self._time_entered_coro = None def _get_time_in_coro(self): t = self._accumulated_time if self._time_entered_coro: t += super().time() - self._time_entered_coro 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) # 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 - self._get_time_in_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())