mirror of
https://github.com/trezor/trezor-firmware.git
synced 2024-12-12 17:38:13 +00:00
571 lines
18 KiB
Python
571 lines
18 KiB
Python
"""
|
|
Implements an event loop with cooperative multitasking and async I/O. Tasks in
|
|
the form of python coroutines (either plain generators or `async` functions) are
|
|
stepped through until completion, and can get asynchronously blocked by
|
|
`yield`ing or `await`ing a syscall.
|
|
|
|
See `schedule`, `run`, and syscalls `sleep`, `wait`, `signal` and `race`.
|
|
"""
|
|
|
|
import utime
|
|
import utimeq
|
|
from typing import TYPE_CHECKING
|
|
|
|
from trezor import io, log
|
|
|
|
if TYPE_CHECKING:
|
|
from typing import Any, Awaitable, Callable, Coroutine, Generator, Union
|
|
|
|
Task = Union[Coroutine, Generator, "wait"]
|
|
AwaitableTask = Task | Awaitable
|
|
Finalizer = Callable[[Task, Any], None]
|
|
|
|
# tasks scheduled for execution in the future
|
|
_queue = utimeq.utimeq(64)
|
|
|
|
# tasks paused on I/O
|
|
_paused: dict[int, set[Task]] = {}
|
|
|
|
# functions to execute after a task is finished
|
|
_finalizers: dict[int, Finalizer] = {}
|
|
|
|
# reference to the task that is currently executing
|
|
this_task: Task | None = None
|
|
|
|
|
|
class TaskClosed(Exception):
|
|
pass
|
|
|
|
|
|
TASK_CLOSED = TaskClosed()
|
|
|
|
|
|
def schedule(
|
|
task: Task,
|
|
value: Any = None,
|
|
deadline: int | None = None,
|
|
finalizer: Finalizer | None = None,
|
|
reschedule: bool = False,
|
|
) -> None:
|
|
"""
|
|
Schedule task to be executed with `value` on given `deadline` (in
|
|
milliseconds). Does not start the event loop itself, see `run`.
|
|
Usually done in very low-level cases, see `race` for more user-friendly
|
|
and correct concept.
|
|
|
|
If `reschedule` is set, updates an existing entry.
|
|
"""
|
|
if reschedule:
|
|
_queue.discard(task)
|
|
if deadline is None:
|
|
deadline = utime.ticks_ms()
|
|
if finalizer is not None:
|
|
_finalizers[id(task)] = finalizer
|
|
_queue.push(deadline, task, value)
|
|
|
|
|
|
def pause(task: Task, iface: int) -> None:
|
|
"""
|
|
Block task on given message interface. Task is resumed when the interface
|
|
is activated. It is most probably wrong to call `pause` from user code,
|
|
see the `wait` syscall for the correct concept.
|
|
"""
|
|
tasks = _paused.get(iface, None)
|
|
if tasks is None:
|
|
tasks = _paused[iface] = set()
|
|
tasks.add(task)
|
|
|
|
|
|
def finalize(task: Task, value: Any) -> None:
|
|
"""Call and remove any finalization callbacks registered for given task."""
|
|
fn = _finalizers.pop(id(task), None)
|
|
if fn is not None:
|
|
fn(task, value)
|
|
|
|
|
|
def close(task: Task) -> None:
|
|
"""
|
|
Unschedule and unblock a task, close it so it can release all resources, and
|
|
call its finalizer.
|
|
"""
|
|
for iface in _paused: # pylint: disable=consider-using-dict-items
|
|
_paused[iface].discard(task)
|
|
for iface in _paused: # pylint: disable=consider-using-dict-items
|
|
if not _paused[iface]:
|
|
del _paused[iface]
|
|
_queue.discard(task)
|
|
task.close()
|
|
finalize(task, GeneratorExit())
|
|
|
|
|
|
def run() -> None:
|
|
"""
|
|
Loop forever, stepping through scheduled tasks and awaiting I/O events
|
|
in between. Use `schedule` first to add a coroutine to the task queue.
|
|
Tasks yield back to the scheduler on any I/O, usually by calling `await` on
|
|
a `Syscall`.
|
|
"""
|
|
task_entry = [0, 0, 0] # deadline, task, value
|
|
msg_entry = [0, 0] # iface | flags, value
|
|
while _queue or _paused:
|
|
# compute the maximum amount of time we can wait for a message
|
|
if _queue:
|
|
delay = utime.ticks_diff(_queue.peektime(), utime.ticks_ms())
|
|
else:
|
|
delay = 1000 # wait for 1 sec maximum if queue is empty
|
|
|
|
if io.poll(_paused, msg_entry, delay):
|
|
# message received, run tasks paused on the interface
|
|
msg_tasks = _paused.pop(msg_entry[0], ())
|
|
for task in msg_tasks:
|
|
_step(task, msg_entry[1])
|
|
else:
|
|
# timeout occurred, run the first scheduled task
|
|
if _queue:
|
|
_queue.pop(task_entry)
|
|
_step(task_entry[1], task_entry[2]) # type: ignore [Argument of type "int" cannot be assigned to parameter "task" of type "Task" in function "_step"]
|
|
# error: Argument 1 to "_step" has incompatible type "int"; expected "Coroutine[Any, Any, Any]"
|
|
# rationale: We use untyped lists here, because that is what the C API supports.
|
|
|
|
|
|
def clear() -> None:
|
|
"""Clear all queue state. Any scheduled or paused tasks will be forgotten."""
|
|
_ = [0, 0, 0]
|
|
while _queue:
|
|
_queue.pop(_)
|
|
_paused.clear()
|
|
_finalizers.clear()
|
|
|
|
|
|
def _step(task: Task, value: Any) -> None:
|
|
"""
|
|
Step through the task by sending value to it. This can result in either:
|
|
1. The task raises an exception:
|
|
a) StopIteration
|
|
- The Task is completed and we call finalize() to finish it.
|
|
b) Exception
|
|
- An error occurred. We still need to call finalize().
|
|
2. Task does not raise exception and returns either:
|
|
a) Syscall
|
|
- Syscall.handle() is called.
|
|
b) None
|
|
- The Task is simply scheduled to continue.
|
|
c) Something else
|
|
- This should not happen - error.
|
|
"""
|
|
global this_task
|
|
this_task = task
|
|
try:
|
|
if isinstance(value, BaseException):
|
|
result = task.throw(value)
|
|
else:
|
|
result = task.send(value)
|
|
except StopIteration as e:
|
|
if __debug__:
|
|
log.debug(__name__, "finish: %s", task)
|
|
finalize(task, e.value)
|
|
except Exception as e:
|
|
if __debug__:
|
|
log.exception(__name__, e)
|
|
finalize(task, e)
|
|
else:
|
|
if isinstance(result, Syscall):
|
|
result.handle(task)
|
|
elif result is None:
|
|
schedule(task)
|
|
else:
|
|
if __debug__:
|
|
log.error(__name__, "unknown syscall: %s", result)
|
|
|
|
|
|
class Syscall:
|
|
"""
|
|
When tasks want to perform any I/O, or do any sort of communication with the
|
|
scheduler, they do so through instances of a class derived from `Syscall`.
|
|
"""
|
|
|
|
def __iter__(self) -> Generator:
|
|
# support `yield from` or `await` on syscalls
|
|
return (yield self)
|
|
|
|
if TYPE_CHECKING:
|
|
|
|
def __await__(self) -> Generator[Any, Any, Any]:
|
|
return self.__iter__()
|
|
|
|
def handle(self, task: Task) -> None:
|
|
pass
|
|
|
|
|
|
class Timeout(Exception):
|
|
pass
|
|
|
|
|
|
_TIMEOUT_ERROR = Timeout()
|
|
|
|
|
|
class sleep(Syscall):
|
|
"""Pause current task and resume it after given delay.
|
|
|
|
Result value is the calculated deadline.
|
|
|
|
Example:
|
|
|
|
>>> planned = await loop.sleep(1000) # sleep for 1s
|
|
>>> print(f"missed by {utime.ticks_diff(utime.ticks_ms(), planned)} ms")
|
|
"""
|
|
|
|
def __init__(self, delay_ms: int) -> None:
|
|
self.delay_ms = delay_ms
|
|
|
|
def handle(self, task: Task) -> None:
|
|
deadline = utime.ticks_add(utime.ticks_ms(), self.delay_ms)
|
|
schedule(task, deadline, deadline)
|
|
|
|
|
|
class wait(Syscall):
|
|
"""
|
|
Pause current task, and resume only after a message on `msg_iface` is
|
|
received. Messages are received either from an USB interface, or the
|
|
touch display. Result value is a tuple of message values.
|
|
|
|
Example:
|
|
|
|
>>> hid_report, = await loop.wait(0xABCD) # await USB HID report
|
|
>>> event, x, y = await loop.wait(io.TOUCH) # await touch event
|
|
"""
|
|
|
|
_DO_NOT_RESCHEDULE = Syscall()
|
|
|
|
def __init__(self, msg_iface: int, timeout_ms: int | None = None) -> None:
|
|
self.msg_iface = msg_iface
|
|
self.timeout_ms = timeout_ms
|
|
self.task: Task | None = None
|
|
|
|
def handle(self, task: Task) -> None:
|
|
self.task = task
|
|
pause(self, self.msg_iface)
|
|
if self.timeout_ms is not None:
|
|
deadline = utime.ticks_add(utime.ticks_ms(), self.timeout_ms)
|
|
schedule(self, _TIMEOUT_ERROR, deadline)
|
|
|
|
def send(self, __value: Any) -> Any:
|
|
assert self.task is not None
|
|
self.close()
|
|
_step(self.task, __value)
|
|
return self._DO_NOT_RESCHEDULE
|
|
|
|
throw = send
|
|
|
|
def close(self) -> None:
|
|
_queue.discard(self)
|
|
if self.msg_iface in _paused:
|
|
_paused[self.msg_iface].discard(self)
|
|
if not _paused[self.msg_iface]:
|
|
del _paused[self.msg_iface]
|
|
|
|
def __iter__(self) -> Generator:
|
|
try:
|
|
return (yield self)
|
|
finally:
|
|
# whichever way we got here, we must be removed from the paused list
|
|
self.close()
|
|
|
|
|
|
_type_gen: type[Generator] = type((lambda: (yield))())
|
|
|
|
|
|
class race(Syscall):
|
|
"""
|
|
Given a list of either children tasks or syscalls, `race` waits until one of
|
|
them completes (tasks are executed in parallel, syscalls are waited upon,
|
|
directly). Return value of `race` is the return value of the child that
|
|
triggered the completion. Other running children are killed (by cancelling
|
|
any pending schedules and raising a `GeneratorExit` by calling `close()`).
|
|
|
|
Example:
|
|
|
|
>>> # async def wait_for_touch(): ...
|
|
>>> # async def animate_logo(): ...
|
|
>>> touch_task = wait_for_touch()
|
|
>>> animation_task = animate_logo()
|
|
>>> racer = loop.race(touch_task, animation_task)
|
|
>>> result = await racer
|
|
|
|
Note: You should not directly `yield` a `race` instance, see logic in
|
|
`race.__iter__` for explanation. Always use `await`.
|
|
"""
|
|
|
|
def __init__(self, *children: AwaitableTask) -> None:
|
|
self.children = children
|
|
self.finished = False
|
|
self.scheduled: list[Task] = [] # scheduled wrapper tasks
|
|
|
|
def handle(self, task: Task) -> None:
|
|
"""
|
|
Schedule all children Tasks and set `task` as callback.
|
|
"""
|
|
finalizer = self._finish
|
|
scheduled = self.scheduled
|
|
self.finished = False
|
|
|
|
self.callback = task
|
|
scheduled.clear()
|
|
|
|
for child in self.children:
|
|
child_task: Task
|
|
if isinstance(child, _type_gen):
|
|
# child is a coroutine/generator
|
|
# i.e., async function, or function using yield (these are identical
|
|
# in micropython)
|
|
child_task = child
|
|
else:
|
|
# child is a layout -- type-wise, it is an Awaitable, but
|
|
# implementation-wise it is an Iterable and we know that its __iter__
|
|
# will return a Generator.
|
|
child_task = child.__iter__() # type: ignore [Cannot access attribute "__iter__" for class "Awaitable[Unknown]";;Cannot access attribute "__iter__" for class "Coroutine[Unknown, Unknown, Unknown]"]
|
|
schedule(child_task, None, None, finalizer)
|
|
scheduled.append(child_task)
|
|
|
|
def exit(self, except_for: Task | None = None) -> None:
|
|
for task in self.scheduled:
|
|
if task != except_for:
|
|
close(task)
|
|
|
|
def _finish(self, task: Task, result: Any) -> None:
|
|
if not self.finished:
|
|
self.finished = True
|
|
self.exit(task)
|
|
schedule(self.callback, result)
|
|
|
|
def __iter__(self) -> Task:
|
|
try:
|
|
return (yield self)
|
|
except: # noqa: E722
|
|
# exception was raised on the waiting task externally with
|
|
# close() or throw(), kill the children tasks and re-raise
|
|
# Make sure finalizers don't continue processing.
|
|
self.finished = True
|
|
self.exit()
|
|
raise
|
|
|
|
|
|
class mailbox(Syscall):
|
|
"""
|
|
Wait to receive a value.
|
|
|
|
In terms of synchronization primitives, this is a condition variable that also
|
|
contains a value. It is a simplification of Go channels, which is one-ended and
|
|
only has a buffer of size 1.
|
|
|
|
The receiving end pauses until a value is received, and then empties the mailbox
|
|
to wait again.
|
|
|
|
The sending end synchronously posts a value. It is impossible to wait until
|
|
the value is consumed. Trying to post a value when the mailbox is full raises
|
|
an error, unless `replace=True` is specified
|
|
|
|
Example:
|
|
|
|
>>> # in task #1:
|
|
>>> box = loop.mailbox()
|
|
>>> while True:
|
|
>>> result = await box
|
|
>>> print("awaited result:", result)
|
|
|
|
>>> # in task #2:
|
|
>>> box.put("Hello from the other task")
|
|
>>> print("put completed")
|
|
|
|
Example Output:
|
|
|
|
put completed
|
|
awaited result: Hello from the other task
|
|
"""
|
|
|
|
_NO_VALUE = object()
|
|
|
|
def __init__(self, initial_value: Any = _NO_VALUE) -> None:
|
|
self.value = initial_value
|
|
self.taker: Task | None = None
|
|
|
|
def is_empty(self) -> bool:
|
|
"""Is the mailbox empty?"""
|
|
return self.value is self._NO_VALUE
|
|
|
|
def clear(self) -> None:
|
|
"""Empty the mailbox."""
|
|
assert self.taker is None
|
|
self.value = self._NO_VALUE
|
|
|
|
def put(self, value: Any, replace: bool = False) -> None:
|
|
"""Put a value into the mailbox.
|
|
|
|
If there is another task waiting for the value, it will be scheduled to resume.
|
|
Otherwise, the mailbox will hold the value until someone consumes it.
|
|
|
|
It is an error to call `put()` when there is a value already held, unless
|
|
`replace` is set to `True`. In such case, the held value is replaced with
|
|
the new one.
|
|
"""
|
|
if not self.is_empty() and not replace:
|
|
raise ValueError("mailbox already has a value")
|
|
|
|
self.value = value
|
|
if self.taker is not None:
|
|
self._take(self.taker)
|
|
|
|
def _take(self, task: Task) -> None:
|
|
"""Take a value and schedule the taker."""
|
|
self.taker = None
|
|
schedule(task, self.value)
|
|
self.clear()
|
|
|
|
def handle(self, task: Task) -> None:
|
|
assert self.taker is None
|
|
if not self.is_empty():
|
|
self._take(task)
|
|
else:
|
|
self.taker = task
|
|
|
|
def __iter__(self) -> Generator:
|
|
assert self.taker is None
|
|
|
|
# short-circuit if there is a value already
|
|
if not self.is_empty():
|
|
value = self.value
|
|
self.clear()
|
|
return value
|
|
|
|
# otherwise, wait for a value
|
|
try:
|
|
return (yield self)
|
|
finally:
|
|
# Clear the taker even in case of exception. This way stale takers don't
|
|
# blow up someone calling `maybe_close()`
|
|
self.taker = None
|
|
|
|
def maybe_close(self) -> None:
|
|
"""Shut down the taker if possible."""
|
|
taker = self.taker
|
|
self.taker = None
|
|
if taker is not None and taker is not this_task:
|
|
taker.close()
|
|
|
|
|
|
class spawn(Syscall):
|
|
"""Spawn a task asynchronously and get an awaitable reference to it.
|
|
|
|
Abstraction over `loop.schedule` and `loop.close`. Useful when you need to start
|
|
a task in the background, but want to be able to kill it from the outside.
|
|
|
|
Examples:
|
|
|
|
1. Spawn a background task, get its result later.
|
|
|
|
>>> wire_read = loop.spawn(read_from_wire())
|
|
>>> long_result = await long_running_operation()
|
|
>>> wire_result = await wire_read
|
|
|
|
2. Allow the user to kill a long-running operation:
|
|
|
|
>>> try:
|
|
>>> operation = loop.spawn(long_running_operation())
|
|
>>> result = await operation
|
|
>>> print("finished with result", result)
|
|
>>> except loop.TaskClosed:
|
|
>>> print("task was closed before it could finish")
|
|
>>>
|
|
>>> # meanwhile, on the other side of town...
|
|
>>> controller.close()
|
|
|
|
Task is spawned only once. Multiple attempts to `await spawned_object` will return
|
|
the original return value (or raise the original exception).
|
|
"""
|
|
|
|
def __init__(self, task: Task) -> None:
|
|
self.task = task
|
|
self.callback: Task | None = None
|
|
self.finalizer_callback: Callable[["spawn"], None] | None = None
|
|
self.finished = False
|
|
self.return_value: Any = None
|
|
|
|
# schedule task immediately
|
|
if __debug__:
|
|
log.debug(__name__, "spawn new task: %s", task)
|
|
|
|
assert isinstance(task, _type_gen)
|
|
schedule(task, finalizer=self._finalize)
|
|
|
|
def _finalize(self, task: Task, value: Any) -> None:
|
|
# sanity check: make sure finalizer is for our task
|
|
assert task is self.task
|
|
# sanity check: make sure finalizer is not called more than once
|
|
assert self.finished is False
|
|
|
|
# now we are truly finished
|
|
self.finished = True
|
|
if isinstance(value, GeneratorExit):
|
|
# coerce GeneratorExit to a catchable TaskClosed
|
|
self.return_value = TASK_CLOSED
|
|
else:
|
|
self.return_value = value
|
|
|
|
if self.callback is not None:
|
|
schedule(self.callback, self.return_value)
|
|
self.callback = None
|
|
if self.finalizer_callback is not None:
|
|
self.finalizer_callback(self)
|
|
|
|
def __iter__(self) -> Task:
|
|
if self.finished:
|
|
# exit immediately if we already have a return value
|
|
if isinstance(self.return_value, BaseException):
|
|
raise self.return_value
|
|
else:
|
|
return self.return_value
|
|
|
|
try:
|
|
return (yield self)
|
|
except BaseException:
|
|
# Clear out the callback. Otherwise we would raise the exception into it,
|
|
# AND schedule it with the closing value of the child task.
|
|
self.callback = None
|
|
assert self.task is not this_task # closing parent from child :(
|
|
close(self.task)
|
|
raise
|
|
|
|
def handle(self, caller: Task) -> None:
|
|
# the same spawn should not be awaited multiple times
|
|
assert self.callback is None
|
|
self.callback = caller
|
|
|
|
def close(self) -> None:
|
|
"""Shut down the spawned task.
|
|
|
|
If another caller is awaiting its result it will get a TaskClosed exception.
|
|
If the task was already finished, the call has no effect.
|
|
"""
|
|
if not self.finished:
|
|
if __debug__:
|
|
log.debug(__name__, "close spawned task: %s", self.task)
|
|
close(self.task)
|
|
|
|
def set_finalizer(self, finalizer_callback: Callable[["spawn"], None]) -> None:
|
|
"""Register a finalizer callback.
|
|
|
|
The provided function is executed synchronously when the spawned task ends,
|
|
with the spawn object as an argument.
|
|
"""
|
|
if self.finished:
|
|
finalizer_callback(self)
|
|
self.finalizer_callback = finalizer_callback
|
|
|
|
def is_running(self) -> bool:
|
|
"""Check if the caller is executing from the spawned task.
|
|
|
|
Useful for checking if it is OK to call `task.close()`. If `task.is_running()`
|
|
is True, it would be calling close on self, which will result in a ValueError.
|
|
"""
|
|
return self.task is this_task
|