""" aiosched: Yet another Python async scheduler Copyright (C) 2022 nocturn9x 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 https: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. """ import signal import itertools from collections import deque from functools import partial from aiosched.task import Task, TaskState from timeit import default_timer from aiosched.internals.queue import TimeQueue from aiosched.util.debugging import BaseDebugger from typing import Callable, Any, Coroutine from aiosched.errors import ( InternalError, ResourceBusy, Cancelled, ResourceClosed, ResourceBroken, ) from aiosched.context import TaskContext from selectors import DefaultSelector, BaseSelector, EVENT_READ, EVENT_WRITE class FIFOKernel: """ An asynchronous event loop implementation with a FIFO scheduling policy. :param clock: The function used to keep track of time. Defaults to timeit.default_timer :param debugger: A subclass of aiosched.util.BaseDebugger or None if no debugging output is desired :type debugger: :class: aiosched.util.debugging.BaseDebugger, optional :param selector: The selector to use for I/O multiplexing, defaults to selectors.DefaultSelector :type selector: :class: selectors.DefaultSelector """ def __init__( self, clock: Callable[[], float] = default_timer, debugger: BaseDebugger | None = None, selector: BaseSelector = DefaultSelector(), ): """ Public constructor """ self.clock = clock if debugger and not issubclass(type(debugger), BaseDebugger): raise InternalError( "The debugger must be a subclass of aiosched.util.debugging.BaseDebugger" ) # The debugger object. If it is none we create a dummy object that immediately returns an empty # lambda which in turn returns None every time we access any of its attributes to avoid lots of # if self.debugger clauses self.debugger = ( debugger or type( "DumbDebugger", (object,), {"__getattr__": lambda *_: lambda *_: None}, )() ) # Abstraction layer over low-level OS # primitives for asynchronous I/O self.selector: BaseSelector = selector # Tasks that are ready to run self.run_ready: deque[Task] = deque() # Tasks that are paused and waiting # for some deadline to expire self.paused: TimeQueue = TimeQueue(self.clock) # Data that is to be sent back to coroutines self.data: dict[Task, Any] = {} # The currently running task self.current_task: Task | None = None # The loop's entry point self.entry_point: Task | None = None # Did we receive a Ctrl+C? self._sigint_handled: bool = False # Are we executing any task code? self._running: bool = False def __repr__(self): """ Returns repr(self) """ fields = { "debugger", "run_ready", "selector", "clock", "data", "paused", "current_task", } data = ", ".join( name + "=" + str(value) for name, value in zip(fields, (getattr(self, field) for field in fields)) ) return f"{type(self).__name__}({data})" def _sigint_handler(self, *_args): """ Handles SIGINT :return: """ self._sigint_handled = True def done(self) -> bool: """ Returns whether the loop has no more work to do """ if self.current_task and not self.current_task.done(): # Current task isn't done yet! return False if any([self.paused, self.run_ready]): # There's tasks sleeping and/or on the # ready queue! return False if self.get_active_io_count(): # We don't just do any([self.paused, self.run_ready, self.selector.get_map()]) # because we don't want to just know if there's any resources we're waiting on, # but if there's at least one non-terminated task that owns a resource we're # waiting on. This avoids issues such as the event loop never exiting if the # user forgets to close a socket, for example return False return True def close(self, force: bool = False): """ Closes the event loop. If force equals False, which is the default, raises an InternalError exception. If force equals True, cancels all tasks """ if not self.done() and not force: self.current_task.throw( InternalError("cannot shut down a running event loop") ) for task in self.all(): self.cancel(task) def all(self) -> Task: """ Yields all the tasks the event loop is keeping track of """ for task in itertools.chain(self.run_ready, self.paused): task: Task yield task def shutdown(self): """ Shuts down the event loop """ for task in self.all(): self.io_release_task(task) self.selector.close() self.close() def wait_io(self): """ Waits for I/O and schedules tasks when their associated resource is ready to be used """ self._running = False before_time = self.clock() # Used for the debugger timeout = 0.0 if self.run_ready: # If there is work to do immediately (tasks to run) we # can't wait. # TODO: This could cause I/O starvation in highly concurrent # environments: maybe a more convoluted scheduling strategy # where I/O timeouts can only be skipped n times before a # mandatory x-second timeout occurs is needed? It should of # course take deadlines into account so that timeouts are # always delivered in a timely manner and tasks awake from # sleeping at the right moment timeout = 0.0 elif self.paused: # If there are asleep tasks or deadlines, wait until the closest date timeout = self.paused.get_closest_deadline() - self.clock() self.debugger.before_io(timeout) # Get sockets that are ready and schedule their tasks for key, _ in self.selector.select(timeout): key.data: dict[int, Task] for task in key.data.values(): # We don't reschedule a task that wasn't # blocking on I/O before: this way if a # task waits on a socket and then goes to # sleep, it won't be woken up early if the # resource becomes available before its # deadline expires self.run_ready.append(task) # Resource ready? Schedule its task self.debugger.after_io(self.clock() - before_time) def awake_tasks(self): """ Reschedules paused tasks if their deadline has elapsed """ self._running = False while self.paused and self.paused.get_closest_deadline() <= self.clock(): # Reschedules tasks when their deadline has elapsed task, _ = self.paused.get() slept = self.clock() - task.paused_when self.run_ready.append(task) task.paused_when = 0 task.next_deadline = 0 self.debugger.after_sleep(task, slept) def reschedule_running(self): """ Reschedules the currently running task """ self.run_ready.append(self.current_task) def schedule(self, task: Task): """ Schedules a task that was previously suspended """ self.run_ready.append(task) self.reschedule_running() def suspend(self): """ Suspends execution of the current task. This is basically a do-nothing method, since it will not reschedule the task before returning. The task will stay suspended as long as something else outside the loop reschedules it (possibly forever) """ self.current_task.state = TaskState.PAUSED def run_task_step(self): """ Runs a single step for the current task. A step ends when the task awaits any of our primitives or async methods. Note that this method does NOT catch any errors arising from tasks, nor does it take StopIteration or Cancelled exceptions into account """ # Sets the currently running task self.current_task = self.run_ready.popleft() while self.current_task.done(): # We make sure not to schedule # any terminated tasks. Might want # to eventually get rid of this code, # but for now it does the job if not self.run_ready: # We'll let run() handle the I/O # or the shutdown if necessary, as # there are no more runnable tasks return self.current_task = self.run_ready.popleft() # We nullify the exception object just in case the # entry point raised and caught an error so that # self.start() doesn't raise it again at the end self.current_task.exc = None self.debugger.before_task_step(self.current_task) self._running = True # Some debugging and internal chatter here self.current_task.state = TaskState.RUN self.current_task.steps += 1 if self._sigint_handled: self._sigint_handled = False self.current_task.throw(KeyboardInterrupt()) elif self.current_task.pending_cancellation: # We perform the deferred cancellation # if it was previously scheduled self.cancel(self.current_task) else: # Run a single step with the calculation (i.e. until a yield # somewhere) method, args, kwargs = self.current_task.run( self.data.pop(self.current_task, None) ) if not hasattr(self, method) or not callable(getattr(self, method)): # This if block is meant to be triggered by other async # libraries, which most likely have different method names and behaviors # compared to us. If you get this exception, and you're 100% sure you're # not mixing async primitives from other libraries, then it's a bug! self.current_task.throw( InternalError( "Uh oh! Something bad just happened: did you try to mix primitives from other async libraries?" ) ) # Sneaky method call, thanks to David Beazley for this ;) getattr(self, method)(*args, **kwargs) self.debugger.after_task_step(self.current_task) def run(self): """ The event loop's runner function. This method drives execution for the entire framework and orchestrates I/O, events, sleeping, cancellations and deadlines, but the actual functionality for all of that is implemented in object wrappers. This keeps the size of this module to a minimum while allowing anyone to replace it with their own, as long as the system calls required by higher-level object wrappers are implemented. If you want to add features to the library, don't add them here, but take inspiration from the current API (i.e. not depending on any implementation detail from the loop aside from system calls) """ while True: if self.done(): # If we're done, which means there are # both no paused tasks and no running tasks, we # simply tear us down and return to self.start self.shutdown() break elif self._sigint_handled: # We got Ctrl+C-ed while not running a task! Pick a # random task and blow it up with a KeyboardInterrupt # exception: our existing error handling machinery will # deal with it accordingly task: Task | None = None if self.selector.get_map(): # Pretty convoluted, huh? Sorry, but I wanted this on one line ;) task = next(iter(next(iter(self.selector.get_map().values())).data.values())) elif self.paused: # TODO task = None self.run_ready.append(task) self.handle_errors(self.run_task_step) elif not self.run_ready: # If there are no actively running tasks, we start by # checking for I/O. This method will wait for I/O until # the closest deadline to avoid starving sleeping tasks # or missing deadlines if self.selector.get_map(): self.wait_io() if self.paused: # Next we check for deadlines self.awake_tasks() else: # Otherwise, while there are tasks ready to run, we run them! self.handle_errors(self.run_task_step) def start( self, func: Callable[..., Coroutine[Any, Any, Any]], *args, **kwargs ) -> Any: """ Starts the event loop from a synchronous context """ signal.signal(signal.SIGINT, self._sigint_handler) self.entry_point = Task(func.__name__ or str(func), func(*args, **kwargs)) self.run_ready.append(self.entry_point) self.debugger.on_start() try: self.run() finally: self.debugger.on_exit() if ( self.entry_point.exc and self.entry_point.context is None and self.entry_point.propagate ): # Contexts already manage exceptions for us, # no need to raise it manually. If a context # is not used, *then* we can raise the error raise self.entry_point.exc return self.entry_point.result def io_release(self, resource): """ Releases the given resource from our selector :param resource: The resource to be released """ if resource in self.selector.get_map(): self.selector.unregister(resource) self.debugger.on_io_unschedule(resource) if resource is self.current_task.last_io[1]: self.current_task.last_io = () self.reschedule_running() def io_release_task(self, task: Task): """ Calls self.io_release in a loop for each I/O resource the given task owns """ for key in filter( lambda k: task in k.data.values(), dict(self.selector.get_map()).values() ): self.notify_closing(key.fileobj, broken=True) self.selector.unregister(key.fileobj) task.last_io = () def get_active_io_count(self) -> int: """ Returns the number of streams that are currently being used by any active task """ result = 0 for key in (self.selector.get_map() or {}).values(): key.data: dict[int, Task] for task in key.data.values(): if task.done(): continue result += 1 return result def notify_closing(self, stream, broken: bool = False): """ Notifies paused tasks that a stream is about to be closed. The stream itself is not touched and must be closed by the caller """ if not broken: exc = ResourceClosed("stream has been closed") else: exc = ResourceBroken("stream might be corrupted") for k in filter( lambda o: o.fileobj == stream, dict(self.selector.get_map()).values(), ): for task in k.data.values(): if task is not self.current_task: # We don't want to raise an error inside # the task that's trying to close the stream! self.handle_errors(partial(k.data.throw, exc), k.data) self.reschedule_running() def cancel(self, task: Task): """ Attempts to cancel the given task or schedules cancellation for later if it fails """ self.handle_errors(partial(task.throw, Cancelled(task)), task) if task.state != TaskState.CANCELLED: task.pending_cancellation = True self.io_release_task(task) self.paused.discard(task) self.reschedule_running() def handle_errors(self, func: Callable, task: Task | None = None): """ Convenience method for handling various exceptions from tasks """ try: func() except StopIteration as ret: # We re-define it because we call run_task_step # with this method and that changes the current # task task = task or self.current_task # At the end of the day, coroutines are generator functions with # some tricky behaviors, and this is one of them. When a coroutine # hits a return statement (either explicit or implicit), it raises # a StopIteration exception, which has an attribute named value that # represents the return value of the coroutine, if it has one. Of course # this exception is not an error, and we should happily keep going after it: # most of this code below is just useful for internal/debugging purposes task.state = TaskState.FINISHED task.result = ret.value self.wait(task) except Cancelled: # When a task needs to be cancelled, aiosched tries to do it gracefully # first: if the task is paused in either I/O or sleeping, that's perfect. # But we also need to cancel a task if it was not sleeping or waiting on # any I/O because it could never do so (therefore blocking everything # forever). So, when cancellation can't be done right away, we schedule # it for the next execution step of the task. aiosched will also make sure # to re-raise cancellations at every checkpoint until the task lets the # exception propagate into us, because we *really* want the task to be # cancelled task = task or self.current_task task.state = TaskState.CANCELLED task.pending_cancellation = False self.debugger.after_cancel(task) self.wait(task) except Exception as err: # Any other exception is caught here task = task or self.current_task task.exc = err task.state = TaskState.CRASHED self.debugger.on_exception_raised(task, err) self.wait(task) def sleep(self, seconds: int | float): """ Puts the current task to sleep for a given amount of seconds """ if seconds: self.debugger.before_sleep(self.current_task, seconds) self.paused.put(self.current_task, seconds) else: # When we're called with a timeout of 0, this method acts as a checkpoint # that allows aiosched to kick in and to its job without pausing the task's # execution for too long. It is recommended to put a couple of checkpoints # like these in your code if you see degraded concurrent performance in parts # of your code that block the loop self.reschedule_running() def wait(self, task: Task): """ Makes the current task wait for completion of the given one by only rescheduling it once the given task has finished executing """ if task.done(): self.paused.discard(task) self.io_release_task(task) self.run_ready.extend(task.joiners) def join(self, task: Task): """ Tells the event loop that the current task wants to wait on the given one, but without actually waiting for its completion """ task.joiners.add(self.current_task) self.reschedule_running() def spawn(self, func: Callable[..., Coroutine[Any, Any, Any]], *args, **kwargs): """ Spawns a task from a coroutine function. All positional and keyword arguments besides the coroutine function itself are passed to the newly created coroutine """ task = Task(func.__name__ or repr(func), func(*args, **kwargs)) self.data[self.current_task] = task self.run_ready.append(task) self.reschedule_running() self.debugger.on_task_spawn(task) def set_context(self, ctx: TaskContext): """ Sets the current task context. This is implemented as simply wrapping the current task inside the context and replacing the Task object with the TaskContext one. This may also wrap another task context into a new one, but the loop doesn't need to care about that: the API is designed exactly for this """ ctx.entry_point = self.current_task ctx.tasks.append(ctx.entry_point) self.current_task.context = ctx self.current_task = ctx self.debugger.on_context_creation(ctx) self.reschedule_running() def close_context(self, ctx: TaskContext): """ Closes the given context """ self.debugger.on_context_exit(ctx) task = ctx.entry_point task.context = None self.current_task = task self.reschedule_running() def get_current_task(self): """ Returns the current task to an asynchronous caller """ self.data[self.current_task] = self.current_task self.reschedule_running() def perform_io(self, resource, evt_type: int): """ Registers the given resource inside our selector to perform I/O multiplexing :param resource: The resource on which a read or write operation has to be performed :param evt_type: The type of event to perform on the given socket, either selectors.EVENT_READ or selectors.EVENT_WRITE :type evt_type: int """ self.current_task.state = TaskState.IO if self.current_task.last_io: # Since most of the time tasks will perform multiple # I/O operations on a given resource, unregistering them # every time isn't a sensible approach. A quick and # easy optimization to address this problem is to # store the last I/O operation that the task performed, # together with the resource itself, inside the task # object. If the task then tries to perform the same # operation on the same resource again, this method then # returns immediately as the resource is already being watched # by the selector. If the resource is the same, but the # event type has changed, then we modify the resource's # associated event. Only if the resource is different from # the last one used then this method will register a new # one if self.current_task.last_io == (evt_type, resource): # Selector is already listening for that event on # this resource return elif self.current_task.last_io[1] == resource: # If the event to listen for has changed we just modify it self.selector.modify(resource, evt_type, {evt_type: self.current_task}) self.current_task.last_io = (evt_type, resource) self.debugger.on_io_schedule(resource, evt_type) elif not self.current_task.last_io or self.current_task.last_io[1] != resource: # The task has either registered a new resource or is doing # I/O for the first time self.current_task.last_io = evt_type, resource try: self.selector.register(resource, evt_type, {evt_type: self.current_task}) self.debugger.on_io_schedule(resource, evt_type) except KeyError: # The stream is already being used key = self.selector.get_key(resource) if key.data[key.events] == self.current_task: # If the task that registered the stream # changed their mind on what they want # to do with it, who are we to deny their # request? self.selector.modify(resource, key.events | evt_type, {EVENT_READ: self.current_task, EVENT_WRITE: self.current_task}) self.debugger.on_io_schedule(resource, evt_type) elif key.events != evt_type: # We also modify the event in # our selector so that one task can read # off a given stream while another one is # writing to it self.selector.modify(resource, key.events | evt_type, {evt_type: self.current_task, key.events: list(key.data.values())[0]}) else: # One task reading and one writing on the same # resource is fine (think producer-consumer), # but having two tasks reading/writing at the # same time can't lead to anything good, better # disallow it self.current_task.throw(ResourceBusy(f"The resource is being read from/written by another task"))