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