1 use crate::sync::task::AtomicWaker;
2
3 use futures::future::poll_fn;
4 use loom::future::block_on;
5 use loom::sync::atomic::AtomicUsize;
6 use loom::thread;
7 use std::sync::atomic::Ordering::Relaxed;
8 use std::sync::Arc;
9 use std::task::Poll::{Pending, Ready};
10
11 struct Chan {
12 num: AtomicUsize,
13 task: AtomicWaker,
14 }
15
16 #[test]
basic_notification()17 fn basic_notification() {
18 const NUM_NOTIFY: usize = 2;
19
20 loom::model(|| {
21 let chan = Arc::new(Chan {
22 num: AtomicUsize::new(0),
23 task: AtomicWaker::new(),
24 });
25
26 for _ in 0..NUM_NOTIFY {
27 let chan = chan.clone();
28
29 thread::spawn(move || {
30 chan.num.fetch_add(1, Relaxed);
31 chan.task.wake();
32 });
33 }
34
35 block_on(poll_fn(move |cx| {
36 chan.task.register_by_ref(cx.waker());
37
38 if NUM_NOTIFY == chan.num.load(Relaxed) {
39 return Ready(());
40 }
41
42 Pending
43 }));
44 });
45 }
46
47 #[test]
test_panicky_waker()48 fn test_panicky_waker() {
49 use std::panic;
50 use std::ptr;
51 use std::task::{RawWaker, RawWakerVTable, Waker};
52
53 static PANICKING_VTABLE: RawWakerVTable =
54 RawWakerVTable::new(|_| panic!("clone"), |_| (), |_| (), |_| ());
55
56 let panicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &PANICKING_VTABLE)) };
57
58 // If you're working with this test (and I sure hope you never have to!),
59 // uncomment the following section because there will be a lot of panics
60 // which would otherwise log.
61 //
62 // We can't however leaved it uncommented, because it's global.
63 // panic::set_hook(Box::new(|_| ()));
64
65 const NUM_NOTIFY: usize = 2;
66
67 loom::model(move || {
68 let chan = Arc::new(Chan {
69 num: AtomicUsize::new(0),
70 task: AtomicWaker::new(),
71 });
72
73 for _ in 0..NUM_NOTIFY {
74 let chan = chan.clone();
75
76 thread::spawn(move || {
77 chan.num.fetch_add(1, Relaxed);
78 chan.task.wake();
79 });
80 }
81
82 // Note: this panic should have no effect on the overall state of the
83 // waker and it should proceed as normal.
84 //
85 // A thread above might race to flag a wakeup, and a WAKING state will
86 // be preserved if this expected panic races with that so the below
87 // procedure should be allowed to continue uninterrupted.
88 let _ = panic::catch_unwind(|| chan.task.register_by_ref(&panicking));
89
90 block_on(poll_fn(move |cx| {
91 chan.task.register_by_ref(cx.waker());
92
93 if NUM_NOTIFY == chan.num.load(Relaxed) {
94 return Ready(());
95 }
96
97 Pending
98 }));
99 });
100 }
101