src/liballoc/task.rs

   1 #![unstable(feature = "wake_trait", issue = "69912")]
   2 //! Types and Traits for working with asynchronous tasks.
   3 use core::mem::{self, 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         let waker: Arc<W> = Arc::from_raw(waker as *const W);
  64         mem::forget(Arc::clone(&waker));
  65         raw_waker(waker)
  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: Arc<W> = 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: ManuallyDrop<Arc<W>> = 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         mem::drop(Arc::from_raw(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 }