ntex_service/
fn_shutdown.rs

1use std::{cell::Cell, convert::Infallible, fmt, marker::PhantomData};
2
3use crate::{Ctx, Service, ServiceFactory};
4
5/// Function that can act as a `on_shutdown` callback.
6pub struct FnShutdown<F, Err> {
7    f_shutdown: Cell<Option<F>>,
8    err: PhantomData<Err>,
9}
10
11impl<F, Err> FnShutdown<F, Err> {
12    pub fn new<St>(f: F) -> Self
13    where
14        F: AsyncFnOnce(&St),
15    {
16        Self {
17            f_shutdown: Cell::new(Some(f)),
18            err: PhantomData,
19        }
20    }
21}
22
23impl<F, Err> Clone for FnShutdown<F, Err>
24where
25    F: Clone,
26{
27    #[inline]
28    fn clone(&self) -> Self {
29        let f = self.f_shutdown.take();
30        self.f_shutdown.set(f.clone());
31        Self {
32            f_shutdown: Cell::new(f),
33            err: PhantomData,
34        }
35    }
36}
37
38impl<F, Err> fmt::Debug for FnShutdown<F, Err> {
39    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
40        f.debug_struct("FnShutdown")
41            .field("fn", &std::any::type_name::<F>())
42            .finish()
43    }
44}
45
46impl<F, St, Req, Err> Service<St, Req> for FnShutdown<F, Err>
47where
48    F: AsyncFnOnce(&St),
49{
50    type Res = Req;
51    type Error = Err;
52
53    #[inline]
54    async fn shutdown(&self, ctx: Ctx<'_, Self, St>) {
55        if let Some(f) = self.f_shutdown.take() {
56            (f)(ctx.st()).await;
57        }
58    }
59
60    #[inline]
61    async fn call(&self, req: Req, _: Ctx<'_, Self, St>) -> Result<Req, Err> {
62        Ok(req)
63    }
64}
65
66impl<F, St, Req, Err> ServiceFactory<St, Req> for FnShutdown<F, Err>
67where
68    F: AsyncFnOnce(&St) + Clone,
69{
70    type Res = Req;
71    type Error = Err;
72
73    type Service = FnShutdown<F, Err>;
74    type InitError = Infallible;
75
76    #[inline]
77    async fn create(&self, _: &St) -> Result<Self::Service, Self::InitError> {
78        if let Some(f) = self.f_shutdown.take() {
79            self.f_shutdown.set(Some(f.clone()));
80            Ok(FnShutdown {
81                f_shutdown: Cell::new(Some(f)),
82                err: PhantomData,
83            })
84        } else {
85            panic!("FnShutdown was used already");
86        }
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use std::{future::poll_fn, rc::Rc};
93
94    use crate::{Pipeline, factory, service};
95
96    use super::*;
97
98    #[ntex::test]
99    async fn test_fn_shutdown() {
100        // Service factory
101        let is_called = Rc::new(Cell::new(false));
102        let is_called2 = is_called.clone();
103        let fac = factory(|()| async { Ok::<_, ()>("pipe") }).shutdown(async move |()| {
104            is_called2.set(true);
105        });
106        let _ = format!("{fac:?}");
107
108        let pipe = Pipeline::new((), fac.clone().create(&()).await.unwrap());
109
110        let res = pipe.call(()).await;
111        assert_eq!(pipe.ready().await, Ok(()));
112        assert!(res.is_ok());
113        assert_eq!(res.unwrap(), "pipe");
114        assert!(!pipe.is_shutdown());
115        pipe.shutdown().await;
116        assert!(is_called.get());
117        assert!(pipe.is_shutdown());
118
119        poll_fn(|cx| pipe.poll_shutdown(cx)).await;
120        assert!(pipe.is_shutdown());
121
122        // Service
123        let is_called = Rc::new(Cell::new(false));
124        let is_called2 = is_called.clone();
125        let svc = service(|()| async { Ok::<_, ()>("pipe") }).shutdown(async move |()| {
126            is_called2.set(true);
127        });
128        let _ = format!("{fac:?}");
129
130        let pipe = Pipeline::new((), svc);
131
132        let res = pipe.call(()).await;
133        assert_eq!(pipe.ready().await, Ok(()));
134        assert!(res.is_ok());
135        assert_eq!(res.unwrap(), "pipe");
136        assert!(!pipe.is_shutdown());
137        pipe.shutdown().await;
138        assert!(is_called.get());
139        assert!(pipe.is_shutdown());
140
141        poll_fn(|cx| pipe.poll_shutdown(cx)).await;
142        assert!(pipe.is_shutdown());
143    }
144}
ntex_service/fn_shutdown.rs

ntex_service/
fn_shutdown.rs
