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