1use futures_core::task::Waker;
2use futures_util::task::{self, ArcWake};
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::sync::Arc;
5
6/// Number of times the waker was awoken.
7///
8/// See [`new_count_waker`] for usage.
9#[derive(Debug)]
10pub struct AwokenCount {
11 inner: Arc<WakerInner>,
12}
13
14impl AwokenCount {
15 /// Get the current count.
16 pub fn get(&self) -> usize {
17 self.inner.count.load(Ordering::SeqCst)
18 }
19}
20
21impl PartialEq<usize> for AwokenCount {
22 fn eq(&self, other: &usize) -> bool {
23 self.get() == *other
24 }
25}
26
27#[derive(Debug)]
28struct WakerInner {
29 count: AtomicUsize,
30}
31
32impl ArcWake for WakerInner {
33 fn wake_by_ref(arc_self: &Arc<Self>) {
34 let _ = arc_self.count.fetch_add(1, Ordering::SeqCst);
35 }
36}
37
38/// Create a new [`Waker`] that counts the number of times it's awoken.
39///
40/// [`Waker`]: futures_core::task::Waker
41///
42/// # Examples
43///
44/// ```
45/// use futures_test::task::new_count_waker;
46///
47/// let (waker, count) = new_count_waker();
48///
49/// assert_eq!(count, 0);
50///
51/// waker.wake_by_ref();
52/// waker.wake();
53///
54/// assert_eq!(count, 2);
55/// ```
56pub fn new_count_waker() -> (Waker, AwokenCount) {
57 let inner = Arc::new(WakerInner { count: AtomicUsize::new(0) });
58 (task::waker(inner.clone()), AwokenCount { inner })
59}
60