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