ntex_service/
fn_shutdown.rs

1use std::{cell::Cell, fmt, marker::PhantomData};
2
3use crate::{Service, ServiceCtx};
4
5#[inline]
6/// Create `FnShutdown` for function that can act as a `on_shutdown` callback.
7pub fn fn_shutdown<Req, Err, F>(f: F) -> FnShutdown<Req, Err, F>
8where
9    F: FnOnce(),
10{
11    FnShutdown::new(f)
12}
13
14pub struct FnShutdown<Req, Err, F> {
15    f_shutdown: Cell<Option<F>>,
16    _t: PhantomData<(Req, Err)>,
17}
18
19impl<Req, Err, F> FnShutdown<Req, Err, F> {
20    pub(crate) fn new(f: F) -> Self {
21        Self {
22            f_shutdown: Cell::new(Some(f)),
23            _t: PhantomData,
24        }
25    }
26}
27
28impl<Req, Err, F> Clone for FnShutdown<Req, Err, F>
29where
30    F: Clone,
31{
32    #[inline]
33    fn clone(&self) -> Self {
34        let f = self.f_shutdown.take();
35        self.f_shutdown.set(f.clone());
36        Self {
37            f_shutdown: Cell::new(f),
38            _t: PhantomData,
39        }
40    }
41}
42
43impl<Req, Err, F> fmt::Debug for FnShutdown<Req, Err, F> {
44    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45        f.debug_struct("FnShutdown")
46            .field("fn", &std::any::type_name::<F>())
47            .finish()
48    }
49}
50
51impl<Req, Err, F> Service<Req> for FnShutdown<Req, Err, F>
52where
53    F: FnOnce(),
54{
55    type Response = Req;
56    type Error = Err;
57
58    #[inline]
59    async fn shutdown(&self) {
60        if let Some(f) = self.f_shutdown.take() {
61            (f)()
62        }
63    }
64
65    #[inline]
66    async fn call(&self, req: Req, _: ServiceCtx<'_, Self>) -> Result<Req, Err> {
67        Ok(req)
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use std::rc::Rc;
74
75    use crate::{chain, fn_service, Pipeline};
76
77    use super::*;
78
79    #[ntex::test]
80    async fn test_fn_shutdown() {
81        let is_called = Rc::new(Cell::new(false));
82        let srv = fn_service(|_| async { Ok::<_, ()>("pipe") });
83        let is_called2 = is_called.clone();
84        let on_shutdown = fn_shutdown(|| {
85            is_called2.set(true);
86        });
87
88        let pipe = Pipeline::new(chain(srv).and_then(on_shutdown).clone());
89
90        let res = pipe.call(()).await;
91        assert_eq!(pipe.ready().await, Ok(()));
92        assert!(res.is_ok());
93        assert_eq!(res.unwrap(), "pipe");
94        pipe.shutdown().await;
95        assert!(is_called.get());
96
97        let _ = format!("{:?}", pipe);
98    }
99}