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::{Ctx, Middleware, Service};
5
6use crate::task::LocalWaker;
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,
18            ready: Cell::new(true),
19            waker: LocalWaker::new(),
20        }
21    }
22}
23
24/// Service wrapper that allows only one call to run at a time.
25///
26/// This type is intentionally not cloneable. Cloning its readiness state would
27/// create another independent gate and allow calls to overlap.
28#[derive(Debug)]
29pub struct OneRequestService<S> {
30    waker: LocalWaker,
31    service: S,
32    ready: Cell<bool>,
33}
34
35impl<S> OneRequestService<S> {
36    /// Wraps a service so that concurrent callers wait for the active call.
37    pub fn new<St, Req>(service: S) -> Self
38    where
39        S: Service<St, Req>,
40    {
41        Self {
42            service,
43            ready: Cell::new(true),
44            waker: LocalWaker::new(),
45        }
46    }
47}
48
49impl<S: Service<St, Req>, St, Req> Service<St, Req> for OneRequestService<S> {
50    type Res = S::Res;
51    type Error = S::Error;
52
53    #[inline]
54    async fn ready(&self, ctx: Ctx<'_, Self, St>) -> Result<(), S::Error> {
55        if !self.ready.get() {
56            poll_fn(|cx| {
57                self.waker.register(cx.waker());
58                if self.ready.get() {
59                    Poll::Ready(())
60                } else {
61                    Poll::Pending
62                }
63            })
64            .await;
65        }
66        ctx.ready(&self.service).await
67    }
68
69    #[inline]
70    async fn call(&self, req: Req, ctx: Ctx<'_, Self, St>) -> Result<S::Res, S::Error> {
71        self.ready.set(false);
72
73        let result = ctx.call(&self.service, req).await;
74        self.ready.set(true);
75        self.waker.wake();
76        result
77    }
78
79    ntex_service::forward_shutdown!(St, service);
80}
81
82#[cfg(test)]
83mod tests {
84    use ntex_service::{Pipeline, apply, fn_factory};
85    use std::{cell::RefCell, time::Duration};
86
87    use super::*;
88    use crate::{channel::oneshot, future::lazy};
89
90    struct SleepService(oneshot::Receiver<()>);
91
92    impl Service<(), ()> for SleepService {
93        type Res = ();
94        type Error = ();
95
96        async fn call(&self, _r: (), _: Ctx<'_, Self>) -> Result<(), ()> {
97            let _ = self.0.recv().await;
98            Ok::<_, ()>(())
99        }
100    }
101
102    #[ntex::test]
103    async fn test_oneshot() {
104        let (tx, rx) = oneshot::channel();
105
106        let srv = Pipeline::new((), OneRequestService::new(SleepService(rx)));
107        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
108
109        let srv2 = srv.bind();
110        ntex::rt::spawn(async move {
111            let _ = srv2.call(()).await;
112        });
113        crate::time::sleep(Duration::from_millis(25)).await;
114        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
115
116        let _ = tx.send(());
117        crate::time::sleep(Duration::from_millis(25)).await;
118        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
119        srv.shutdown().await;
120    }
121
122    #[ntex::test]
123    async fn test_middleware() {
124        assert_eq!(format!("{OneRequest:?}"), "OneRequest");
125
126        let (tx, rx) = oneshot::channel();
127        let rx = RefCell::new(Some(rx));
128        let sf = apply(
129            OneRequest,
130            fn_factory(move |(): &()| {
131                let rx = rx.borrow_mut().take().unwrap();
132                async move { Ok::<_, ()>(SleepService(rx)) }
133            }),
134        );
135
136        let srv = sf.pipeline(()).await.unwrap();
137        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
138
139        let srv1 = srv.bind();
140        ntex::rt::spawn(async move {
141            let _ = srv1.call(()).await;
142        });
143        crate::time::sleep(Duration::from_millis(25)).await;
144        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
145
146        let _ = tx.send(());
147        crate::time::sleep(Duration::from_millis(25)).await;
148        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
149    }
150
151    #[ntex::test]
152    async fn test_middleware2() {
153        assert_eq!(format!("{OneRequest:?}"), "OneRequest");
154
155        let (tx, rx) = oneshot::channel();
156        let rx = RefCell::new(Some(rx));
157        let sf = apply(
158            OneRequest,
159            fn_factory(move |(): &()| {
160                let rx = rx.borrow_mut().take().unwrap();
161                async move { Ok::<_, ()>(SleepService(rx)) }
162            }),
163        );
164
165        let srv = sf.pipeline(()).await.unwrap();
166        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
167
168        let srv1 = srv.bind();
169        ntex::rt::spawn(async move {
170            let _ = srv1.call(()).await;
171        });
172        crate::time::sleep(Duration::from_millis(25)).await;
173        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
174
175        let _ = tx.send(());
176        crate::time::sleep(Duration::from_millis(25)).await;
177        assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
178    }
179}
ntex_util/services/onerequest.rs

ntex_util/services/
onerequest.rs
