1 use std::marker::PhantomData;
2 use std::mem::ManuallyDrop;
3 use std::ops::Deref;
4 use std::sync::Arc;
5 use std::task::{RawWaker, RawWakerVTable, Waker};
6
7 /// Simplified waking interface based on Arcs
8 pub(crate) trait Wake: Send + Sync {
9 /// Wake by value
wake(self: Arc<Self>)10 fn wake(self: Arc<Self>);
11
12 /// Wake by reference
wake_by_ref(arc_self: &Arc<Self>)13 fn wake_by_ref(arc_self: &Arc<Self>);
14 }
15
16 /// A `Waker` that is only valid for a given lifetime.
17 #[derive(Debug)]
18 pub(crate) struct WakerRef<'a> {
19 waker: ManuallyDrop<Waker>,
20 _p: PhantomData<&'a ()>,
21 }
22
23 impl Deref for WakerRef<'_> {
24 type Target = Waker;
25
deref(&self) -> &Waker26 fn deref(&self) -> &Waker {
27 &self.waker
28 }
29 }
30
31 /// Creates a reference to a `Waker` from a reference to `Arc<impl Wake>`.
waker_ref<W: Wake>(wake: &Arc<W>) -> WakerRef<'_>32 pub(crate) fn waker_ref<W: Wake>(wake: &Arc<W>) -> WakerRef<'_> {
33 let ptr = &**wake as *const _ as *const ();
34
35 let waker = unsafe { Waker::from_raw(RawWaker::new(ptr, waker_vtable::<W>())) };
36
37 WakerRef {
38 waker: ManuallyDrop::new(waker),
39 _p: PhantomData,
40 }
41 }
42
waker_vtable<W: Wake>() -> &'static RawWakerVTable43 fn waker_vtable<W: Wake>() -> &'static RawWakerVTable {
44 &RawWakerVTable::new(
45 clone_arc_raw::<W>,
46 wake_arc_raw::<W>,
47 wake_by_ref_arc_raw::<W>,
48 drop_arc_raw::<W>,
49 )
50 }
51
inc_ref_count<T: Wake>(data: *const ())52 unsafe fn inc_ref_count<T: Wake>(data: *const ()) {
53 // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
54 let arc = ManuallyDrop::new(Arc::<T>::from_raw(data as *const T));
55
56 // Now increase refcount, but don't drop new refcount either
57 let arc_clone: ManuallyDrop<_> = arc.clone();
58
59 // Drop explicitly to avoid clippy warnings
60 drop(arc);
61 drop(arc_clone);
62 }
63
clone_arc_raw<T: Wake>(data: *const ()) -> RawWaker64 unsafe fn clone_arc_raw<T: Wake>(data: *const ()) -> RawWaker {
65 inc_ref_count::<T>(data);
66 RawWaker::new(data, waker_vtable::<T>())
67 }
68
wake_arc_raw<T: Wake>(data: *const ())69 unsafe fn wake_arc_raw<T: Wake>(data: *const ()) {
70 let arc: Arc<T> = Arc::from_raw(data as *const T);
71 Wake::wake(arc);
72 }
73
74 // used by `waker_ref`
wake_by_ref_arc_raw<T: Wake>(data: *const ())75 unsafe fn wake_by_ref_arc_raw<T: Wake>(data: *const ()) {
76 // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
77 let arc = ManuallyDrop::new(Arc::<T>::from_raw(data as *const T));
78 Wake::wake_by_ref(&arc);
79 }
80
drop_arc_raw<T: Wake>(data: *const ())81 unsafe fn drop_arc_raw<T: Wake>(data: *const ()) {
82 drop(Arc::<T>::from_raw(data as *const T))
83 }
84