Run un-readiness check in separate task (#464)

ntex-service/CHANGES.md

@@ -1,5 +1,9 @@
 # Changes
 
+## [3.3.3] - 2024-11-10
+
+* Add Pipeline::is_shutdown() helper
+
 ## [3.3.2] - 2024-11-10
 
 * Fix un-needed wakeups for unready future

ntex-service/Cargo.toml

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

ntex-service/src/ctx.rs

@@ -12,6 +12,7 @@ pub struct ServiceCtx<'a, S: ?Sized> {
 #[derive(Debug)]
 pub(crate) struct WaitersRef {
     cur: cell::Cell<u32>,
+    shutdown: cell::Cell<bool>,
     wakers: cell::UnsafeCell<Vec<u32>>,
     indexes: cell::UnsafeCell<slab::Slab<Option<Waker>>>,
 }
@@ -19,10 +20,15 @@ pub(crate) struct WaitersRef {
 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 {
                 cur: cell::Cell::new(u32::MAX),
+                shutdown: cell::Cell::new(false),
                 indexes: cell::UnsafeCell::new(waiters),
                 wakers: cell::UnsafeCell::new(Vec::default()),
             },
@@ -62,6 +68,18 @@ 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() {
@@ -90,6 +108,14 @@ impl WaitersRef {
             false
         }
     }
+
+    pub(crate) fn shutdown(&self) {
+        self.shutdown.set(true);
+    }
+
+    pub(crate) fn is_shutdown(&self) -> bool {
+        self.shutdown.get()
+    }
 }
 
 impl<'a, S> ServiceCtx<'a, S> {

ntex-service/src/pipeline.rs

@@ -110,6 +110,12 @@ impl<S> Pipeline<S> {
         }
     }
 
+    #[inline]
+    /// Check if shutdown is initiated.
+    pub fn is_shutdown(&self) -> bool {
+        self.state.waiters.is_shutdown()
+    }
+
     #[inline]
     /// Shutdown enclosed service.
     pub async fn shutdown<R>(&self)
@@ -202,6 +208,12 @@ where
         &self.pl.state.svc
     }
 
+    #[inline]
+    /// Get pipeline
+    pub fn pipeline(&self) -> Pipeline<S> {
+        self.pl.clone()
+    }
+
     #[inline]
     /// Returns `Ready` when the pipeline is able to process requests.
     ///
@@ -263,6 +275,8 @@ where
                 // `self` is alive
                 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),
@@ -300,6 +314,12 @@ where
         }
     }
 
+    #[inline]
+    /// Check if shutdown is initiated.
+    pub fn is_shutdown(&self) -> bool {
+        self.pl.state.waiters.is_shutdown()
+    }
+
     #[inline]
     /// Shutdown enclosed service.
     pub async fn shutdown(&self) {
@@ -449,14 +469,14 @@ struct CheckUnReadiness<S: 'static, F, Fut> {
 
 impl<S, F, Fut> Unpin for CheckUnReadiness<S, F, Fut> {}
 
-impl<T, S, F, Fut> Future 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 = T>,
+    Fut: Future<Output = ()>,
 {
-    type Output = T;
+    type Output = ();
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
         let mut slf = self.as_mut();
 
         if slf.fut.is_none() {
@@ -464,11 +484,18 @@ where
         }
         let fut = slf.fut.as_mut().unwrap();
         match unsafe { Pin::new_unchecked(fut) }.poll(cx) {
-            Poll::Pending => Poll::Pending,
-            Poll::Ready(res) => {
+            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(res)
+                Poll::Ready(())
             }
         }
     }