1 //! Creates a Waker that can be observed from tests.
2 
3 use std::mem::forget;
4 use std::sync::atomic::{AtomicU32, Ordering};
5 use std::sync::Arc;
6 use std::task::{RawWaker, RawWakerVTable, Waker};
7 
8 #[derive(Default)]
9 pub struct WakerHandle {
10     clone_count: AtomicU32,
11     drop_count: AtomicU32,
12     wake_count: AtomicU32,
13 }
14 
15 impl WakerHandle {
clone_count(&self) -> u3216     pub fn clone_count(&self) -> u32 {
17         self.clone_count.load(Ordering::Relaxed)
18     }
19 
drop_count(&self) -> u3220     pub fn drop_count(&self) -> u32 {
21         self.drop_count.load(Ordering::Relaxed)
22     }
23 
wake_count(&self) -> u3224     pub fn wake_count(&self) -> u32 {
25         self.wake_count.load(Ordering::Relaxed)
26     }
27 }
28 
waker() -> (Waker, Arc<WakerHandle>)29 pub fn waker() -> (Waker, Arc<WakerHandle>) {
30     let waker_handle = Arc::new(WakerHandle::default());
31     let waker_handle_ptr = Arc::into_raw(waker_handle.clone());
32     let raw_waker = RawWaker::new(waker_handle_ptr as *const _, waker_vtable());
33     (unsafe { Waker::from_raw(raw_waker) }, waker_handle)
34 }
35 
waker_vtable() -> &'static RawWakerVTable36 pub(super) fn waker_vtable() -> &'static RawWakerVTable {
37     &RawWakerVTable::new(clone_raw, wake_raw, wake_by_ref_raw, drop_raw)
38 }
39 
clone_raw(data: *const ()) -> RawWaker40 unsafe fn clone_raw(data: *const ()) -> RawWaker {
41     let handle: Arc<WakerHandle> = Arc::from_raw(data as *const _);
42     handle.clone_count.fetch_add(1, Ordering::Relaxed);
43     forget(handle.clone());
44     forget(handle);
45     RawWaker::new(data, waker_vtable())
46 }
47 
wake_raw(data: *const ())48 unsafe fn wake_raw(data: *const ()) {
49     let handle: Arc<WakerHandle> = Arc::from_raw(data as *const _);
50     handle.wake_count.fetch_add(1, Ordering::Relaxed);
51     handle.drop_count.fetch_add(1, Ordering::Relaxed);
52 }
53 
wake_by_ref_raw(data: *const ())54 unsafe fn wake_by_ref_raw(data: *const ()) {
55     let handle: Arc<WakerHandle> = Arc::from_raw(data as *const _);
56     handle.wake_count.fetch_add(1, Ordering::Relaxed);
57     forget(handle)
58 }
59 
drop_raw(data: *const ())60 unsafe fn drop_raw(data: *const ()) {
61     let handle: Arc<WakerHandle> = Arc::from_raw(data as *const _);
62     handle.drop_count.fetch_add(1, Ordering::Relaxed);
63     drop(handle)
64 }
65