library/alloc/src/task.rs

   1 #![stable(feature = "wake_trait", since = "1.51.0")]
   2
   3 //! Types and Traits for working with asynchronous tasks.
   4 //!
   5 //! **Note**: This module is only available on platforms that support atomic
   6 //! loads and stores of pointers. This may be detected at compile time using
   7 //! `#[cfg(target_has_atomic = "ptr")]`.
   8
   9 use core::mem::ManuallyDrop;
  10 use core::task::{RawWaker, RawWakerVTable, Waker};
  11
  12 use crate::sync::Arc;
  13
  14 /// The implementation of waking a task on an executor.
  15 ///
  16 /// This trait can be used to create a [`Waker`]. An executor can define an
  17 /// implementation of this trait, and use that to construct a Waker to pass
  18 /// to the tasks that are executed on that executor.
  19 ///
  20 /// This trait is a memory-safe and ergonomic alternative to constructing a
  21 /// [`RawWaker`]. It supports the common executor design in which the data used
  22 /// to wake up a task is stored in an [`Arc`]. Some executors (especially
  23 /// those for embedded systems) cannot use this API, which is why [`RawWaker`]
  24 /// exists as an alternative for those systems.
  25 ///
  26 /// [arc]: ../../std/sync/struct.Arc.html
  27 ///
  28 /// # Examples
  29 ///
  30 /// A basic `block_on` function that takes a future and runs it to completion on
  31 /// the current thread.
  32 ///
  33 /// **Note:** This example trades correctness for simplicity. In order to prevent
  34 /// deadlocks, production-grade implementations will also need to handle
  35 /// intermediate calls to `thread::unpark` as well as nested invocations.
  36 ///
  37 /// ```rust
  38 /// use std::future::Future;
  39 /// use std::sync::Arc;
  40 /// use std::task::{Context, Poll, Wake};
  41 /// use std::thread::{self, Thread};
  42 ///
  43 /// /// A waker that wakes up the current thread when called.
  44 /// struct ThreadWaker(Thread);
  45 ///
  46 /// impl Wake for ThreadWaker {
  47 ///     fn wake(self: Arc<Self>) {
  48 ///         self.0.unpark();
  49 ///     }
  50 /// }
  51 ///
  52 /// /// Run a future to completion on the current thread.
  53 /// fn block_on<T>(fut: impl Future<Output = T>) -> T {
  54 ///     // Pin the future so it can be polled.
  55 ///     let mut fut = Box::pin(fut);
  56 ///
  57 ///     // Create a new context to be passed to the future.
  58 ///     let t = thread::current();
  59 ///     let waker = Arc::new(ThreadWaker(t)).into();
  60 ///     let mut cx = Context::from_waker(&waker);
  61 ///
  62 ///     // Run the future to completion.
  63 ///     loop {
  64 ///         match fut.as_mut().poll(&mut cx) {
  65 ///             Poll::Ready(res) => return res,
  66 ///             Poll::Pending => thread::park(),
  67 ///         }
  68 ///     }
  69 /// }
  70 ///
  71 /// block_on(async {
  72 ///     println!("Hi from inside a future!");
  73 /// });
  74 /// ```
  75 #[stable(feature = "wake_trait", since = "1.51.0")]
  76 pub trait Wake {
  77     /// Wake this task.
  78     #[stable(feature = "wake_trait", since = "1.51.0")]
  79     fn wake(self: Arc<Self>);
  80
  81     /// Wake this task without consuming the waker.
  82     ///
  83     /// If an executor supports a cheaper way to wake without consuming the
  84     /// waker, it should override this method. By default, it clones the
  85     /// [`Arc`] and calls [`wake`] on the clone.
  86     ///
  87     /// [`wake`]: Wake::wake
  88     #[stable(feature = "wake_trait", since = "1.51.0")]
  89     fn wake_by_ref(self: &Arc<Self>) {
  90         self.clone().wake();
  91     }
  92 }
  93
  94 #[stable(feature = "wake_trait", since = "1.51.0")]
  95 impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for Waker {
  96     /// Use a `Wake`-able type as a `Waker`.
  97     ///
  98     /// No heap allocations or atomic operations are used for this conversion.
  99     fn from(waker: Arc<W>) -> Waker {
 100         // SAFETY: This is safe because raw_waker safely constructs
 101         // a RawWaker from Arc<W>.
 102         unsafe { Waker::from_raw(raw_waker(waker)) }
 103     }
 104 }
 105
 106 #[stable(feature = "wake_trait", since = "1.51.0")]
 107 impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for RawWaker {
 108     /// Use a `Wake`-able type as a `RawWaker`.
 109     ///
 110     /// No heap allocations or atomic operations are used for this conversion.
 111     fn from(waker: Arc<W>) -> RawWaker {
 112         raw_waker(waker)
 113     }
 114 }
 115
 116 // NB: This private function for constructing a RawWaker is used, rather than
 117 // inlining this into the `From<Arc<W>> for RawWaker` impl, to ensure that
 118 // the safety of `From<Arc<W>> for Waker` does not depend on the correct
 119 // trait dispatch - instead both impls call this function directly and
 120 // explicitly.
 121 #[inline(always)]
 122 fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker {
 123     // Increment the reference count of the arc to clone it.
 124     unsafe fn clone_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) -> RawWaker {
 125         unsafe { Arc::increment_strong_count(waker as *const W) };
 126         RawWaker::new(
 127             waker as *const (),
 128             &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
 129         )
 130     }
 131
 132     // Wake by value, moving the Arc into the Wake::wake function
 133     unsafe fn wake<W: Wake + Send + Sync + 'static>(waker: *const ()) {
 134         let waker = unsafe { Arc::from_raw(waker as *const W) };
 135         <W as Wake>::wake(waker);
 136     }
 137
 138     // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it
 139     unsafe fn wake_by_ref<W: Wake + Send + Sync + 'static>(waker: *const ()) {
 140         let waker = unsafe { ManuallyDrop::new(Arc::from_raw(waker as *const W)) };
 141         <W as Wake>::wake_by_ref(&waker);
 142     }
 143
 144     // Decrement the reference count of the Arc on drop
 145     unsafe fn drop_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) {
 146         unsafe { Arc::decrement_strong_count(waker as *const W) };
 147     }
 148
 149     RawWaker::new(
 150         Arc::into_raw(waker) as *const (),
 151         &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
 152     )
 153 }