library/alloc/src/task.rs

   1 #![unstable(feature = "wake_trait", issue = "69912")]
   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 #[unstable(feature = "wake_trait", issue = "69912")]
  20 pub trait Wake {
  21     /// Wake this task.
  22     #[unstable(feature = "wake_trait", issue = "69912")]
  23     fn wake(self: Arc<Self>);
  24
  25     /// Wake this task without consuming the waker.
  26     ///
  27     /// If an executor supports a cheaper way to wake without consuming the
  28     /// waker, it should override this method. By default, it clones the
  29     /// [`Arc`] and calls `wake` on the clone.
  30     #[unstable(feature = "wake_trait", issue = "69912")]
  31     fn wake_by_ref(self: &Arc<Self>) {
  32         self.clone().wake();
  33     }
  34 }
  35
  36 #[unstable(feature = "wake_trait", issue = "69912")]
  37 impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for Waker {
  38     fn from(waker: Arc<W>) -> Waker {
  39         // SAFETY: This is safe because raw_waker safely constructs
  40         // a RawWaker from Arc<W>.
  41         unsafe { Waker::from_raw(raw_waker(waker)) }
  42     }
  43 }
  44
  45 #[unstable(feature = "wake_trait", issue = "69912")]
  46 impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for RawWaker {
  47     fn from(waker: Arc<W>) -> RawWaker {
  48         raw_waker(waker)
  49     }
  50 }
  51
  52 // NB: This private function for constructing a RawWaker is used, rather than
  53 // inlining this into the `From<Arc<W>> for RawWaker` impl, to ensure that
  54 // the safety of `From<Arc<W>> for Waker` does not depend on the correct
  55 // trait dispatch - instead both impls call this function directly and
  56 // explicitly.
  57 #[inline(always)]
  58 fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker {
  59     // Increment the reference count of the arc to clone it.
  60     unsafe fn clone_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) -> RawWaker {
  61         unsafe { Arc::incr_strong_count(waker as *const W) };
  62         RawWaker::new(
  63             waker as *const (),
  64             &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
  65         )
  66     }
  67
  68     // Wake by value, moving the Arc into the Wake::wake function
  69     unsafe fn wake<W: Wake + Send + Sync + 'static>(waker: *const ()) {
  70         let waker = unsafe { Arc::from_raw(waker as *const W) };
  71         <W as Wake>::wake(waker);
  72     }
  73
  74     // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it
  75     unsafe fn wake_by_ref<W: Wake + Send + Sync + 'static>(waker: *const ()) {
  76         let waker = unsafe { ManuallyDrop::new(Arc::from_raw(waker as *const W)) };
  77         <W as Wake>::wake_by_ref(&waker);
  78     }
  79
  80     // Decrement the reference count of the Arc on drop
  81     unsafe fn drop_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) {
  82         unsafe { Arc::decr_strong_count(waker as *const W) };
  83     }
  84
  85     RawWaker::new(
  86         Arc::into_raw(waker) as *const (),
  87         &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
  88     )
  89 }