ntex_util/services/
onerequest.rs

1//! Middleware that permits one service call at a time.
2use std::{cell::Cell, future::poll_fn, task::Poll};
3
4use ntex_service::{BindCtx, Ctx, Middleware, ReadyNotifier, Service, ServiceBinding};
5
6use crate::{channel::condition::Condition, future::OnDropFn};
7
8/// Middleware that serializes calls to its wrapped service.
9#[derive(Copy, Clone, Default, Debug)]
10pub struct OneRequest;
11
12impl<S, St> Middleware<S, St> for OneRequest {
13    type Service = OneRequestService<S>;
14
15    fn create(&self, _: &St, service: S) -> Self::Service {
16        OneRequestService {
17            service: ServiceBinding::new(service),
18            ready: Cell::new(true),
19            waiters: Condition::new(),
20            changed: ReadyNotifier::new(),
21        }
22    }
23}
24
25/// Service wrapper that allows only one call to run at a time.
26///
27/// This type is intentionally not cloneable. Cloning its readiness state would
28/// create another independent gate and allow calls to overlap.
29///
30/// Acquiring or releasing the call slot invalidates this service's readiness.
31/// The wrapped service has an independent readiness cache and must notify when
32/// its readiness changes.
33#[derive(Debug)]
34pub struct OneRequestService<S> {
35    waiters: Condition,
36    service: ServiceBinding<S>,
37    ready: Cell<bool>,
38    changed: ReadyNotifier,
39}
40
41impl<S> OneRequestService<S> {
42    /// Wraps a service so that concurrent callers wait for the active call.
43    pub fn new<St, Req>(service: S) -> Self
44    where
45        S: Service<St, Req>,
46    {
47        Self {
48            service: ServiceBinding::new(service),
49            ready: Cell::new(true),
50            waiters: Condition::new(),
51            changed: ReadyNotifier::new(),
52        }
53    }
54}
55
56impl<S> OneRequestService<S> {
57    /// Waits until no call is active.
58    async fn acquire(&self) {
59        if self.ready.get() {
60            return;
61        }
62
63        let waiter = self.waiters.wait();
64        poll_fn(|cx| {
65            // the waiter stays registered after a notification, another task
66            // may have started a call before this one was polled
67            let _ = waiter.poll_ready(cx);
68            if self.ready.get() {
69                Poll::Ready(())
70            } else {
71                Poll::Pending
72            }
73        })
74        .await;
75    }
76}
77
78/// Releases the call slot even if the call future is dropped or panics.
79struct Release<'a, S>(&'a OneRequestService<S>);
80
81impl<S> Drop for Release<'_, S> {
82    fn drop(&mut self) {
83        self.0.ready.set(true);
84        let _wake = OnDropFn::new(|| self.0.waiters.notify(()));
85        self.0.changed.notify();
86    }
87}
88
89impl<S: Service<St, Req>, St, Req> Service<St, Req> for OneRequestService<S> {
90    type Res = S::Res;
91    type Error = S::Error;
92
93    fn bind(&self, ctx: BindCtx<'_>) {
94        ctx.subscribe(&self.changed);
95        ctx.bind::<_, St, Req>(&self.service);
96    }
97
98    #[inline]
99    async fn ready(&self, ctx: Ctx<'_, Self, St>) -> Result<(), S::Error> {
100        self.acquire().await;
101        ctx.ready(&self.service).await
102    }
103
104    #[inline]
105    async fn call(&self, req: Req, ctx: Ctx<'_, Self, St>) -> Result<S::Res, S::Error> {
106        // several callers can observe readiness before any of them calls
107        self.acquire().await;
108        self.ready.set(false);
109        let _release = Release(self);
110        self.changed.notify();
111
112        ctx.call(&self.service, req).await
113    }
114
115    ntex_service::forward_shutdown!(St, service);
116}
117
118#[cfg(test)]
119mod tests {
120    use ntex_service::{Pipeline, apply, fn_factory};
121    use std::{cell::RefCell, rc::Rc, time::Duration};
122
123    use super::*;
124    use crate::{channel::oneshot, future::lazy};
125
126    struct SleepService(oneshot::Receiver<()>);
127
128    impl Service<(), ()> for SleepService {
129        type Res = ();
130        type Error = ();
131
132        async fn call(&self, _r: (), _: Ctx<'_, Self>) -> Result<(), ()> {
133            let _ = self.0.recv().await;
134            Ok::<_, ()>(())
135        }
136    }
137
138    #[ntex::test]
139    async fn test_oneshot() {
140        let (tx, rx) = oneshot::channel();
141
142        let srv = Pipeline::new((), OneRequestService::new(SleepService(rx)));
143        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
144
145        let srv2 = srv.bind();
146        ntex::rt::spawn(async move {
147            let _ = srv2.call(()).await;
148        });
149        crate::time::sleep(Duration::from_millis(25)).await;
150        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
151
152        let _ = tx.send(());
153        crate::time::sleep(Duration::from_millis(25)).await;
154        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
155        srv.shutdown().await;
156    }
157
158    #[ntex::test]
159    async fn test_middleware() {
160        assert_eq!(format!("{OneRequest:?}"), "OneRequest");
161
162        let (tx, rx) = oneshot::channel();
163        let rx = RefCell::new(Some(rx));
164        let sf = apply(
165            OneRequest,
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 srv1 = srv.bind();
176        ntex::rt::spawn(async move {
177            let _ = srv1.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(());
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!(format!("{OneRequest:?}"), "OneRequest");
190
191        let (tx, rx) = oneshot::channel();
192        let rx = RefCell::new(Some(rx));
193        let sf = apply(
194            OneRequest,
195            fn_factory(move |(): &()| {
196                let rx = rx.borrow_mut().take().unwrap();
197                async move { Ok::<_, ()>(SleepService(rx)) }
198            }),
199        );
200
201        let srv = sf.pipeline(()).await.unwrap();
202        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
203
204        let srv1 = srv.bind();
205        ntex::rt::spawn(async move {
206            let _ = srv1.call(()).await;
207        });
208        crate::time::sleep(Duration::from_millis(25)).await;
209        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
210
211        let _ = tx.send(());
212        crate::time::sleep(Duration::from_millis(25)).await;
213        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
214    }
215
216    #[ntex::test]
217    async fn test_cancelled_call_releases() {
218        let (_tx, rx) = oneshot::channel();
219        let srv = Pipeline::new((), OneRequestService::new(SleepService(rx)));
220
221        let res = crate::time::timeout(Duration::from_millis(25), srv.call(())).await;
222        assert!(res.is_err());
223        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
224    }
225
226    struct CountService {
227        active: Rc<Cell<usize>>,
228        max: Rc<Cell<usize>>,
229    }
230
231    impl Service<(), ()> for CountService {
232        type Res = ();
233        type Error = ();
234
235        async fn call(&self, _r: (), _: Ctx<'_, Self>) -> Result<(), ()> {
236            self.active.set(self.active.get() + 1);
237            self.max.set(self.max.get().max(self.active.get()));
238            crate::time::sleep(Duration::from_millis(10)).await;
239            self.active.set(self.active.get() - 1);
240            Ok(())
241        }
242    }
243
244    fn count_srv() -> (Pipeline<(), (), ()>, Rc<Cell<usize>>) {
245        let max = Rc::new(Cell::new(0));
246        let srv = Pipeline::new(
247            (),
248            OneRequestService::new(CountService {
249                active: Rc::new(Cell::new(0)),
250                max: max.clone(),
251            }),
252        );
253        (srv, max)
254    }
255
256    #[ntex::test]
257    async fn test_many_waiters() {
258        let (srv, max) = count_srv();
259        let done = Rc::new(Cell::new(0));
260        for _ in 0..4 {
261            let (srv, done) = (srv.bind(), done.clone());
262            ntex::rt::spawn(async move {
263                srv.call(()).await.unwrap();
264                done.set(done.get() + 1);
265            });
266        }
267
268        // calls run one after another, wait for all of them with a generous
269        // deadline, timers on loaded CI runners overshoot a lot
270        for _ in 0..100 {
271            if done.get() == 4 {
272                break;
273            }
274            crate::time::sleep(Duration::from_millis(50)).await;
275        }
276        assert_eq!(done.get(), 4);
277        assert_eq!(max.get(), 1);
278    }
279
280    #[ntex::test]
281    async fn test_no_overlap_after_ready() {
282        let (srv, max) = count_srv();
283
284        srv.ready().await.unwrap();
285        srv.ready().await.unwrap();
286        let (r1, r2) = crate::future::join(srv.call_static(()), srv.call_static(())).await;
287        assert!(r1.is_ok() && r2.is_ok());
288        assert_eq!(max.get(), 1);
289    }
290}
ntex_util/services/onerequest.rs

ntex_util/services/
onerequest.rs
