From 807d0d532b82df9d5c2b0691b584a0f70e937c06 Mon Sep 17 00:00:00 2001 From: Sean McArthur Date: Mon, 29 Jun 2026 11:33:21 -0400 Subject: [PATCH 1/2] fix(pool): preserve Cache readiness with clones --- src/client/pool/cache.rs | 194 +++++++++++++++++++++++++++++++++++++-- 1 file changed, 185 insertions(+), 9 deletions(-) diff --git a/src/client/pool/cache.rs b/src/client/pool/cache.rs index ca15ec6e..cf42598b 100644 --- a/src/client/pool/cache.rs +++ b/src/client/pool/cache.rs @@ -55,6 +55,7 @@ mod internal { connector: M, shared: Arc>>, events: Ev, + ready: Ready, } /// A builder to configure a `Cache`. @@ -86,6 +87,12 @@ mod internal { // todo: on_idle } + #[derive(Debug)] + enum Ready { + None, + Cached(S), + } + pub enum CacheFuture where M: Service, @@ -148,6 +155,7 @@ mod internal { Cache { connector, events: self.events, + ready: Ready::None, shared: Arc::new(Mutex::new(Shared { services: Vec::new(), waiters: VecDeque::new(), @@ -167,12 +175,19 @@ mod internal { where F: FnMut(&mut M::Response) -> bool, { + let mut predicate = predicate; + if let Ready::Cached(svc) = &mut self.ready { + if !predicate(svc) { + self.ready = Ready::None; + } + } + self.shared.lock().unwrap().services.retain_mut(predicate); } /// Check whether this cache has no cached services. pub fn is_empty(&self) -> bool { - self.shared.lock().unwrap().services.is_empty() + matches!(self.ready, Ready::None) && self.shared.lock().unwrap().services.is_empty() } } @@ -188,23 +203,38 @@ mod internal { type Future = CacheFuture; fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll> { - if !self.shared.lock().unwrap().services.is_empty() { - Poll::Ready(Ok(())) - } else { - self.connector.poll_ready(cx) + match self.ready { + Ready::Cached(_) => return Poll::Ready(Ok(())), + Ready::None => {} } + + if let Some(svc) = self.shared.lock().unwrap().take() { + self.ready = Ready::Cached(svc); + return Poll::Ready(Ok(())); + } + + self.connector.poll_ready(cx) } fn call(&mut self, target: Dst) -> Self::Future { // 1. If already cached, easy! - let waiter = { - let mut locked = self.shared.lock().unwrap(); - if let Some(found) = locked.take() { + match std::mem::replace(&mut self.ready, Ready::None) { + Ready::Cached(svc) => { return CacheFuture::Cached { - svc: Some(Cached::new(found, Arc::downgrade(&self.shared))), + svc: Some(Cached::new(svc, Arc::downgrade(&self.shared))), }; } + Ready::None => { + if let Some(svc) = self.shared.lock().unwrap().take() { + return CacheFuture::Cached { + svc: Some(Cached::new(svc, Arc::downgrade(&self.shared))), + }; + } + } + } + let waiter = { + let mut locked = self.shared.lock().unwrap(); let (tx, rx) = oneshot::channel(); locked.waiters.push_back(tx); rx @@ -230,6 +260,20 @@ mod internal { connector: self.connector.clone(), events: self.events.clone(), shared: self.shared.clone(), + ready: Ready::None, + } + } + } + + impl Drop for Cache + where + M: Service, + { + fn drop(&mut self) { + if let Ready::Cached(svc) = std::mem::replace(&mut self.ready, Ready::None) { + if let Ok(mut shared) = self.shared.lock() { + shared.put(svc); + } } } } @@ -441,6 +485,13 @@ mod events { #[cfg(test)] mod tests { + use std::convert::Infallible; + use std::sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }; + use std::task::{self, Poll}; + use futures_util::future; use tower_service::Service; use tower_test::assert_request_eq; @@ -560,4 +611,129 @@ mod tests { Poll::Pending => panic!("last waiter was not woken"), } } + + #[tokio::test] + async fn clone_readiness_reserves_idle_service() { + let connector = StrictConnector::default(); + let poll_ready_count = connector.poll_ready_count.clone(); + let calls = connector.calls.clone(); + let mut cache = super::builder().build(connector); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + let cached = cache.call(1).await.unwrap(); + assert_eq!(*cached.inner(), 0); + drop(cached); + + let mut a = cache.clone(); + let mut b = cache.clone(); + + std::future::poll_fn(|cx| a.poll_ready(cx)).await.unwrap(); + assert_eq!(poll_ready_count.load(Ordering::SeqCst), 1); + assert!(!a.is_empty()); + + std::future::poll_fn(|cx| b.poll_ready(cx)).await.unwrap(); + assert_eq!(poll_ready_count.load(Ordering::SeqCst), 2); + + let a_cached = a.call(10).await.unwrap(); + assert_eq!(*a_cached.inner(), 0); + + let b_cached = b.call(20).await.unwrap(); + assert_eq!(*b_cached.inner(), 1); + + assert_eq!(*calls.lock().unwrap(), vec![1, 20]); + } + + #[tokio::test] + async fn dropped_ready_slot_returns_idle_service() { + let connector = StrictConnector::default(); + let poll_ready_count = connector.poll_ready_count.clone(); + let mut cache = super::builder().build(connector); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + let cached = cache.call(1).await.unwrap(); + drop(cached); + + let mut clone = cache.clone(); + std::future::poll_fn(|cx| clone.poll_ready(cx)) + .await + .unwrap(); + drop(clone); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + assert_eq!(poll_ready_count.load(Ordering::SeqCst), 1); + + let cached = cache.call(2).await.unwrap(); + assert_eq!(*cached.inner(), 0); + } + + #[tokio::test] + async fn retain_checks_ready_slot() { + let connector = StrictConnector::default(); + let poll_ready_count = connector.poll_ready_count.clone(); + let mut cache = super::builder().build(connector); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + let cached = cache.call(1).await.unwrap(); + drop(cached); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + assert!(!cache.is_empty()); + + cache.retain(|svc| *svc != 0); + assert!(cache.is_empty()); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + assert_eq!(poll_ready_count.load(Ordering::SeqCst), 2); + } + + #[derive(Default)] + struct StrictConnector { + poll_ready_count: Arc, + next: Arc, + calls: Arc>>, + ready: bool, + } + + impl Clone for StrictConnector { + fn clone(&self) -> Self { + StrictConnector { + poll_ready_count: self.poll_ready_count.clone(), + next: self.next.clone(), + calls: self.calls.clone(), + ready: false, + } + } + } + + impl Service for StrictConnector { + type Response = usize; + type Error = Infallible; + type Future = std::future::Ready>; + + fn poll_ready(&mut self, _cx: &mut task::Context<'_>) -> Poll> { + self.ready = true; + self.poll_ready_count.fetch_add(1, Ordering::SeqCst); + Poll::Ready(Ok(())) + } + + fn call(&mut self, target: usize) -> Self::Future { + assert!(self.ready, "connector called without poll_ready"); + self.ready = false; + self.calls.lock().unwrap().push(target); + let id = self.next.fetch_add(1, Ordering::SeqCst); + std::future::ready(Ok(id)) + } + } } From d351715d4b053f07ed191420ca28aba7c17db3f7 Mon Sep 17 00:00:00 2001 From: Sean McArthur Date: Mon, 6 Jul 2026 14:22:48 -0400 Subject: [PATCH 2/2] register a waker in the shared idle list from poll_ready --- src/client/pool/cache.rs | 336 ++++++++++++++++++++++++++++++--------- 1 file changed, 262 insertions(+), 74 deletions(-) diff --git a/src/client/pool/cache.rs b/src/client/pool/cache.rs index cf42598b..4499a9f1 100644 --- a/src/client/pool/cache.rs +++ b/src/client/pool/cache.rs @@ -23,10 +23,8 @@ mod internal { use std::future::Future; use std::pin::Pin; use std::sync::{Arc, Mutex, Weak}; - use std::task::{self, Poll, ready}; + use std::task::{self, Poll, Waker, ready}; - use futures_util::future; - use tokio::sync::oneshot; use tower_service::Service; use super::events; @@ -56,6 +54,7 @@ mod internal { shared: Arc>>, events: Ev, ready: Ready, + ready_waiter: Option, } /// A builder to configure a `Cache`. @@ -99,14 +98,10 @@ mod internal { { Racing { shared: Arc>>, - select: future::Select, M::Future>, + waiter: WaiterId, + future: Option, events: Ev, }, - Connecting { - // TODO: could be Weak even here... - shared: Arc>>, - future: M::Future, - }, Cached { svc: Option>, }, @@ -116,7 +111,18 @@ mod internal { #[derive(Debug)] pub struct Shared { services: Vec, - waiters: VecDeque>, + waiters: VecDeque, + reservations: Vec<(WaiterId, S)>, + next_waiter: usize, + } + + #[derive(Clone, Copy, Debug, Eq, PartialEq)] + pub struct WaiterId(usize); + + #[derive(Debug)] + struct Waiter { + id: WaiterId, + waker: Option, } // impl Builder @@ -156,9 +162,12 @@ mod internal { connector, events: self.events, ready: Ready::None, + ready_waiter: None, shared: Arc::new(Mutex::new(Shared { services: Vec::new(), waiters: VecDeque::new(), + reservations: Vec::new(), + next_waiter: 0, })), } } @@ -208,12 +217,34 @@ mod internal { Ready::None => {} } - if let Some(svc) = self.shared.lock().unwrap().take() { - self.ready = Ready::Cached(svc); - return Poll::Ready(Ok(())); + { + let mut shared = self.shared.lock().unwrap(); + if let Some(id) = self.ready_waiter { + if let Some(svc) = shared.take_reserved(id) { + self.ready_waiter = None; + self.ready = Ready::Cached(svc); + return Poll::Ready(Ok(())); + } + } else if let Some(svc) = shared.take_available() { + self.ready = Ready::Cached(svc); + return Poll::Ready(Ok(())); + } + + let id = *self + .ready_waiter + .get_or_insert_with(|| shared.push_waiter()); + shared.store_waker(id, cx.waker()); } - self.connector.poll_ready(cx) + match self.connector.poll_ready(cx) { + Poll::Ready(result) => { + if let Some(id) = self.ready_waiter.take() { + self.shared.lock().unwrap().cancel_waiter(id); + } + Poll::Ready(result) + } + Poll::Pending => Poll::Pending, + } } fn call(&mut self, target: Dst) -> Self::Future { @@ -225,7 +256,16 @@ mod internal { }; } Ready::None => { - if let Some(svc) = self.shared.lock().unwrap().take() { + if let Some(id) = self.ready_waiter.take() { + let mut shared = self.shared.lock().unwrap(); + if let Some(svc) = shared.take_reserved(id) { + return CacheFuture::Cached { + svc: Some(Cached::new(svc, Arc::downgrade(&self.shared))), + }; + } + shared.cancel_waiter(id); + } + if let Some(svc) = self.shared.lock().unwrap().take_available() { return CacheFuture::Cached { svc: Some(Cached::new(svc, Arc::downgrade(&self.shared))), }; @@ -235,16 +275,15 @@ mod internal { let waiter = { let mut locked = self.shared.lock().unwrap(); - let (tx, rx) = oneshot::channel(); - locked.waiters.push_back(tx); - rx + locked.push_waiter() }; // 2. Otherwise, we start a new connect, and also listen for // any newly idle. CacheFuture::Racing { shared: self.shared.clone(), - select: future::select(waiter, self.connector.call(target)), + waiter, + future: Some(self.connector.call(target)), events: self.events.clone(), } } @@ -261,6 +300,7 @@ mod internal { events: self.events.clone(), shared: self.shared.clone(), ready: Ready::None, + ready_waiter: None, } } } @@ -275,6 +315,24 @@ mod internal { shared.put(svc); } } + if let Some(id) = self.ready_waiter.take() { + if let Ok(mut shared) = self.shared.lock() { + shared.cancel_waiter(id); + } + } + } + } + + impl Drop for CacheFuture + where + M: Service, + { + fn drop(&mut self) { + if let CacheFuture::Racing { shared, waiter, .. } = self { + if let Ok(mut shared) = shared.lock() { + shared.cancel_waiter(*waiter); + } + } } } @@ -288,50 +346,40 @@ mod internal { type Output = Result, M::Error>; fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll { - loop { - match &mut *self.as_mut() { - CacheFuture::Racing { - shared, - select, - events, - } => { - match ready!(Pin::new(select).poll(cx)) { - future::Either::Left((Err(_pool_closed), connecting)) => { - // pool was dropped, so we'll never get it from a waiter, - // but if this future still exists, then the user still - // wants a connection. just wait for the connecting - *self = CacheFuture::Connecting { - shared: shared.clone(), - future: connecting, - }; - } - future::Either::Left((Ok(pool_got), connecting)) => { - events.on_race_lost(BackgroundConnect { - future: connecting, - shared: Arc::downgrade(&shared), - }); - return Poll::Ready(Ok(Cached::new( - pool_got, - Arc::downgrade(&shared), - ))); - } - future::Either::Right((connected, _waiter)) => { - let inner = connected?; - return Poll::Ready(Ok(Cached::new( - inner, - Arc::downgrade(&shared), - ))); - } + match &mut *self.as_mut() { + CacheFuture::Racing { + shared, + waiter, + future, + events, + } => { + { + let mut locked = shared.lock().unwrap(); + if let Some(pool_got) = locked.take_reserved(*waiter) { + events.on_race_lost(BackgroundConnect { + future: future.take().expect("racing future polled after done"), + shared: Arc::downgrade(&shared), + }); + return Poll::Ready(Ok(Cached::new(pool_got, Arc::downgrade(&shared)))); } + locked.store_waker(*waiter, cx.waker()); } - CacheFuture::Connecting { shared, future } => { - let inner = ready!(Pin::new(future).poll(cx))?; - return Poll::Ready(Ok(Cached::new(inner, Arc::downgrade(&shared)))); - } - CacheFuture::Cached { svc } => { - return Poll::Ready(Ok(svc.take().unwrap())); - } + + let connected = match ready!( + Pin::new(future.as_mut().expect("racing future polled after done")) + .poll(cx) + ) { + Ok(inner) => inner, + Err(err) => { + shared.lock().unwrap().cancel_waiter(*waiter); + return Poll::Ready(Err(err)); + } + }; + + shared.lock().unwrap().cancel_waiter(*waiter); + Poll::Ready(Ok(Cached::new(connected, Arc::downgrade(&shared)))) } + CacheFuture::Cached { svc } => Poll::Ready(Ok(svc.take().unwrap())), } } } @@ -407,26 +455,61 @@ mod internal { impl Shared { fn put(&mut self, val: V) { - let mut val = Some(val); - while let Some(tx) = self.waiters.pop_front() { - if !tx.is_closed() { - match tx.send(val.take().unwrap()) { - Ok(()) => break, - Err(v) => { - val = Some(v); - } - } + if let Some(mut waiter) = self.waiters.pop_front() { + self.reservations.push((waiter.id, val)); + if let Some(waker) = waiter.waker.take() { + waker.wake(); } + return; } - if let Some(val) = val { - self.services.push(val); + self.services.push(val); + } + + fn take_available(&mut self) -> Option { + if self.waiters.is_empty() { + self.services.pop() + } else { + None } } - fn take(&mut self) -> Option { - // TODO: take in a loop - self.services.pop() + fn push_waiter(&mut self) -> WaiterId { + let id = WaiterId(self.next_waiter); + self.next_waiter = self.next_waiter.wrapping_add(1); + self.waiters.push_back(Waiter { id, waker: None }); + id + } + + fn store_waker(&mut self, id: WaiterId, waker: &Waker) { + if let Some(waiter) = self.waiters.iter_mut().find(|waiter| waiter.id == id) { + if waiter + .waker + .as_ref() + .is_none_or(|current| !current.will_wake(waker)) + { + waiter.waker = Some(waker.clone()); + } + } + } + + fn take_reserved(&mut self, id: WaiterId) -> Option { + let index = self + .reservations + .iter() + .position(|(reserved_id, _)| *reserved_id == id)?; + Some(self.reservations.remove(index).1) + } + + fn cancel_waiter(&mut self, id: WaiterId) { + if let Some(index) = self.waiters.iter().position(|waiter| waiter.id == id) { + self.waiters.remove(index); + return; + } + + if let Some(svc) = self.take_reserved(id) { + self.put(svc); + } } } @@ -612,6 +695,48 @@ mod tests { } } + #[tokio::test] + async fn dropped_racing_future_cancels_waiter() { + use std::future::Future; + use std::task::{Context, Poll, Waker}; + + let (mock, mut handle) = tower_test::mock::pair::(); + let mut cache = super::builder().build(mock); + handle.allow(16); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + let held = future::join(cache.call(0), async { + assert_request_eq!(handle, 0).send_response("conn"); + }) + .await + .0 + .expect("call"); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + let mut dropped = Box::pin(cache.call(1)); + let mut cx = Context::from_waker(Waker::noop()); + assert!(dropped.as_mut().poll(&mut cx).is_pending()); + drop(dropped); + + drop(held); + + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + let mut reused = Box::pin(cache.call(2)); + match reused.as_mut().poll(&mut cx) { + Poll::Ready(Ok(cached)) => { + assert_eq!(*cached.inner(), "conn"); + } + Poll::Ready(Err(err)) => panic!("unexpected error: {err}"), + Poll::Pending => panic!("dropped waiter blocked idle reuse"), + } + } + #[tokio::test] async fn clone_readiness_reserves_idle_service() { let connector = StrictConnector::default(); @@ -698,6 +823,69 @@ mod tests { assert_eq!(poll_ready_count.load(Ordering::SeqCst), 2); } + #[tokio::test] + async fn idle_return_wakes_pending_poll_ready() { + use std::future::Future; + use std::sync::atomic::AtomicBool; + use std::task::{Context, Waker}; + + let connector = PendingConnector::default(); + let allow_ready = connector.allow_ready.clone(); + let mut cache = super::builder().build(connector); + + allow_ready.store(true, Ordering::SeqCst); + std::future::poll_fn(|cx| cache.poll_ready(cx)) + .await + .unwrap(); + let held = cache.call(1).await.unwrap(); + assert_eq!(*held.inner(), 0); + + let mut ready = Box::pin(std::future::poll_fn(|cx| cache.poll_ready(cx))); + let mut cx = Context::from_waker(Waker::noop()); + assert!(ready.as_mut().poll(&mut cx).is_pending()); + + drop(held); + + match ready.as_mut().poll(&mut cx) { + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(err)) => match err {}, + Poll::Pending => panic!("idle return did not wake pending poll_ready"), + } + drop(ready); + + let cached = cache.call(2).await.unwrap(); + assert_eq!(*cached.inner(), 0); + + #[derive(Default)] + struct PendingConnector { + allow_ready: Arc, + next: Arc, + ready: bool, + } + + impl Service for PendingConnector { + type Response = usize; + type Error = Infallible; + type Future = std::future::Ready>; + + fn poll_ready(&mut self, _cx: &mut task::Context<'_>) -> Poll> { + if self.allow_ready.swap(false, Ordering::SeqCst) { + self.ready = true; + Poll::Ready(Ok(())) + } else { + Poll::Pending + } + } + + fn call(&mut self, _target: usize) -> Self::Future { + assert!(self.ready, "connector called without poll_ready"); + self.ready = false; + let id = self.next.fetch_add(1, Ordering::SeqCst); + std::future::ready(Ok(id)) + } + } + } + #[derive(Default)] struct StrictConnector { poll_ready_count: Arc,