ntex-util 4.2.0

//! Task wake-up and scheduling utilities.
use std::{cell::Cell, fmt, marker::PhantomData, rc, task::Waker};

/// A single-threaded synchronization primitive for task wake-up.
///
/// Sometimes the task interested in a given event will change over time.
/// A `LocalWaker` can coordinate 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`.
#[derive(Default)]
pub struct LocalWaker {
    waker: Cell<Option<Waker>>,
    _t: PhantomData<rc::Rc<()>>,
}

impl LocalWaker {
    /// Creates an empty `LocalWaker`.
    pub fn new() -> Self {
        LocalWaker::with(None)
    }

    /// Creates a `LocalWaker` with an optional initial waker.
    pub fn with(waker: Option<Waker>) -> Self {
        LocalWaker {
            waker: Cell::new(waker),
            _t: PhantomData,
        }
    }

    #[inline]
    /// Registers the waker to be notified on calls to `wake`.
    ///
    /// Returns `true` if a waker was already registered.
    pub fn register(&self, waker: &Waker) -> bool {
        match self.waker.take() {
            Some(prev) if prev.will_wake(waker) => {
                self.waker.set(Some(prev));
                true
            }
            prev => {
                self.waker.set(Some(waker.clone()));
                prev.is_some()
            }
        }
    }

    #[inline]
    /// Calls `wake` on the last `Waker` passed to `register`.
    ///
    /// The waker is consumed, so a later `wake` does nothing until `register`
    /// is called again. If no waker is registered, this does nothing.
    pub fn wake(&self) {
        if let Some(waker) = self.take() {
            waker.wake();
        }
    }

    #[inline]
    /// Calls `wake` on the last `Waker` passed to `register`.
    ///
    /// Returns `false` if no waker is registered. Like [`wake`](Self::wake),
    /// this consumes the registered waker.
    pub fn wake_checked(&self) -> bool {
        if let Some(waker) = self.take() {
            waker.wake();
            true
        } else {
            false
        }
    }

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

    #[doc(hidden)]
    /// Returns `true` if a waker is registered.
    pub fn is_set(&self) -> bool {
        let waker = self.waker.take();
        let set = waker.is_some();
        self.waker.set(waker);
        set
    }
}

/// Cloning creates an empty `LocalWaker`; the registered waker is not copied.
impl Clone for LocalWaker {
    fn clone(&self) -> Self {
        LocalWaker::new()
    }
}

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

/// Yields execution back to the current runtime.
pub async fn yield_to() {
    use std::{future::Future, pin::Pin, task::Context, task::Poll};

    struct Yield {
        completed: bool,
    }

    impl Future for Yield {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
            if self.completed {
                return Poll::Ready(());
            }

            self.completed = true;
            cx.waker().wake_by_ref();

            Poll::Pending
        }
    }

    Yield { completed: false }.await;
}

#[cfg(test)]
mod test {
    use super::*;

    #[ntex::test]
    async fn yield_test() {
        yield_to().await;
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::task::{RawWaker, RawWakerVTable};

    use super::*;

    static CLONES: AtomicUsize = AtomicUsize::new(0);

    static VTABLE: RawWakerVTable = RawWakerVTable::new(
        |p| {
            CLONES.fetch_add(1, Ordering::Relaxed);
            RawWaker::new(p, &VTABLE)
        },
        |_| {},
        |_| {},
        |_| {},
    );

    #[test]
    fn test_register_same_waker() {
        static A: u8 = 0;
        static B: u8 = 0;
        let a = unsafe { Waker::from_raw(RawWaker::new((&raw const A).cast(), &VTABLE)) };
        let b = unsafe { Waker::from_raw(RawWaker::new((&raw const B).cast(), &VTABLE)) };

        let w = LocalWaker::new();
        assert!(!w.register(&a));
        assert!(w.register(&a));
        assert!(w.register(&a));
        assert_eq!(CLONES.load(Ordering::Relaxed), 1);

        assert!(w.register(&b));
        assert_eq!(CLONES.load(Ordering::Relaxed), 2);
        assert!(w.take().unwrap().will_wake(&b));
    }

    #[test]
    fn local_waker_clone_is_empty() {
        let waker = LocalWaker::new();
        waker.register(std::task::Waker::noop());
        assert!(waker.is_set());
        let cloned = waker.clone();
        assert!(!cloned.is_set());
        assert!(!cloned.wake_checked());
        assert_eq!(format!("{waker:?}"), "LocalWaker");
    }
}