Add waker_fn (#18)
* Add waker_fn
* Add a waker_fn test
* Double sleep times
* More benches
* Prohibit recursive block_on calls
* Reformat code
diff --git a/src/lib.rs b/src/lib.rs
index 5fe858a..1bd7ca8 100644
--- a/src/lib.rs
+++ b/src/lib.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 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;
diff --git a/src/raw.rs b/src/raw.rs
index 1440e7e..ed3ee97 100644
--- a/src/raw.rs
+++ b/src/raw.rs
@@ -15,9 +15,6 @@
/// 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 @@
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 @@
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 @@
/// 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 @@
abort();
}
- RawWaker::new(ptr, raw_waker_vtable)
+ RawWaker::new(ptr, &Self::RAW_WAKER_VTABLE)
}
/// Drops a waker.
@@ -464,10 +461,7 @@
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);
diff --git a/src/waker_fn.rs b/src/waker_fn.rs
new file mode 100644
index 0000000..37105f1
--- /dev/null
+++ b/src/waker_fn.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));
+ }
+}