AtomicWaker test

Author: stepancheg

futures-util/tests/atomic_waker.rs

@@ -0,0 +1,54 @@
+#![feature(futures_api)]
+
+extern crate futures_util;
+
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+use std::sync::Arc;
+use std::thread;
+
+use futures_core::Poll;
+use futures_executor::block_on;
+use futures_util::future::poll_fn;
+use futures_util::task::AtomicWaker;
+
+#[test]
+fn basic() {
+    let atomic_waker = Arc::new(AtomicWaker::new());
+    let atomic_waker_copy = atomic_waker.clone();
+
+    let returned_pending = Arc::new(AtomicUsize::new(0));
+    let returned_pending_copy = returned_pending.clone();
+
+    let woken = Arc::new(AtomicUsize::new(0));
+    let woken_copy = woken.clone();
+
+    let t = thread::spawn(move || {
+        let mut pending_count = 0;
+
+        block_on(poll_fn(move |lw| {
+            if woken_copy.load(Ordering::Relaxed) == 1 {
+                Poll::Ready(())
+            } else {
+                // Assert we return pending exactly once
+                assert_eq!(0, pending_count);
+                pending_count += 1;
+                atomic_waker_copy.register(lw);
+
+                returned_pending_copy.store(1, Ordering::Relaxed);
+
+                Poll::Pending
+            }
+        }))
+    });
+
+    while returned_pending.load(Ordering::Relaxed) == 0 {}
+
+    // give spawned thread some time to sleep in `block_on`
+    thread::yield_now();
+
+    woken.store(1, Ordering::Relaxed);
+    atomic_waker.wake();
+
+    t.join().unwrap();
+}