1 // Copyright 2020 The Chromium OS Authors. All rights reserved.
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, ManuallyDrop};
6 use std::sync::Weak;
7 use std::task::{RawWaker, RawWakerVTable, Waker};
8
9 /// Wrapper around a usize used as a token to uniquely identify a pending waker.
10 #[derive(Debug)]
11 pub(crate) struct WakerToken(pub(crate) usize);
12
13 /// Like `futures::task::ArcWake` but uses `Weak<T>` instead of `Arc<T>`.
14 pub(crate) trait WeakWake: Send + Sync {
wake_by_ref(weak_self: &Weak<Self>)15 fn wake_by_ref(weak_self: &Weak<Self>);
16
wake(weak_self: Weak<Self>)17 fn wake(weak_self: Weak<Self>) {
18 Self::wake_by_ref(&weak_self)
19 }
20 }
21
waker_vtable<W: WeakWake>() -> &'static RawWakerVTable22 fn waker_vtable<W: WeakWake>() -> &'static RawWakerVTable {
23 &RawWakerVTable::new(
24 clone_weak_raw::<W>,
25 wake_weak_raw::<W>,
26 wake_by_ref_weak_raw::<W>,
27 drop_weak_raw::<W>,
28 )
29 }
30
clone_weak_raw<W: WeakWake>(data: *const ()) -> RawWaker31 unsafe fn clone_weak_raw<W: WeakWake>(data: *const ()) -> RawWaker {
32 // Get a handle to the Weak<T> but wrap it in a ManuallyDrop so that we don't reduce the
33 // refcount at the end of this function.
34 let weak = ManuallyDrop::new(Weak::<W>::from_raw(data as *const W));
35
36 // Now increase the weak count and keep it in a ManuallyDrop so that it doesn't get decreased
37 // at the end of this function.
38 let _weak_clone: ManuallyDrop<_> = weak.clone();
39
40 RawWaker::new(data, waker_vtable::<W>())
41 }
42
wake_weak_raw<W: WeakWake>(data: *const ())43 unsafe fn wake_weak_raw<W: WeakWake>(data: *const ()) {
44 let weak: Weak<W> = Weak::from_raw(data as *const W);
45
46 WeakWake::wake(weak)
47 }
48
wake_by_ref_weak_raw<W: WeakWake>(data: *const ())49 unsafe fn wake_by_ref_weak_raw<W: WeakWake>(data: *const ()) {
50 // Get a handle to the Weak<T> but wrap it in a ManuallyDrop so that we don't reduce the
51 // refcount at the end of this function.
52 let weak = ManuallyDrop::new(Weak::<W>::from_raw(data as *const W));
53
54 WeakWake::wake_by_ref(&weak)
55 }
56
drop_weak_raw<W: WeakWake>(data: *const ())57 unsafe fn drop_weak_raw<W: WeakWake>(data: *const ()) {
58 drop(Weak::from_raw(data as *const W))
59 }
60
new_waker<W: WeakWake>(w: Weak<W>) -> Waker61 pub(crate) fn new_waker<W: WeakWake>(w: Weak<W>) -> Waker {
62 unsafe {
63 Waker::from_raw(RawWaker::new(
64 w.into_raw() as *const (),
65 waker_vtable::<W>(),
66 ))
67 }
68 }
69