1use crate::runtime::task::{Header, RawTask, Schedule};
2
3use std::marker::PhantomData;
4use std::mem::ManuallyDrop;
5use std::ops;
6use std::ptr::NonNull;
7use std::task::{RawWaker, RawWakerVTable, Waker};
8
9pub(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.
16pub(super) fn waker_ref<S>(header: &NonNull<Header>) -> WakerRef<'_, S>
17where
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: ManuallyDrop = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker(*header))) };
29
30 WakerRef {
31 waker,
32 _p: PhantomData,
33 }
34}
35
36impl<S> ops::Deref for WakerRef<'_, S> {
37 type Target = Waker;
38
39 fn deref(&self) -> &Waker {
40 &self.waker
41 }
42}
43
44cfg_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
58cfg_not_trace! {
59 macro_rules! trace {
60 ($header:expr, $op:expr) => {
61 // noop
62 let _ = &$header;
63 }
64 }
65}
66
67unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
68 let header: NonNull
= 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
74unsafe fn drop_waker(ptr: *const ()) {
75 let ptr: NonNull
= NonNull::new_unchecked(ptr as *mut Header);
76 trace!(ptr, "waker.drop");
77 let raw: RawTask = RawTask::from_raw(ptr);
78 raw.drop_reference();
79}
80
81unsafe fn wake_by_val(ptr: *const ()) {
82 let ptr: NonNull
= NonNull::new_unchecked(ptr as *mut Header);
83 trace!(ptr, "waker.wake");
84 let raw: RawTask = RawTask::from_raw(ptr);
85 raw.wake_by_val();
86}
87
88// Wake without consuming the waker
89unsafe fn wake_by_ref(ptr: *const ()) {
90 let ptr: NonNull
= NonNull::new_unchecked(ptr as *mut Header);
91 trace!(ptr, "waker.wake_by_ref");
92 let raw: RawTask = RawTask::from_raw(ptr);
93 raw.wake_by_ref();
94}
95
96static WAKER_VTABLE: RawWakerVTable =
97 RawWakerVTable::new(clone_waker, wake_by_val, wake_by_ref, drop_waker);
98
99fn raw_waker(header: NonNull<Header>) -> RawWaker {
100 let ptr: *const () = header.as_ptr() as *const ();
101 RawWaker::new(data:ptr, &WAKER_VTABLE)
102}
103