ntex_util/services/
inflight.rs

1//! Service that limits number of in-flight async requests.
2use ntex_service::{Ctx, Middleware, Service};
3
4use super::counter::Counter;
5
6/// `InFlight` - service factory for service that can limit number of in-flight
7/// async requests.
8///
9/// Default number of in-flight requests is 15
10#[derive(Copy, Clone, Debug)]
11pub struct InFlight {
12    max_inflight: usize,
13}
14
15impl InFlight {
16    pub fn new(max: usize) -> Self {
17        Self { max_inflight: max }
18    }
19}
20
21impl Default for InFlight {
22    fn default() -> Self {
23        Self::new(15)
24    }
25}
26
27impl<S, St> Middleware<S, St> for InFlight {
28    type Service = InFlightService<S>;
29
30    fn create(&self, _: &St, service: S) -> Self::Service {
31        InFlightService {
32            service,
33            count: Counter::new(self.max_inflight),
34        }
35    }
36}
37
38#[derive(Debug)]
39pub struct InFlightService<S> {
40    count: Counter,
41    service: S,
42}
43
44impl<S> InFlightService<S> {
45    pub fn new(max: usize, service: S) -> Self {
46        Self {
47            service,
48            count: Counter::new(max),
49        }
50    }
51}
52
53impl<S, St, Req> Service<St, Req> for InFlightService<S>
54where
55    S: Service<St, Req>,
56{
57    type Res = S::Res;
58    type Error = S::Error;
59
60    #[inline]
61    async fn ready(&self, ctx: Ctx<'_, Self, St>) -> Result<(), S::Error> {
62        if self.count.is_available() {
63            ctx.ready(&self.service).await
64        } else {
65            crate::future::join(self.count.available(), ctx.ready(&self.service))
66                .await
67                .1
68        }
69    }
70
71    #[inline]
72    async fn call(&self, req: Req, ctx: Ctx<'_, Self, St>) -> Result<S::Res, S::Error> {
73        ctx.ready(self).await?;
74        let _guard = self.count.get();
75        ctx.call(&self.service, req).await
76    }
77
78    ntex_service::forward_shutdown!(St, service);
79}
80
81#[cfg(test)]
82mod tests {
83    use std::{cell::Cell, cell::RefCell, rc::Rc, task::Poll, time::Duration};
84
85    use async_channel as mpmc;
86    use ntex_service::{Pipeline, apply, fn_factory};
87
88    use super::*;
89    use crate::{channel::oneshot, future::lazy};
90
91    struct SleepService(mpmc::Receiver<()>);
92
93    impl Service<(), ()> for SleepService {
94        type Res = ();
95        type Error = ();
96
97        async fn call(&self, _r: (), _: Ctx<'_, Self>) -> Result<(), ()> {
98            let _ = self.0.recv().await;
99            Ok(())
100        }
101    }
102
103    #[ntex::test]
104    async fn test_service() {
105        let (tx, rx) = mpmc::unbounded();
106        let counter = Rc::new(Cell::new(0));
107
108        let srv = Pipeline::new((), InFlightService::new(1, SleepService(rx)));
109        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
110
111        let counter2 = counter.clone();
112        let fut = srv.call_nowait(());
113        ntex::rt::spawn(async move {
114            let _ = fut.await;
115            counter2.set(counter2.get() + 1);
116        });
117        crate::time::sleep(Duration::from_millis(25)).await;
118        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
119
120        let counter2 = counter.clone();
121        let fut = srv.call_nowait(());
122        ntex::rt::spawn(async move {
123            let _ = fut.await;
124            counter2.set(counter2.get() + 1);
125        });
126        crate::time::sleep(Duration::from_millis(25)).await;
127        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
128
129        let counter2 = counter.clone();
130        let fut = srv.call_static(());
131        let (stx, srx) = oneshot::channel::<()>();
132        ntex::rt::spawn(async move {
133            let _ = fut.await;
134            counter2.set(counter2.get() + 1);
135            let _ = stx.send(());
136        });
137        crate::time::sleep(Duration::from_millis(25)).await;
138        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
139
140        let _ = tx.send(()).await;
141        crate::time::sleep(Duration::from_millis(25)).await;
142        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
143
144        let _ = tx.send(()).await;
145        crate::time::sleep(Duration::from_millis(25)).await;
146        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
147
148        let _ = tx.send(()).await;
149        let _ = srx.recv().await;
150        assert_eq!(counter.get(), 3);
151        srv.shutdown().await;
152    }
153
154    #[ntex::test]
155    async fn test_middleware() {
156        assert_eq!(InFlight::default().max_inflight, 15);
157        assert_eq!(
158            format!("{:?}", InFlight::new(1)),
159            "InFlight { max_inflight: 1 }"
160        );
161
162        let (tx, rx) = mpmc::unbounded();
163        let rx = RefCell::new(Some(rx));
164        let sf = apply(
165            InFlight::new(1),
166            fn_factory(move |(): &()| {
167                let rx = rx.borrow_mut().take().unwrap();
168                async move { Ok::<_, ()>(SleepService(rx)) }
169            }),
170        );
171
172        let srv = sf.pipeline(()).await.unwrap();
173        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
174
175        let srv2 = srv.bind();
176        ntex::rt::spawn(async move {
177            let _ = srv2.call(()).await;
178        });
179        crate::time::sleep(Duration::from_millis(25)).await;
180        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
181
182        let _ = tx.send(()).await;
183        crate::time::sleep(Duration::from_millis(25)).await;
184        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
185    }
186
187    #[ntex::test]
188    async fn test_middleware2() {
189        assert_eq!(InFlight::default().max_inflight, 15);
190        assert_eq!(
191            format!("{:?}", InFlight::new(1)),
192            "InFlight { max_inflight: 1 }"
193        );
194
195        let (tx, rx) = mpmc::unbounded();
196        let rx = RefCell::new(Some(rx));
197        let sf = apply(
198            InFlight::new(1),
199            fn_factory(move |(): &()| {
200                let rx = rx.borrow_mut().take().unwrap();
201                async move { Ok::<_, ()>(SleepService(rx)) }
202            }),
203        );
204
205        let srv = sf.pipeline(()).await.unwrap();
206        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
207
208        let srv2 = srv.bind();
209        ntex::rt::spawn(async move {
210            let _ = srv2.call(()).await;
211        });
212        crate::time::sleep(Duration::from_millis(25)).await;
213        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
214
215        let _ = tx.send(()).await;
216        crate::time::sleep(Duration::from_millis(25)).await;
217        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
218    }
219}
ntex_util/services/inflight.rs

ntex_util/services/
inflight.rs
