1 use futures_core::task::Waker;
2 use futures_util::task::{self, ArcWake};
3 use std::sync::atomic::{AtomicUsize, Ordering};
4 use std::sync::Arc;
5
6 /// Number of times the waker was awoken.
7 ///
8 /// See [`new_count_waker`] for usage.
9 #[derive(Debug)]
10 pub struct AwokenCount {
11 inner: Arc<WakerInner>,
12 }
13
14 impl AwokenCount {
15 /// Get the current count.
get(&self) -> usize16 pub fn get(&self) -> usize {
17 self.inner.count.load(Ordering::SeqCst)
18 }
19 }
20
21 impl PartialEq<usize> for AwokenCount {
eq(&self, other: &usize) -> bool22 fn eq(&self, other: &usize) -> bool {
23 self.get() == *other
24 }
25 }
26
27 #[derive(Debug)]
28 struct WakerInner {
29 count: AtomicUsize,
30 }
31
32 impl ArcWake for WakerInner {
wake_by_ref(arc_self: &Arc<Self>)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 /// ```
new_count_waker() -> (Waker, AwokenCount)56 pub 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