diff --git a/src/liballoc/lib.rs b/src/liballoc/lib.rs index d877ac6ac5c79..d55a1a3b63584 100644 --- a/src/liballoc/lib.rs +++ b/src/liballoc/lib.rs @@ -161,6 +161,8 @@ pub mod str; pub mod string; #[cfg(target_has_atomic = "ptr")] pub mod sync; +#[cfg(target_has_atomic = "ptr")] +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 0000000000000..981095302c730 --- /dev/null +++ b/src/liballoc/task.rs @@ -0,0 +1,87 @@ +#![unstable(feature = "wake_trait", issue = "69912")] +//! Types and Traits for working with asynchronous tasks. +use core::mem::{self, 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. +#[unstable(feature = "wake_trait", issue = "69912")] +pub trait Wake { + /// Wake this task. + #[unstable(feature = "wake_trait", issue = "69912")] + fn wake(self: Arc); + + /// 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 = "69912")] + fn wake_by_ref(self: &Arc) { + self.clone().wake(); + } +} + +#[unstable(feature = "wake_trait", issue = "69912")] +impl From> for Waker { + fn from(waker: Arc) -> Waker { + // SAFETY: This is safe because raw_waker safely constructs + // a RawWaker from Arc. + unsafe { Waker::from_raw(raw_waker(waker)) } + } +} + +#[unstable(feature = "wake_trait", issue = "69912")] +impl From> for RawWaker { + fn from(waker: Arc) -> RawWaker { + raw_waker(waker) + } +} + +// NB: This private function for constructing a RawWaker is used, rather than +// inlining this into the `From> for RawWaker` impl, to ensure that +// the safety of `From> for Waker` does not depend on the correct +// trait dispatch - instead both impls call this function directly and +// explicitly. +#[inline(always)] +fn raw_waker(waker: Arc) -> RawWaker { + // Increment the reference count of the arc to clone it. + unsafe fn clone_waker(waker: *const ()) -> RawWaker { + let waker: Arc = Arc::from_raw(waker as *const W); + mem::forget(Arc::clone(&waker)); + raw_waker(waker) + } + + // Wake by value, moving the Arc into the Wake::wake function + unsafe fn wake(waker: *const ()) { + let waker: Arc = Arc::from_raw(waker as *const W); + ::wake(waker); + } + + // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it + unsafe fn wake_by_ref(waker: *const ()) { + let waker: ManuallyDrop> = ManuallyDrop::new(Arc::from_raw(waker as *const W)); + ::wake_by_ref(&waker); + } + + // Decrement the reference count of the Arc on drop + unsafe fn drop_waker(waker: *const ()) { + mem::drop(Arc::from_raw(waker as *const W)); + } + + RawWaker::new( + Arc::into_raw(waker) as *const (), + &RawWakerVTable::new(clone_waker::, wake::, wake_by_ref::, drop_waker::), + ) +} diff --git a/src/libstd/lib.rs b/src/libstd/lib.rs index 1f122b02b6a29..e5dad307a209a 100644 --- a/src/libstd/lib.rs +++ b/src/libstd/lib.rs @@ -310,6 +310,7 @@ #![feature(untagged_unions)] #![feature(unwind_attributes)] #![feature(vec_into_raw_parts)] +#![feature(wake_trait)] // NB: the above list is sorted to minimize merge conflicts. #![default_lib_allocator] @@ -463,9 +464,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 = "69912")] + pub use alloc::task::*; } #[stable(feature = "futures_api", since = "1.36.0")]