Skip to main content

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, ReadyBinding, ReadyNotifier, Service};
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: ReadyBinding::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: ReadyBinding<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: ReadyBinding::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(&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    #[inline]
116    async fn shutdown(&self, ctx: Ctx<'_, Self, St>) {
117        ctx.shutdown(&self.service).await;
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use ntex_service::{Pipeline, apply, fn_factory};
124    use std::{cell::RefCell, rc::Rc, time::Duration};
125
126    use super::*;
127    use crate::{channel::oneshot, future::lazy};
128
129    struct SleepService(oneshot::Receiver<()>);
130
131    impl Service<(), ()> for SleepService {
132        type Res = ();
133        type Error = ();
134
135        async fn call(&self, _r: (), _: Ctx<'_, Self>) -> Result<(), ()> {
136            let _ = self.0.recv().await;
137            Ok::<_, ()>(())
138        }
139    }
140
141    /// STATEMENT: An active request blocks readiness until it completes.
142    ///
143    /// HOW:
144    /// - Start a channel-blocked request after checking initial readiness.
145    /// - Observe pending readiness, release the request, and verify readiness returns.
146    #[ntex::test]
147    async fn test_oneshot() {
148        let (tx, rx) = oneshot::channel();
149
150        let srv = Pipeline::new((), OneRequestService::new(SleepService(rx)));
151        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
152
153        let srv2 = srv.bind();
154        ntex::rt::spawn(async move {
155            let _ = srv2.call(()).await;
156        });
157        crate::time::sleep(Duration::from_millis(25)).await;
158        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
159
160        let _ = tx.send(());
161        crate::time::sleep(Duration::from_millis(25)).await;
162        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
163        srv.shutdown().await;
164    }
165
166    /// STATEMENT: Factory-applied `OneRequest` blocks readiness during an active call.
167    ///
168    /// HOW:
169    /// - Create middleware around a channel-blocked service and check its debug name.
170    /// - Verify readiness before, during, and after releasing a spawned request.
171    #[ntex::test]
172    async fn test_middleware() {
173        assert_eq!(format!("{OneRequest:?}"), "OneRequest");
174
175        let (tx, rx) = oneshot::channel();
176        let rx = RefCell::new(Some(rx));
177        let sf = apply(
178            OneRequest,
179            fn_factory(move |(): &()| {
180                let rx = rx.borrow_mut().take().unwrap();
181                async move { Ok::<_, ()>(SleepService(rx)) }
182            }),
183        );
184
185        let srv = sf.pipeline(()).await.unwrap();
186        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
187
188        let srv1 = srv.bind();
189        ntex::rt::spawn(async move {
190            let _ = srv1.call(()).await;
191        });
192        crate::time::sleep(Duration::from_millis(25)).await;
193        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
194
195        let _ = tx.send(());
196        crate::time::sleep(Duration::from_millis(25)).await;
197        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
198    }
199
200    /// STATEMENT: A `OneRequest` middleware pipeline becomes ready again after completion.
201    ///
202    /// HOW:
203    /// - Build the middleware through a factory and start a channel-blocked request.
204    /// - Check readiness becomes pending, then send completion and check readiness recovers.
205    #[ntex::test]
206    async fn test_middleware2() {
207        assert_eq!(format!("{OneRequest:?}"), "OneRequest");
208
209        let (tx, rx) = oneshot::channel();
210        let rx = RefCell::new(Some(rx));
211        let sf = apply(
212            OneRequest,
213            fn_factory(move |(): &()| {
214                let rx = rx.borrow_mut().take().unwrap();
215                async move { Ok::<_, ()>(SleepService(rx)) }
216            }),
217        );
218
219        let srv = sf.pipeline(()).await.unwrap();
220        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
221
222        let srv1 = srv.bind();
223        ntex::rt::spawn(async move {
224            let _ = srv1.call(()).await;
225        });
226        crate::time::sleep(Duration::from_millis(25)).await;
227        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
228
229        let _ = tx.send(());
230        crate::time::sleep(Duration::from_millis(25)).await;
231        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
232    }
233
234    /// STATEMENT: Canceling a request releases the single-request readiness gate.
235    ///
236    /// HOW:
237    /// - Time out a call whose inner service is waiting on a channel.
238    /// - Verify the timeout and immediately poll the pipeline as ready.
239    #[ntex::test]
240    async fn test_cancelled_call_releases() {
241        let (_tx, rx) = oneshot::channel();
242        let srv = Pipeline::new((), OneRequestService::new(SleepService(rx)));
243
244        let res = crate::time::timeout(Duration::from_millis(25), srv.call(())).await;
245        assert!(res.is_err());
246        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
247    }
248
249    struct CountService {
250        active: Rc<Cell<usize>>,
251        max: Rc<Cell<usize>>,
252    }
253
254    impl Service<(), ()> for CountService {
255        type Res = ();
256        type Error = ();
257
258        async fn call(&self, _r: (), _: Ctx<'_, Self>) -> Result<(), ()> {
259            self.active.set(self.active.get() + 1);
260            self.max.set(self.max.get().max(self.active.get()));
261            crate::time::sleep(Duration::from_millis(10)).await;
262            self.active.set(self.active.get() - 1);
263            Ok(())
264        }
265    }
266
267    fn count_srv() -> (Pipeline<(), (), ()>, Rc<Cell<usize>>) {
268        let max = Rc::new(Cell::new(0));
269        let srv = Pipeline::new(
270            (),
271            OneRequestService::new(CountService {
272                active: Rc::new(Cell::new(0)),
273                max: max.clone(),
274            }),
275        );
276        (srv, max)
277    }
278
279    /// STATEMENT: Multiple waiting requests complete without overlapping inner calls.
280    ///
281    /// HOW:
282    /// - Spawn four requests against a service that counts active calls.
283    /// - Wait for all completions and verify the maximum concurrency is one.
284    #[ntex::test]
285    async fn test_many_waiters() {
286        let (srv, max) = count_srv();
287        let done = Rc::new(Cell::new(0));
288        for _ in 0..4 {
289            let (srv, done) = (srv.bind(), done.clone());
290            ntex::rt::spawn(async move {
291                srv.call(()).await.unwrap();
292                done.set(done.get() + 1);
293            });
294        }
295
296        // calls run one after another, wait for all of them with a generous
297        // deadline, timers on loaded CI runners overshoot a lot
298        for _ in 0..100 {
299            if done.get() == 4 {
300                break;
301            }
302            crate::time::sleep(Duration::from_millis(50)).await;
303        }
304        assert_eq!(done.get(), 4);
305        assert_eq!(max.get(), 1);
306    }
307
308    /// STATEMENT: Repeated readiness checks do not permit overlapping requests.
309    ///
310    /// HOW:
311    /// - Check readiness twice, then join two concurrent calls.
312    /// - Verify both succeed while the inner service's maximum concurrency stays at one.
313    #[ntex::test]
314    async fn test_no_overlap_after_ready() {
315        let (srv, max) = count_srv();
316
317        srv.ready().await.unwrap();
318        srv.ready().await.unwrap();
319        let (r1, r2) = crate::future::join(srv.call_static(()), srv.call_static(())).await;
320        assert!(r1.is_ok() && r2.is_ok());
321        assert_eq!(max.get(), 1);
322    }
323}