Add Wake trait for safe construction of Wakers.

Currently, constructing a waker requires calling the unsafe
`Waker::from_raw` API. This API requires the user to manually construct
a vtable for the waker themself - which is both cumbersome and very
error prone. This API would provide an ergonomic, straightforward and
guaranteed memory-safe way of constructing a waker.

It has been our longstanding intention that the `Waker` type essentially
function as an `Arc<dyn Wake>`, with a `Wake` trait as defined here. Two
considerations prevented the original API from being shipped as simply
an `Arc<dyn Wake>`:

- We want to support futures on embedded systems, which may not have an
  allocator, and in optimized executors for which this API may not be
  best-suited. Therefore, we have always explicitly supported the
  maximally-flexible (but also memory-unsafe) `RawWaker` API, and
  `Waker` has always lived in libcore.
- Because `Waker` lives in libcore and `Arc` lives in liballoc, it has
  not been feasible to provide a constructor for `Waker` from `Arc<dyn
  Wake>`.

Therefore, the Wake trait was left out of the initial version of the
task waker API.

However, as Rust 1.41, it is possible under the more flexible orphan
rules to implement `From<Arc<W>> for Waker where W: Wake` in liballoc.
Therefore, we can now define this constructor even though `Waker` lives
in libcore.

This PR adds these APIs:

- A `Wake` trait, which contains two methods
    - A required method `wake`, which is called by `Waker::wake`
    - A provided method `wake_by_ref`, which is called by
      `Waker::wake_by_ref` and which implementors can override if they
      can optimize this use case.
- An implementation of `From<Arc<W>> for Waker where W: Wake + Send +
  Sync + 'static`
- A similar implementation of `From<Arc<W>> for RawWaker`.

3 files changed, 97 insertions, 0 deletions

diff --git a/src/liballoc/lib.rs b/src/liballoc/lib.rs
index d877ac6ac5c..4d9e8cd3460 100644
--- a/src/liballoc/lib.rs
+++ b/src/liballoc/lib.rs
@@ -161,6 +161,7 @@ pub mod str;
 pub mod string;
 #[cfg(target_has_atomic = "ptr")]
 pub mod sync;
+pub mod task;
 #[cfg(test)]
 mod tests;
 pub mod vec;
diff --git a/src/liballoc/task.rs b/src/liballoc/task.rs
new file mode 100644
index 00000000000..8e0466d9a2f
--- /dev/null
+++ b/src/liballoc/task.rs
@@ -0,0 +1,91 @@
+#![unstable(feature = "wake_trait", issue = "0")]
+//! Types and Traits for working with asynchronous tasks.
+use core::mem;
+use core::task::{Waker, RawWaker, RawWakerVTable};
+
+use crate::sync::Arc;
+
+/// The implementation of waking a task on an executor.
+/// 
+/// This trait can be used to create a [`Waker`]. An executor can define an
+/// implementation of this trait, and use that to construct a Waker to pass
+/// to the tasks that are executed on that executor.
+///
+/// This trait is a memory-safe and ergonomic alternative to constructing a
+/// [`RawWaker`]. It supports the common executor design in which the data
+/// used to wake up a task is stored in an [`Arc`]. Some executors (especially
+/// those for embedded systems) cannot use this API, which is way [`RawWaker`]
+/// exists as an alternative for those systems.
+#[unstable(feature = "wake_trait", issue = "0")]
+pub trait Wake {
+    /// Wake this task.
+    #[unstable(feature = "wake_trait", issue = "0")]
+    fn wake(self: Arc<Self>);
+
+    /// Wake this task without consuming the waker.
+    ///
+    /// If an executor supports a cheaper way to wake without consuming the
+    /// waker, it should override this method. By default, it clones the
+    /// [`Arc`] and calls `wake` on the clone.
+    #[unstable(feature = "wake_trait", issue = "0")]
+    fn wake_by_ref(self: &Arc<Self>) {
+        self.clone().wake();
+    }
+}
+
+#[unstable(feature = "wake_trait", issue = "0")]
+impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for Waker {
+    fn from(waker: Arc<W>) -> Waker {
+        unsafe {
+            Waker::from_raw(raw_waker(waker))
+        }
+    }
+}
+
+#[unstable(feature = "wake_trait", issue = "0")]
+impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for RawWaker {
+    fn from(waker: Arc<W>) -> RawWaker {
+        raw_waker(waker)
+    }
+}
+
+// NB: This private function for constructing a RawWaker is used, rather than
+// inlining this into the `From<Arc<W>> for RawWaker` impl, to ensure that
+// the safety of `From<Arc<W>> for Waker` does not depend on the correct
+// trait dispatch - instead both impls call this function directly and
+// explicitly.
+#[inline(always)]
+fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker {
+
+    // Increment the reference count of the arc to clone it.
+    unsafe fn clone_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) -> RawWaker {
+        let waker: Arc<W> = Arc::from_raw(waker as *const W);
+        mem::forget(waker.clone());
+        raw_waker(waker)
+    }
+
+    // Wake by value, moving the Arc into the Wake::wake function
+    unsafe fn wake<W: Wake + Send + Sync + 'static>(waker: *const ()) {
+        let waker: Arc<W> = Arc::from_raw(waker as *const W);
+        Wake::wake(waker);
+    }
+
+    // Wake by reference, forgetting the Arc to avoid decrementing the reference count
+    unsafe fn wake_by_ref<W: Wake + Send + Sync + 'static>(waker: *const ()) {
+        let waker: Arc<W> = Arc::from_raw(waker as *const W);
+        Wake::wake_by_ref(&waker);
+        mem::forget(waker);
+    }
+
+    // Decrement the reference count of the Arc on drop
+    unsafe fn drop_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) {
+        mem::drop(Arc::from_raw(waker as *const W));
+    }
+
+    RawWaker::new(Arc::into_raw(waker) as *const (), &RawWakerVTable::new(
+        clone_waker::<W>,
+        wake::<W>,
+        wake_by_ref::<W>,
+        drop_waker::<W>,
+    ))
+}
diff --git a/src/libstd/lib.rs b/src/libstd/lib.rs
index 1f122b02b6a..ac6701e6119 100644
--- a/src/libstd/lib.rs
+++ b/src/libstd/lib.rs
@@ -463,9 +463,14 @@ pub mod time;
 #[stable(feature = "futures_api", since = "1.36.0")]
 pub mod task {
     //! Types and Traits for working with asynchronous tasks.
+
     #[doc(inline)]
     #[stable(feature = "futures_api", since = "1.36.0")]
     pub use core::task::*;
+
+    #[doc(inline)]
+    #[unstable(feature = "wake_trait", issue = "0")]
+    pub use alloc::task::*;
 }
 
 #[stable(feature = "futures_api", since = "1.36.0")]