1use futures::task::{self, ArcWake, Waker};
2use std::panic;
3use std::sync::{Arc, Mutex};
4
5struct CountingWaker {
6 nr_wake: Mutex<i32>,
7}
8
9impl CountingWaker {
10 fn new() -> Self {
11 Self { nr_wake: Mutex::new(0) }
12 }
13
14 fn wakes(&self) -> i32 {
15 *self.nr_wake.lock().unwrap()
16 }
17}
18
19impl ArcWake for CountingWaker {
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]
27fn 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]
48fn 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]
60fn 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