#
local_waker/
lib.rs

    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

    //! A synchronization primitive for thread-local task wakeup.
    //!
    //! See docs for [`LocalWaker`].

    #![no_std]
    #![deny(rust_2018_idioms, nonstandard_style)]
    #![warn(future_incompatible, missing_docs)]

    use core::{cell::Cell, fmt, marker::PhantomData, task::Waker};

    /// A synchronization primitive for task wakeup.
    ///
    /// Sometimes the task interested in a given event will change over time. A `LocalWaker` can
    /// coordinate concurrent notifications with the consumer, potentially "updating" the underlying
    /// task to wake up. This is useful in scenarios where a computation completes in another task and
    /// wants to notify the consumer, but the consumer is in the process of being migrated to a new
    /// logical task.
    ///
    /// Consumers should call [`register`] before checking the result of a computation and producers
    /// should call [`wake`] after producing the computation (this differs from the usual `thread::park`
    /// pattern). It is also permitted for [`wake`] to be called _before_ [`register`]. This results in
    /// a no-op.
    ///
    /// A single `LocalWaker` may be reused for any number of calls to [`register`] or [`wake`].
    ///
    /// [`register`]: LocalWaker::register
    /// [`wake`]: LocalWaker::wake
    #[derive(Default)]
    pub struct LocalWaker {
        pub(crate) waker: Cell<Option<Waker>>,
        // mark LocalWaker as a !Send type.
        _phantom: PhantomData<*const ()>,
    }

    impl LocalWaker {
        /// Creates a new, empty `LocalWaker`.
        pub fn new() -> Self {
            LocalWaker::default()
        }

        /// Registers the waker to be notified on calls to `wake`.
        ///
        /// Returns `true` if waker was registered before.
        #[inline]
        pub fn register(&self, waker: &Waker) -> bool {
            let last_waker = self.waker.replace(Some(waker.clone()));
            last_waker.is_some()
        }

        /// Calls `wake` on the last `Waker` passed to `register`.
        ///
        /// If `register` has not been called yet, then this does nothing.
        #[inline]
        pub fn wake(&self) {
            if let Some(waker) = self.take() {
                waker.wake();
            }
        }

        /// Returns the last `Waker` passed to `register`, so that the user can wake it.
        ///
        /// If a waker has not been registered, this returns `None`.
        #[inline]
        pub fn take(&self) -> Option<Waker> {
            self.waker.take()
        }
    }

    impl fmt::Debug for LocalWaker {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "LocalWaker")
        }
    }