Add waker_fn (#18)

* Add waker_fn

* Add a waker_fn test

* Double sleep times

* More benches

* Prohibit recursive block_on calls

* Reformat code

Author: Stjepan Glavina

Cargo.toml

@@ -23,3 +23,4 @@ crossbeam = "0.7.3"
 crossbeam-utils = "0.7.0"
 futures = "0.3.1"
 lazy_static = "1.4.0"
+pin-utils = "0.1.0-alpha.4"

benches/bench.rs renamed to benches/spawn.rs

@@ -0,0 +1,89 @@
+#![feature(test)]
+
+extern crate test;
+
+use std::cell::RefCell;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll, Waker};
+
+use crossbeam::sync::Parker;
+use test::Bencher;
+
+/// Runs a future to completion on the current thread.
+fn block_on<F: Future>(future: F) -> F::Output {
+    // Pin the future on the stack.
+    pin_utils::pin_mut!(future);
+
+    thread_local! {
+        // Parker and waker associated with the current thread.
+        static CACHE: RefCell<(Parker, Waker)> = {
+            let parker = Parker::new();
+            let unparker = parker.unparker().clone();
+            let waker = async_task::waker_fn(move || unparker.unpark());
+            RefCell::new((parker, waker))
+        };
+    }
+
+    CACHE.with(|cache| {
+        // Panic if `block_on()` is called recursively.
+        let (parker, waker) = &mut *cache.try_borrow_mut().ok().expect("recursive `block_on`");
+
+        // Create the task context.
+        let cx = &mut Context::from_waker(&waker);
+
+        // Keep polling the future until completion.
+        loop {
+            match future.as_mut().poll(cx) {
+                Poll::Ready(output) => return output,
+                Poll::Pending => parker.park(),
+            }
+        }
+    })
+}
+
+#[bench]
+fn custom_block_on_0_yields(b: &mut Bencher) {
+    b.iter(|| block_on(Yields(0)));
+}
+
+#[bench]
+fn custom_block_on_10_yields(b: &mut Bencher) {
+    b.iter(|| block_on(Yields(10)));
+}
+
+#[bench]
+fn custom_block_on_50_yields(b: &mut Bencher) {
+    b.iter(|| block_on(Yields(50)));
+}
+
+#[bench]
+fn futures_block_on_0_yields(b: &mut Bencher) {
+    b.iter(|| futures::executor::block_on(Yields(0)));
+}
+
+#[bench]
+fn futures_block_on_10_yields(b: &mut Bencher) {
+    b.iter(|| futures::executor::block_on(Yields(10)));
+}
+
+#[bench]
+fn futures_block_on_50_yields(b: &mut Bencher) {
+    b.iter(|| futures::executor::block_on(Yields(50)));
+}
+
+struct Yields(u32);
+
+impl Future for Yields {
+    type Output = ();
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        if self.0 == 0 {
+            Poll::Ready(())
+        } else {
+            self.0 -= 1;
+            cx.waker().wake_by_ref();
+            Poll::Pending
+        }
+    }
+}

benches/waker_fn.rs

@@ -0,0 +1,60 @@
+//! A simple implementation of `block_on`.
+
+use std::cell::RefCell;
+use std::future::Future;
+use std::task::{Context, Poll, Waker};
+use std::thread;
+use std::time::Duration;
+
+use crossbeam::sync::Parker;
+use futures::channel::oneshot;
+
+/// Runs a future to completion on the current thread.
+fn block_on<F: Future>(future: F) -> F::Output {
+    // Pin the future on the stack.
+    pin_utils::pin_mut!(future);
+
+    thread_local! {
+        // Parker and waker associated with the current thread.
+        static CACHE: RefCell<(Parker, Waker)> = {
+            let parker = Parker::new();
+            let unparker = parker.unparker().clone();
+            let waker = async_task::waker_fn(move || unparker.unpark());
+            RefCell::new((parker, waker))
+        };
+    }
+
+    CACHE.with(|cache| {
+        // Panic if `block_on()` is called recursively.
+        let (parker, waker) = &mut *cache.try_borrow_mut().ok().expect("recursive block_on()");
+
+        // Create the task context.
+        let cx = &mut Context::from_waker(&waker);
+
+        // Keep polling the future until completion.
+        loop {
+            match future.as_mut().poll(cx) {
+                Poll::Ready(output) => return output,
+                Poll::Pending => parker.park(),
+            }
+        }
+    })
+}
+
+fn main() {
+    let (s, r) = oneshot::channel();
+
+    // Spawn a thread that will send a message through the channel.
+    thread::spawn(move || {
+        thread::sleep(Duration::from_secs(1));
+        s.send("Hello, world!").unwrap();
+    });
+
+    // Block until the message is received.
+    let msg = block_on(async {
+        println!("Awaiting...");
+        r.await.unwrap()
+    });
+
+    println!("{}", msg);
+}

examples/block.rs

@@ -1,13 +1,11 @@
 //! Task abstraction for building executors.
 //!
+//! # Spawning
+//!
 //! To spawn a future onto an executor, we first need to allocate it on the heap and keep some
 //! state alongside it. The state indicates whether the future is ready for polling, waiting to be
 //! woken up, or completed. Such a future is called a *task*.
 //!
-//! This crate helps with task allocation and polling its future to completion.
-//!
-//! # Spawning
-//!
 //! All executors have some kind of queue that holds runnable tasks:
 //!
 //! ```
@@ -89,13 +87,31 @@
 //! Task construction incurs a single allocation that holds its state, the schedule function, and
 //! the future or the result of the future if completed.
 //!
-//! The layout of a task is equivalent to 4 words followed by the schedule function, and then by a
-//! union of the future and its output.
+//! The layout of a task is equivalent to 4 `usize`s followed by the schedule function, and then by
+//! a union of the future and its output.
+//!
+//! # Waking
+//!
+//! The handy [`waker_fn`] constructor converts any function into a [`Waker`]. Every time it is
+//! woken, the function gets called:
+//!
+//! ```
+//! let waker = async_task::waker_fn(|| println!("Wake!"));
+//!
+//! // Prints "Wake!" twice.
+//! waker.wake_by_ref();
+//! waker.wake_by_ref();
+//! ```
+//!
+//! This is useful for implementing single-future executors like [`block_on`].
 //!
 //! [`spawn`]: fn.spawn.html
 //! [`spawn_local`]: fn.spawn_local.html
+//! [`waker_fn`]: fn.waker_fn.html
 //! [`Task`]: struct.Task.html
 //! [`JoinHandle`]: struct.JoinHandle.html
+//! [`Waker`]: https://doc.rust-lang.org/std/task/struct.Waker.html
+//! [`block_on`]: https://github.com/async-rs/async-task/blob/master/examples/block.rs
 
 #![no_std]
 #![warn(missing_docs, missing_debug_implementations, rust_2018_idioms)]
@@ -110,6 +126,8 @@ mod raw;
 mod state;
 mod task;
 mod utils;
+mod waker_fn;
 
 pub use crate::join_handle::JoinHandle;
 pub use crate::task::{spawn, spawn_local, Task};
+pub use crate::waker_fn::waker_fn;

src/lib.rs

@@ -15,9 +15,6 @@ use crate::Task;
 
 /// The vtable for a task.
 pub(crate) struct TaskVTable {
-    /// The raw waker vtable.
-    pub(crate) raw_waker_vtable: RawWakerVTable,
-
     /// Schedules the task.
     pub(crate) schedule: unsafe fn(*const ()),
 
@@ -101,6 +98,13 @@ where
     F: Future<Output = R> + 'static,
     S: Fn(Task<T>) + Send + Sync + 'static,
 {
+    const RAW_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
+        Self::clone_waker,
+        Self::wake,
+        Self::wake_by_ref,
+        Self::drop_waker,
+    );
+
     /// Allocates a task with the given `future` and `schedule` function.
     ///
     /// It is assumed that initially only the `Task` reference and the `JoinHandle` exist.
@@ -122,12 +126,6 @@ where
                 state: AtomicUsize::new(SCHEDULED | HANDLE | REFERENCE),
                 awaiter: UnsafeCell::new(None),
                 vtable: &TaskVTable {
-                    raw_waker_vtable: RawWakerVTable::new(
-                        Self::clone_waker,
-                        Self::wake,
-                        Self::wake_by_ref,
-                        Self::drop_waker,
-                    ),
                     schedule: Self::schedule,
                     drop_future: Self::drop_future,
                     get_output: Self::get_output,
@@ -335,7 +333,6 @@ where
     /// Clones a waker.
     unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
         let raw = Self::from_ptr(ptr);
-        let raw_waker_vtable = &(*raw.header).vtable.raw_waker_vtable;
 
         // Increment the reference count. With any kind of reference-counted data structure,
         // relaxed ordering is appropriate when incrementing the counter.
@@ -346,7 +343,7 @@ where
             abort();
         }
 
-        RawWaker::new(ptr, raw_waker_vtable)
+        RawWaker::new(ptr, &Self::RAW_WAKER_VTABLE)
     }
 
     /// Drops a waker.
@@ -464,10 +461,7 @@ where
         let raw = Self::from_ptr(ptr);
 
         // Create a context from the raw task pointer and the vtable inside the its header.
-        let waker = ManuallyDrop::new(Waker::from_raw(RawWaker::new(
-            ptr,
-            &(*raw.header).vtable.raw_waker_vtable,
-        )));
+        let waker = ManuallyDrop::new(Waker::from_raw(RawWaker::new(ptr, &Self::RAW_WAKER_VTABLE)));
         let cx = &mut Context::from_waker(&waker);
 
         let mut state = (*raw.header).state.load(Ordering::Acquire);

src/raw.rs

@@ -0,0 +1,43 @@
+use alloc::sync::Arc;
+use core::mem::{self, ManuallyDrop};
+use core::task::{RawWaker, RawWakerVTable, Waker};
+
+/// Creates a waker from a wake function.
+///
+/// The function gets called every time the waker is woken.
+pub fn waker_fn<F: Fn() + Send + Sync + 'static>(f: F) -> Waker {
+    let raw = Arc::into_raw(Arc::new(f)) as *const ();
+    let vtable = &Helper::<F>::VTABLE;
+    unsafe { Waker::from_raw(RawWaker::new(raw, vtable)) }
+}
+
+struct Helper<F>(F);
+
+impl<F: Fn() + Send + Sync + 'static> Helper<F> {
+    const VTABLE: RawWakerVTable = RawWakerVTable::new(
+        Self::clone_waker,
+        Self::wake,
+        Self::wake_by_ref,
+        Self::drop_waker,
+    );
+
+    unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
+        let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const F));
+        mem::forget(arc.clone());
+        RawWaker::new(ptr, &Self::VTABLE)
+    }
+
+    unsafe fn wake(ptr: *const ()) {
+        let arc = Arc::from_raw(ptr as *const F);
+        (arc)();
+    }
+
+    unsafe fn wake_by_ref(ptr: *const ()) {
+        let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const F));
+        (arc)();
+    }
+
+    unsafe fn drop_waker(ptr: *const ()) {
+        drop(Arc::from_raw(ptr as *const F));
+    }
+}