1use 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#[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#[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 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 async fn acquire(&self) {
59 if self.ready.get() {
60 return;
61 }
62
63 let waiter = self.waiters.wait();
64 poll_fn(|cx| {
65 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
78struct 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 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 #[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 #[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 #[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 #[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 #[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 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 #[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}