tokio/util/
wake.rs

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

tokio/util/
wake.rs
