hyper_util/client/pool/
cache.rs

1//! A cache of services
2//!
3//! The cache is a single list of cached services, bundled with a `MakeService`.
4//! Calling the cache returns either an existing service, or makes a new one.
5//! The returned `impl Service` can be used to send requests, and when dropped,
6//! it will try to be returned back to the cache.
7
8pub use self::internal::builder;
9
10#[cfg(docsrs)]
11pub use self::internal::Builder;
12#[cfg(docsrs)]
13pub use self::internal::Cache;
14#[cfg(docsrs)]
15pub use self::internal::Cached;
16
17// For now, nothing else in this module is nameable. We can always make things
18// more public, but we can't change type shapes (generics) once things are
19// public.
20mod internal {
21    use std::collections::VecDeque;
22    use std::fmt;
23    use std::pin::Pin;
24    use std::sync::{Arc, Mutex, Weak};
25    use std::task::{self, Poll, Waker, ready};
26
27    use tower_service::Service;
28
29    use super::events;
30
31    /// Start a builder to construct a `Cache` pool.
32    pub fn builder() -> Builder<events::Ignore> {
33        Builder {
34            events: events::Ignore,
35        }
36    }
37
38    /// A cache pool of services from the inner make service.
39    ///
40    /// Created with [`builder()`].
41    ///
42    /// # Unnameable
43    ///
44    /// This type is normally unnameable, forbidding naming of the type within
45    /// code. The type is exposed in the documentation to show which methods
46    /// can be publicly called.
47    #[derive(Debug)]
48    pub struct Cache<M, Dst, Ev>
49    where
50        M: Service<Dst>,
51    {
52        connector: M,
53        shared: Arc<Mutex<Shared<M::Response>>>,
54        events: Ev,
55        ready: Ready<M::Response>,
56        ready_waiter: Option<WaiterId>,
57    }
58
59    /// A builder to configure a `Cache`.
60    ///
61    /// # Unnameable
62    ///
63    /// This type is normally unnameable, forbidding naming of the type within
64    /// code. The type is exposed in the documentation to show which methods
65    /// can be publicly called.
66    #[derive(Debug)]
67    pub struct Builder<Ev> {
68        events: Ev,
69    }
70
71    /// A cached service returned from a [`Cache`].
72    ///
73    /// Implements `Service` by delegating to the inner service. Once dropped,
74    /// tries to reinsert into the `Cache`.
75    ///
76    /// # Unnameable
77    ///
78    /// This type is normally unnameable, forbidding naming of the type within
79    /// code. The type is exposed in the documentation to show which methods
80    /// can be publicly called.
81    pub struct Cached<S> {
82        is_closed: bool,
83        inner: Option<S>,
84        shared: Weak<Mutex<Shared<S>>>,
85        // todo: on_idle
86    }
87
88    #[derive(Debug)]
89    enum Ready<S> {
90        None,
91        Cached(S),
92    }
93
94    pub enum CacheFuture<M, Dst, Ev>
95    where
96        M: Service<Dst>,
97    {
98        Racing {
99            shared: Arc<Mutex<Shared<M::Response>>>,
100            waiter: WaiterId,
101            future: Option<M::Future>,
102            events: Ev,
103        },
104        Cached {
105            svc: Option<Cached<M::Response>>,
106        },
107    }
108
109    // shouldn't be pub
110    #[derive(Debug)]
111    pub struct Shared<S> {
112        services: Vec<S>,
113        waiters: VecDeque<Waiter>,
114        reservations: Vec<(WaiterId, S)>,
115        next_waiter: usize,
116    }
117
118    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
119    pub struct WaiterId(usize);
120
121    #[derive(Debug)]
122    struct Waiter {
123        id: WaiterId,
124        waker: Option<Waker>,
125    }
126
127    // impl Builder
128
129    impl<Ev> Builder<Ev> {
130        /// Provide a `Future` executor to be used by the `Cache`.
131        ///
132        /// The executor is used handle some optional background tasks that
133        /// can improve the behavior of the cache, such as reducing connection
134        /// thrashing when a race is won. If not configured with an executor,
135        /// the default behavior is to ignore any of these optional background
136        /// tasks.
137        ///
138        /// The executor should implmenent [`hyper::rt::Executor`].
139        ///
140        /// # Example
141        ///
142        /// ```rust
143        /// # #[cfg(feature = "tokio")]
144        /// # fn run() {
145        /// let builder = hyper_util::client::pool::cache::builder()
146        ///     .executor(hyper_util::rt::TokioExecutor::new());
147        /// # }
148        /// ```
149        pub fn executor<E>(self, exec: E) -> Builder<events::WithExecutor<E>> {
150            Builder {
151                events: events::WithExecutor(exec),
152            }
153        }
154
155        /// Build a `Cache` pool around the `connector`.
156        pub fn build<M, Dst>(self, connector: M) -> Cache<M, Dst, Ev>
157        where
158            M: Service<Dst>,
159        {
160            Cache {
161                connector,
162                events: self.events,
163                ready: Ready::None,
164                ready_waiter: None,
165                shared: Arc::new(Mutex::new(Shared {
166                    services: Vec::new(),
167                    waiters: VecDeque::new(),
168                    reservations: Vec::new(),
169                    next_waiter: 0,
170                })),
171            }
172        }
173    }
174
175    // impl Cache
176
177    impl<M, Dst, Ev> Cache<M, Dst, Ev>
178    where
179        M: Service<Dst>,
180    {
181        /// Retain all cached services indicated by the predicate.
182        pub fn retain<F>(&mut self, predicate: F)
183        where
184            F: FnMut(&mut M::Response) -> bool,
185        {
186            let mut predicate = predicate;
187            if let Ready::Cached(svc) = &mut self.ready {
188                if !predicate(svc) {
189                    self.ready = Ready::None;
190                }
191            }
192
193            self.shared.lock().unwrap().services.retain_mut(predicate);
194        }
195
196        /// Check whether this cache has no cached services.
197        pub fn is_empty(&self) -> bool {
198            matches!(self.ready, Ready::None) && self.shared.lock().unwrap().services.is_empty()
199        }
200    }
201
202    impl<M, Dst, Ev> Service<Dst> for Cache<M, Dst, Ev>
203    where
204        M: Service<Dst>,
205        M::Future: Unpin,
206        M::Response: Unpin,
207        Ev: events::Events<BackgroundConnect<M::Future, M::Response>> + Clone + Unpin,
208    {
209        type Response = Cached<M::Response>;
210        type Error = M::Error;
211        type Future = CacheFuture<M, Dst, Ev>;
212
213        fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
214            match self.ready {
215                Ready::Cached(_) => return Poll::Ready(Ok(())),
216                Ready::None => {}
217            }
218
219            {
220                let mut shared = self.shared.lock().unwrap();
221                if let Some(id) = self.ready_waiter {
222                    if let Some(svc) = shared.take_reserved(id) {
223                        self.ready_waiter = None;
224                        self.ready = Ready::Cached(svc);
225                        return Poll::Ready(Ok(()));
226                    }
227                } else if let Some(svc) = shared.take_available() {
228                    self.ready = Ready::Cached(svc);
229                    return Poll::Ready(Ok(()));
230                }
231
232                let id = *self
233                    .ready_waiter
234                    .get_or_insert_with(|| shared.push_waiter());
235                shared.store_waker(id, cx.waker());
236            }
237
238            match self.connector.poll_ready(cx) {
239                Poll::Ready(result) => {
240                    if let Some(id) = self.ready_waiter.take() {
241                        self.shared.lock().unwrap().cancel_waiter(id);
242                    }
243                    Poll::Ready(result)
244                }
245                Poll::Pending => Poll::Pending,
246            }
247        }
248
249        fn call(&mut self, target: Dst) -> Self::Future {
250            // 1. If already cached, easy!
251            match std::mem::replace(&mut self.ready, Ready::None) {
252                Ready::Cached(svc) => {
253                    return CacheFuture::Cached {
254                        svc: Some(Cached::new(svc, Arc::downgrade(&self.shared))),
255                    };
256                }
257                Ready::None => {
258                    if let Some(id) = self.ready_waiter.take() {
259                        let mut shared = self.shared.lock().unwrap();
260                        if let Some(svc) = shared.take_reserved(id) {
261                            return CacheFuture::Cached {
262                                svc: Some(Cached::new(svc, Arc::downgrade(&self.shared))),
263                            };
264                        }
265                        shared.cancel_waiter(id);
266                    }
267                    if let Some(svc) = self.shared.lock().unwrap().take_available() {
268                        return CacheFuture::Cached {
269                            svc: Some(Cached::new(svc, Arc::downgrade(&self.shared))),
270                        };
271                    }
272                }
273            }
274
275            let waiter = {
276                let mut locked = self.shared.lock().unwrap();
277                locked.push_waiter()
278            };
279
280            // 2. Otherwise, we start a new connect, and also listen for
281            //    any newly idle.
282            CacheFuture::Racing {
283                shared: self.shared.clone(),
284                waiter,
285                future: Some(self.connector.call(target)),
286                events: self.events.clone(),
287            }
288        }
289    }
290
291    impl<M, Dst, Ev> Clone for Cache<M, Dst, Ev>
292    where
293        M: Service<Dst> + Clone,
294        Ev: Clone,
295    {
296        fn clone(&self) -> Self {
297            Self {
298                connector: self.connector.clone(),
299                events: self.events.clone(),
300                shared: self.shared.clone(),
301                ready: Ready::None,
302                ready_waiter: None,
303            }
304        }
305    }
306
307    impl<M, Dst, Ev> Drop for Cache<M, Dst, Ev>
308    where
309        M: Service<Dst>,
310    {
311        fn drop(&mut self) {
312            if let Ready::Cached(svc) = std::mem::replace(&mut self.ready, Ready::None) {
313                if let Ok(mut shared) = self.shared.lock() {
314                    shared.put(svc);
315                }
316            }
317            if let Some(id) = self.ready_waiter.take() {
318                if let Ok(mut shared) = self.shared.lock() {
319                    shared.cancel_waiter(id);
320                }
321            }
322        }
323    }
324
325    impl<M, Dst, Ev> Drop for CacheFuture<M, Dst, Ev>
326    where
327        M: Service<Dst>,
328    {
329        fn drop(&mut self) {
330            if let CacheFuture::Racing { shared, waiter, .. } = self {
331                if let Ok(mut shared) = shared.lock() {
332                    shared.cancel_waiter(*waiter);
333                }
334            }
335        }
336    }
337
338    impl<M, Dst, Ev> Future for CacheFuture<M, Dst, Ev>
339    where
340        M: Service<Dst>,
341        M::Future: Unpin,
342        M::Response: Unpin,
343        Ev: events::Events<BackgroundConnect<M::Future, M::Response>> + Unpin,
344    {
345        type Output = Result<Cached<M::Response>, M::Error>;
346
347        fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
348            match &mut *self.as_mut() {
349                CacheFuture::Racing {
350                    shared,
351                    waiter,
352                    future,
353                    events,
354                } => {
355                    {
356                        let mut locked = shared.lock().unwrap();
357                        if let Some(pool_got) = locked.take_reserved(*waiter) {
358                            events.on_race_lost(BackgroundConnect {
359                                future: future.take().expect("racing future polled after done"),
360                                shared: Arc::downgrade(&shared),
361                            });
362                            return Poll::Ready(Ok(Cached::new(pool_got, Arc::downgrade(&shared))));
363                        }
364                        locked.store_waker(*waiter, cx.waker());
365                    }
366
367                    let connected = match ready!(
368                        Pin::new(future.as_mut().expect("racing future polled after done"))
369                            .poll(cx)
370                    ) {
371                        Ok(inner) => inner,
372                        Err(err) => {
373                            shared.lock().unwrap().cancel_waiter(*waiter);
374                            return Poll::Ready(Err(err));
375                        }
376                    };
377
378                    shared.lock().unwrap().cancel_waiter(*waiter);
379                    Poll::Ready(Ok(Cached::new(connected, Arc::downgrade(&shared))))
380                }
381                CacheFuture::Cached { svc } => Poll::Ready(Ok(svc.take().unwrap())),
382            }
383        }
384    }
385
386    // impl Cached
387
388    impl<S> Cached<S> {
389        fn new(inner: S, shared: Weak<Mutex<Shared<S>>>) -> Self {
390            Cached {
391                is_closed: false,
392                inner: Some(inner),
393                shared,
394            }
395        }
396
397        // TODO: inner()? looks like `tower` likes `get_ref()` and `get_mut()`.
398
399        /// Get a reference to the inner service.
400        pub fn inner(&self) -> &S {
401            self.inner.as_ref().expect("inner only taken in drop")
402        }
403
404        /// Get a mutable reference to the inner service.
405        pub fn inner_mut(&mut self) -> &mut S {
406            self.inner.as_mut().expect("inner only taken in drop")
407        }
408    }
409
410    impl<S, Req> Service<Req> for Cached<S>
411    where
412        S: Service<Req>,
413    {
414        type Response = S::Response;
415        type Error = S::Error;
416        type Future = S::Future;
417
418        fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
419            self.inner.as_mut().unwrap().poll_ready(cx).map_err(|err| {
420                self.is_closed = true;
421                err
422            })
423        }
424
425        fn call(&mut self, req: Req) -> Self::Future {
426            self.inner.as_mut().unwrap().call(req)
427        }
428    }
429
430    impl<S> Drop for Cached<S> {
431        fn drop(&mut self) {
432            if self.is_closed {
433                return;
434            }
435            if let Some(value) = self.inner.take() {
436                if let Some(shared) = self.shared.upgrade() {
437                    if let Ok(mut shared) = shared.lock() {
438                        shared.put(value);
439                    }
440                }
441            }
442        }
443    }
444
445    impl<S: fmt::Debug> fmt::Debug for Cached<S> {
446        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
447            f.debug_tuple("Cached")
448                .field(self.inner.as_ref().unwrap())
449                .finish()
450        }
451    }
452
453    // impl Shared
454
455    impl<V> Shared<V> {
456        fn put(&mut self, val: V) {
457            if let Some(mut waiter) = self.waiters.pop_front() {
458                self.reservations.push((waiter.id, val));
459                if let Some(waker) = waiter.waker.take() {
460                    waker.wake();
461                }
462                return;
463            }
464
465            self.services.push(val);
466        }
467
468        fn take_available(&mut self) -> Option<V> {
469            if self.waiters.is_empty() {
470                self.services.pop()
471            } else {
472                None
473            }
474        }
475
476        fn push_waiter(&mut self) -> WaiterId {
477            let id = WaiterId(self.next_waiter);
478            self.next_waiter = self.next_waiter.wrapping_add(1);
479            self.waiters.push_back(Waiter { id, waker: None });
480            id
481        }
482
483        fn store_waker(&mut self, id: WaiterId, waker: &Waker) {
484            if let Some(waiter) = self.waiters.iter_mut().find(|waiter| waiter.id == id) {
485                if waiter
486                    .waker
487                    .as_ref()
488                    .is_none_or(|current| !current.will_wake(waker))
489                {
490                    waiter.waker = Some(waker.clone());
491                }
492            }
493        }
494
495        fn take_reserved(&mut self, id: WaiterId) -> Option<V> {
496            let index = self
497                .reservations
498                .iter()
499                .position(|(reserved_id, _)| *reserved_id == id)?;
500            Some(self.reservations.remove(index).1)
501        }
502
503        fn cancel_waiter(&mut self, id: WaiterId) {
504            if let Some(index) = self.waiters.iter().position(|waiter| waiter.id == id) {
505                self.waiters.remove(index);
506                return;
507            }
508
509            if let Some(svc) = self.take_reserved(id) {
510                self.put(svc);
511            }
512        }
513    }
514
515    pub struct BackgroundConnect<CF, S> {
516        future: CF,
517        shared: Weak<Mutex<Shared<S>>>,
518    }
519
520    impl<CF, S, E> Future for BackgroundConnect<CF, S>
521    where
522        CF: Future<Output = Result<S, E>> + Unpin,
523    {
524        type Output = ();
525
526        fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
527            match ready!(Pin::new(&mut self.future).poll(cx)) {
528                Ok(svc) => {
529                    if let Some(shared) = self.shared.upgrade() {
530                        if let Ok(mut locked) = shared.lock() {
531                            locked.put(svc);
532                        }
533                    }
534                    Poll::Ready(())
535                }
536                Err(_e) => Poll::Ready(()),
537            }
538        }
539    }
540}
541
542mod events {
543    #[derive(Clone, Debug)]
544    #[non_exhaustive]
545    pub struct Ignore;
546
547    #[derive(Clone, Debug)]
548    pub struct WithExecutor<E>(pub(super) E);
549
550    pub trait Events<CF> {
551        fn on_race_lost(&self, fut: CF);
552    }
553
554    impl<CF> Events<CF> for Ignore {
555        fn on_race_lost(&self, _fut: CF) {}
556    }
557
558    impl<E, CF> Events<CF> for WithExecutor<E>
559    where
560        E: hyper::rt::Executor<CF>,
561    {
562        fn on_race_lost(&self, fut: CF) {
563            self.0.execute(fut);
564        }
565    }
566}
567
568#[cfg(test)]
569mod tests {
570    use std::convert::Infallible;
571    use std::sync::{
572        Arc, Mutex,
573        atomic::{AtomicUsize, Ordering},
574    };
575    use std::task::{self, Poll};
576
577    use futures_util::future;
578    use tower_service::Service;
579    use tower_test::assert_request_eq;
580
581    #[tokio::test]
582    async fn test_makes_svc_when_empty() {
583        let (mock, mut handle) = tower_test::mock::pair();
584        let mut cache = super::builder().build(mock);
585        handle.allow(1);
586
587        std::future::poll_fn(|cx| cache.poll_ready(cx))
588            .await
589            .unwrap();
590
591        let f = cache.call(1);
592
593        future::join(f, async move {
594            assert_request_eq!(handle, 1).send_response("one");
595        })
596        .await
597        .0
598        .expect("call");
599    }
600
601    #[tokio::test]
602    async fn test_reuses_after_idle() {
603        let (mock, mut handle) = tower_test::mock::pair();
604        let mut cache = super::builder().build(mock);
605
606        // only 1 connection should ever be made
607        handle.allow(1);
608
609        std::future::poll_fn(|cx| cache.poll_ready(cx))
610            .await
611            .unwrap();
612        let f = cache.call(1);
613        let cached = future::join(f, async {
614            assert_request_eq!(handle, 1).send_response("one");
615        })
616        .await
617        .0
618        .expect("call");
619        drop(cached);
620
621        std::future::poll_fn(|cx| cache.poll_ready(cx))
622            .await
623            .unwrap();
624        let f = cache.call(1);
625        let cached = f.await.expect("call");
626        drop(cached);
627    }
628
629    // A returned connection is handed to waiters in the order they parked
630    // (FIFO), so a waiter cannot be starved by later arrivals.
631    #[tokio::test]
632    async fn test_waiters_woken_in_fifo_order() {
633        use std::task::{Context, Poll, Waker};
634
635        let (mock, mut handle) = tower_test::mock::pair::<u32, &'static str>();
636        let mut cache = super::builder().build(mock);
637        handle.allow(16);
638
639        // Establish one connection and hold it, so the next checkouts find no
640        // idle service and park.
641        std::future::poll_fn(|cx| cache.poll_ready(cx))
642            .await
643            .unwrap();
644        let held = future::join(cache.call(0), async {
645            assert_request_eq!(handle, 0).send_response("conn");
646        })
647        .await
648        .0
649        .expect("call");
650
651        // Park three checkouts in order. Each misses and starts a connect, but
652        // the connect is never completed, so each parks on its waiter.
653        let mut cx = Context::from_waker(Waker::noop());
654        std::future::poll_fn(|cx| cache.poll_ready(cx))
655            .await
656            .unwrap();
657        let mut first = Box::pin(cache.call(1));
658        assert!(first.as_mut().poll(&mut cx).is_pending());
659        std::future::poll_fn(|cx| cache.poll_ready(cx))
660            .await
661            .unwrap();
662        let mut second = Box::pin(cache.call(2));
663        assert!(second.as_mut().poll(&mut cx).is_pending());
664        std::future::poll_fn(|cx| cache.poll_ready(cx))
665            .await
666            .unwrap();
667        let mut third = Box::pin(cache.call(3));
668        assert!(third.as_mut().poll(&mut cx).is_pending());
669
670        // Returning the connection wakes the oldest waiter first.
671        drop(held);
672        let first = match first.as_mut().poll(&mut cx) {
673            Poll::Ready(r) => r.expect("first"),
674            Poll::Pending => panic!("oldest waiter was not woken first"),
675        };
676        assert!(second.as_mut().poll(&mut cx).is_pending());
677        assert!(third.as_mut().poll(&mut cx).is_pending());
678
679        // Returning it again wakes the next-oldest, then the last.
680        drop(first);
681        let second = match second.as_mut().poll(&mut cx) {
682            Poll::Ready(r) => r.expect("second"),
683            Poll::Pending => panic!("second waiter was not woken next"),
684        };
685        assert!(third.as_mut().poll(&mut cx).is_pending());
686
687        drop(second);
688        match third.as_mut().poll(&mut cx) {
689            Poll::Ready(r) => {
690                r.expect("third");
691            }
692            Poll::Pending => panic!("last waiter was not woken"),
693        }
694    }
695
696    #[tokio::test]
697    async fn dropped_racing_future_cancels_waiter() {
698        use std::task::{Context, Poll, Waker};
699
700        let (mock, mut handle) = tower_test::mock::pair::<u32, &'static str>();
701        let mut cache = super::builder().build(mock);
702        handle.allow(16);
703
704        std::future::poll_fn(|cx| cache.poll_ready(cx))
705            .await
706            .unwrap();
707        let held = future::join(cache.call(0), async {
708            assert_request_eq!(handle, 0).send_response("conn");
709        })
710        .await
711        .0
712        .expect("call");
713
714        std::future::poll_fn(|cx| cache.poll_ready(cx))
715            .await
716            .unwrap();
717        let mut dropped = Box::pin(cache.call(1));
718        let mut cx = Context::from_waker(Waker::noop());
719        assert!(dropped.as_mut().poll(&mut cx).is_pending());
720        drop(dropped);
721
722        drop(held);
723
724        std::future::poll_fn(|cx| cache.poll_ready(cx))
725            .await
726            .unwrap();
727        let mut reused = Box::pin(cache.call(2));
728        match reused.as_mut().poll(&mut cx) {
729            Poll::Ready(Ok(cached)) => {
730                assert_eq!(*cached.inner(), "conn");
731            }
732            Poll::Ready(Err(err)) => panic!("unexpected error: {err}"),
733            Poll::Pending => panic!("dropped waiter blocked idle reuse"),
734        }
735    }
736
737    #[tokio::test]
738    async fn clone_readiness_reserves_idle_service() {
739        let connector = StrictConnector::default();
740        let poll_ready_count = connector.poll_ready_count.clone();
741        let calls = connector.calls.clone();
742        let mut cache = super::builder().build(connector);
743
744        std::future::poll_fn(|cx| cache.poll_ready(cx))
745            .await
746            .unwrap();
747        let cached = cache.call(1).await.unwrap();
748        assert_eq!(*cached.inner(), 0);
749        drop(cached);
750
751        let mut a = cache.clone();
752        let mut b = cache.clone();
753
754        std::future::poll_fn(|cx| a.poll_ready(cx)).await.unwrap();
755        assert_eq!(poll_ready_count.load(Ordering::SeqCst), 1);
756        assert!(!a.is_empty());
757
758        std::future::poll_fn(|cx| b.poll_ready(cx)).await.unwrap();
759        assert_eq!(poll_ready_count.load(Ordering::SeqCst), 2);
760
761        let a_cached = a.call(10).await.unwrap();
762        assert_eq!(*a_cached.inner(), 0);
763
764        let b_cached = b.call(20).await.unwrap();
765        assert_eq!(*b_cached.inner(), 1);
766
767        assert_eq!(*calls.lock().unwrap(), vec![1, 20]);
768    }
769
770    #[tokio::test]
771    async fn dropped_ready_slot_returns_idle_service() {
772        let connector = StrictConnector::default();
773        let poll_ready_count = connector.poll_ready_count.clone();
774        let mut cache = super::builder().build(connector);
775
776        std::future::poll_fn(|cx| cache.poll_ready(cx))
777            .await
778            .unwrap();
779        let cached = cache.call(1).await.unwrap();
780        drop(cached);
781
782        let mut clone = cache.clone();
783        std::future::poll_fn(|cx| clone.poll_ready(cx))
784            .await
785            .unwrap();
786        drop(clone);
787
788        std::future::poll_fn(|cx| cache.poll_ready(cx))
789            .await
790            .unwrap();
791        assert_eq!(poll_ready_count.load(Ordering::SeqCst), 1);
792
793        let cached = cache.call(2).await.unwrap();
794        assert_eq!(*cached.inner(), 0);
795    }
796
797    #[tokio::test]
798    async fn retain_checks_ready_slot() {
799        let connector = StrictConnector::default();
800        let poll_ready_count = connector.poll_ready_count.clone();
801        let mut cache = super::builder().build(connector);
802
803        std::future::poll_fn(|cx| cache.poll_ready(cx))
804            .await
805            .unwrap();
806        let cached = cache.call(1).await.unwrap();
807        drop(cached);
808
809        std::future::poll_fn(|cx| cache.poll_ready(cx))
810            .await
811            .unwrap();
812        assert!(!cache.is_empty());
813
814        cache.retain(|svc| *svc != 0);
815        assert!(cache.is_empty());
816
817        std::future::poll_fn(|cx| cache.poll_ready(cx))
818            .await
819            .unwrap();
820        assert_eq!(poll_ready_count.load(Ordering::SeqCst), 2);
821    }
822
823    #[tokio::test]
824    async fn idle_return_wakes_pending_poll_ready() {
825        use std::sync::atomic::AtomicBool;
826        use std::task::{Context, Waker};
827
828        let connector = PendingConnector::default();
829        let allow_ready = connector.allow_ready.clone();
830        let mut cache = super::builder().build(connector);
831
832        allow_ready.store(true, Ordering::SeqCst);
833        std::future::poll_fn(|cx| cache.poll_ready(cx))
834            .await
835            .unwrap();
836        let held = cache.call(1).await.unwrap();
837        assert_eq!(*held.inner(), 0);
838
839        let mut ready = Box::pin(std::future::poll_fn(|cx| cache.poll_ready(cx)));
840        let mut cx = Context::from_waker(Waker::noop());
841        assert!(ready.as_mut().poll(&mut cx).is_pending());
842
843        drop(held);
844
845        match ready.as_mut().poll(&mut cx) {
846            Poll::Ready(Ok(())) => {}
847            Poll::Ready(Err(err)) => match err {},
848            Poll::Pending => panic!("idle return did not wake pending poll_ready"),
849        }
850        drop(ready);
851
852        let cached = cache.call(2).await.unwrap();
853        assert_eq!(*cached.inner(), 0);
854
855        #[derive(Default)]
856        struct PendingConnector {
857            allow_ready: Arc<AtomicBool>,
858            next: Arc<AtomicUsize>,
859            ready: bool,
860        }
861
862        impl Service<usize> for PendingConnector {
863            type Response = usize;
864            type Error = Infallible;
865            type Future = std::future::Ready<Result<usize, Infallible>>;
866
867            fn poll_ready(&mut self, _cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
868                if self.allow_ready.swap(false, Ordering::SeqCst) {
869                    self.ready = true;
870                    Poll::Ready(Ok(()))
871                } else {
872                    Poll::Pending
873                }
874            }
875
876            fn call(&mut self, _target: usize) -> Self::Future {
877                assert!(self.ready, "connector called without poll_ready");
878                self.ready = false;
879                let id = self.next.fetch_add(1, Ordering::SeqCst);
880                std::future::ready(Ok(id))
881            }
882        }
883    }
884
885    #[derive(Default)]
886    struct StrictConnector {
887        poll_ready_count: Arc<AtomicUsize>,
888        next: Arc<AtomicUsize>,
889        calls: Arc<Mutex<Vec<usize>>>,
890        ready: bool,
891    }
892
893    impl Clone for StrictConnector {
894        fn clone(&self) -> Self {
895            StrictConnector {
896                poll_ready_count: self.poll_ready_count.clone(),
897                next: self.next.clone(),
898                calls: self.calls.clone(),
899                ready: false,
900            }
901        }
902    }
903
904    impl Service<usize> for StrictConnector {
905        type Response = usize;
906        type Error = Infallible;
907        type Future = std::future::Ready<Result<usize, Infallible>>;
908
909        fn poll_ready(&mut self, _cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
910            self.ready = true;
911            self.poll_ready_count.fetch_add(1, Ordering::SeqCst);
912            Poll::Ready(Ok(()))
913        }
914
915        fn call(&mut self, target: usize) -> Self::Future {
916            assert!(self.ready, "connector called without poll_ready");
917            self.ready = false;
918            self.calls.lock().unwrap().push(target);
919            let id = self.next.fetch_add(1, Ordering::SeqCst);
920            std::future::ready(Ok(id))
921        }
922    }
923}
hyper_util/client/pool/cache.rs

hyper_util/client/pool/
cache.rs
