[FEATURE] FnShutdown service with on_shutdown callback (#163)

* FnShutdown service with on_shutdown callback

ntex-service/CHANGES.md

@@ -1,5 +1,9 @@
 # Changes
 
+## [1.0.1] - 2023-1-23
+
+* Add `FnShutdown` service to provide on_shutdown callback
+
 ## [1.0.0-beta.0] - 2022-12-28
 
 * Rename Transform to Middleware

ntex-service/src/fn_shutdown.rs

@@ -0,0 +1,94 @@
+use std::cell::Cell;
+use std::future::{ready, Ready};
+use std::marker::PhantomData;
+use std::task::{Context, Poll};
+
+use crate::Service;
+
+#[inline]
+/// Create `FnShutdown` for function that can act as a `on_shutdown` callback.
+pub fn fn_shutdown<Req, Err, F>(f: F) -> FnShutdown<Req, Err, F>
+where
+    F: FnOnce(),
+{
+    FnShutdown::new(f)
+}
+
+pub struct FnShutdown<Req, Err, F> {
+    f_shutdown: Cell<Option<F>>,
+    _t: PhantomData<(Req, Err)>,
+}
+
+impl<Req, Err, F> FnShutdown<Req, Err, F> {
+    pub(crate) fn new(f: F) -> Self {
+        Self {
+            f_shutdown: Cell::new(Some(f)),
+            _t: PhantomData,
+        }
+    }
+}
+
+impl<Req, Err, F> Clone for FnShutdown<Req, Err, F>
+where
+    F: Clone,
+{
+    #[inline]
+    fn clone(&self) -> Self {
+        let f = self.f_shutdown.take();
+        self.f_shutdown.set(f.clone());
+        Self {
+            f_shutdown: Cell::new(f),
+            _t: PhantomData,
+        }
+    }
+}
+
+impl<Req, Err, F> Service<Req> for FnShutdown<Req, Err, F>
+where
+    F: FnOnce(),
+{
+    type Response = Req;
+    type Error = Err;
+    type Future<'f> = Ready<Result<Req, Err>> where Self: 'f, Req: 'f;
+
+    #[inline]
+    fn poll_shutdown(&self, _: &mut Context<'_>) -> Poll<()> {
+        if let Some(f) = self.f_shutdown.take() {
+            (f)()
+        }
+        Poll::Ready(())
+    }
+
+    #[inline]
+    fn call(&self, req: Req) -> Self::Future<'_> {
+        ready(Ok(req))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use ntex_util::future::lazy;
+    use std::task::Poll;
+
+    use crate::{fn_service, pipeline};
+
+    use super::*;
+
+    #[ntex::test]
+    async fn test_fn_shutdown() {
+        let mut is_called = false;
+        let srv = fn_service(|_| async { Ok::<_, ()>("pipe") });
+        let on_shutdown = fn_shutdown(|| {
+            is_called = true;
+        });
+
+        let pipe = pipeline(srv).and_then(on_shutdown);
+
+        let res = pipe.call(()).await;
+        assert_eq!(lazy(|cx| pipe.poll_ready(cx)).await, Poll::Ready(Ok(())));
+        assert!(res.is_ok());
+        assert_eq!(res.unwrap(), "pipe");
+        assert_eq!(lazy(|cx| pipe.poll_shutdown(cx)).await, Poll::Ready(()));
+        assert!(is_called);
+    }
+}

ntex-service/src/lib.rs

@@ -11,6 +11,7 @@ mod and_then;
 mod apply;
 pub mod boxed;
 mod fn_service;
+mod fn_shutdown;
 mod macros;
 mod map;
 mod map_config;
@@ -22,6 +23,7 @@ mod then;
 
 pub use self::apply::{apply_fn, apply_fn_factory};
 pub use self::fn_service::{fn_factory, fn_factory_with_config, fn_service};
+pub use self::fn_shutdown::fn_shutdown;
 pub use self::map_config::{map_config, unit_config};
 pub use self::middleware::{apply, Identity, Middleware, Stack};
 pub use self::pipeline::{pipeline, pipeline_factory, Pipeline, PipelineFactory};
@@ -365,6 +367,7 @@ pub mod dev {
     pub use crate::fn_service::{
         FnService, FnServiceConfig, FnServiceFactory, FnServiceNoConfig,
     };
+    pub use crate::fn_shutdown::FnShutdown;
     pub use crate::map::{Map, MapFactory};
     pub use crate::map_config::{MapConfig, UnitConfig};
     pub use crate::map_err::{MapErr, MapErrFactory};