226 lines
6.2 KiB
Python
226 lines
6.2 KiB
Python
"""
|
|
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 inspect
|
|
from types import coroutine
|
|
from typing import Any, Callable, Coroutine
|
|
from aiosched.task import Task, TaskState
|
|
from aiosched.errors import SchedulerError
|
|
from selectors import EVENT_READ, EVENT_WRITE
|
|
|
|
|
|
@coroutine
|
|
def syscall(method: str, *args, **kwargs) -> Any | None:
|
|
"""
|
|
Lowest-level primitive to interact with the event loop:
|
|
calls a loop method with the provided arguments. This
|
|
function should not be used directly, but through abstraction
|
|
layers. All positional and keyword arguments are passed to
|
|
the method itself and its return value is provided once the
|
|
loop yields control back to us
|
|
|
|
:param method: The loop method to call
|
|
:type method: str
|
|
:returns: The result of the method call, if any
|
|
"""
|
|
|
|
result = yield method, args, kwargs
|
|
return result
|
|
|
|
|
|
async def schedule(task: Task):
|
|
"""
|
|
Reschedules a task that had been
|
|
previously suspended
|
|
"""
|
|
|
|
await syscall("schedule", task)
|
|
|
|
|
|
async def spawn(func: Callable[..., Coroutine[Any, Any, Any]], *args, **kwargs) -> Task:
|
|
"""
|
|
Spawns a task from a coroutine and returns it. The coroutine
|
|
is put on the running queue and is executed as soon as possible.
|
|
Any positional and keyword arguments are passed along to the coroutine
|
|
|
|
:param func: The coroutine function to instantiate. Note that this should NOT
|
|
be a coroutine object: the arguments to the coroutine should be passed to
|
|
spawn as well
|
|
:return:
|
|
"""
|
|
|
|
if inspect.iscoroutine(func):
|
|
raise TypeError(
|
|
"Looks like you tried to call spawn(your_func(arg1, arg2, ...)), that is wrong!"
|
|
"\nWhat you wanna do, instead, is this: spawn(your_func, arg1, arg2, ...)"
|
|
)
|
|
elif inspect.iscoroutinefunction(func):
|
|
return await syscall("spawn", func, *args, **kwargs)
|
|
else:
|
|
raise TypeError("func must be a coroutine function")
|
|
|
|
|
|
async def sleep(delay: int | float):
|
|
"""
|
|
Puts the calling task to sleep for the
|
|
given amount of time. If the delay is equal
|
|
to zero, this call acts as a checkpoint to
|
|
perform task switching and is useful to release
|
|
pressure from the scheduler in highly concurrent
|
|
environments
|
|
|
|
:param delay: The amount of time (in seconds) that
|
|
the task has to be put to sleep for. Must be
|
|
greater than zero
|
|
:type delay: int | float
|
|
"""
|
|
|
|
await syscall("sleep", delay)
|
|
|
|
|
|
async def checkpoint():
|
|
"""
|
|
Shorthand for sleep(0)
|
|
"""
|
|
|
|
await sleep(0)
|
|
|
|
|
|
async def suspend():
|
|
"""
|
|
Suspends the calling task indefinitely.
|
|
The task can be unsuspended by a timer,
|
|
an event or an incoming I/O operation
|
|
"""
|
|
|
|
await syscall("suspend")
|
|
|
|
|
|
async def current_task() -> Task:
|
|
"""
|
|
Returns the currently running
|
|
task object
|
|
"""
|
|
|
|
return await syscall("get_current_task")
|
|
|
|
|
|
async def join(task: Task):
|
|
"""
|
|
Tells the event loop that the current task
|
|
wants to wait on the given one, but without
|
|
waiting for its completion. This is a low
|
|
level trap and should not be used on its
|
|
own
|
|
"""
|
|
|
|
await syscall("join", task)
|
|
|
|
|
|
async def wait(task: Task) -> Any | None:
|
|
"""
|
|
Waits for the completion of a
|
|
given task and returns its
|
|
return value. Can be called
|
|
multiple times by multiple tasks.
|
|
Returns immediately if the task has
|
|
completed already, but exceptions are
|
|
propagated only once. Returns the task's
|
|
return value, if it has one (returned once
|
|
for each call).
|
|
|
|
:param task: The task to wait for
|
|
:type task: :class: Task
|
|
:returns: The task's return value, if any
|
|
"""
|
|
|
|
if task == await current_task():
|
|
# We don't do an "x is y" check because
|
|
# tasks and task contexts can compare equal
|
|
# despite having different memory addresses
|
|
raise SchedulerError("a task cannot join itself")
|
|
await syscall("wait", task)
|
|
if task.exc and task.state != TaskState.CANCELLED and task.propagate:
|
|
# The task raised an error that wasn't directly caused by a cancellation:
|
|
# raise it, but do so only the first time wait was called
|
|
task.propagate = False
|
|
raise task.exc
|
|
return task.result
|
|
|
|
|
|
async def cancel(task: Task, block: bool = False):
|
|
"""
|
|
Cancels the given task. Note that
|
|
cancellations may not happen immediately
|
|
if the task is blocked in an uninterruptible
|
|
state. If block equals False, the default,
|
|
this function returns immediately, otherwise
|
|
it waits for the task to receive the cancellation
|
|
|
|
:param task: The task to wait for
|
|
:type task: :class: Task
|
|
:param block: Whether to wait for the task to be
|
|
actually cancelled or not, defaults to False
|
|
:type block: bool, optional
|
|
"""
|
|
|
|
await syscall("cancel", task)
|
|
if block:
|
|
await wait(task)
|
|
if not task.state == TaskState.CANCELLED:
|
|
raise SchedulerError(f"task {task.name!r} ignored cancellation")
|
|
|
|
|
|
async def closing(stream):
|
|
"""
|
|
Notifies the event loop that the
|
|
given stream is about to be closed,
|
|
causing all callers waiting on it
|
|
to error out with an exception instead
|
|
of blocking forever
|
|
"""
|
|
|
|
await syscall("notify_closing", stream)
|
|
|
|
|
|
async def wait_readable(stream):
|
|
"""
|
|
Waits until the given stream is
|
|
readable
|
|
"""
|
|
|
|
await syscall("perform_io", stream, EVENT_READ)
|
|
|
|
|
|
async def wait_writable(stream):
|
|
"""
|
|
Waits until the given stream is
|
|
writable
|
|
"""
|
|
|
|
await syscall("perform_io", stream, EVENT_WRITE)
|
|
|
|
|
|
async def io_release(stream):
|
|
"""
|
|
Signals to the event loop to
|
|
release a given I/O resource
|
|
"""
|
|
|
|
await syscall("io_release", stream)
|