cros_async/
waker.rs

1// Copyright 2020 The ChromiumOS Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use std::mem::drop;
6use std::mem::ManuallyDrop;
7use std::sync::Weak;
8use std::task::RawWaker;
9use std::task::RawWakerVTable;
10use std::task::Waker;
11
12/// Wrapper around a usize used as a token to uniquely identify a pending waker.
13#[derive(Debug, Hash, PartialEq, Eq, Copy, Clone)]
14pub(crate) struct WakerToken(pub(crate) usize);
15
16/// Like `futures::task::ArcWake` but uses `Weak<T>` instead of `Arc<T>`.
17pub(crate) trait WeakWake: Send + Sync {
18    fn wake_by_ref(weak_self: &Weak<Self>);
19
20    fn wake(weak_self: Weak<Self>) {
21        Self::wake_by_ref(&weak_self)
22    }
23}
24
25fn waker_vtable<W: WeakWake>() -> &'static RawWakerVTable {
26    &RawWakerVTable::new(
27        clone_weak_raw::<W>,
28        wake_weak_raw::<W>,
29        wake_by_ref_weak_raw::<W>,
30        drop_weak_raw::<W>,
31    )
32}
33
34unsafe fn clone_weak_raw<W: WeakWake>(data: *const ()) -> RawWaker {
35    // Get a handle to the Weak<T> but wrap it in a ManuallyDrop so that we don't reduce the
36    // refcount at the end of this function.
37    let weak = ManuallyDrop::new(Weak::<W>::from_raw(data as *const W));
38
39    // Now increase the weak count and keep it in a ManuallyDrop so that it doesn't get decreased
40    // at the end of this function.
41    let _weak_clone: ManuallyDrop<_> = weak.clone();
42
43    RawWaker::new(data, waker_vtable::<W>())
44}
45
46unsafe fn wake_weak_raw<W: WeakWake>(data: *const ()) {
47    let weak: Weak<W> = Weak::from_raw(data as *const W);
48
49    WeakWake::wake(weak)
50}
51
52unsafe fn wake_by_ref_weak_raw<W: WeakWake>(data: *const ()) {
53    // Get a handle to the Weak<T> but wrap it in a ManuallyDrop so that we don't reduce the
54    // refcount at the end of this function.
55    let weak = ManuallyDrop::new(Weak::<W>::from_raw(data as *const W));
56
57    WeakWake::wake_by_ref(&weak)
58}
59
60unsafe fn drop_weak_raw<W: WeakWake>(data: *const ()) {
61    drop(Weak::from_raw(data as *const W))
62}
63
64pub(crate) fn new_waker<W: WeakWake>(w: Weak<W>) -> Waker {
65    // TODO(b/315998194): Add safety comment
66    #[allow(clippy::undocumented_unsafe_blocks)]
67    unsafe {
68        Waker::from_raw(RawWaker::new(
69            w.into_raw() as *const (),
70            waker_vtable::<W>(),
71        ))
72    }
73}
cros_async/waker.rs

cros_async/
waker.rs
