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