1#![allow(clippy::type_complexity)]
3use std::cell::{Cell, RefCell};
4use std::task::Poll;
5use std::{collections::VecDeque, fmt, future, future::Future, marker, mem, pin, rc::Rc};
6
7use ntex_service::{BindCtx, Ctx, Middleware, ReadyNotifier, Service, pipeline::PipelineState};
8
9use crate::{channel::oneshot, future::OnDropFn, task::LocalWaker};
10
11#[derive(Copy, Clone, Debug)]
12pub struct Buffer<St: Clone, Req, Res, Err> {
18 buf_size: usize,
19 cancel_on_shutdown: bool,
20 st: marker::PhantomData<fn(St, Req) -> Result<Res, Err>>,
21}
22
23impl<St: Clone, Req, Res, Err> Buffer<St, Req, Res, Err> {
24 #[must_use]
28 pub fn buf_size(mut self, size: usize) -> Self {
29 self.buf_size = size;
30 self
31 }
32
33 #[must_use]
37 pub fn cancel_on_shutdown(mut self) -> Self {
38 self.cancel_on_shutdown = true;
39 self
40 }
41}
42
43impl<St: Clone, Req, Res, Err> Default for Buffer<St, Req, Res, Err> {
44 fn default() -> Self {
45 Self {
46 buf_size: 16,
47 cancel_on_shutdown: false,
48 st: marker::PhantomData,
49 }
50 }
51}
52
53impl<S, St, Req, Res, Err> Middleware<S, St> for Buffer<St, Req, Res, Err>
54where
55 S: Service<St, Req, Res = Res, Error = Err> + 'static,
56 St: Clone + 'static,
57 Req: 'static,
58 Res: 'static,
59 Err: 'static,
60{
61 type Service = BufferService<St, Req, Res, Err>;
62
63 fn create(&self, _: &St, service: S) -> Self::Service {
64 let service = BufferService::new(self.buf_size, PipelineState::new(service));
65 if self.cancel_on_shutdown {
66 service.cancel_on_shutdown()
67 } else {
68 service
69 }
70 }
71}
72
73#[derive(Clone, Copy, Debug, PartialEq, Eq)]
75pub enum BufferServiceError<E> {
76 Service(E),
78 RequestCanceled,
80}
81
82impl<E> From<E> for BufferServiceError<E> {
83 fn from(err: E) -> Self {
84 BufferServiceError::Service(err)
85 }
86}
87
88impl<E: fmt::Display> fmt::Display for BufferServiceError<E> {
89 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
90 match self {
91 BufferServiceError::Service(e) => fmt::Display::fmt(e, f),
92 BufferServiceError::RequestCanceled => f.write_str("buffer service request canceled"),
93 }
94 }
95}
96
97impl<E: fmt::Display + fmt::Debug> std::error::Error for BufferServiceError<E> {}
98
99pub struct BufferService<St, Req, Res, Err> {
121 shared: Rc<Shared>,
122 service: PipelineState<St, Req, Res, Err>,
123 cancel_on_shutdown: bool,
124}
125
126struct Shared {
127 size: usize,
128 buf: RefCell<VecDeque<oneshot::Sender<Token>>>,
129 has_token: Cell<bool>,
131 in_flight: Cell<usize>,
133 canceled: Cell<bool>,
134 notifier: ReadyNotifier,
135 token_waker: LocalWaker,
136 drained: LocalWaker,
137}
138
139impl Shared {
140 fn has_room(&self) -> bool {
141 self.buf.borrow().len() < self.size
142 }
143
144 fn is_drained(&self) -> bool {
145 !self.has_token.get() && self.in_flight.get() == 0 && self.buf.borrow().is_empty()
146 }
147
148 fn release_next(self: &Rc<Self>) {
150 loop {
151 let (sender, freed) = {
152 let mut buf = self.buf.borrow_mut();
153 let full = buf.len() >= self.size;
154 let sender = buf.pop_front();
155 (sender, full && buf.len() < self.size)
156 };
157 let Some(sender) = sender else {
158 self.has_token.set(false);
159 self.drained.wake();
160 return;
161 };
162 let token = Token(Some(self.clone()));
163 if freed {
164 self.notifier.notify();
165 }
166 match sender.send(token) {
167 Ok(()) => return,
168 Err(token) => token.disarm(),
170 }
171 }
172 }
173
174 fn prune(&self) {
176 let freed = {
177 let mut buf = self.buf.borrow_mut();
178 let full = buf.len() >= self.size;
179 buf.retain(|sender| !sender.is_canceled());
180 full && buf.len() < self.size
181 };
182 if freed {
183 self.notifier.notify();
184 }
185 }
186}
187
188struct Token(Option<Rc<Shared>>);
190
191impl Token {
192 fn disarm(mut self) {
193 self.0 = None;
194 }
195}
196
197impl Drop for Token {
198 fn drop(&mut self) {
199 if let Some(shared) = self.0.take() {
200 shared.release_next();
201 }
202 }
203}
204
205struct Waiter<'a> {
207 shared: &'a Shared,
208 rx: Option<oneshot::Receiver<Token>>,
209}
210
211impl Drop for Waiter<'_> {
212 fn drop(&mut self) {
213 if let Some(rx) = self.rx.take() {
214 drop(rx);
216 self.shared.prune();
217 }
218 }
219}
220
221struct InFlight<'a>(&'a Shared);
222
223impl Drop for InFlight<'_> {
224 fn drop(&mut self) {
225 self.0.in_flight.set(self.0.in_flight.get() - 1);
226 self.0.drained.wake();
227 }
228}
229
230impl<St, Req, Res, Err> BufferService<St, Req, Res, Err>
231where
232 St: Clone + 'static,
233{
234 #[must_use]
235 pub fn new(size: usize, service: PipelineState<St, Req, Res, Err>) -> Self {
237 let size = size.max(1);
238 Self {
239 shared: Rc::new(Shared {
240 size,
241 buf: RefCell::new(VecDeque::with_capacity(size)),
242 has_token: Cell::new(false),
243 in_flight: Cell::new(0),
244 canceled: Cell::new(false),
245 notifier: ReadyNotifier::new(),
246 token_waker: LocalWaker::new(),
247 drained: LocalWaker::new(),
248 }),
249 service,
250 cancel_on_shutdown: false,
251 }
252 }
253
254 #[must_use]
255 pub fn cancel_on_shutdown(self) -> Self {
257 Self {
258 cancel_on_shutdown: true,
259 ..self
260 }
261 }
262}
263
264impl<St, Req, Res, Err> fmt::Debug for BufferService<St, Req, Res, Err> {
265 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
266 f.debug_struct("BufferService")
267 .field("size", &self.shared.size)
268 .field("cancel_on_shutdown", &self.cancel_on_shutdown)
269 .field("buffered", &self.shared.buf.borrow().len())
270 .field("in_flight", &self.shared.in_flight.get())
271 .field("service", &self.service)
272 .finish()
273 }
274}
275
276impl<St, Req, Res, Err> Service<St, Req> for BufferService<St, Req, Res, Err>
277where
278 St: Clone + 'static,
279 Req: 'static,
280 Res: 'static,
281 Err: 'static,
282{
283 type Res = Res;
284 type Error = BufferServiceError<Err>;
285
286 fn bind(&self, ctx: BindCtx<'_>) {
287 ctx.subscribe(&self.shared.notifier);
289 }
290
291 async fn ready(&self, _: Ctx<'_, Self, St>) -> Result<(), Self::Error> {
292 if self.shared.has_room() {
293 Ok(())
294 } else {
295 log::trace!("Buffer limit exceeded");
296 future::pending().await
298 }
299 }
300
301 async fn shutdown(&self, ctx: Ctx<'_, Self, St>) {
302 let shared = &self.shared;
303 if self.cancel_on_shutdown {
304 shared.canceled.set(true);
305 let buf = mem::take(&mut *shared.buf.borrow_mut());
306 let freed = buf.len() >= shared.size;
307 let _cancel = OnDropFn::new(|| {
308 shared.token_waker.wake();
309 drop(buf);
311 });
312 if freed {
313 shared.notifier.notify();
314 }
315 }
316
317 future::poll_fn(|cx| {
318 if shared.is_drained() {
319 return Poll::Ready(());
320 }
321 shared.drained.register(cx.waker());
322 if shared.is_drained() {
323 Poll::Ready(())
324 } else {
325 Poll::Pending
326 }
327 })
328 .await;
329
330 self.service.shutdown(ctx.st()).await;
331 }
332
333 async fn call(&self, req: Req, ctx: Ctx<'_, Self, St>) -> Result<Res, Self::Error> {
334 let shared = &self.shared;
335 let idle = !shared.has_token.get() && shared.buf.borrow().is_empty();
336 if idle && self.service.is_ready() {
337 return Ok(self.service.call(req, ctx.st()).await?);
338 }
339
340 let (token, mut buffered) = if idle {
341 shared.has_token.set(true);
342 (Token(Some(shared.clone())), false)
343 } else {
344 let (tx, rx) = oneshot::channel();
345 let full = {
346 let mut buf = shared.buf.borrow_mut();
347 buf.push_back(tx);
348 buf.len() == shared.size
349 };
350 let mut waiter = Waiter {
351 shared,
352 rx: Some(rx),
353 };
354 if full {
355 shared.notifier.notify();
356 }
357 let result = future::poll_fn(|cx| waiter.rx.as_ref().unwrap().poll_recv(cx)).await;
358 waiter.rx = None;
359 let token = result.map_err(|_| {
360 log::trace!("Buffered service request canceled");
361 BufferServiceError::RequestCanceled
362 })?;
363 (token, true)
364 };
365
366 let mut ready = pin::pin!(self.service.ready(ctx.st()));
368 future::poll_fn(|cx| {
369 if shared.canceled.get() {
370 return Poll::Ready(Err(BufferServiceError::RequestCanceled));
371 }
372 if let Poll::Ready(result) = ready.as_mut().poll(cx) {
373 Poll::Ready(result.map_err(BufferServiceError::Service))
374 } else {
375 buffered = true;
376 shared.token_waker.register(cx.waker());
377 Poll::Pending
378 }
379 })
380 .await?;
381
382 let _in_flight = buffered.then(|| {
383 shared.in_flight.set(shared.in_flight.get() + 1);
384 InFlight(shared)
385 });
386 drop(token);
387
388 Ok(self.service.call(req, ctx.st()).await?)
389 }
390}
391
392#[cfg(test)]
393mod tests {
394 #![allow(clippy::unused_async_trait_impl)]
395 use ntex_service::{Pipeline, ReadyCallback, apply, fn_factory};
396 use std::{cell::RefCell, pin::Pin, time::Duration};
397 use std::{panic::AssertUnwindSafe, panic::catch_unwind, task::Context};
398
399 use super::*;
400 use crate::future::lazy;
401
402 #[derive(Debug, Clone)]
403 struct TestService(Rc<Inner>);
404
405 #[derive(Debug, Default)]
406 struct Inner {
407 ready: Cell<bool>,
408 waker: LocalWaker,
409 changed: ReadyNotifier,
410 count: Cell<usize>,
411 }
412
413 impl Inner {
414 fn new(ready: bool) -> Rc<Self> {
415 let inner = Rc::new(Self::default());
416 inner.ready.set(ready);
417 inner
418 }
419
420 fn set_ready(&self) {
421 self.ready.set(true);
422 self.waker.wake();
423 }
424 }
425
426 impl Service<(), ()> for TestService {
427 type Res = ();
428 type Error = ();
429
430 fn bind(&self, ctx: BindCtx<'_>) {
431 ctx.subscribe(&self.0.changed);
432 }
433
434 async fn ready(&self, _: Ctx<'_, Self, ()>) -> Result<(), Self::Error> {
435 future::poll_fn(|cx| {
436 self.0.waker.register(cx.waker());
437 if self.0.ready.get() {
438 Poll::Ready(Ok(()))
439 } else {
440 Poll::Pending
441 }
442 })
443 .await
444 }
445
446 async fn call(&self, _r: (), _: Ctx<'_, Self, ()>) -> Result<(), ()> {
447 self.0.ready.set(false);
449 self.0.changed.notify();
450 self.0.count.set(self.0.count.get() + 1);
451 Ok(())
452 }
453 }
454
455 fn spawn_call(
456 srv: &Pipeline<(), (), BufferServiceError<()>>,
457 ) -> Rc<RefCell<Option<Result<(), BufferServiceError<()>>>>> {
458 let result = Rc::new(RefCell::new(None));
459 let result2 = result.clone();
460 let srv = srv.bind();
461 ntex::rt::spawn(async move {
462 *result2.borrow_mut() = Some(srv.call(()).await);
463 });
464 result
465 }
466
467 async fn settle() {
468 crate::time::sleep(Duration::from_millis(25)).await;
469 }
470
471 #[ntex::test]
477 async fn test_service() {
478 let inner = Inner::new(false);
479 let svc = BufferService::new(2, PipelineState::new(TestService(inner.clone())));
480 assert!(format!("{svc:?}").contains("BufferService"));
481
482 let srv = Pipeline::new((), svc);
483 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
484
485 let _r1 = spawn_call(&srv);
487 let _r2 = spawn_call(&srv);
488 settle().await;
489 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
490 let _r3 = spawn_call(&srv);
491 settle().await;
492 assert_eq!(inner.count.get(), 0);
493 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
494
495 inner.set_ready();
497 settle().await;
498 assert_eq!(inner.count.get(), 1);
499 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
500
501 inner.set_ready();
502 settle().await;
503 assert_eq!(inner.count.get(), 2);
504
505 inner.set_ready();
506 settle().await;
507 assert_eq!(inner.count.get(), 3);
508 assert!(lazy(|cx| srv.poll_shutdown(cx)).await.is_ready());
509
510 let err = BufferServiceError::from("test");
511 assert!(format!("{err}").contains("test"));
512 assert!(format!("{:?}", Buffer::<(), (), (), ()>::default()).contains("Buffer"));
513 }
514
515 #[ntex::test]
520 async fn sequential_calls_with_ready_service() {
521 let inner = Rc::new(PendingInner::default());
522 let srv = Pipeline::new(
523 (),
524 BufferService::new(2, PipelineState::new(ReadyService(inner.clone()))),
525 );
526
527 for _ in 0..3 {
528 crate::time::timeout(Duration::from_secs(1), srv.call(()))
529 .await
530 .unwrap()
531 .unwrap();
532 }
533 assert_eq!(inner.count.get(), 3);
534 }
535
536 #[ntex::test]
542 #[allow(clippy::redundant_clone)]
543 async fn test_middleware() {
544 let inner = Inner::new(false);
545 let inner2 = inner.clone();
546
547 let srv = apply(
548 Buffer::default().buf_size(1),
549 fn_factory(async move |(): &()| Ok::<_, ()>(TestService(inner2.clone()))),
550 );
551
552 let srv = srv.pipeline(()).await.unwrap();
553 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
554
555 let srv1 = srv.bind();
556 ntex::rt::spawn(async move {
557 let _ = srv1.call(()).await;
558 });
559 let srv1 = srv.bind();
560 ntex::rt::spawn(async move {
561 let _ = srv1.call(()).await;
562 });
563 settle().await;
564 assert_eq!(inner.count.get(), 0);
565 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
566
567 inner.set_ready();
568 settle().await;
569 assert_eq!(inner.count.get(), 1);
570 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
571
572 inner.set_ready();
573 settle().await;
574 assert_eq!(inner.count.get(), 2);
575 }
576
577 #[ntex::test]
583 async fn middleware_cancels_buffered_requests_on_shutdown() {
584 let inner = Inner::new(false);
585 let inner2 = inner.clone();
586
587 let srv = apply(
588 Buffer::default().buf_size(1).cancel_on_shutdown(),
589 fn_factory(async move |(): &()| Ok::<_, ()>(TestService(inner2.clone()))),
590 )
591 .pipeline(())
592 .await
593 .unwrap();
594 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
595
596 let r1 = spawn_call(&srv);
598 let r2 = spawn_call(&srv);
599 settle().await;
600 crate::time::timeout(Duration::from_secs(1), srv.shutdown())
601 .await
602 .unwrap();
603 settle().await;
604
605 assert!(matches!(
606 *r1.borrow(),
607 Some(Err(BufferServiceError::RequestCanceled))
608 ));
609 assert!(matches!(
610 *r2.borrow(),
611 Some(Err(BufferServiceError::RequestCanceled))
612 ));
613 assert_eq!(inner.count.get(), 0);
614 }
615
616 #[derive(Default)]
617 struct PendingInner {
618 started: Cell<bool>,
619 completed: Cell<bool>,
620 shutdown: Cell<bool>,
621 count: Cell<usize>,
622 }
623
624 struct ReadyService(Rc<PendingInner>);
625
626 impl Service<(), ()> for ReadyService {
627 type Res = ();
628 type Error = ();
629
630 async fn call(&self, (): (), _: Ctx<'_, Self, ()>) -> Result<(), ()> {
631 self.0.count.set(self.0.count.get() + 1);
632 Ok(())
633 }
634 }
635
636 struct PendingService {
637 inner: Rc<PendingInner>,
638 release: RefCell<Option<oneshot::Receiver<()>>>,
639 }
640
641 impl Service<(), ()> for PendingService {
642 type Res = ();
643 type Error = ();
644
645 async fn call(&self, (): (), _: Ctx<'_, Self, ()>) -> Result<(), ()> {
646 self.inner.started.set(true);
647 let release = self.release.borrow_mut().take().unwrap();
648 release.recv().await.unwrap();
649 self.inner.completed.set(true);
650 Ok(())
651 }
652
653 async fn shutdown(&self, _: Ctx<'_, Self, ()>) {
654 self.inner.shutdown.set(true);
655 }
656 }
657
658 #[ntex::test]
664 async fn shutdown_with_in_flight_request() {
665 let inner = Rc::new(PendingInner::default());
666 let (release_tx, release_rx) = oneshot::channel();
667 let service = PendingService {
668 inner: inner.clone(),
669 release: RefCell::new(Some(release_rx)),
670 };
671 let srv = Pipeline::new((), BufferService::new(1, PipelineState::new(service)));
672
673 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
674
675 let srv2 = srv.bind();
676 ntex::rt::spawn(async move {
677 srv2.call(()).await.unwrap();
678 });
679 settle().await;
680
681 assert!(inner.started.get());
683 assert!(!inner.completed.get());
684 assert!(lazy(|cx| srv.poll_shutdown(cx)).await.is_ready());
685 assert!(inner.shutdown.get());
686 assert!(!inner.completed.get());
687
688 release_tx.send(()).unwrap();
689 settle().await;
690
691 assert!(inner.completed.get());
692 }
693
694 #[derive(Default)]
695 struct DrainInner {
696 ready: Cell<bool>,
697 waker: LocalWaker,
698 active: Cell<usize>,
699 count: Cell<usize>,
700 active_on_shutdown: Cell<Option<usize>>,
701 }
702
703 struct DrainService(Rc<DrainInner>);
704
705 impl Service<(), ()> for DrainService {
706 type Res = ();
707 type Error = ();
708
709 async fn ready(&self, _: Ctx<'_, Self, ()>) -> Result<(), ()> {
710 future::poll_fn(|cx| {
711 self.0.waker.register(cx.waker());
712 if self.0.ready.get() {
713 Poll::Ready(Ok(()))
714 } else {
715 Poll::Pending
716 }
717 })
718 .await
719 }
720
721 async fn call(&self, (): (), _: Ctx<'_, Self, ()>) -> Result<(), ()> {
722 let inner = &self.0;
723 inner.active.set(inner.active.get() + 1);
724 crate::time::sleep(Duration::from_millis(25)).await;
725 inner.active.set(inner.active.get() - 1);
726 inner.count.set(inner.count.get() + 1);
727 Ok(())
728 }
729
730 async fn shutdown(&self, _: Ctx<'_, Self, ()>) {
731 self.0.active_on_shutdown.set(Some(self.0.active.get()));
732 }
733 }
734
735 #[ntex::test]
741 async fn shutdown_drains_buffer() {
742 let inner = Rc::new(DrainInner::default());
743 let srv = Pipeline::new(
744 (),
745 BufferService::new(4, PipelineState::new(DrainService(inner.clone()))),
746 );
747
748 for _ in 0..2 {
749 srv.ready().await.unwrap();
750 let srv = srv.bind();
751 ntex::rt::spawn(async move {
752 srv.call(()).await.unwrap();
753 });
754 }
755 crate::time::sleep(Duration::from_millis(10)).await;
756 assert_eq!(inner.count.get(), 0);
757
758 inner.ready.set(true);
759 inner.waker.wake();
760 srv.shutdown().await;
761
762 assert_eq!(inner.count.get(), 2);
763 assert_eq!(inner.active_on_shutdown.get(), Some(0));
764 }
765
766 #[test]
771 fn request_canceled_display() {
772 let err = BufferServiceError::<&str>::RequestCanceled;
773 assert_eq!(err.to_string(), "buffer service request canceled");
774 }
775
776 #[ntex::test]
782 async fn dropped_requests_are_skipped() {
783 let inner = Inner::new(false);
784 let srv = Pipeline::new(
785 (),
786 BufferService::new(1, PipelineState::new(TestService(inner.clone()))),
787 );
788 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
789
790 let mut fut1 = srv.call_static(());
792 assert!(lazy(|cx| Pin::new(&mut fut1).poll(cx)).await.is_pending());
793 let mut fut2 = srv.call_static(());
794 assert!(lazy(|cx| Pin::new(&mut fut2).poll(cx)).await.is_pending());
795 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Pending);
796
797 drop(fut2);
799 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
800 let r3 = spawn_call(&srv);
801 settle().await;
802
803 drop(fut1);
805 inner.set_ready();
806 settle().await;
807 assert!(matches!(*r3.borrow(), Some(Ok(()))));
808 assert_eq!(inner.count.get(), 1);
809
810 let mut fut4 = srv.call_static(());
812 assert!(lazy(|cx| Pin::new(&mut fut4).poll(cx)).await.is_pending());
813 let mut fut5 = srv.call_static(());
814 assert!(lazy(|cx| Pin::new(&mut fut5).poll(cx)).await.is_pending());
815 drop(fut4);
816 drop(fut5);
817 inner.set_ready();
818 crate::time::timeout(Duration::from_secs(1), srv.call(()))
819 .await
820 .unwrap()
821 .unwrap();
822 assert_eq!(inner.count.get(), 2);
823 crate::time::timeout(Duration::from_secs(1), srv.shutdown())
824 .await
825 .unwrap();
826 }
827
828 #[ntex::test]
834 async fn shutdown_skips_dropped_buffered_request() {
835 let inner = Inner::new(false);
836 let srv = Pipeline::new(
837 (),
838 BufferService::new(2, PipelineState::new(TestService(inner.clone()))),
839 );
840 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
841
842 let mut fut1 = srv.call_static(());
843 assert!(lazy(|cx| Pin::new(&mut fut1).poll(cx)).await.is_pending());
844 let mut fut2 = srv.call_static(());
845 assert!(lazy(|cx| Pin::new(&mut fut2).poll(cx)).await.is_pending());
846 drop(fut2);
847 drop(fut1);
848
849 crate::time::timeout(Duration::from_secs(1), srv.shutdown())
850 .await
851 .unwrap();
852 assert_eq!(inner.count.get(), 0);
853 }
854
855 #[derive(Default)]
856 struct FailInner {
857 fail: Cell<bool>,
858 changed: ReadyNotifier,
859 }
860
861 struct FailService(Rc<FailInner>);
862
863 impl Service<(), ()> for FailService {
864 type Res = ();
865 type Error = ();
866
867 fn bind(&self, ctx: BindCtx<'_>) {
868 ctx.subscribe(&self.0.changed);
869 }
870
871 async fn ready(&self, _: Ctx<'_, Self, ()>) -> Result<(), ()> {
872 if self.0.fail.get() {
873 Err(())
874 } else {
875 future::pending().await
876 }
877 }
878
879 async fn call(&self, (): (), _: Ctx<'_, Self, ()>) -> Result<(), ()> {
880 unreachable!("FailService is never ready")
881 }
882 }
883
884 #[ntex::test]
890 async fn inner_ready_error_is_returned_to_requests() {
891 let inner = Rc::new(FailInner::default());
892 let srv = Pipeline::new(
893 (),
894 BufferService::new(2, PipelineState::new(FailService(inner.clone()))),
895 );
896 assert_eq!(lazy(|cx| srv.poll_ready(cx)).await, Poll::Ready(Ok(())));
897
898 let r1 = spawn_call(&srv);
899 let r2 = spawn_call(&srv);
900 settle().await;
901 assert!(r1.borrow().is_none());
902
903 inner.fail.set(true);
904 inner.changed.notify();
905 crate::time::timeout(Duration::from_secs(1), srv.shutdown())
906 .await
907 .unwrap();
908 settle().await;
909
910 assert!(matches!(
911 *r1.borrow(),
912 Some(Err(BufferServiceError::Service(())))
913 ));
914 assert!(matches!(
915 *r2.borrow(),
916 Some(Err(BufferServiceError::Service(())))
917 ));
918 }
919
920 #[ntex::test]
926 async fn shared_buffer_capacity_notifies_all_roots() {
927 let inner = Inner::new(false);
928 let buffer = Rc::new(BufferService::new(
929 1,
930 PipelineState::new(TestService(inner.clone())),
931 ));
932 let first = Pipeline::new((), buffer.clone());
933 let second = Pipeline::new((), buffer);
934 first.ready().await.unwrap();
935 second.ready().await.unwrap();
936 let mut cx = Context::from_waker(std::task::Waker::noop());
937
938 let mut active = Box::pin(first.call(()));
939 assert!(active.as_mut().poll(&mut cx).is_pending());
940 let mut queued = Box::pin(second.call(()));
941 assert!(queued.as_mut().poll(&mut cx).is_pending());
942 assert!(first.poll_ready(&mut cx).is_pending());
943 assert!(second.poll_ready(&mut cx).is_pending());
944
945 drop(queued);
946 assert_eq!(first.poll_ready(&mut cx), Poll::Ready(Ok(())));
947 assert_eq!(second.poll_ready(&mut cx), Poll::Ready(Ok(())));
948 inner.set_ready();
949 assert_eq!(active.as_mut().poll(&mut cx), Poll::Ready(Ok(())));
950 }
951
952 #[ntex::test]
958 async fn notification_panics_do_not_strand_buffered_requests() {
959 let inner = Inner::new(false);
960 let buffer = Rc::new(BufferService::new(
961 1,
962 PipelineState::new(TestService(inner.clone())),
963 ));
964 let srv = Pipeline::new((), buffer.clone());
965 let panic_once = Rc::new(Cell::new(false));
966 let flag = panic_once.clone();
967 srv.on_notify(&ReadyCallback::new(move |_| {
968 assert!(!flag.replace(false), "notification panic");
969 }));
970 let mut cx = Context::from_waker(std::task::Waker::noop());
971 let mut active = Box::pin(srv.call(()));
972 assert!(active.as_mut().poll(&mut cx).is_pending());
973
974 panic_once.set(true);
975 let mut queued = Box::pin(srv.call(()));
976 assert!(catch_unwind(AssertUnwindSafe(|| queued.as_mut().poll(&mut cx))).is_err());
977 drop(queued);
978 assert!(buffer.shared.buf.borrow().is_empty());
979 assert_eq!(srv.poll_ready(&mut cx), Poll::Ready(Ok(())));
980
981 let mut queued = Box::pin(srv.call(()));
982 assert!(queued.as_mut().poll(&mut cx).is_pending());
983 panic_once.set(true);
984 assert!(catch_unwind(AssertUnwindSafe(|| drop(active))).is_err());
985 assert_eq!(
986 queued.as_mut().poll(&mut cx),
987 Poll::Ready(Err(BufferServiceError::RequestCanceled))
988 );
989 assert!(buffer.shared.is_drained());
990
991 inner.set_ready();
992 assert_eq!(srv.call(()).await, Ok(()));
993 }
994
995 #[ntex::test]
1001 async fn shutdown_notification_panic_still_cancels_waiters() {
1002 let inner = Inner::new(false);
1003 let buffer = Rc::new(
1004 BufferService::new(1, PipelineState::new(TestService(inner))).cancel_on_shutdown(),
1005 );
1006 let srv = Pipeline::new((), buffer.clone());
1007 let mut cx = Context::from_waker(std::task::Waker::noop());
1008 let mut active = Box::pin(srv.call(()));
1009 assert!(active.as_mut().poll(&mut cx).is_pending());
1010 let mut queued = Box::pin(srv.call(()));
1011 assert!(queued.as_mut().poll(&mut cx).is_pending());
1012 srv.on_notify(&ReadyCallback::new(|_| panic!("shutdown notification")));
1013
1014 assert!(catch_unwind(AssertUnwindSafe(|| srv.poll_shutdown(&mut cx))).is_err());
1015 assert_eq!(
1016 active.as_mut().poll(&mut cx),
1017 Poll::Ready(Err(BufferServiceError::RequestCanceled))
1018 );
1019 assert_eq!(
1020 queued.as_mut().poll(&mut cx),
1021 Poll::Ready(Err(BufferServiceError::RequestCanceled))
1022 );
1023 assert!(buffer.shared.is_drained());
1024 }
1025}