ntex_service/

ctx.rs

use std::{cell, fmt, future, marker, ops, pin, rc::Rc, task::Context, task::Poll, task::Waker};

use crate::Service;

/// Context provided to [`Service`] lifecycle methods.
///
/// A context gives a service access to pipeline state and coordinates calls to
/// inner services with the pipeline's readiness and shutdown machinery.
pub struct Ctx<'a, Svc: ?Sized, St = ()> {
    idx: u32,
    st: &'a St,
    waiters: &'a WaitersRef,
    _t: marker::PhantomData<Rc<Svc>>,
}

#[derive(Debug)]
pub(crate) struct WaitersRef {
    cur: cell::Cell<u32>,
    running: cell::Cell<bool>,
    shutdown: cell::Cell<bool>,
    ready: cell::Cell<bool>,
    wakers: cell::UnsafeCell<Vec<u32>>,
    indexes: cell::UnsafeCell<slab::Slab<Option<Waker>>>,
}

impl WaitersRef {
    pub(crate) fn new() -> Self {
        let mut waiters = slab::Slab::with_capacity(16);
        waiters.insert(None);
        WaitersRef {
            cur: cell::Cell::new(u32::MAX),
            running: cell::Cell::new(false),
            shutdown: cell::Cell::new(false),
            ready: cell::Cell::new(false),
            indexes: cell::UnsafeCell::new(waiters),
            wakers: cell::UnsafeCell::new(Vec::default()),
        }
    }

    #[allow(clippy::mut_from_ref)]
    pub(crate) fn get(&self) -> &mut slab::Slab<Option<Waker>> {
        unsafe { &mut *self.indexes.get() }
    }

    #[allow(clippy::mut_from_ref)]
    pub(crate) fn get_wakers(&self) -> &mut Vec<u32> {
        unsafe { &mut *self.wakers.get() }
    }

    pub(crate) fn insert(&self) -> u32 {
        self.get().insert(None) as u32
    }

    pub(crate) fn remove(&self, idx: u32) {
        self.get().remove(idx as usize);

        if self.cur.get() == idx {
            self.notify();
        }
    }

    pub(crate) fn notify(&self) {
        let wakers = self.get_wakers();
        if !wakers.is_empty() {
            let indexes = self.get();
            for idx in wakers.drain(..) {
                if let Some(item) = indexes.get_mut(idx as usize)
                    && let Some(waker) = item.take()
                {
                    waker.wake();
                }
            }
        }

        self.cur.set(u32::MAX);
    }

    pub(crate) fn run<F, R>(&self, idx: u32, cx: &mut Context<'_>, f: F) -> Poll<R>
    where
        F: FnOnce(&mut Context<'_>) -> Poll<R>,
    {
        // Ctx::poll_xxx() methods requires current waker always available,
        // nested readiness checks poll with the same waker
        let slot = &mut self.get()[idx as usize];
        if !slot.as_ref().is_some_and(|w| w.will_wake(cx.waker())) {
            *slot = Some(cx.waker().clone());
        }

        // calculate owner for readiness check
        let cur = self.cur.get();
        let can_check = if cur == idx {
            true
        } else if cur == u32::MAX {
            self.cur.set(idx);
            true
        } else {
            false
        };

        if can_check {
            // only one readiness check can manage waiters
            let initial_run = !self.running.get();
            if initial_run {
                self.running.set(true);
            }

            let result = f(cx);

            if initial_run {
                if result.is_pending() {
                    self.get_wakers().push(idx);
                } else {
                    self.notify();
                }
                self.running.set(false);
            }
            result
        } else {
            // Another pipeline binding owns the readiness check.
            self.get_wakers().push(idx);
            Poll::Pending
        }
    }

    pub(crate) fn shutdown(&self) {
        self.shutdown.set(true);
        self.ready.set(false);
    }

    /// Records the result of a pipeline readiness check
    pub(crate) fn set_ready(&self, ready: bool) {
        self.ready.set(ready);
    }

    /// Consumes the result of the last successful pipeline readiness check
    pub(crate) fn take_ready(&self) -> bool {
        self.ready.replace(false)
    }

    pub(crate) fn is_shutdown(&self) -> bool {
        self.shutdown.get()
    }
}

impl<'a, Svc, St> Ctx<'a, Svc, St> {
    pub(crate) fn new(idx: u32, waiters: &'a WaitersRef, st: &'a St) -> Self {
        Self {
            idx,
            waiters,
            st,
            _t: marker::PhantomData,
        }
    }

    pub(crate) fn inner(self) -> (u32, &'a WaitersRef, &'a St) {
        (self.idx, self.waiters, self.st)
    }

    #[inline]
    /// Returns the identifier of the current pipeline binding.
    pub fn id(&self) -> u32 {
        self.idx
    }

    #[inline]
    /// Returns the pipeline state.
    pub fn st(&'a self) -> &'a St {
        self.st
    }

    /// Waits until `svc` is ready to process a request.
    pub async fn ready<S, Req>(&self, svc: &'a S) -> Result<(), S::Error>
    where
        S: Service<St, Req>,
    {
        // check readiness and notify waiters
        ReadyCall {
            completed: false,
            fut: svc.ready(Ctx {
                st: self.st,
                idx: self.idx,
                waiters: self.waiters,
                _t: marker::PhantomData,
            }),
            idx: self.idx,
            waiters: self.waiters,
        }
        .await
    }

    #[inline]
    /// Waits for `svc` to become ready, then calls it.
    pub async fn call<S, Req>(&self, svc: &'a S, req: Req) -> Result<S::Res, S::Error>
    where
        S: Service<St, Req>,
    {
        self.ready(svc).await?;

        svc.call(
            req,
            Ctx {
                idx: self.idx,
                st: self.st,
                waiters: self.waiters,
                _t: marker::PhantomData,
            },
        )
        .await
    }

    #[inline]
    /// Calls `svc` without checking readiness.
    ///
    /// The caller must ensure that `svc` is ready.
    pub async fn call_nowait<S, Req>(&self, svc: &'a S, req: Req) -> Result<S::Res, S::Error>
    where
        S: Service<St, Req>,
    {
        svc.call(
            req,
            Ctx {
                st: self.st,
                idx: self.idx,
                waiters: self.waiters,
                _t: marker::PhantomData,
            },
        )
        .await
    }

    #[inline]
    /// Polls a closure until completion using the dispatcher's task [`Context`].
    ///
    /// The [`Context`] uses the waker of the task that currently owns the readiness
    /// check. Outside of a readiness check a no-op waker is used, so a `Pending`
    /// result is not woken. If the current owner is not the main runner, the main
    /// runner is also scheduled for wake-up.
    pub async fn poll_fn<F, R>(&'a self, f: F) -> R
    where
        F: Fn(&mut Context<'a>) -> Poll<R>,
    {
        future::poll_fn(move |_| {
            let wakers = self.waiters.get();
            let idx = self.waiters.cur.get() as usize;
            let mut ctx = if let Some(w) = wakers.get(idx).and_then(|w| w.as_ref()) {
                Context::from_waker(w)
            } else {
                Context::from_waker(Waker::noop())
            };

            // if the current runner is not the main one, wake up the main task
            // so it re-registers all required wakers
            if idx != 0 {
                self.waiters.get_wakers().push(0);
            }

            f(&mut ctx)
        })
        .await
    }

    #[inline]
    /// Calls a closure once with the dispatcher's task [`Context`].
    ///
    /// The [`Context`] uses the waker of the task that currently owns the readiness
    /// check. Outside of a readiness check a no-op waker is used, so a `Pending`
    /// result is not woken. If the current owner is not the main runner, the main
    /// runner is also scheduled for wake-up.
    pub fn poll_once<F, R>(&'a self, f: F) -> R
    where
        F: FnOnce(&mut Context<'a>) -> R,
    {
        let wakers = self.waiters.get();
        let idx = self.waiters.cur.get() as usize;
        let mut ctx = if let Some(w) = wakers.get(idx).and_then(|w| w.as_ref()) {
            Context::from_waker(w)
        } else {
            Context::from_waker(Waker::noop())
        };

        // if the current runner is not the main one, wake up the main task
        // so it re-registers all required wakers
        if idx != 0 {
            self.waiters.get_wakers().push(0);
        }

        f(&mut ctx)
    }

    #[inline]
    /// Shuts down `svc`.
    pub async fn shutdown<S, Req>(&self, svc: &'a S)
    where
        S: Service<St, Req>,
    {
        svc.shutdown(Ctx {
            idx: self.idx,
            st: self.st,
            waiters: self.waiters,
            _t: marker::PhantomData,
        })
        .await;
    }

    #[inline]
    /// Returns a context that uses `st` as its state.
    pub fn map_state<NewSt>(&'a self, st: &'a NewSt) -> Ctx<'a, Self, NewSt> {
        Ctx {
            st,
            idx: self.idx,
            waiters: self.waiters,
            _t: marker::PhantomData,
        }
    }
}

impl<S, St> Copy for Ctx<'_, S, St> {}

impl<S, St> Clone for Ctx<'_, S, St> {
    #[inline]
    fn clone(&self) -> Self {
        *self
    }
}

impl<S, St> ops::Deref for Ctx<'_, S, St> {
    type Target = St;

    #[inline]
    fn deref(&self) -> &St {
        self.st
    }
}

impl<S, St> fmt::Debug for Ctx<'_, S, St> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Ctx")
            .field("idx", &self.idx)
            .field("waiters", &self.waiters.get().len())
            .finish()
    }
}

struct ReadyCall<'a, F: future::Future> {
    completed: bool,
    fut: F,
    idx: u32,
    waiters: &'a WaitersRef,
}

impl<F: future::Future> Drop for ReadyCall<'_, F> {
    fn drop(&mut self) {
        if !self.completed && self.waiters.cur.get() == self.idx {
            self.waiters.notify();
        }
    }
}

impl<F: future::Future> Unpin for ReadyCall<'_, F> {}

impl<F: future::Future> future::Future for ReadyCall<'_, F> {
    type Output = F::Output;

    fn poll(mut self: pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.waiters.run(self.idx, cx, |cx| {
            // SAFETY: `fut` never moves
            let result = unsafe { pin::Pin::new_unchecked(&mut self.as_mut().fut).poll(cx) };
            if result.is_ready() {
                self.completed = true;
            }
            result
        })
    }
}

#[cfg(test)]
mod tests {
    use std::{cell::Cell, cell::RefCell, future::poll_fn};

    use ntex::channel::{condition, oneshot};
    use ntex::{rt::spawn, time, util::lazy, util::select};

    use super::*;
    use crate::Pipeline;

    struct Srv(Rc<Cell<usize>>, condition::Waiter);

    impl Service<(), &'static str> for Srv {
        type Res = &'static str;
        type Error = ();

        async fn ready(&self, _: Ctx<'_, Self>) -> Result<(), Self::Error> {
            self.0.set(self.0.get() + 1);
            self.1.ready().await;
            Ok(())
        }

        async fn call(
            &self,
            req: &'static str,
            ctx: Ctx<'_, Self>,
        ) -> Result<Self::Res, Self::Error> {
            let _ = format!("{ctx:?}");
            let _ = format!("{:?}", ctx.id());
            let () = *ctx;
            #[allow(clippy::clone_on_copy)]
            let _ = ctx.clone();
            Ok(req)
        }
    }

    #[ntex::test]
    async fn test_ready() {
        let cnt = Rc::new(Cell::new(0));
        let con = condition::Condition::new();

        let srv = Pipeline::new((), Srv(cnt.clone(), con.wait()));
        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Pending);
        assert_eq!(cnt.get(), 1);

        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Pending);
        assert_eq!(cnt.get(), 1);

        con.notify(());
        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Ready(Ok(())));
        assert_eq!(cnt.get(), 1);

        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Pending);
        assert_eq!(cnt.get(), 2);

        con.notify(());
        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Ready(Ok(())));
        assert_eq!(cnt.get(), 2);

        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Pending);
        assert_eq!(cnt.get(), 3);
    }

    struct WakeCounter {
        clones: Cell<usize>,
        wakes: Cell<usize>,
    }

    fn counting_waker() -> (&'static WakeCounter, Waker) {
        use std::task::{RawWaker, RawWakerVTable};

        static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake, drop);

        unsafe fn clone(data: *const ()) -> RawWaker {
            let cnt = unsafe { &*data.cast::<WakeCounter>() };
            cnt.clones.set(cnt.clones.get() + 1);
            RawWaker::new(data, &VTABLE)
        }
        unsafe fn wake(data: *const ()) {
            let cnt = unsafe { &*data.cast::<WakeCounter>() };
            cnt.wakes.set(cnt.wakes.get() + 1);
        }
        unsafe fn drop(_: *const ()) {}

        let cnt: &'static WakeCounter = Box::leak(Box::new(WakeCounter {
            clones: Cell::new(0),
            wakes: Cell::new(0),
        }));
        let raw = RawWaker::new(std::ptr::from_ref(cnt).cast(), &VTABLE);
        (cnt, unsafe { Waker::from_raw(raw) })
    }

    /// Wakes the dispatcher's task context during the readiness check
    struct WakeSrv;

    impl Service<(), &'static str> for WakeSrv {
        type Res = &'static str;
        type Error = ();

        async fn ready(&self, ctx: Ctx<'_, Self>) -> Result<(), Self::Error> {
            ctx.poll_once(|cx| cx.waker().wake_by_ref());
            Ok(())
        }

        async fn call(&self, req: &'static str, _: Ctx<'_, Self>) -> Result<&'static str, ()> {
            Ok(req)
        }
    }

    struct Nested<S>(S);

    impl<S: Service<(), &'static str>> Service<(), &'static str> for Nested<S> {
        type Res = S::Res;
        type Error = S::Error;

        async fn ready(&self, ctx: Ctx<'_, Self>) -> Result<(), Self::Error> {
            ctx.ready(&self.0).await
        }

        async fn call(
            &self,
            req: &'static str,
            ctx: Ctx<'_, Self>,
        ) -> Result<Self::Res, Self::Error> {
            ctx.call(&self.0, req).await
        }
    }

    #[ntex::test]
    async fn test_ready_waker_not_recloned() {
        let srv = Pipeline::new((), Nested(Nested(WakeSrv)));

        let (cnt1, waker1) = counting_waker();
        let mut cx = Context::from_waker(&waker1);
        for _ in 0..4 {
            assert_eq!(srv.poll_ready(&mut cx), Poll::Ready(Ok(())));
        }
        assert_eq!(cnt1.clones.get(), 1);
        assert_eq!(cnt1.wakes.get(), 4);

        // a different waker replaces the stored one
        let (cnt2, waker2) = counting_waker();
        let mut cx = Context::from_waker(&waker2);
        assert_eq!(srv.poll_ready(&mut cx), Poll::Ready(Ok(())));
        assert_eq!(cnt2.clones.get(), 1);
        assert_eq!(cnt2.wakes.get(), 1);
        assert_eq!(cnt1.wakes.get(), 4);
    }

    #[ntex::test]
    async fn test_ready_on_drop() {
        let cnt = Rc::new(Cell::new(0));
        let con = condition::Condition::new();
        let srv = Pipeline::new((), Srv(cnt.clone(), con.wait()));

        let srv1 = srv.bind();
        let (tx, rx) = oneshot::channel();
        spawn(async move {
            select(rx, srv1.ready()).await;
            time::sleep(time::Millis(25000)).await;
        });
        time::sleep(time::Millis(250)).await;

        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Pending);

        let _ = tx.send(());
        time::sleep(time::Millis(250)).await;

        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Pending);

        con.notify(());
        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Ready(Ok(())));
    }

    #[ntex::test]
    async fn test_ready_after_shutdown() {
        let cnt = Rc::new(Cell::new(0));
        let con = condition::Condition::new();
        let srv = Pipeline::new((), Srv(cnt.clone(), con.wait()));

        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Pending);

        let (tx, rx) = oneshot::channel();
        let (tx2, rx2) = oneshot::channel();
        spawn(async move {
            select(rx, srv.ready()).await;
            srv.shutdown().await;
            let _ = tx2.send(srv);
        });
        time::sleep(time::Millis(250)).await;

        let _ = tx.send(());
        let srv = rx2.await.unwrap();

        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Ready(Ok(())));

        con.notify(());
        let res = lazy(|cx| srv.poll_ready(cx)).await;
        assert_eq!(res, Poll::Ready(Ok(())));
    }

    #[ntex::test]
    async fn test_pipeline_binding_after_shutdown() {
        let cnt = Rc::new(Cell::new(0));
        let con = condition::Condition::new();
        let srv = Pipeline::new((), Srv(cnt.clone(), con.wait()));
        poll_fn(|cx| srv.poll_shutdown(cx)).await;
        let _ = poll_fn(|cx| srv.poll_ready(cx)).await;
    }

    #[ntex::test]
    async fn test_shared_call() {
        let data = Rc::new(RefCell::new(Vec::new()));

        let cnt = Rc::new(Cell::new(0));
        let con = condition::Condition::new();

        let srv = Pipeline::new((), Srv(cnt.clone(), con.wait()));

        let srv1 = srv.bind();
        let data1 = data.clone();
        ntex::rt::spawn(async move {
            let _ = srv1.ready().await;
            let fut = srv1.call_static("srv1");
            assert!(format!("{fut:?}").contains("PipelineCall"));
            let i = fut.await.unwrap();
            data1.borrow_mut().push(i);
        });

        let srv2 = srv.bind();
        let data2 = data.clone();
        ntex::rt::spawn(async move {
            let i = srv2.call("srv2").await.unwrap();
            data2.borrow_mut().push(i);
        });
        time::sleep(time::Millis(50)).await;

        con.notify(());
        time::sleep(time::Millis(150)).await;

        assert_eq!(cnt.get(), 2);
        assert_eq!(&*data.borrow(), &["srv1"]);

        con.notify(());
        time::sleep(time::Millis(150)).await;

        assert_eq!(cnt.get(), 2);
        assert_eq!(&*data.borrow(), &["srv1", "srv2"]);
    }
}