diff options
Diffstat (limited to 'library/alloc/src/task.rs')
-rw-r--r-- | library/alloc/src/task.rs | 147 |
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>), + ) +} |