tokio/runtime/task/
waker.rs

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 = 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    /// # Safety
46    ///
47    /// `$header` must be a valid pointer to a [`Header`].
48    macro_rules! trace {
49        ($header:expr, $op:expr) => {
50            if let Some(id) = Header::get_tracing_id(&$header) {
51                tracing::trace!(
52                    target: "tokio::task::waker",
53                    op = $op,
54                    task.id = id.into_u64(),
55                );
56            }
57        }
58    }
59}
60
61cfg_not_trace! {
62    macro_rules! trace {
63        ($header:expr, $op:expr) => {
64            // noop
65            let _ = &$header;
66        }
67    }
68}
69
70unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
71    // Safety: `ptr` was created from a `Header` pointer in function `waker_ref`.
72    let header = unsafe { NonNull::new_unchecked(ptr as *mut Header) };
73    #[cfg_attr(not(all(tokio_unstable, feature = "tracing")), allow(unused_unsafe))]
74    unsafe {
75        trace!(header, "waker.clone");
76    }
77    unsafe { header.as_ref() }.state.ref_inc();
78    raw_waker(header)
79}
80
81unsafe fn drop_waker(ptr: *const ()) {
82    // Safety: `ptr` was created from a `Header` pointer in function `waker_ref`.
83    let ptr = unsafe { NonNull::new_unchecked(ptr as *mut Header) };
84    // TODO; replace to #[expect(unused_unsafe)] after bumping MSRV to 1.81.0.
85    #[cfg_attr(not(all(tokio_unstable, feature = "tracing")), allow(unused_unsafe))]
86    unsafe {
87        trace!(ptr, "waker.drop");
88    }
89    let raw = unsafe { RawTask::from_raw(ptr) };
90    raw.drop_reference();
91}
92
93unsafe fn wake_by_val(ptr: *const ()) {
94    // Safety: `ptr` was created from a `Header` pointer in function `waker_ref`.
95    let ptr = unsafe { NonNull::new_unchecked(ptr as *mut Header) };
96    // TODO; replace to #[expect(unused_unsafe)] after bumping MSRV to 1.81.0.
97    #[cfg_attr(not(all(tokio_unstable, feature = "tracing")), allow(unused_unsafe))]
98    unsafe {
99        trace!(ptr, "waker.wake");
100    }
101    let raw = unsafe { RawTask::from_raw(ptr) };
102    raw.wake_by_val();
103}
104
105// Wake without consuming the waker
106unsafe fn wake_by_ref(ptr: *const ()) {
107    // Safety: `ptr` was created from a `Header` pointer in function `waker_ref`.
108    let ptr = unsafe { NonNull::new_unchecked(ptr as *mut Header) };
109    // TODO; replace to #[expect(unused_unsafe)] after bumping MSRV to 1.81.0.
110    #[cfg_attr(not(all(tokio_unstable, feature = "tracing")), allow(unused_unsafe))]
111    unsafe {
112        trace!(ptr, "waker.wake_by_ref");
113    }
114    let raw = unsafe { RawTask::from_raw(ptr) };
115    raw.wake_by_ref();
116}
117
118static WAKER_VTABLE: RawWakerVTable =
119    RawWakerVTable::new(clone_waker, wake_by_val, wake_by_ref, drop_waker);
120
121fn raw_waker(header: NonNull<Header>) -> RawWaker {
122    let ptr = header.as_ptr() as *const ();
123    RawWaker::new(ptr, &WAKER_VTABLE)
124}
tokio/runtime/task/waker.rs

tokio/runtime/task/
waker.rs
