shuttle/runtime/task/
waker.rs

1use crate::runtime::execution::ExecutionState;
2use crate::runtime::task::TaskId;
3use std::task::{RawWaker, RawWakerVTable, Waker};
4
5// Safety: the `RawWaker` interface is unsafe because it requires manually enforcing resource
6// management contracts on each method in the vtable:
7// * `clone` should create an additional RawWaker, including creating all the resources required
8// * `wake` should consume the waker it was invoked on and release its resources
9// * `wake_by_ref` is like `wake` but does not consume or release the resources
10// * `drop` releases all the resources associated with a waker
11// Our wakers don't have any resources associated with them -- the `data` pointer's bits are just
12// the task ID -- so all these safety requirements are trivial.
13
14/// Create a `Waker` that will make the given `task_id` runnable when invoked.
15pub(crate) fn make_waker(task_id: TaskId) -> Waker {
16    // We stash the task ID into the bits of the `data` pointer that all the vtable method below
17    // receive as an argument.
18    let data = task_id.0 as *const ();
19    // Safety: see above
20    unsafe { Waker::from_raw(RawWaker::new(data, &RAW_WAKER_VTABLE)) }
21}
22
23unsafe fn raw_waker_clone(data: *const ()) -> RawWaker {
24    // No resources associated with our wakers, so just duplicate the pointer
25    RawWaker::new(data, &RAW_WAKER_VTABLE)
26}
27
28unsafe fn raw_waker_wake(data: *const ()) {
29    let task_id = TaskId::from(data as usize);
30    ExecutionState::with(|state| {
31        if state.is_finished() {
32            return;
33        }
34
35        let waiter = state.get_mut(task_id);
36
37        if waiter.finished() {
38            return;
39        }
40
41        waiter.wake();
42    });
43}
44
45unsafe fn raw_waker_wake_by_ref(data: *const ()) {
46    // Our wakers have no resources associated with then, so `wake` and `wake_by_ref` are the same
47    raw_waker_wake(data);
48}
49
50unsafe fn raw_waker_drop(_data: *const ()) {
51    // No resources associated with our wakers, so nothing to do on drop
52}
53
54const RAW_WAKER_VTABLE: RawWakerVTable =
55    RawWakerVTable::new(raw_waker_clone, raw_waker_wake, raw_waker_wake_by_ref, raw_waker_drop);