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