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
5 use std::mem::drop;
6 use std::mem::ManuallyDrop;
7 use std::sync::Weak;
8 use std::task::RawWaker;
9 use std::task::RawWakerVTable;
10 use 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)]
14 pub(crate) struct WakerToken(pub(crate) usize);
15
16 /// Like `futures::task::ArcWake` but uses `Weak<T>` instead of `Arc<T>`.
17 pub(crate) trait WeakWake: Send + Sync {
wake_by_ref(weak_self: &Weak<Self>)18 fn wake_by_ref(weak_self: &Weak<Self>);
19
wake(weak_self: Weak<Self>)20 fn wake(weak_self: Weak<Self>) {
21 Self::wake_by_ref(&weak_self)
22 }
23 }
24
waker_vtable<W: WeakWake>() -> &'static RawWakerVTable25 fn 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
clone_weak_raw<W: WeakWake>(data: *const ()) -> RawWaker34 unsafe 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
wake_weak_raw<W: WeakWake>(data: *const ())46 unsafe 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
wake_by_ref_weak_raw<W: WeakWake>(data: *const ())52 unsafe 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
drop_weak_raw<W: WeakWake>(data: *const ())60 unsafe fn drop_weak_raw<W: WeakWake>(data: *const ()) {
61 drop(Weak::from_raw(data as *const W))
62 }
63
new_waker<W: WeakWake>(w: Weak<W>) -> Waker64 pub(crate) fn new_waker<W: WeakWake>(w: Weak<W>) -> Waker {
65 unsafe {
66 Waker::from_raw(RawWaker::new(
67 w.into_raw() as *const (),
68 waker_vtable::<W>(),
69 ))
70 }
71 }
72