1 use futures::task::{self, ArcWake, Waker};
2 use std::panic;
3 use std::sync::{Arc, Mutex};
4
5 struct CountingWaker {
6 nr_wake: Mutex<i32>,
7 }
8
9 impl CountingWaker {
new() -> Self10 fn new() -> Self {
11 Self { nr_wake: Mutex::new(0) }
12 }
13
wakes(&self) -> i3214 fn wakes(&self) -> i32 {
15 *self.nr_wake.lock().unwrap()
16 }
17 }
18
19 impl ArcWake for CountingWaker {
wake_by_ref(arc_self: &Arc<Self>)20 fn wake_by_ref(arc_self: &Arc<Self>) {
21 let mut lock = arc_self.nr_wake.lock().unwrap();
22 *lock += 1;
23 }
24 }
25
26 #[test]
create_from_arc()27 fn create_from_arc() {
28 let some_w = Arc::new(CountingWaker::new());
29
30 let w1: Waker = task::waker(some_w.clone());
31 assert_eq!(2, Arc::strong_count(&some_w));
32 w1.wake_by_ref();
33 assert_eq!(1, some_w.wakes());
34
35 let w2 = w1.clone();
36 assert_eq!(3, Arc::strong_count(&some_w));
37
38 w2.wake_by_ref();
39 assert_eq!(2, some_w.wakes());
40
41 drop(w2);
42 assert_eq!(2, Arc::strong_count(&some_w));
43 drop(w1);
44 assert_eq!(1, Arc::strong_count(&some_w));
45 }
46
47 #[test]
ref_wake_same()48 fn ref_wake_same() {
49 let some_w = Arc::new(CountingWaker::new());
50
51 let w1: Waker = task::waker(some_w.clone());
52 let w2 = task::waker_ref(&some_w);
53 let w3 = w2.clone();
54
55 assert!(w1.will_wake(&w2));
56 assert!(w2.will_wake(&w3));
57 }
58
59 #[test]
proper_refcount_on_wake_panic()60 fn proper_refcount_on_wake_panic() {
61 struct PanicWaker;
62
63 impl ArcWake for PanicWaker {
64 fn wake_by_ref(_arc_self: &Arc<Self>) {
65 panic!("WAKE UP");
66 }
67 }
68
69 let some_w = Arc::new(PanicWaker);
70
71 let w1: Waker = task::waker(some_w.clone());
72 assert_eq!(
73 "WAKE UP",
74 *panic::catch_unwind(|| w1.wake_by_ref()).unwrap_err().downcast::<&str>().unwrap()
75 );
76 assert_eq!(2, Arc::strong_count(&some_w)); // some_w + w1
77 drop(w1);
78 assert_eq!(1, Arc::strong_count(&some_w)); // some_w
79 }
80