library/alloc/src/task.rs

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