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);