util: Add task::Wake

portable-atomic-util/README.md

@@ -9,6 +9,7 @@
 Synchronization primitives built with [portable-atomic].
 
 - Provide `Arc`. (optional, requires the `std` or `alloc` feature)
+- Provide `task::Wake`. (optional, requires the `std` or `alloc` feature)
 <!-- - Provide generic `Atomic<T>` type. (optional, requires the `generic` feature) -->
 
 See [#1] for other primitives being considered for addition to this crate.

portable-atomic-util/build.rs

@@ -33,6 +33,10 @@ fn main() {
     if !version.probe(36, 2019, 4, 14) {
         println!("cargo:rustc-cfg=portable_atomic_no_alloc");
     }
+    // std::{future,task} stabilized in Rust 1.36 (nightly-2019-04-25) https://github.com/rust-lang/rust/pull/59739
+    if !version.probe(36, 2019, 4, 24) {
+        println!("cargo:rustc-cfg=portable_atomic_no_futures_api");
+    }
     // Layout::{align_to,pad_to_align,extend,array} stabilized in Rust 1.44 (nightly-2020-04-22) https://github.com/rust-lang/rust/pull/69362
     if !version.probe(44, 2020, 4, 21) {
         println!("cargo:rustc-cfg=portable_atomic_no_alloc_layout_extras");

portable-atomic-util/src/arc.rs

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
-// The module is based on alloc::sync::Arc, licensed under "Apache-2.0 OR MIT".
+// This module is based on alloc::sync::Arc.
 //
 // The code has been adjusted to work with stable Rust and to
 // avoid UBs (https://github.com/rust-lang/rust/issues/119241).
@@ -77,10 +77,10 @@ macro_rules! acquire {
 
 /// A thread-safe, strongly reference counted pointer.
 ///
-/// This is an equivalent to [`std::sync::Arc`], but using [`portable-atomic`] for synchronization.
+/// This is an equivalent to [`std::sync::Arc`], but using [portable-atomic] for synchronization.
 /// See the documentation for [`std::sync::Arc`] for more details.
 ///
-/// [`portable-atomic`]: https://crates.io/crates/portable-atomic
+/// [portable-atomic]: https://crates.io/crates/portable-atomic
 ///
 /// # Examples
 ///
@@ -126,10 +126,10 @@ impl<T: ?Sized> Arc<T> {
 
 /// A weakly reference counted pointer.
 ///
-/// This is an equivalent to [`std::sync::Weak`], but using [`portable-atomic`] for synchronization.
+/// This is an equivalent to [`std::sync::Weak`], but using [portable-atomic] for synchronization.
 /// See the documentation for [`std::sync::Weak`] for more details.
 ///
-/// [`portable-atomic`]: https://crates.io/crates/portable-atomic
+/// [portable-atomic]: https://crates.io/crates/portable-atomic
 ///
 /// # Examples
 ///
@@ -2389,10 +2389,6 @@ impl<T: std::error::Error + ?Sized> std::error::Error for Arc<T> {
 //   - https://doc.rust-lang.org/nightly/alloc/sync/struct.Arc.html#impl-From%3C%26CStr%3E-for-Arc%3CCStr%3E
 //   - https://doc.rust-lang.org/nightly/alloc/sync/struct.Arc.html#impl-From%3CCString%3E-for-Arc%3CCStr%3E
 //   - Currently, we cannot implement these since CStr layout is not stable.
-// - alloc::task
-//   - https://doc.rust-lang.org/nightly/alloc/sync/struct.Arc.html#impl-From%3CArc%3CW%3E%3E-for-RawWaker
-//   - https://doc.rust-lang.org/nightly/alloc/sync/struct.Arc.html#impl-From%3CArc%3CW%3E%3E-for-Waker
-//   - https://doc.rust-lang.org/nightly/alloc/task/trait.Wake.html
 // - std::os
 //   - https://doc.rust-lang.org/nightly/std/sync/struct.Arc.html#impl-AsFd-for-Arc%3CT%3E
 //   - https://doc.rust-lang.org/nightly/std/sync/struct.Arc.html#impl-AsHandle-for-Arc%3CT%3E

portable-atomic-util/src/lib.rs

@@ -5,6 +5,7 @@
 Synchronization primitives built with [portable-atomic].
 
 - Provide `Arc`. (optional, requires the `std` or `alloc` feature)
+- Provide `task::Wake`. (optional, requires the `std` or `alloc` feature)
 <!-- - Provide generic `Atomic<T>` type. (optional, requires the `generic` feature) -->
 
 See [#1] for other primitives being considered for addition to this crate.
@@ -73,3 +74,8 @@ extern crate std as alloc;
 mod arc;
 #[cfg(any(all(feature = "alloc", not(portable_atomic_no_alloc)), feature = "std"))]
 pub use arc::{Arc, Weak};
+
+#[cfg(not(portable_atomic_no_futures_api))]
+#[cfg(any(all(feature = "alloc", not(portable_atomic_no_alloc)), feature = "std"))]
+#[cfg_attr(portable_atomic_doc_cfg, doc(cfg(any(feature = "alloc", feature = "std"))))]
+pub mod task;

portable-atomic-util/src/task.rs

@@ -0,0 +1,156 @@
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+//! Types and Traits for working with asynchronous tasks.
+
+// This module is based on alloc::task::Wake.
+//
+// Source: https://github.com/rust-lang/rust/blob/5151b8c42712c473e7da56e213926b929d0212ef/library/alloc/src/task.rs.
+//
+// Copyright & License of the original code:
+// - https://github.com/rust-lang/rust/blob/5151b8c42712c473e7da56e213926b929d0212ef/COPYRIGHT
+// - https://github.com/rust-lang/rust/blob/5151b8c42712c473e7da56e213926b929d0212ef/LICENSE-APACHE
+// - https://github.com/rust-lang/rust/blob/5151b8c42712c473e7da56e213926b929d0212ef/LICENSE-MIT
+
+use core::{
+    mem::ManuallyDrop,
+    task::{RawWaker, RawWakerVTable, Waker},
+};
+
+use crate::Arc;
+
+/// The implementation of waking a task on an executor.
+///
+/// This is an equivalent to [`std::task::Wake`], but using [`portable_atomic_util::Arc`](crate::Arc)
+/// as a reference-counted pointer. See the documentation for [`std::task::Wake`] for more details.
+///
+/// **Note:** Unlike `std::task::Wake`, all methods take `this:` instead of `self:`.
+/// This is because using `portable_atomic_util::Arc` as a receiver requires the
+/// [unstable `arbitrary_self_types` feature](https://github.com/rust-lang/rust/issues/44874).
+///
+/// # 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.
+///
+/// ```
+/// use portable_atomic_util::{task::Wake, Arc};
+/// use std::{
+///     future::Future,
+///     task::{Context, Poll},
+///     thread::{self, Thread},
+/// };
+///
+/// /// A waker that wakes up the current thread when called.
+/// struct ThreadWaker(Thread);
+///
+/// impl Wake for ThreadWaker {
+///     fn wake(this: Arc<Self>) {
+///         this.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!");
+/// });
+/// ```
+pub trait Wake {
+    /// Wake this task.
+    fn wake(this: 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
+    fn wake_by_ref(this: &Arc<Self>) {
+        Self::wake(this.clone());
+    }
+}
+
+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>) -> Self {
+        // SAFETY: This is safe because raw_waker safely constructs
+        // a RawWaker from Arc<W>.
+        unsafe { Self::from_raw(raw_waker(waker)) }
+    }
+}
+
+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>) -> Self {
+        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 {
+        // SAFETY: the caller must uphold the safety contract.
+        unsafe { Arc::increment_strong_count(waker as *const W) };
+        RawWaker::new(
+            waker,
+            &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 ()) {
+        // SAFETY: the caller must uphold the safety contract.
+        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 ()) {
+        // SAFETY: the caller must uphold the safety contract.
+        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 ()) {
+        // SAFETY: the caller must uphold the safety contract.
+        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>),
+    )
+}