1 use crate::runtime::task::{Header, RawTask, Schedule};
2 
3 use std::marker::PhantomData;
4 use std::mem::ManuallyDrop;
5 use std::ops;
6 use std::ptr::NonNull;
7 use std::task::{RawWaker, RawWakerVTable, Waker};
8 
9 pub(super) struct WakerRef<'a, S: 'static> {
10     waker: ManuallyDrop<Waker>,
11     _p: PhantomData<(&'a Header, S)>,
12 }
13 
14 /// Returns a `WakerRef` which avoids having to preemptively increase the
15 /// refcount if there is no need to do so.
waker_ref<S>(header: &NonNull<Header>) -> WakerRef<'_, S> where S: Schedule,16 pub(super) fn waker_ref<S>(header: &NonNull<Header>) -> WakerRef<'_, S>
17 where
18     S: Schedule,
19 {
20     // `Waker::will_wake` uses the VTABLE pointer as part of the check. This
21     // means that `will_wake` will always return false when using the current
22     // task's waker. (discussion at rust-lang/rust#66281).
23     //
24     // To fix this, we use a single vtable. Since we pass in a reference at this
25     // point and not an *owned* waker, we must ensure that `drop` is never
26     // called on this waker instance. This is done by wrapping it with
27     // `ManuallyDrop` and then never calling drop.
28     let waker = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker(*header))) };
29 
30     WakerRef {
31         waker,
32         _p: PhantomData,
33     }
34 }
35 
36 impl<S> ops::Deref for WakerRef<'_, S> {
37     type Target = Waker;
38 
deref(&self) -> &Waker39     fn deref(&self) -> &Waker {
40         &self.waker
41     }
42 }
43 
44 cfg_trace! {
45     macro_rules! trace {
46         ($header:expr, $op:expr) => {
47             if let Some(id) = Header::get_tracing_id(&$header) {
48                 tracing::trace!(
49                     target: "tokio::task::waker",
50                     op = $op,
51                     task.id = id.into_u64(),
52                 );
53             }
54         }
55     }
56 }
57 
58 cfg_not_trace! {
59     macro_rules! trace {
60         ($header:expr, $op:expr) => {
61             // noop
62             let _ = &$header;
63         }
64     }
65 }
66 
clone_waker(ptr: *const ()) -> RawWaker67 unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
68     let header = NonNull::new_unchecked(ptr as *mut Header);
69     trace!(header, "waker.clone");
70     header.as_ref().state.ref_inc();
71     raw_waker(header)
72 }
73 
drop_waker(ptr: *const ())74 unsafe fn drop_waker(ptr: *const ()) {
75     let ptr = NonNull::new_unchecked(ptr as *mut Header);
76     trace!(ptr, "waker.drop");
77     let raw = RawTask::from_raw(ptr);
78     raw.drop_reference();
79 }
80 
wake_by_val(ptr: *const ())81 unsafe fn wake_by_val(ptr: *const ()) {
82     let ptr = NonNull::new_unchecked(ptr as *mut Header);
83     trace!(ptr, "waker.wake");
84     let raw = RawTask::from_raw(ptr);
85     raw.wake_by_val();
86 }
87 
88 // Wake without consuming the waker
wake_by_ref(ptr: *const ())89 unsafe fn wake_by_ref(ptr: *const ()) {
90     let ptr = NonNull::new_unchecked(ptr as *mut Header);
91     trace!(ptr, "waker.wake_by_ref");
92     let raw = RawTask::from_raw(ptr);
93     raw.wake_by_ref();
94 }
95 
96 static WAKER_VTABLE: RawWakerVTable =
97     RawWakerVTable::new(clone_waker, wake_by_val, wake_by_ref, drop_waker);
98 
raw_waker(header: NonNull<Header>) -> RawWaker99 fn raw_waker(header: NonNull<Header>) -> RawWaker {
100     let ptr = header.as_ptr() as *const ();
101     RawWaker::new(ptr, &WAKER_VTABLE)
102 }
103