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