1 | use futures::executor::block_on; |
2 | use futures::future::poll_fn; |
3 | use futures::task::{AtomicWaker, Poll}; |
4 | use std::sync::atomic::AtomicUsize; |
5 | use std::sync::atomic::Ordering; |
6 | use std::sync::Arc; |
7 | use std::thread; |
8 | |
9 | #[test] |
10 | fn basic() { |
11 | let atomic_waker = Arc::new(AtomicWaker::new()); |
12 | let atomic_waker_copy = atomic_waker.clone(); |
13 | |
14 | let returned_pending = Arc::new(AtomicUsize::new(0)); |
15 | let returned_pending_copy = returned_pending.clone(); |
16 | |
17 | let woken = Arc::new(AtomicUsize::new(0)); |
18 | let woken_copy = woken.clone(); |
19 | |
20 | let t = thread::spawn(move || { |
21 | let mut pending_count = 0; |
22 | |
23 | block_on(poll_fn(move |cx| { |
24 | if woken_copy.load(Ordering::Relaxed) == 1 { |
25 | Poll::Ready(()) |
26 | } else { |
27 | // Assert we return pending exactly once |
28 | assert_eq!(0, pending_count); |
29 | pending_count += 1; |
30 | atomic_waker_copy.register(cx.waker()); |
31 | |
32 | returned_pending_copy.store(1, Ordering::Relaxed); |
33 | |
34 | Poll::Pending |
35 | } |
36 | })) |
37 | }); |
38 | |
39 | while returned_pending.load(Ordering::Relaxed) == 0 {} |
40 | |
41 | // give spawned thread some time to sleep in `block_on` |
42 | thread::yield_now(); |
43 | |
44 | woken.store(1, Ordering::Relaxed); |
45 | atomic_waker.wake(); |
46 | |
47 | t.join().unwrap(); |
48 | } |
49 | |