1 use crate::runtime::task::harness::Harness;
2 use crate::runtime::task::{Header, Schedule};
3
4 use std::future::Future;
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
clone_waker<T, S>(ptr: *const ()) -> RawWaker where T: Future, S: Schedule,47 unsafe fn clone_waker<T, S>(ptr: *const ()) -> RawWaker
48 where
49 T: Future,
50 S: Schedule,
51 {
52 let header = ptr as *const Header;
53 (*header).state.ref_inc();
54 raw_waker::<T, S>(header)
55 }
56
drop_waker<T, S>(ptr: *const ()) where T: Future, S: Schedule,57 unsafe fn drop_waker<T, S>(ptr: *const ())
58 where
59 T: Future,
60 S: Schedule,
61 {
62 let ptr = NonNull::new_unchecked(ptr as *mut Header);
63 let harness = Harness::<T, S>::from_raw(ptr);
64 harness.drop_reference();
65 }
66
wake_by_val<T, S>(ptr: *const ()) where T: Future, S: Schedule,67 unsafe fn wake_by_val<T, S>(ptr: *const ())
68 where
69 T: Future,
70 S: Schedule,
71 {
72 let ptr = NonNull::new_unchecked(ptr as *mut Header);
73 let harness = Harness::<T, S>::from_raw(ptr);
74 harness.wake_by_val();
75 }
76
77 // Wake without consuming the waker
wake_by_ref<T, S>(ptr: *const ()) where T: Future, S: Schedule,78 unsafe fn wake_by_ref<T, S>(ptr: *const ())
79 where
80 T: Future,
81 S: Schedule,
82 {
83 let ptr = NonNull::new_unchecked(ptr as *mut Header);
84 let harness = Harness::<T, S>::from_raw(ptr);
85 harness.wake_by_ref();
86 }
87
raw_waker<T, S>(header: *const Header) -> RawWaker where T: Future, S: Schedule,88 fn raw_waker<T, S>(header: *const Header) -> RawWaker
89 where
90 T: Future,
91 S: Schedule,
92 {
93 let ptr = header as *const ();
94 let vtable = &RawWakerVTable::new(
95 clone_waker::<T, S>,
96 wake_by_val::<T, S>,
97 wake_by_ref::<T, S>,
98 drop_waker::<T, S>,
99 );
100 RawWaker::new(ptr, vtable)
101 }
102