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 #[cfg_attr(bootstrap, allow(rustc::ineffective_unstable_trait_impl))]
  89 #[cfg_attr(not(bootstrap), allow(ineffective_unstable_trait_impl))]
  90 #[stable(feature = "wake_trait", since = "1.51.0")]
  91 impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for Waker {
  92     fn from(waker: Arc<W>) -> Waker {
  93         // SAFETY: This is safe because raw_waker safely constructs
  94         // a RawWaker from Arc<W>.
  95         unsafe { Waker::from_raw(raw_waker(waker)) }
  96     }
  97 }
  98
  99 #[cfg_attr(bootstrap, allow(rustc::ineffective_unstable_trait_impl))]
 100 #[cfg_attr(not(bootstrap), allow(ineffective_unstable_trait_impl))]
 101 #[stable(feature = "wake_trait", since = "1.51.0")]
 102 impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for RawWaker {
 103     fn from(waker: Arc<W>) -> RawWaker {
 104         raw_waker(waker)
 105     }
 106 }
 107
 108 // NB: This private function for constructing a RawWaker is used, rather than
 109 // inlining this into the `From<Arc<W>> for RawWaker` impl, to ensure that
 110 // the safety of `From<Arc<W>> for Waker` does not depend on the correct
 111 // trait dispatch - instead both impls call this function directly and
 112 // explicitly.
 113 #[inline(always)]
 114 fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker {
 115     // Increment the reference count of the arc to clone it.
 116     unsafe fn clone_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) -> RawWaker {
 117         unsafe { Arc::increment_strong_count(waker as *const W) };
 118         RawWaker::new(
 119             waker as *const (),
 120             &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
 121         )
 122     }
 123
 124     // Wake by value, moving the Arc into the Wake::wake function
 125     unsafe fn wake<W: Wake + Send + Sync + 'static>(waker: *const ()) {
 126         let waker = unsafe { Arc::from_raw(waker as *const W) };
 127         <W as Wake>::wake(waker);
 128     }
 129
 130     // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it
 131     unsafe fn wake_by_ref<W: Wake + Send + Sync + 'static>(waker: *const ()) {
 132         let waker = unsafe { ManuallyDrop::new(Arc::from_raw(waker as *const W)) };
 133         <W as Wake>::wake_by_ref(&waker);
 134     }
 135
 136     // Decrement the reference count of the Arc on drop
 137     unsafe fn drop_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) {
 138         unsafe { Arc::decrement_strong_count(waker as *const W) };
 139     }
 140
 141     RawWaker::new(
 142         Arc::into_raw(waker) as *const (),
 143         &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
 144     )
 145 }