1 files changed, 147 insertions, 0 deletions
diff --git a/library/alloc/src/task.rs b/library/alloc/src/task.rs
new file mode 100644
index 000000000..528ee4ff1
--- /dev/null
+++ b/library/alloc/src/task.rs
@@ -0,0 +1,147 @@
+#![stable(feature = "wake_trait", since = "1.51.0")]
+//! Types and Traits for working with asynchronous tasks.
+use core::mem::ManuallyDrop;
+use core::task::{RawWaker, RawWakerVTable, Waker};
+
+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 why [`RawWaker`]
+/// exists as an alternative for those systems.
+///
+/// [arc]: ../../std/sync/struct.Arc.html
+///
+/// # Examples
+///
+/// A basic `block_on` function that takes a future and runs it to completion on
+/// the current thread.
+///
+/// **Note:** This example trades correctness for simplicity. In order to prevent
+/// deadlocks, production-grade implementations will also need to handle
+/// intermediate calls to `thread::unpark` as well as nested invocations.
+///
+/// ```rust
+/// use std::future::Future;
+/// use std::sync::Arc;
+/// use std::task::{Context, Poll, Wake};
+/// use std::thread::{self, Thread};
+///
+/// /// A waker that wakes up the current thread when called.
+/// struct ThreadWaker(Thread);
+///
+/// impl Wake for ThreadWaker {
+///     fn wake(self: Arc<Self>) {
+///         self.0.unpark();
+///     }
+/// }
+///
+/// /// Run a future to completion on the current thread.
+/// fn block_on<T>(fut: impl Future<Output = T>) -> T {
+///     // Pin the future so it can be polled.
+///     let mut fut = Box::pin(fut);
+///
+///     // Create a new context to be passed to the future.
+///     let t = thread::current();
+///     let waker = Arc::new(ThreadWaker(t)).into();
+///     let mut cx = Context::from_waker(&waker);
+///
+///     // Run the future to completion.
+///     loop {
+///         match fut.as_mut().poll(&mut cx) {
+///             Poll::Ready(res) => return res,
+///             Poll::Pending => thread::park(),
+///         }
+///     }
+/// }
+///
+/// block_on(async {
+///     println!("Hi from inside a future!");
+/// });
+/// ```
+#[stable(feature = "wake_trait", since = "1.51.0")]
+pub trait Wake {
+    /// Wake this task.
+    #[stable(feature = "wake_trait", since = "1.51.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.
+    ///
+    /// [`wake`]: Wake::wake
+    #[stable(feature = "wake_trait", since = "1.51.0")]
+    fn wake_by_ref(self: &Arc<Self>) {
+        self.clone().wake();
+    }
+}
+
+#[stable(feature = "wake_trait", since = "1.51.0")]
+impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for Waker {
+    /// Use a `Wake`-able type as a `Waker`.
+    ///
+    /// No heap allocations or atomic operations are used for this conversion.
+    fn from(waker: Arc<W>) -> Waker {
+        // SAFETY: This is safe because raw_waker safely constructs
+        // a RawWaker from Arc<W>.
+        unsafe { Waker::from_raw(raw_waker(waker)) }
+    }
+}
+
+#[stable(feature = "wake_trait", since = "1.51.0")]
+impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for RawWaker {
+    /// Use a `Wake`-able type as a `RawWaker`.
+    ///
+    /// No heap allocations or atomic operations are used for this conversion.
+    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 {
+        unsafe { Arc::increment_strong_count(waker as *const W) };
+        RawWaker::new(
+            waker as *const (),
+            &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
+        )
+    }
+
+    // Wake by value, moving the Arc into the Wake::wake function
+    unsafe fn wake<W: Wake + Send + Sync + 'static>(waker: *const ()) {
+        let waker = unsafe { Arc::from_raw(waker as *const W) };
+        <W as Wake>::wake(waker);
+    }
+
+    // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it
+    unsafe fn wake_by_ref<W: Wake + Send + Sync + 'static>(waker: *const ()) {
+        let waker = unsafe { ManuallyDrop::new(Arc::from_raw(waker as *const W)) };
+        <W as Wake>::wake_by_ref(&waker);
+    }
+
+    // Decrement the reference count of the Arc on drop
+    unsafe fn drop_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) {
+        unsafe { Arc::decrement_strong_count(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>),
+    )
+}