ntex_util/services/

keepalive.rs

use std::{cell::Cell, convert::Infallible, fmt, task::Poll, time};

use ntex_service::{Ctx, Service, ServiceFactory};

use crate::time::{Millis, Sleep, now, sleep};

/// Middleware that fails readiness once the service has been idle too long.
///
/// Every call records the current time and passes the request through
/// unchanged. When no call has been made for the keep-alive duration,
/// [`ready`](Service::ready) fails with the error produced by `f`.
pub struct KeepAlive<F, E>
where
    F: Fn() -> E + Clone,
{
    f: F,
    ka: Millis,
}

impl<F, E> KeepAlive<F, E>
where
    F: Fn() -> E + Clone,
{
    /// Creates keep-alive middleware.
    ///
    /// `ka` is the maximum idle time between calls, and `f` creates the error
    /// returned once it is exceeded.
    pub fn new(ka: Millis, f: F) -> Self {
        KeepAlive { f, ka }
    }
}

impl<F, E> Clone for KeepAlive<F, E>
where
    F: Fn() -> E + Clone,
{
    fn clone(&self) -> Self {
        KeepAlive {
            f: self.f.clone(),
            ka: self.ka,
        }
    }
}

impl<F, E> fmt::Debug for KeepAlive<F, E>
where
    F: Fn() -> E + Clone,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("KeepAlive")
            .field("ka", &self.ka)
            .field("f", &std::any::type_name::<F>())
            .finish()
    }
}

impl<F, E, St, Req> ServiceFactory<St, Req> for KeepAlive<F, E>
where
    F: Fn() -> E + Clone,
{
    type Res = Req;
    type Error = E;

    type Service = KeepAliveService<F, E>;
    type InitError = Infallible;

    #[inline]
    async fn create(&self, _: &St) -> Result<Self::Service, Self::InitError> {
        Ok(KeepAliveService::new(self.ka, self.f.clone()))
    }
}

/// A service that fails readiness once it has been idle too long.
///
/// See [`KeepAlive`] for details.
pub struct KeepAliveService<F, E>
where
    F: Fn() -> E,
{
    f: F,
    dur: time::Duration,
    sleep: Sleep,
    expire: Cell<time::Instant>,
}

impl<F, E> KeepAliveService<F, E>
where
    F: Fn() -> E,
{
    /// Creates a keep-alive service with the maximum idle time `dur`.
    ///
    /// `f` creates the error returned once the service has been idle longer
    /// than `dur`.
    pub fn new(dur: Millis, f: F) -> Self {
        let expire = Cell::new(now());

        KeepAliveService {
            f,
            expire,
            sleep: sleep(dur),
            dur: time::Duration::from(dur),
        }
    }
}

impl<F, E> fmt::Debug for KeepAliveService<F, E>
where
    F: Fn() -> E,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("KeepAliveService")
            .field("dur", &self.dur)
            .field("expire", &self.expire)
            .field("f", &std::any::type_name::<F>())
            .finish()
    }
}

impl<F, E, St, Req> Service<St, Req> for KeepAliveService<F, E>
where
    F: Fn() -> E,
{
    type Res = Req;
    type Error = E;

    async fn ready(&self, ctx: Ctx<'_, Self, St>) -> Result<(), Self::Error> {
        let expire = self.expire.get() + self.dur;
        if expire <= now() {
            Err((self.f)())
        } else {
            ctx.poll_once(|cx| {
                loop {
                    match self.sleep.poll_elapsed(cx) {
                        Poll::Ready(()) => {
                            let now = now();
                            let expire = self.expire.get() + self.dur;
                            if expire <= now {
                                return Err((self.f)());
                            }
                            let expire = expire - now;

                            // sleep must be reset to non zero duration,
                            // otherwise it stays in elapsed state and waker
                            // never gets registered
                            let expire: u32 = expire.as_millis().try_into().unwrap_or(u32::MAX);
                            self.sleep.reset(Millis(expire.max(1)));
                        }
                        Poll::Pending => return Ok(()),
                    }
                }
            })
        }
    }

    #[inline]
    async fn call(&self, req: Req, _: Ctx<'_, Self, St>) -> Result<Req, E> {
        self.expire.set(now());
        Ok(req)
    }
}

#[cfg(test)]
mod tests {
    use std::{pin::Pin, task::Context, task::Poll, task::ready};

    use ntex_service::{Pipeline, boxed, factory};

    use super::*;
    use crate::{channel::oneshot, spawn};

    #[derive(Debug, PartialEq)]
    struct TestErr;

    struct Dispatcher {
        p: Pipeline<usize, usize, TestErr>,
        tx: Option<oneshot::Sender<()>>,
    }

    impl Future for Dispatcher {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            let mut this = self.as_mut();
            if ready!(this.p.poll_ready(cx)).is_err() {
                if let Some(tx) = this.tx.take() {
                    let _ = tx.send(());
                }
                Poll::Ready(())
            } else {
                Poll::Pending
            }
        }
    }

    #[ntex::test]
    async fn test_ka() {
        // the keep-alive is much longer than the sleep between calls, so a
        // delayed wakeup does not expire it
        let factory = factory::<_, (), usize>(KeepAlive::new(Millis(500), || TestErr));
        assert!(format!("{factory:?}").contains("KeepAlive"));
        let _ = factory.clone();

        let svc = factory.create(&()).await.unwrap();
        assert!(format!("{svc:?}").contains("KeepAliveService"));

        let p = Pipeline::new((), boxed::service(svc));
        assert_eq!(p.call(1usize).await, Ok(1usize));
        let svc = p.bind();

        let (tx, rx) = oneshot::channel();
        spawn(Dispatcher { p, tx: Some(tx) }).detach();

        sleep(Millis(25)).await;
        assert_eq!(svc.call(1usize).await, Ok(1usize));
        sleep(Millis(100)).await;

        let res = rx.await;
        assert_eq!(res, Ok(()));
        assert_eq!(svc.ready().await, Err(TestErr));
    }

    #[ntex::test]
    async fn test_ka_sub_millis() {
        let svc = std::rc::Rc::new(KeepAliveService::new(Millis(100), || TestErr));

        // less than millisecond is left before expiration
        svc.expire
            .set(now().checked_sub(svc.dur).unwrap() + time::Duration::from_micros(500));
        svc.sleep.elapse();

        let p = Pipeline::<usize, usize, TestErr>::new((), svc.clone()).bind();
        assert_eq!(p.ready().await, Ok(()));

        // timer has to be re-armed, otherwise waker never gets registered
        // and service readiness never resolves
        assert!(!svc.sleep.is_elapsed());
    }
}