Skip to main content

ntex_service/
apply.rs

1use std::{fmt, marker};
2
3use crate::ctx::{Ctx, WaitersRef};
4use crate::{IntoService, IntoServiceFactory, Service, ServiceFactory};
5use crate::{ServiceCaller, ServiceChain, ServiceChainFactory};
6
7/// Applies an asynchronous middleware function to a service.
8///
9/// The function receives an input request and an [`ApplyCtx`] that can call the
10/// wrapped service.
11pub fn apply_fn<S, St, Req, F, In, Out, Err>(
12    service: impl IntoService<S, St, Req>,
13    f: F,
14) -> ServiceChain<Apply<S, St, Req, F, In, Out, Err>, St, In>
15where
16    S: Service<St, Req>,
17    F: AsyncFn(In, &ApplyCtx<'_, S, St, Req>) -> Result<Out, Err>,
18    Err: From<S::Error>,
19{
20    crate::service(Apply::new(service.into_service(), f))
21}
22
23/// Applies an asynchronous middleware function to every service from a factory.
24pub fn apply_fn_factory<Sf, St, Req, F, In, Out, Err>(
25    service: impl IntoServiceFactory<Sf, St, Req>,
26    f: F,
27) -> ServiceChainFactory<ApplyFactory<F, Sf, St, Req, In, Out, Err>, St, In>
28where
29    Sf: ServiceFactory<St, Req>,
30    F: AsyncFn(In, &ApplyCtx<'_, Sf::Service, St, Req>) -> Result<Out, Err> + Clone,
31    Err: From<Sf::Error>,
32{
33    crate::factory(ApplyFactory::new(service.into_factory(), f))
34}
35
36#[derive(Debug)]
37/// Context passed to middleware functions created by [`apply_fn`].
38pub struct ApplyCtx<'a, S, St, Req> {
39    idx: u32,
40    waiters: &'a WaitersRef,
41    service: &'a S,
42    st: &'a St,
43    r: marker::PhantomData<Req>,
44}
45
46impl<S: Service<St, Req>, St, Req> ApplyCtx<'_, S, St, Req> {
47    /// Returns the pipeline state.
48    #[inline]
49    pub fn st(&self) -> &St {
50        self.st
51    }
52
53    /// Waits for the wrapped service to become ready, then calls it.
54    #[inline]
55    pub async fn call(&self, req: Req) -> Result<S::Res, S::Error> {
56        Ctx::<S, St>::new(self.idx, self.waiters, self.st)
57            .call(&self.service, req)
58            .await
59    }
60}
61
62impl<S: Service<St, Req>, St, Req> ServiceCaller<Req, S::Res, S::Error>
63    for ApplyCtx<'_, S, St, Req>
64{
65    #[inline]
66    async fn call_service(&self, req: Req) -> Result<S::Res, S::Error> {
67        Ctx::<S, St>::new(self.idx, self.waiters, self.st)
68            .call(&self.service, req)
69            .await
70    }
71}
72
73/// Service produced by [`apply_fn`].
74pub struct Apply<S, St, Req, F, In, Out, Err> {
75    svc: S,
76    f: F,
77    r: marker::PhantomData<fn(St, Req) -> (In, Out, Err)>,
78}
79
80impl<S, St, Req, F, In, Out, Err> Apply<S, St, Req, F, In, Out, Err>
81where
82    F: AsyncFn(In, &ApplyCtx<'_, S, St, Req>) -> Result<Out, Err>,
83{
84    pub(crate) fn new(svc: S, f: F) -> Self {
85        Apply {
86            f,
87            svc,
88            r: marker::PhantomData,
89        }
90    }
91}
92
93impl<S, St, Req, F, In, Out, Err> Clone for Apply<S, St, Req, F, In, Out, Err>
94where
95    S: Clone,
96    F: Clone,
97{
98    fn clone(&self) -> Self {
99        Apply {
100            svc: self.svc.clone(),
101            f: self.f.clone(),
102            r: marker::PhantomData,
103        }
104    }
105}
106
107impl<S, St, Req, F, In, Out, Err> fmt::Debug for Apply<S, St, Req, F, In, Out, Err>
108where
109    S: fmt::Debug,
110{
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        f.debug_struct("Apply")
113            .field("svc", &self.svc)
114            .field("map", &std::any::type_name::<F>())
115            .finish()
116    }
117}
118
119impl<S, St, Req, F, In, Out, Err> Service<St, In> for Apply<S, St, Req, F, In, Out, Err>
120where
121    S: Service<St, Req>,
122    F: AsyncFn(In, &ApplyCtx<'_, S, St, Req>) -> Result<Out, Err>,
123    Err: From<S::Error>,
124{
125    type Res = Out;
126    type Error = Err;
127
128    #[inline]
129    async fn ready(&self, ctx: Ctx<'_, Self, St>) -> Result<(), Err> {
130        ctx.ready(&self.svc).await.map_err(From::from)
131    }
132
133    #[inline]
134    async fn call(&self, req: In, ctx: Ctx<'_, Self, St>) -> Result<Out, Err> {
135        let (idx, waiters, st) = ctx.inner();
136
137        let ctx = ApplyCtx {
138            idx,
139            waiters,
140            st,
141            service: &self.svc,
142            r: marker::PhantomData,
143        };
144        (self.f)(req, &ctx).await
145    }
146
147    crate::forward_shutdown!(St, svc);
148}
149
150/// Service factory produced by [`apply_fn_factory`].
151pub struct ApplyFactory<F, Sf, St, Req, In, Out, Err>
152where
153    F: AsyncFn(In, &ApplyCtx<'_, Sf::Service, St, Req>) -> Result<Out, Err> + Clone,
154    Sf: ServiceFactory<St, Req>,
155{
156    f: F,
157    sf: Sf,
158    r: marker::PhantomData<fn(St, Req) -> (In, Out)>,
159}
160
161impl<F, Sf, St, Req, In, Out, Err> ApplyFactory<F, Sf, St, Req, In, Out, Err>
162where
163    F: AsyncFn(In, &ApplyCtx<'_, Sf::Service, St, Req>) -> Result<Out, Err> + Clone,
164    Sf: ServiceFactory<St, Req>,
165{
166    /// Create new `ApplyFactory` new service instance
167    pub(crate) fn new(sf: Sf, f: F) -> Self
168    where
169        Sf: ServiceFactory<St, Req>,
170        Err: From<Sf::Error>,
171    {
172        Self {
173            f,
174            sf,
175            r: marker::PhantomData,
176        }
177    }
178}
179
180impl<F, Sf, St, Req, In, Out, Err> Clone for ApplyFactory<F, Sf, St, Req, In, Out, Err>
181where
182    F: AsyncFn(In, &ApplyCtx<'_, Sf::Service, St, Req>) -> Result<Out, Err> + Clone,
183    Sf: ServiceFactory<St, Req> + Clone,
184{
185    fn clone(&self) -> Self {
186        Self {
187            f: self.f.clone(),
188            sf: self.sf.clone(),
189            r: marker::PhantomData,
190        }
191    }
192}
193
194impl<F, Sf, St, Req, In, Out, Err> fmt::Debug for ApplyFactory<F, Sf, St, Req, In, Out, Err>
195where
196    F: AsyncFn(In, &ApplyCtx<'_, Sf::Service, St, Req>) -> Result<Out, Err> + Clone,
197    Sf: ServiceFactory<St, Req> + fmt::Debug,
198{
199    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
200        f.debug_struct("ApplyFactory")
201            .field("factory", &self.sf)
202            .field("map", &std::any::type_name::<F>())
203            .finish()
204    }
205}
206
207impl<F, Sf, St, Req, In, Out, Err> ServiceFactory<St, In>
208    for ApplyFactory<F, Sf, St, Req, In, Out, Err>
209where
210    F: AsyncFn(In, &ApplyCtx<'_, Sf::Service, St, Req>) -> Result<Out, Err> + Clone,
211    Sf: ServiceFactory<St, Req>,
212    Err: From<Sf::Error>,
213{
214    type Res = Out;
215    type Error = Err;
216
217    type Service = Apply<Sf::Service, St, Req, F, In, Out, Err>;
218    type InitError = Sf::InitError;
219
220    #[inline]
221    async fn create(&self, st: &St) -> Result<Self::Service, Self::InitError> {
222        self.sf.create(st).await.map(|svc| Apply {
223            svc,
224            f: self.f.clone(),
225            r: marker::PhantomData,
226        })
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use std::{cell::Cell, rc::Rc};
233
234    use super::*;
235    use crate::{factory, fn_factory, service};
236
237    #[derive(Debug, Default, Clone)]
238    struct Srv(Rc<Cell<usize>>);
239
240    impl Service<(), ()> for Srv {
241        type Res = ();
242        type Error = ();
243
244        async fn call(&self, _r: (), _: Ctx<'_, Self>) -> Result<(), ()> {
245            Ok(())
246        }
247
248        async fn ready(&self, _: Ctx<'_, Self>) -> Result<(), ()> {
249            self.0.set(self.0.get() + 1);
250            Ok(())
251        }
252
253        async fn shutdown(&self, _: Ctx<'_, Self, ()>) {
254            self.0.set(self.0.get() + 1);
255        }
256    }
257
258    #[derive(Debug, PartialEq, Eq)]
259    struct Err;
260
261    impl From<()> for Err {
262        fn from(_e: ()) -> Self {
263            Err
264        }
265    }
266
267    #[ntex::test]
268    async fn test_call() {
269        let cnt_sht = Rc::new(Cell::new(0));
270        let srv = service(
271            apply_fn(Srv(cnt_sht.clone()), async move |req: &'static str, svc| {
272                svc.call(()).await.unwrap();
273                Ok((req, ()))
274            })
275            .clone(),
276        )
277        .pipeline(());
278
279        assert_eq!(srv.ready().await, Ok::<_, Err>(()));
280
281        srv.shutdown().await;
282        assert_eq!(cnt_sht.get(), 2);
283
284        let res = srv.call("srv").await;
285        assert!(res.is_ok());
286        assert_eq!(res.unwrap(), ("srv", ()));
287    }
288
289    #[ntex::test]
290    async fn test_call_svc() {
291        let cnt_sht = Rc::new(Cell::new(0));
292        let srv = service(Srv(cnt_sht.clone()))
293            .apply_fn(async move |req: &'static str, svc| {
294                svc.st();
295                svc.call(()).await.unwrap();
296                Ok((req, ()))
297            })
298            .clone();
299        let s = format!("{srv:?}");
300        assert!(s.contains("Apply"), "{}", s);
301
302        let srv = srv.pipeline(());
303        assert_eq!(srv.ready().await, Ok::<_, Err>(()));
304
305        srv.shutdown().await;
306        assert_eq!(cnt_sht.get(), 2);
307
308        let res = srv.call("srv").await;
309        assert!(res.is_ok());
310        assert_eq!(res.unwrap(), ("srv", ()));
311        let _ = format!("{srv:?}");
312    }
313
314    #[ntex::test]
315    async fn test_create() {
316        let new_srv = factory(apply_fn_factory(
317            fn_factory(|(): &()| async { Ok::<_, ()>(Srv::default()) }),
318            async move |req: &'static str, srv| {
319                srv.call(()).await.unwrap();
320                Ok((req, ()))
321            },
322        ));
323
324        let srv = new_srv.pipeline(()).await.unwrap();
325
326        assert_eq!(srv.ready().await, Ok::<_, Err>(()));
327
328        let res = srv.call("srv").await;
329        assert!(res.is_ok());
330        assert_eq!(res.unwrap(), ("srv", ()));
331        assert_eq!(Err, Err::from(()));
332    }
333
334    #[ntex::test]
335    async fn test_create_chain() {
336        let new_srv = factory(fn_factory(|(): &()| async { Ok::<_, ()>(Srv::default()) }))
337            .apply_fn(async move |req: &'static str, srv| {
338                srv.call(()).await.unwrap();
339                Ok((req, ()))
340            })
341            .clone();
342
343        let srv = new_srv.pipeline(()).await.unwrap();
344
345        assert_eq!(srv.ready().await, Ok::<_, Err>(()));
346
347        let res = srv.call("srv").await;
348        assert!(res.is_ok());
349        assert_eq!(res.unwrap(), ("srv", ()));
350        let _ = format!("{new_srv:?}");
351    }
352}