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 { |

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 | |

19 | **impl** 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] |

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] |

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] |

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 | |