|
| 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 |
| 16 | +/// used 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 | + let waker: Arc<W> = Arc::from_raw(waker as *const W); |
| 62 | + mem::forget(Arc::clone(&waker)); |
| 63 | + raw_waker(waker) |
| 64 | + } |
| 65 | + |
| 66 | + // Wake by value, moving the Arc into the Wake::wake function |
| 67 | + unsafe fn wake<W: Wake + Send + Sync + 'static>(waker: *const ()) { |
| 68 | + let waker: Arc<W> = Arc::from_raw(waker as *const W); |
| 69 | + <W as Wake>::wake(waker); |
| 70 | + } |
| 71 | + |
| 72 | + // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it |
| 73 | + unsafe fn wake_by_ref<W: Wake + Send + Sync + 'static>(waker: *const ()) { |
| 74 | + let waker: ManuallyDrop<Arc<W>> = ManuallyDrop::new(Arc::from_raw(waker as *const W)); |
| 75 | + <W as Wake>::wake_by_ref(&waker); |
| 76 | + } |
| 77 | + |
| 78 | + // Decrement the reference count of the Arc on drop |
| 79 | + unsafe fn drop_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) { |
| 80 | + mem::drop(Arc::from_raw(waker as *const W)); |
| 81 | + } |
| 82 | + |
| 83 | + RawWaker::new( |
| 84 | + Arc::into_raw(waker) as *const (), |
| 85 | + &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>), |
| 86 | + ) |
| 87 | +} |
0 commit comments