Added Service::poll() method (#481)

ntex-service/CHANGES.md

@@ -1,5 +1,9 @@
 # Changes
 
+## [3.4.0] - 2024-12-04
+
+* Added Service::poll() method
+
 ## [3.3.3] - 2024-11-10
 
 * Add Pipeline::is_shutdown() helper

ntex-service/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ntex-service"
-version = "3.3.3"
+version = "3.4.0"
 authors = ["ntex contributors <team@ntex.rs>"]
 description = "ntex service"
 keywords = ["network", "framework", "async", "futures"]

ntex-service/src/and_then.rs

@@ -31,8 +31,9 @@ where
     }
 
     #[inline]
-    async fn not_ready(&self) {
-        util::select(self.svc1.not_ready(), self.svc2.not_ready()).await
+    fn poll(&self, cx: &mut std::task::Context<'_>) -> Result<(), Self::Error> {
+        self.svc1.poll(cx)?;
+        self.svc2.poll(cx)
     }
 
     #[inline]
@@ -88,8 +89,8 @@ where
 
 #[cfg(test)]
 mod tests {
-    use ntex_util::time;
-    use std::{cell::Cell, rc::Rc};
+    use ntex_util::future::lazy;
+    use std::{cell::Cell, rc::Rc, task::Context};
 
     use crate::{chain, chain_factory, fn_factory, Service, ServiceCtx};
 
@@ -105,9 +106,9 @@ mod tests {
             Ok(())
         }
 
-        async fn not_ready(&self) {
+        fn poll(&self, _: &mut Context<'_>) -> Result<(), Self::Error> {
             self.0.set(self.0.get() + 1);
-            std::future::pending().await
+            Ok(())
         }
 
         async fn call(
@@ -135,9 +136,9 @@ mod tests {
             Ok(())
         }
 
-        async fn not_ready(&self) {
+        fn poll(&self, _: &mut Context<'_>) -> Result<(), Self::Error> {
             self.0.set(self.0.get() + 1);
-            std::future::pending().await
+            Ok(())
         }
 
         async fn call(
@@ -165,11 +166,7 @@ mod tests {
         assert_eq!(res, Ok(()));
         assert_eq!(cnt.get(), 2);
 
-        let srv2 = srv.clone();
-        ntex::rt::spawn(async move {
-            let _ = srv2.not_ready().await;
-        });
-        time::sleep(time::Millis(25)).await;
+        lazy(|cx| srv.clone().poll(cx)).await.unwrap();
         assert_eq!(cnt.get(), 4);
 
         srv.shutdown().await;

ntex-service/src/apply.rs

@@ -113,7 +113,7 @@ where
         (self.f)(req, self.service.clone()).await
     }
 
-    crate::forward_notready!(service);
+    crate::forward_poll!(service);
     crate::forward_shutdown!(service);
 }
 
@@ -205,7 +205,8 @@ where
 
 #[cfg(test)]
 mod tests {
-    use std::{cell::Cell, rc::Rc};
+    use ntex_util::future::lazy;
+    use std::{cell::Cell, rc::Rc, task::Context};
 
     use super::*;
     use crate::{chain, chain_factory, fn_factory};
@@ -221,8 +222,9 @@ mod tests {
             Ok(())
         }
 
-        async fn not_ready(&self) {
+        fn poll(&self, _: &mut Context<'_>) -> Result<(), Self::Error> {
             self.0.set(self.0.get() + 1);
+            Ok(())
         }
 
         async fn shutdown(&self) {
@@ -253,7 +255,7 @@ mod tests {
 
         assert_eq!(srv.ready().await, Ok::<_, Err>(()));
 
-        srv.not_ready().await;
+        lazy(|cx| srv.poll(cx)).await.unwrap();
         assert_eq!(cnt_sht.get(), 1);
 
         srv.shutdown().await;

ntex-service/src/boxed.rs

@@ -1,4 +1,4 @@
-use std::{fmt, future::Future, pin::Pin};
+use std::{fmt, future::Future, pin::Pin, task::Context};
 
 use crate::ctx::{ServiceCtx, WaitersRef};
 
@@ -54,8 +54,6 @@ trait ServiceObj<Req> {
         waiters: &'a WaitersRef,
     ) -> BoxFuture<'a, (), Self::Error>;
 
-    fn not_ready<'a>(&'a self) -> Pin<Box<dyn Future<Output = ()> + 'a>>;
-
     fn call<'a>(
         &'a self,
         req: Req,
@@ -64,6 +62,8 @@ trait ServiceObj<Req> {
     ) -> BoxFuture<'a, Self::Response, Self::Error>;
 
     fn shutdown<'a>(&'a self) -> Pin<Box<dyn Future<Output = ()> + 'a>>;
+
+    fn poll(&self, cx: &mut Context<'_>) -> Result<(), Self::Error>;
 }
 
 impl<S, Req> ServiceObj<Req> for S
@@ -83,11 +83,6 @@ where
         Box::pin(async move { ServiceCtx::<'a, S>::new(idx, waiters).ready(self).await })
     }
 
-    #[inline]
-    fn not_ready<'a>(&'a self) -> Pin<Box<dyn Future<Output = ()> + 'a>> {
-        Box::pin(crate::Service::not_ready(self))
-    }
-
     #[inline]
     fn shutdown<'a>(&'a self) -> Pin<Box<dyn Future<Output = ()> + 'a>> {
         Box::pin(crate::Service::shutdown(self))
@@ -106,6 +101,11 @@ where
                 .await
         })
     }
+
+    #[inline]
+    fn poll(&self, cx: &mut Context<'_>) -> Result<(), Self::Error> {
+        crate::Service::poll(self, cx)
+    }
 }
 
 trait ServiceFactoryObj<Req, Cfg> {
@@ -158,11 +158,6 @@ where
         self.0.ready(idx, waiters).await
     }
 
-    #[inline]
-    async fn not_ready(&self) {
-        self.0.not_ready().await
-    }
-
     #[inline]
     async fn shutdown(&self) {
         self.0.shutdown().await
@@ -173,6 +168,11 @@ where
         let (idx, waiters) = ctx.inner();
         self.0.call(req, idx, waiters).await
     }
+
+    #[inline]
+    fn poll(&self, cx: &mut Context<'_>) -> Result<(), Self::Error> {
+        self.0.poll(cx)
+    }
 }
 
 impl<C, Req, Res, Err, InitErr> crate::ServiceFactory<Req, C>

ntex-service/src/chain.rs

@@ -171,6 +171,7 @@ impl<Svc: Service<Req>, Req> Service<Req> for ServiceChain<Svc, Req> {
     type Response = Svc::Response;
     type Error = Svc::Error;
 
+    crate::forward_poll!(service);
     crate::forward_ready!(service);
     crate::forward_shutdown!(service);
 

ntex-service/src/ctx.rs

@@ -21,9 +21,6 @@ impl WaitersRef {
     pub(crate) fn new() -> (u32, Self) {
         let mut waiters = slab::Slab::new();
 
-        // first insert for wake ups from services
-        let _ = waiters.insert(None);
-
         (
             waiters.insert(Default::default()) as u32,
             WaitersRef {
@@ -68,18 +65,6 @@ impl WaitersRef {
         self.get()[idx as usize] = Some(cx.waker().clone());
     }
 
-    pub(crate) fn register_unready(&self, cx: &mut Context<'_>) {
-        self.get()[0] = Some(cx.waker().clone());
-    }
-
-    pub(crate) fn notify_unready(&self) {
-        if let Some(item) = self.get().get_mut(0) {
-            if let Some(waker) = item.take() {
-                waker.wake();
-            }
-        }
-    }
-
     pub(crate) fn notify(&self) {
         let wakers = self.get_wakers();
         if !wakers.is_empty() {

ntex-service/src/lib.rs

@@ -6,7 +6,7 @@
     unreachable_pub,
     missing_debug_implementations
 )]
-use std::rc::Rc;
+use std::{rc::Rc, task::Context};
 
 mod and_then;
 mod apply;
@@ -118,7 +118,8 @@ pub trait Service<Req> {
         Ok(())
     }
 
-    #[inline]
+    #[deprecated]
+    #[doc(hidden)]
     /// Returns when the service is not able to process requests.
     ///
     /// Unlike the "ready()" method, the "not_ready()" method returns
@@ -136,6 +137,15 @@ pub trait Service<Req> {
     /// Returns when the service is properly shutdowns.
     async fn shutdown(&self) {}
 
+    #[inline]
+    /// Polls service from the current task.
+    ///
+    /// Service may require to execute asynchronous computation or
+    /// maintain asynchronous state.
+    fn poll(&self, cx: &mut Context<'_>) -> Result<(), Self::Error> {
+        Ok(())
+    }
+
     #[inline]
     /// Map this service's output to a different type, returning a new service of the resulting type.
     ///
@@ -259,8 +269,8 @@ where
     }
 
     #[inline]
-    async fn not_ready(&self) {
-        (**self).not_ready().await
+    fn poll(&self, cx: &mut Context<'_>) -> Result<(), S::Error> {
+        (**self).poll(cx)
     }
 
     #[inline]
@@ -290,11 +300,6 @@ where
         ctx.ready(&**self).await
     }
 
-    #[inline]
-    async fn not_ready(&self) {
-        (**self).not_ready().await
-    }
-
     #[inline]
     async fn shutdown(&self) {
         (**self).shutdown().await
@@ -308,6 +313,11 @@ where
     ) -> Result<Self::Response, Self::Error> {
         ctx.call_nowait(&**self, request).await
     }
+
+    #[inline]
+    fn poll(&self, cx: &mut Context<'_>) -> Result<(), S::Error> {
+        (**self).poll(cx)
+    }
 }
 
 impl<S, Req, Cfg> ServiceFactory<Req, Cfg> for Rc<S>

ntex-service/src/macros.rs

@@ -11,11 +11,6 @@ macro_rules! forward_ready {
                 .await
                 .map_err(::core::convert::Into::into)
         }
-
-        #[inline]
-        async fn not_ready(&self) {
-            self.$field.not_ready().await
-        }
     };
     ($field:ident, $err:expr) => {
         #[inline]
@@ -25,21 +20,28 @@ macro_rules! forward_ready {
         ) -> Result<(), Self::Error> {
             ctx.ready(&self.$field).await.map_err($err)
         }
-
-        #[inline]
-        async fn not_ready(&self) {
-            self.$field.not_ready().await
-        }
     };
 }
 
 /// An implementation of [`not_ready`] that forwards not_ready call to a field.
 #[macro_export]
 macro_rules! forward_notready {
+    ($field:ident) => {};
+}
+
+/// An implementation of [`poll`] that forwards poll call to a field.
+#[macro_export]
+macro_rules! forward_poll {
     ($field:ident) => {
         #[inline]
-        async fn not_ready(&self) {
-            self.$field.not_ready().await
+        fn poll(&self, cx: &mut std::task::Context<'_>) -> Result<(), Self::Error> {
+            self.$field.poll(cx).map_err(From::from)
+        }
+    };
+    ($field:ident, $err:expr) => {
+        #[inline]
+        fn poll(&self, cx: &mut std::task::Context<'_>) -> Result<(), Self::Error> {
+            self.$field.poll(cx).map_err($err)
         }
     };
 }

ntex-service/src/map.rs

@@ -62,6 +62,7 @@ where
     type Error = A::Error;
 
     crate::forward_ready!(service);
+    crate::forward_poll!(service);
     crate::forward_shutdown!(service);
 
     #[inline]

ntex-service/src/map_err.rs

@@ -1,4 +1,4 @@
-use std::{fmt, marker::PhantomData};
+use std::{fmt, marker::PhantomData, task::Context};
 
 use super::{Service, ServiceCtx, ServiceFactory};
 
@@ -67,6 +67,11 @@ where
         ctx.ready(&self.service).await.map_err(&self.f)
     }
 
+    #[inline]
+    fn poll(&self, cx: &mut Context<'_>) -> Result<(), Self::Error> {
+        self.service.poll(cx).map_err(&self.f)
+    }
+
     #[inline]
     async fn call(
         &self,
@@ -77,7 +82,6 @@ where
     }
 
     crate::forward_shutdown!(service);
-    crate::forward_notready!(service);
 }
 
 /// Factory for the `map_err` combinator, changing the type of a new

ntex-service/src/pipeline.rs

@@ -1,4 +1,4 @@
-use std::{cell, fmt, future::Future, pin::Pin, rc::Rc, task::Context, task::Poll};
+use std::{cell, fmt, future::Future, marker, pin::Pin, rc::Rc, task::Context, task::Poll};
 
 use crate::{ctx::WaitersRef, Service, ServiceCtx};
 
@@ -50,13 +50,14 @@ impl<S> Pipeline<S> {
             .await
     }
 
-    #[inline]
+    #[doc(hidden)]
+    #[deprecated]
     /// Returns when the pipeline is not able to process requests.
     pub async fn not_ready<R>(&self)
     where
         S: Service<R>,
     {
-        self.state.svc.not_ready().await
+        std::future::pending().await
     }
 
     #[inline]
@@ -125,6 +126,14 @@ impl<S> Pipeline<S> {
         self.state.svc.shutdown().await
     }
 
+    #[inline]
+    pub fn poll<R>(&self, cx: &mut Context<'_>) -> Result<(), S::Error>
+    where
+        S: Service<R>,
+    {
+        self.state.svc.poll(cx)
+    }
+
     #[inline]
     /// Get current pipeline.
     pub fn bind<R>(self) -> PipelineBinding<S, R>
@@ -175,7 +184,6 @@ where
 {
     pl: Pipeline<S>,
     st: cell::UnsafeCell<State<S::Error>>,
-    not_ready: cell::UnsafeCell<StateNotReady>,
 }
 
 enum State<E> {
@@ -184,11 +192,6 @@ enum State<E> {
     Shutdown(Pin<Box<dyn Future<Output = ()> + 'static>>),
 }
 
-enum StateNotReady {
-    New,
-    Readiness(Pin<Box<dyn Future<Output = ()>>>),
-}
-
 impl<S, R> PipelineBinding<S, R>
 where
     S: Service<R> + 'static,
@@ -198,7 +201,6 @@ where
         PipelineBinding {
             pl,
             st: cell::UnsafeCell::new(State::New),
-            not_ready: cell::UnsafeCell::new(StateNotReady::New),
         }
     }
 
@@ -214,6 +216,11 @@ where
         self.pl.clone()
     }
 
+    #[inline]
+    pub fn poll(&self, cx: &mut Context<'_>) -> Result<(), S::Error> {
+        self.pl.poll(cx)
+    }
+
     #[inline]
     /// Returns `Ready` when the pipeline is able to process requests.
     ///
@@ -230,6 +237,7 @@ where
                 let fut = Box::pin(CheckReadiness {
                     fut: None,
                     f: ready,
+                    _t: marker::PhantomData,
                     pl,
                 });
                 *st = State::Readiness(fut);
@@ -240,27 +248,12 @@ where
         }
     }
 
+    #[doc(hidden)]
+    #[deprecated]
     #[inline]
     /// Returns when the pipeline is not able to process requests.
-    pub fn poll_not_ready(&self, cx: &mut Context<'_>) -> Poll<()> {
-        let st = unsafe { &mut *self.not_ready.get() };
-
-        match st {
-            StateNotReady::New => {
-                // SAFETY: `fut` has same lifetime same as lifetime of `self.pl`.
-                // Pipeline::svc is heap allocated(Rc<S>), and it is being kept alive until
-                // `self` is alive
-                let pl: &'static Pipeline<S> = unsafe { std::mem::transmute(&self.pl) };
-                let fut = Box::pin(CheckUnReadiness {
-                    fut: None,
-                    f: not_ready,
-                    pl,
-                });
-                *st = StateNotReady::Readiness(fut);
-                self.poll_not_ready(cx)
-            }
-            StateNotReady::Readiness(ref mut fut) => Pin::new(fut).poll(cx),
-        }
+    pub fn poll_not_ready(&self, _: &mut Context<'_>) -> Poll<()> {
+        Poll::Pending
     }
 
     #[inline]
@@ -276,7 +269,6 @@ where
                 let pl: &'static Pipeline<S> = unsafe { std::mem::transmute(&self.pl) };
                 *st = State::Shutdown(Box::pin(async move { pl.shutdown().await }));
                 pl.state.waiters.shutdown();
-                pl.state.waiters.notify_unready();
                 self.poll_shutdown(cx)
             }
             State::Shutdown(ref mut fut) => Pin::new(fut).poll(cx),
@@ -345,7 +337,6 @@ where
         Self {
             pl: self.pl.clone(),
             st: cell::UnsafeCell::new(State::New),
-            not_ready: cell::UnsafeCell::new(StateNotReady::New),
         }
     }
 }
@@ -404,23 +395,16 @@ where
         .ready(ServiceCtx::<'_, S>::new(pl.index, pl.state.waiters_ref()))
 }
 
-fn not_ready<S, R>(pl: &'static Pipeline<S>) -> impl Future<Output = ()>
-where
-    S: Service<R>,
-    R: 'static,
-{
-    pl.state.svc.not_ready()
-}
-
-struct CheckReadiness<S: 'static, F, Fut> {
+struct CheckReadiness<S: Service<R> + 'static, R, F, Fut> {
     f: F,
     fut: Option<Fut>,
     pl: &'static Pipeline<S>,
+    _t: marker::PhantomData<R>,
 }
 
-impl<S, F, Fut> Unpin for CheckReadiness<S, F, Fut> {}
+impl<S: Service<R>, R, F, Fut> Unpin for CheckReadiness<S, R, F, Fut> {}
 
-impl<S, F, Fut> Drop for CheckReadiness<S, F, Fut> {
+impl<S: Service<R>, R, F, Fut> Drop for CheckReadiness<S, R, F, Fut> {
     fn drop(&mut self) {
         // future fot dropped during polling, we must notify other waiters
         if self.fut.is_some() {
@@ -429,16 +413,19 @@ impl<S, F, Fut> Drop for CheckReadiness<S, F, Fut> {
     }
 }
 
-impl<T, S, F, Fut> Future for CheckReadiness<S, F, Fut>
+impl<S, R, F, Fut> Future for CheckReadiness<S, R, F, Fut>
 where
+    S: Service<R>,
     F: Fn(&'static Pipeline<S>) -> Fut,
-    Fut: Future<Output = T>,
+    Fut: Future<Output = Result<(), S::Error>>,
 {
-    type Output = T;
+    type Output = Result<(), S::Error>;
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let mut slf = self.as_mut();
 
+        slf.pl.poll(cx)?;
+
         if slf.pl.state.waiters.can_check(slf.pl.index, cx) {
             if slf.fut.is_none() {
                 slf.fut = Some((slf.f)(slf.pl));
@@ -460,43 +447,3 @@ where
         }
     }
 }
-
-struct CheckUnReadiness<S: 'static, F, Fut> {
-    f: F,
-    fut: Option<Fut>,
-    pl: &'static Pipeline<S>,
-}
-
-impl<S, F, Fut> Unpin for CheckUnReadiness<S, F, Fut> {}
-
-impl<S, F, Fut> Future for CheckUnReadiness<S, F, Fut>
-where
-    F: Fn(&'static Pipeline<S>) -> Fut,
-    Fut: Future<Output = ()>,
-{
-    type Output = ();
-
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
-        let mut slf = self.as_mut();
-
-        if slf.fut.is_none() {
-            slf.fut = Some((slf.f)(slf.pl));
-        }
-        let fut = slf.fut.as_mut().unwrap();
-        match unsafe { Pin::new_unchecked(fut) }.poll(cx) {
-            Poll::Pending => {
-                if slf.pl.state.waiters.is_shutdown() {
-                    Poll::Ready(())
-                } else {
-                    slf.pl.state.waiters.register_unready(cx);
-                    Poll::Pending
-                }
-            }
-            Poll::Ready(()) => {
-                let _ = slf.fut.take();
-                slf.pl.state.waiters.notify();
-                Poll::Ready(())
-            }
-        }
-    }
-}

ntex-service/src/then.rs

@@ -31,8 +31,9 @@ where
     }
 
     #[inline]
-    async fn not_ready(&self) {
-        util::select(self.svc1.not_ready(), self.svc2.not_ready()).await
+    fn poll(&self, cx: &mut std::task::Context<'_>) -> Result<(), Self::Error> {
+        self.svc1.poll(cx)?;
+        self.svc2.poll(cx)
     }
 
     #[inline]
@@ -91,8 +92,8 @@ where
 
 #[cfg(test)]
 mod tests {
-    use ntex_util::time;
-    use std::{cell::Cell, rc::Rc};
+    use ntex_util::future::lazy;
+    use std::{cell::Cell, rc::Rc, task::Context};
 
     use crate::{chain, chain_factory, fn_factory, Service, ServiceCtx};
 
@@ -108,9 +109,9 @@ mod tests {
             Ok(())
         }
 
-        async fn not_ready(&self) {
+        fn poll(&self, _: &mut Context<'_>) -> Result<(), Self::Error> {
             self.0.set(self.0.get() + 1);
-            std::future::pending().await
+            Ok(())
         }
 
         async fn call(
@@ -141,9 +142,9 @@ mod tests {
             Ok(())
         }
 
-        async fn not_ready(&self) {
+        fn poll(&self, _: &mut Context<'_>) -> Result<(), Self::Error> {
             self.0.set(self.0.get() + 1);
-            std::future::pending().await
+            Ok(())
         }
 
         async fn call(
@@ -173,11 +174,7 @@ mod tests {
         assert_eq!(res, Ok(()));
         assert_eq!(cnt.get(), 2);
 
-        let srv2 = srv.clone();
-        ntex::rt::spawn(async move {
-            let _ = srv2.not_ready().await;
-        });
-        time::sleep(time::Millis(25)).await;
+        lazy(|cx| srv.clone().poll(cx)).await.unwrap();
         assert_eq!(cnt.get(), 4);
 
         srv.shutdown().await;

ntex-service/src/util.rs

@@ -59,24 +59,3 @@ where
     })
     .await
 }
-
-/// Waits for either one of two differently-typed futures to complete.
-pub(crate) async fn select<A, B>(fut1: A, fut2: B) -> A::Output
-where
-    A: Future,
-    B: Future<Output = A::Output>,
-{
-    let mut fut1 = pin::pin!(fut1);
-    let mut fut2 = pin::pin!(fut2);
-
-    poll_fn(|cx| {
-        if let Poll::Ready(item) = Pin::new(&mut fut1).poll(cx) {
-            return Poll::Ready(item);
-        }
-        if let Poll::Ready(item) = Pin::new(&mut fut2).poll(cx) {
-            return Poll::Ready(item);
-        }
-        Poll::Pending
-    })
-    .await
-}