      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145

    #![stable(feature = "wake_trait", since = "1.51.0")]
    //! Types and Traits for working with asynchronous tasks.
    use core::mem::ManuallyDrop;
    use core::task::{RawWaker, RawWakerVTable, Waker};

    use crate::sync::Arc;

    /// The implementation of waking a task on an executor.
    ///
    /// This trait can be used to create a [`Waker`]. An executor can define an
    /// implementation of this trait, and use that to construct a Waker to pass
    /// to the tasks that are executed on that executor.
    ///
    /// This trait is a memory-safe and ergonomic alternative to constructing a
    /// [`RawWaker`]. It supports the common executor design in which the data used
    /// to wake up a task is stored in an [`Arc`]. Some executors (especially
    /// those for embedded systems) cannot use this API, which is why [`RawWaker`]
    /// exists as an alternative for those systems.
    ///
    /// [arc]: ../../std/sync/struct.Arc.html
    ///
    /// # Examples
    ///
    /// A basic `block_on` function that takes a future and runs it to completion on
    /// the current thread.
    ///
    /// **Note:** This example trades correctness for simplicity. In order to prevent
    /// deadlocks, production-grade implementations will also need to handle
    /// intermediate calls to `thread::unpark` as well as nested invocations.
    ///
    /// ```rust
    /// use std::future::Future;
    /// use std::sync::Arc;
    /// use std::task::{Context, Poll, Wake};
    /// use std::thread::{self, Thread};
    ///
    /// /// A waker that wakes up the current thread when called.
    /// struct ThreadWaker(Thread);
    ///
    /// impl Wake for ThreadWaker {
    ///     fn wake(self: Arc<Self>) {
    ///         self.0.unpark();
    ///     }
    /// }
    ///
    /// /// Run a future to completion on the current thread.
    /// fn block_on<T>(fut: impl Future<Output = T>) -> T {
    ///     // Pin the future so it can be polled.
    ///     let mut fut = Box::pin(fut);
    ///
    ///     // Create a new context to be passed to the future.
    ///     let t = thread::current();
    ///     let waker = Arc::new(ThreadWaker(t)).into();
    ///     let mut cx = Context::from_waker(&waker);
    ///
    ///     // Run the future to completion.
    ///     loop {
    ///         match fut.as_mut().poll(&mut cx) {
    ///             Poll::Ready(res) => return res,
    ///             Poll::Pending => thread::park(),
    ///         }
    ///     }
    /// }
    ///
    /// block_on(async {
    ///     println!("Hi from inside a future!");
    /// });
    /// ```
    #[stable(feature = "wake_trait", since = "1.51.0")]
    pub trait Wake {
        /// Wake this task.
        #[stable(feature = "wake_trait", since = "1.51.0")]
        fn wake(self: Arc<Self>);

        /// Wake this task without consuming the waker.
        ///
        /// If an executor supports a cheaper way to wake without consuming the
        /// waker, it should override this method. By default, it clones the
        /// [`Arc`] and calls [`wake`] on the clone.
        ///
        /// [`wake`]: Wake::wake
        #[stable(feature = "wake_trait", since = "1.51.0")]
        fn wake_by_ref(self: &Arc<Self>) {
            self.clone().wake();
        }
    }

    #[cfg_attr(bootstrap, allow(rustc::ineffective_unstable_trait_impl))]
    #[cfg_attr(not(bootstrap), allow(ineffective_unstable_trait_impl))]
    #[stable(feature = "wake_trait", since = "1.51.0")]
    impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for Waker {
        fn from(waker: Arc<W>) -> Waker {
            // SAFETY: This is safe because raw_waker safely constructs
            // a RawWaker from Arc<W>.
            unsafe { Waker::from_raw(raw_waker(waker)) }
        }
    }

    #[cfg_attr(bootstrap, allow(rustc::ineffective_unstable_trait_impl))]
    #[cfg_attr(not(bootstrap), allow(ineffective_unstable_trait_impl))]
    #[stable(feature = "wake_trait", since = "1.51.0")]
    impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for RawWaker {
        fn from(waker: Arc<W>) -> RawWaker {
            raw_waker(waker)
        }
    }

    // NB: This private function for constructing a RawWaker is used, rather than
    // inlining this into the `From<Arc<W>> for RawWaker` impl, to ensure that
    // the safety of `From<Arc<W>> for Waker` does not depend on the correct
    // trait dispatch - instead both impls call this function directly and
    // explicitly.
    #[inline(always)]
    fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker {
        // Increment the reference count of the arc to clone it.
        unsafe fn clone_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) -> RawWaker {
            unsafe { Arc::increment_strong_count(waker as *const W) };
            RawWaker::new(
                waker as *const (),
                &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
            )
        }

        // Wake by value, moving the Arc into the Wake::wake function
        unsafe fn wake<W: Wake + Send + Sync + 'static>(waker: *const ()) {
            let waker = unsafe { Arc::from_raw(waker as *const W) };
            <W as Wake>::wake(waker);
        }

        // Wake by reference, wrap the waker in ManuallyDrop to avoid dropping it
        unsafe fn wake_by_ref<W: Wake + Send + Sync + 'static>(waker: *const ()) {
            let waker = unsafe { ManuallyDrop::new(Arc::from_raw(waker as *const W)) };
            <W as Wake>::wake_by_ref(&waker);
        }

        // Decrement the reference count of the Arc on drop
        unsafe fn drop_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) {
            unsafe { Arc::decrement_strong_count(waker as *const W) };
        }

        RawWaker::new(
            Arc::into_raw(waker) as *const (),
            &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
        )
    }