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