tokio/util/
wake.rs

1use crate::loom::sync::Arc;
2
3use std::marker::PhantomData;
4use std::mem::ManuallyDrop;
5use std::ops::Deref;
6use std::task::{RawWaker, RawWakerVTable, Waker};
7
8/// Simplified waking interface based on Arcs.
9pub(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)]
19pub(crate) struct WakerRef<'a> {
20    waker: ManuallyDrop<Waker>,
21    _p: PhantomData<&'a ()>,
22}
23
24impl 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>`.
33pub(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
44fn 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
53unsafe 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
58unsafe 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`
64unsafe 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
70unsafe fn drop_arc_raw<T: Wake>(data: *const ()) {
71    drop(Arc::<T>::from_raw(data.cast()));
72}
tokio/util/wake.rs

tokio/util/
wake.rs
