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); + + /// 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.clone().wake(); + } +} + +#[unstable(feature = "wake_trait", issue = "0")] +impl From> for Waker { + fn from(waker: Arc) -> Waker { + unsafe { + Waker::from_raw(raw_waker(waker)) + } + } +} + +#[unstable(feature = "wake_trait", issue = "0")] +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(waker.clone()); + 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::wake(waker); + } + + // Wake by reference, forgetting the Arc to avoid decrementing the reference count + unsafe fn wake_by_ref(waker: *const ()) { + let waker: Arc = 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(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 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")]