Add a loom aware block_on for running futures to completion

src/executor.rs

@@ -0,0 +1,80 @@
+//! An async executor that is loom aware.
+//!
+//! Used to test the correctness of async code under loom.
+
+use crate::sync::{Condvar, Mutex};
+use std::mem;
+use std::sync::Arc;
+use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
+
+/// Runs a future to completion in a loom aware fashion. Can be used to test the correctness
+/// of a `Future` implementation.
+pub fn block_on<F: std::future::Future>(mut f: F) -> F::Output {
+    let mut f = unsafe { std::pin::Pin::new_unchecked(&mut f) };
+    let loom_waker = Arc::new(LoomWaker::default());
+    let waker = {
+        let loom_waker_ptr = Arc::into_raw(loom_waker.clone());
+        let raw_waker = RawWaker::new(loom_waker_ptr as *const _, &VTABLE);
+        unsafe { Waker::from_raw(raw_waker) }
+    };
+    let mut cx = Context::from_waker(&waker);
+
+    loop {
+        match f.as_mut().poll(&mut cx) {
+            Poll::Pending => loom_waker.park(),
+            Poll::Ready(value) => break value,
+        }
+    }
+}
+
+struct LoomWaker(Mutex<bool>, Condvar);
+
+// FIXME: Can be replaced with #[derive(Default)] when #138 is merged.
+impl Default for LoomWaker {
+    fn default() -> Self {
+        Self(Mutex::new(false), Condvar::new())
+    }
+}
+
+impl LoomWaker {
+    /// Used by the `Waker` to unpark the thread blocked in `block_on`.
+    pub fn unpark(&self) {
+        *self.0.lock().unwrap() = true;
+        self.1.notify_one();
+    }
+
+    /// Used in `block_on` to park the thread until the `Waker` wakes it up.
+    pub fn park(&self) {
+        let mut wake = self.0.lock().unwrap();
+        while !*wake {
+            wake = self.1.wait(wake).unwrap();
+        }
+        *wake = false;
+    }
+}
+
+static VTABLE: RawWakerVTable = RawWakerVTable::new(
+    raw_waker_clone,
+    raw_waker_wake,
+    raw_waker_wake_by_ref,
+    raw_waker_drop,
+);
+
+unsafe fn raw_waker_clone(waker_ptr: *const ()) -> RawWaker {
+    let waker = Arc::from_raw(waker_ptr as *const LoomWaker);
+    mem::forget(waker.clone());
+    mem::forget(waker);
+    RawWaker::new(waker_ptr, &VTABLE)
+}
+
+unsafe fn raw_waker_wake(waker_ptr: *const ()) {
+    Arc::from_raw(waker_ptr as *const LoomWaker).unpark()
+}
+
+unsafe fn raw_waker_wake_by_ref(waker_ptr: *const ()) {
+    (&*(waker_ptr as *const LoomWaker)).unpark()
+}
+
+unsafe fn raw_waker_drop(waker_ptr: *const ()) {
+    Arc::from_raw(waker_ptr as *const LoomWaker);
+}

src/lib.rs

@@ -150,6 +150,7 @@ mod rt;
 
 pub mod alloc;
 pub mod cell;
+pub mod executor;
 pub mod lazy_static;
 pub mod model;
 pub mod sync;

tests/executor.rs

@@ -0,0 +1,91 @@
+use loom::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use loom::sync::Arc;
+use loom::thread;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+#[test]
+fn block_on_empty_async_block() {
+    loom::model(|| {
+        let _nothing: () = loom::executor::block_on(async {});
+    })
+}
+
+#[test]
+fn block_on_simple_value() {
+    loom::model(|| {
+        let i: u128 = loom::executor::block_on(async { 95u128 });
+        assert_eq!(i, 95);
+    })
+}
+
+#[test]
+fn block_on_futures_returning_pending() {
+    struct Delay<T> {
+        value: Option<T>,
+        thread_spawned: bool,
+        done: Arc<AtomicBool>,
+    }
+
+    impl<T> Unpin for Delay<T> {}
+
+    impl<T> Delay<T> {
+        pub fn new(value: T) -> Self {
+            Self {
+                value: Some(value),
+                thread_spawned: false,
+                done: Arc::new(AtomicBool::new(false)),
+            }
+        }
+    }
+
+    impl<T> std::future::Future for Delay<T> {
+        type Output = T;
+
+        fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
+            if !self.thread_spawned {
+                self.thread_spawned = true;
+                let done = self.done.clone();
+                let waker = cx.waker().clone();
+                thread::spawn(move || {
+                    done.store(true, Ordering::SeqCst);
+                    waker.wake();
+                });
+            }
+            if self.done.load(Ordering::SeqCst) {
+                Poll::Ready(self.value.take().unwrap())
+            } else {
+                Poll::Pending
+            }
+        }
+    }
+
+    loom::model(|| {
+        let i: u128 = loom::executor::block_on(async {
+            let a = Delay::new(5u128).await;
+            let b = Delay::new(6u128).await;
+            a + b
+        });
+        assert_eq!(i, 11);
+    })
+}
+
+#[test]
+#[should_panic]
+fn buggy_concurrent_inc_future() {
+    loom::model(|| {
+        let num = Arc::new(AtomicUsize::new(0));
+        let num_clone = num.clone();
+
+        let t = thread::spawn(move || {
+            loom::executor::block_on(async {
+                let curr = num_clone.load(Ordering::Acquire);
+                num_clone.store(curr + 1, Ordering::Release);
+            })
+        });
+        num.fetch_add(1, Ordering::Relaxed);
+        t.join().unwrap();
+
+        assert_eq!(2, num.load(Ordering::Relaxed));
+    })
+}