From e3e174f8ea4bfc2d2705d84f074d4279a9d3ea8a Mon Sep 17 00:00:00 2001 From: kdcokenny Date: Fri, 2 Oct 2026 12:07:30 +0000 Subject: [PATCH 1/6] Wake capacity waiters after background connection failure A min_idle attempt can occupy the last capacity slot while get() waits on Notify. On upstream, failure releases the reservation but never wakes that getter to start a replacement. The manually polled regression fails on untouched upstream under both mutexes, and passes with only the added failure notification. Notify after releasing the pool lock. Require Tokio 1.26 for atomic notify_waiters across batched wakeups. This fixes an existing background-failure wakeup gap independently of caller-owned attempts. --- bb8/Cargo.toml | 2 +- bb8/src/inner.rs | 2 ++ bb8/tests/background.rs | 76 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 79 insertions(+), 1 deletion(-) create mode 100644 bb8/tests/background.rs diff --git a/bb8/Cargo.toml b/bb8/Cargo.toml index aab5878..18a3a93 100644 --- a/bb8/Cargo.toml +++ b/bb8/Cargo.toml @@ -12,7 +12,7 @@ readme = "../README.md" [dependencies] futures-util = { version = "0.3.2", default-features = false, features = ["alloc"] } parking_lot = { version = "0.12", optional = true } -tokio = { version = "1.0", features = ["rt", "sync", "time"] } +tokio = { version = "1.26", features = ["rt", "sync", "time"] } [target.'cfg(not(target_has_atomic = "64"))'.dependencies] portable-atomic = "1" diff --git a/bb8/src/inner.rs b/bb8/src/inner.rs index d7ad1d5..a2e539b 100644 --- a/bb8/src/inner.rs +++ b/bb8/src/inner.rs @@ -217,6 +217,8 @@ impl PoolInner { { let mut locked = shared.internals.lock(); locked.connect_failed(approval); + drop(locked); + self.inner.notify.notify_waiters(); return Err(e); } else { self.inner.forward_error(e); diff --git a/bb8/tests/background.rs b/bb8/tests/background.rs new file mode 100644 index 0000000..8058e68 --- /dev/null +++ b/bb8/tests/background.rs @@ -0,0 +1,76 @@ +use bb8::{ManageConnection, Pool}; +use std::{ + future::Future, + sync::Arc, + task::{Context, Wake, Waker}, + time::Duration, +}; +use tokio::sync::{mpsc, oneshot}; + +#[derive(Debug)] +struct Failure; + +struct Controlled { + attempts: mpsc::UnboundedSender>>, +} + +impl ManageConnection for Controlled { + type Connection = (); + type Error = Failure; + + async fn connect(&self) -> Result<(), Failure> { + let (tx, rx) = oneshot::channel(); + self.attempts.send(tx).unwrap(); + rx.await.unwrap() + } + + async fn is_valid(&self, _: &mut ()) -> Result<(), Failure> { + Ok(()) + } + + fn has_broken(&self, _: &mut ()) -> bool { + false + } +} + +#[derive(Default)] +struct Wakes(std::sync::atomic::AtomicUsize); + +impl Wake for Wakes { + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + } +} + +#[tokio::test] +async fn failed_background_attempt_wakes_capacity_waiter() { + let (attempts, mut rx) = mpsc::unbounded_channel(); + let pool = Pool::builder() + .max_size(1) + .min_idle(1) + .retry_connection(false) + .connection_timeout(Duration::from_secs(1)) + .build_unchecked(Controlled { attempts }); + let background = rx.recv().await.unwrap(); + let wakes = Arc::new(Wakes::default()); + let waker = Waker::from(wakes.clone()); + let mut get = Box::pin(pool.get()); + assert!(get + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending()); + let before = wakes.0.load(std::sync::atomic::Ordering::SeqCst); + background.send(Err(Failure)).unwrap(); + tokio::task::yield_now().await; + assert!(wakes.0.load(std::sync::atomic::Ordering::SeqCst) > before); + assert!(get + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending()); + rx.recv().await.unwrap().send(Ok(())).unwrap(); + get.await.unwrap(); +} From 1bfa3e18d3087f0cddb79f0bb25fe47d945fc7fd Mon Sep 17 00:00:00 2001 From: kdcokenny Date: Fri, 2 Oct 2026 11:52:14 +0000 Subject: [PATCH 2/6] Separate connection retries from background pool insertion Extract the existing connect/customize/backoff loop so it returns an unpooled connection. Background insertion still consumes its approval, records Created, and releases capacity with a wakeup on failure. Keep retry delays and the existing elapsed-time check unchanged. Removing the weak upgrade is safe because this method already borrows a PoolInner owning the strong Arc. This mechanical step makes the later ownership change smaller without changing error routing or cancellation. --- bb8/src/inner.rs | 38 +++++++++++++++++++++----------------- 1 file changed, 21 insertions(+), 17 deletions(-) diff --git a/bb8/src/inner.rs b/bb8/src/inner.rs index a2e539b..18deba1 100644 --- a/bb8/src/inner.rs +++ b/bb8/src/inner.rs @@ -186,16 +186,30 @@ impl PoolInner { // Outside of Pool to avoid borrow splitting issues on self async fn add_connection(&self, approval: Approval) -> Result<(), M::Error> { - let new_shared = Arc::downgrade(&self.inner); - let shared = match new_shared.upgrade() { - None => return Ok(()), - Some(shared) => shared, - }; + match self.connect_with_retry().await { + Ok(conn) => { + self.inner.internals.lock().put( + Conn::new(conn), + Some(approval), + self.inner.clone(), + ); + self.inner.statistics.record(StatsKind::Created); + Ok(()) + } + Err(e) => { + self.inner.internals.lock().connect_failed(approval); + self.inner.notify.notify_waiters(); + Err(e) + } + } + } + async fn connect_with_retry(&self) -> Result { let start = Instant::now(); let mut delay = Duration::from_secs(0); loop { - let conn = shared + let conn = self + .inner .manager .connect() .and_then(|mut c| async { self.on_acquire_connection(&mut c).await.map(|_| c) }) @@ -203,22 +217,12 @@ impl PoolInner { match conn { Ok(conn) => { - let conn = Conn::new(conn); - shared - .internals - .lock() - .put(conn, Some(approval), self.inner.clone()); - self.inner.statistics.record(StatsKind::Created); - return Ok(()); + return Ok(conn); } Err(e) => { if !self.inner.statics.retry_connection || Instant::now() - start > self.inner.statics.connection_timeout { - let mut locked = shared.internals.lock(); - locked.connect_failed(approval); - drop(locked); - self.inner.notify.notify_waiters(); return Err(e); } else { self.inner.forward_error(e); From 9e2e516354ccd9c1b1bf75b049851f8a5d176bcf Mon Sep 17 00:00:00 2001 From: kdcokenny Date: Fri, 2 Oct 2026 11:53:15 +0000 Subject: [PATCH 3/6] Let managers classify non-retryable connection errors Add the default-false error_is_fatal hook at the end of ManageConnection, as requested in review. Apply it to connect and customization errors in the shared retry loop. Fatal build failures return immediately; fatal background failures go to the sink once and do not disable the pool. Exercise both initialization error sources in one test before changing which attempts get() owns. Existing managers retain their retry behavior. --- bb8/src/api.rs | 14 ++++++++++-- bb8/src/inner.rs | 1 + bb8/tests/background.rs | 47 ++++++++++++++++++++++++++++++++++++++++- 3 files changed, 59 insertions(+), 3 deletions(-) diff --git a/bb8/src/api.rs b/bb8/src/api.rs index 6a80e70..e37f4b6 100644 --- a/bb8/src/api.rs +++ b/bb8/src/api.rs @@ -369,6 +369,8 @@ impl Builder { /// or intermittent network failures. Some applications however are smart enough to /// know that the server is down and retries won't help (and could actually hurt recovery). /// In that case, it's better to disable retries here and let the pool error out. + /// Errors classified as fatal by [`ManageConnection::error_is_fatal`] are + /// never retried. /// /// Defaults to enabled. #[must_use] @@ -463,6 +465,14 @@ pub trait ManageConnection: Sized + Send + Sync + 'static { ) -> impl Future> + Send; /// Synchronously determine if the connection is no longer usable, if possible. fn has_broken(&self, conn: &mut Self::Connection) -> bool; + + /// Classifies an error from [`Self::connect`] or connection customization as non-retryable. + /// + /// Defaults to false. Fatal failures stop retries for that attempt without + /// disabling the pool. + fn error_is_fatal(&self, _error: &Self::Error) -> bool { + false + } } /// A trait which provides functionality to initialize a connection @@ -472,8 +482,8 @@ pub trait CustomizeConnection: /// Called with connections immediately after they are returned from /// `ManageConnection::connect`. /// - /// The default implementation simply returns `Ok(())`. If this method returns an - /// error, it will be forwarded to the configured error sink. + /// The default implementation simply returns `Ok(())`. Errors follow the + /// same retry and reporting behavior as [`ManageConnection::connect`]. fn on_acquire<'a>( &'a self, _connection: &'a mut C, diff --git a/bb8/src/inner.rs b/bb8/src/inner.rs index 18deba1..2cb18a6 100644 --- a/bb8/src/inner.rs +++ b/bb8/src/inner.rs @@ -221,6 +221,7 @@ impl PoolInner { } Err(e) => { if !self.inner.statics.retry_connection + || self.inner.manager.error_is_fatal(&e) || Instant::now() - start > self.inner.statics.connection_timeout { return Err(e); diff --git a/bb8/tests/background.rs b/bb8/tests/background.rs index 8058e68..b402bd2 100644 --- a/bb8/tests/background.rs +++ b/bb8/tests/background.rs @@ -12,6 +12,7 @@ struct Failure; struct Controlled { attempts: mpsc::UnboundedSender>>, + fatal: bool, } impl ManageConnection for Controlled { @@ -31,6 +32,10 @@ impl ManageConnection for Controlled { fn has_broken(&self, _: &mut ()) -> bool { false } + + fn error_is_fatal(&self, _: &Failure) -> bool { + self.fatal + } } #[derive(Default)] @@ -54,7 +59,10 @@ async fn failed_background_attempt_wakes_capacity_waiter() { .min_idle(1) .retry_connection(false) .connection_timeout(Duration::from_secs(1)) - .build_unchecked(Controlled { attempts }); + .build_unchecked(Controlled { + attempts, + fatal: false, + }); let background = rx.recv().await.unwrap(); let wakes = Arc::new(Wakes::default()); let waker = Waker::from(wakes.clone()); @@ -74,3 +82,40 @@ async fn failed_background_attempt_wakes_capacity_waiter() { rx.recv().await.unwrap().send(Ok(())).unwrap(); get.await.unwrap(); } + +#[derive(Debug)] +struct FailingCustomizer; + +impl bb8::CustomizeConnection<(), Failure> for FailingCustomizer { + fn on_acquire<'a>( + &'a self, + _: &'a mut (), + ) -> std::pin::Pin> + Send + 'a>> { + Box::pin(async { Err(Failure) }) + } +} + +#[tokio::test] +async fn fatal_connect_and_customizer_errors_stop_build_retries() { + // Fail either connect() or the customizer after connect() succeeds. + for result in [Err(Failure), Ok(())] { + let (attempts, mut rx) = mpsc::unbounded_channel(); + let build = tokio::spawn( + Pool::builder() + .min_idle(1) + .connection_timeout(Duration::from_secs(2)) + .connection_customizer(Box::new(FailingCustomizer)) + .build(Controlled { + attempts, + fatal: true, + }), + ); + rx.recv().await.unwrap().send(result).unwrap(); + tokio::time::timeout(Duration::from_millis(100), build) + .await + .expect("fatal initialization error must return before retry backoff") + .unwrap() + .unwrap_err(); + assert!(rx.try_recv().is_err()); + } +} From ad7a00cb67297551ec07986cbb46f97049f47c2a Mon Sep 17 00:00:00 2001 From: kdcokenny Date: Fri, 2 Oct 2026 12:10:39 +0000 Subject: [PATCH 4/6] Let each checkout own its demand connection attempt Return the exact non-retried manager/customizer error to the checkout reserving the attempt, without Clone or Sync bounds or shared stale errors. Race the owned attempt against idle availability under one checkout deadline while preserving transient retry backoff and sink reporting. Keep capacity reservations, cancellation drop, runtime-handle min_idle refill, wakeups, and all demand regressions together. The immediate-validation timeout and slow-validation notification regressions pass on untouched upstream under both mutexes. Their safeguards belong in this architectural change: inline ready connects make a validation loop able to starve the deadline, and registering Notify before shared-state inspection requires dropping that registration before validation. Strengthen the timeout regression with an explicit first-poll yield assertion, using its existing validator. The independently reproduced upstream background-failure wakeup fix precedes this commit. The retry extraction and fatal classifier can be reviewed independently, but this commit remains the principal architectural review burden. The sequence does not establish that its scope is acceptable to the maintainer. --- bb8/Cargo.toml | 2 +- bb8/src/api.rs | 11 + bb8/src/inner.rs | 89 ++++++-- bb8/src/internals.rs | 69 ++++++- bb8/tests/checkout.rs | 141 +++++++++++++ bb8/tests/demand.rs | 468 ++++++++++++++++++++++++++++++++++++++++++ bb8/tests/test.rs | 2 +- 7 files changed, 755 insertions(+), 27 deletions(-) create mode 100644 bb8/tests/checkout.rs create mode 100644 bb8/tests/demand.rs diff --git a/bb8/Cargo.toml b/bb8/Cargo.toml index 18a3a93..77bb8d3 100644 --- a/bb8/Cargo.toml +++ b/bb8/Cargo.toml @@ -18,7 +18,7 @@ tokio = { version = "1.26", features = ["rt", "sync", "time"] } portable-atomic = "1" [dev-dependencies] -tokio = { version = "1.0", features = ["macros"] } +tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] } [features] parking_lot = ["dep:parking_lot", "tokio/parking_lot"] diff --git a/bb8/src/api.rs b/bb8/src/api.rs index e37f4b6..659cdeb 100644 --- a/bb8/src/api.rs +++ b/bb8/src/api.rs @@ -22,6 +22,15 @@ impl Pool { } /// Retrieves a connection from the pool. + /// + /// A connection attempt reserved by this call returns [`RunError::User`] + /// when retries are disabled or [`ManageConnection::error_is_fatal`] returns + /// true. Retried and background errors go to the configured error sink. + /// This call is bounded by [`Builder::connection_timeout`]. + /// + /// Cancelling or timing out this call drops its connection and customization + /// futures, releasing reserved capacity. Background attempts to maintain + /// [`Builder::min_idle`] are independent of this call. pub async fn get(&self) -> Result, RunError> { self.inner.get().await } @@ -30,6 +39,7 @@ impl Pool { /// /// Using an owning `PooledConnection` makes it easier to leak the connection pool. Therefore, [`Pool::get`] /// (which stores a lifetime-bound reference to the pool) should be preferred whenever possible. + /// Error and cancellation behavior is the same as [`Pool::get`]. pub async fn get_owned(&self) -> Result, RunError> { Ok(PooledConnection { conn: self.get().await?.take(), @@ -458,6 +468,7 @@ pub trait ManageConnection: Sized + Send + Sync + 'static { /// Attempts to create a new connection. fn connect(&self) -> impl Future> + Send; + /// Determines if the connection is still connected to the database. fn is_valid( &self, diff --git a/bb8/src/inner.rs b/bb8/src/inner.rs index 2cb18a6..4d1d6ee 100644 --- a/bb8/src/inner.rs +++ b/bb8/src/inner.rs @@ -4,6 +4,7 @@ use std::future::Future; use std::sync::{Arc, Weak}; use std::time::{Duration, Instant}; +use futures_util::future::{select, Either}; use futures_util::stream::{FuturesUnordered, StreamExt}; use futures_util::TryFutureExt; use tokio::spawn; @@ -12,7 +13,9 @@ use tokio::time::{interval_at, sleep, timeout, Interval}; use crate::api::{ AddError, Builder, ConnectionState, ManageConnection, PooledConnection, RunError, State, }; -use crate::internals::{Approval, ApprovalIter, Conn, SharedPool, StatsGetKind, StatsKind}; +use crate::internals::{ + Approval, ApprovalIter, Conn, GettingResult, SharedPool, StatsGetKind, StatsKind, +}; pub(crate) struct PoolInner { inner: Arc>, @@ -46,6 +49,10 @@ impl PoolInner { Ok(()) } + pub(crate) fn replenish_after_demand(inner: Arc>) { + Self { inner }.spawn_start_connections(); + } + pub(crate) fn spawn_start_connections(&self) { let mut locked = self.inner.internals.lock(); self.spawn_replenishing_approvals(locked.wanted(&self.inner.statics)); @@ -86,23 +93,64 @@ impl PoolInner { let future = async { let getting = self.inner.start_get(); + let mut connecting = None; loop { - let (conn, approvals) = getting.get(); - self.spawn_replenishing_approvals(approvals); - - // Cancellation safety: make sure to wrap the connection in a `PooledConnection` - // before allowing the code to hit an `await`, so we don't lose the connection. - - let mut conn = match conn { - Some(conn) => PooledConnection::new(self, conn), - None => { - wait_time_start = Some(Instant::now()); - kind = StatsGetKind::Waited; - self.inner.notify.notified().await; - continue; + // A validator must not consume another getter's notification. + let mut conn = { + // Register before inspecting shared state: multiple notifications + // must not collapse while independent getters check the queue. + let notified = self.inner.notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + let GettingResult { + conn, + approvals, + demand, + } = getting.get(connecting.is_some()); + self.spawn_replenishing_approvals(approvals); + if let Some(demand) = demand { + debug_assert!(connecting.is_none()); + connecting = Some(Box::pin(async move { + let conn = self.connect_with_retry(false).await?; + let (conn, approvals) = demand.checkout(conn); + self.spawn_replenishing_approvals(approvals); + Ok::<_, M::Error>(conn) + })); + } + + // Cancellation safety: make sure to wrap the connection in a `PooledConnection` + // before allowing the code to hit an `await`, so we don't lose the connection. + + match conn { + Some(conn) => PooledConnection::new(self, conn), + None => { + wait_time_start = Some(Instant::now()); + kind = StatsGetKind::Waited; + match connecting.as_mut() { + Some(attempt) => { + match select(attempt.as_mut(), notified.as_mut()).await { + Either::Left((result, _)) => { + connecting = None; + PooledConnection::new( + self, + result.map_err(RunError::User)?, + ) + } + Either::Right(_) => continue, + } + } + None => { + notified.await; + continue; + } + } + } } }; + // Release a superseded demand reservation before validation can block. + drop(connecting.take()); + if !self.inner.statics.test_on_check_out { return Ok(conn); } @@ -113,6 +161,9 @@ impl PoolInner { self.inner.statistics.record(StatsKind::ClosedInvalid); self.inner.forward_error(e); conn.state = ConnectionState::Invalid; + drop(conn); + // An immediately failing validator must still let the timeout run. + tokio::task::yield_now().await; continue; } } @@ -186,7 +237,7 @@ impl PoolInner { // Outside of Pool to avoid borrow splitting issues on self async fn add_connection(&self, approval: Approval) -> Result<(), M::Error> { - match self.connect_with_retry().await { + match self.connect_with_retry(true).await { Ok(conn) => { self.inner.internals.lock().put( Conn::new(conn), @@ -204,7 +255,7 @@ impl PoolInner { } } - async fn connect_with_retry(&self) -> Result { + async fn connect_with_retry(&self, background: bool) -> Result { let start = Instant::now(); let mut delay = Duration::from_secs(0); loop { @@ -217,12 +268,16 @@ impl PoolInner { match conn { Ok(conn) => { + if !background { + self.inner.statistics.record(StatsKind::Created); + } return Ok(conn); } Err(e) => { if !self.inner.statics.retry_connection || self.inner.manager.error_is_fatal(&e) - || Instant::now() - start > self.inner.statics.connection_timeout + || (background + && Instant::now() - start > self.inner.statics.connection_timeout) { return Err(e); } else { diff --git a/bb8/src/internals.rs b/bb8/src/internals.rs index b98e990..dfcc348 100644 --- a/bb8/src/internals.rs +++ b/bb8/src/internals.rs @@ -225,18 +225,31 @@ pub(crate) struct Getting { } impl Getting { - pub(crate) fn get(&self) -> (Option>, ApprovalIter) { + pub(crate) fn get(&self, connecting: bool) -> GettingResult { let mut locked = self.inner.internals.lock(); if let Some(IdleConn { conn, .. }) = locked.conns.pop_front() { - return (Some(conn), locked.wanted(&self.inner.statics)); + return GettingResult { + conn: Some(conn), + approvals: locked.wanted(&self.inner.statics), + demand: None, + }; } - let approvals = match locked.in_flight > locked.pending_conns { - true => 1, - false => 0, - }; - - (None, locked.approvals(&self.inner.statics, approvals)) + let wanted = !connecting && locked.in_flight > locked.pending_conns; + let approval = locked + .approvals(&self.inner.statics, u32::from(wanted)) + .next(); + drop(locked); + let demand = approval.map(|approval| Demand { + inner: self.inner.clone(), + approval: Some(approval), + runtime: tokio::runtime::Handle::current(), + }); + GettingResult { + conn: None, + approvals: ApprovalIter { num: 0 }, + demand, + } } } @@ -257,6 +270,46 @@ impl Drop for Getting { } } +pub(crate) struct GettingResult { + pub(crate) conn: Option>, + pub(crate) approvals: ApprovalIter, + pub(crate) demand: Option>, +} + +pub(crate) struct Demand { + inner: Arc>, + approval: Option, + runtime: tokio::runtime::Handle, +} + +impl Demand { + pub(crate) fn checkout(mut self, conn: M::Connection) -> (Conn, ApprovalIter) { + let approvals = { + let mut locked = self.inner.internals.lock(); + let _ = self.approval.take().expect("live reservation"); + locked.pending_conns -= 1; + locked.num_conns += 1; + locked.wanted(&self.inner.statics) + }; + (Conn::new(conn), approvals) + } +} + +impl Drop for Demand { + fn drop(&mut self) { + if let Some(approval) = self.approval.take() { + self.inner.internals.lock().connect_failed(approval); + // The canceled reservation may have counted toward min_idle. + { + // Drop can run outside the runtime that first polled get(). + let _entered = self.runtime.enter(); + crate::inner::PoolInner::replenish_after_demand(self.inner.clone()); + } + self.inner.notify.notify_waiters(); + } + } +} + #[derive(Debug)] pub(crate) struct Conn { pub(crate) conn: C, diff --git a/bb8/tests/checkout.rs b/bb8/tests/checkout.rs new file mode 100644 index 0000000..f4dc7f5 --- /dev/null +++ b/bb8/tests/checkout.rs @@ -0,0 +1,141 @@ +use bb8::{ManageConnection, Pool, RunError}; +use std::{ + future::Future, + sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + Arc, Mutex, OnceLock, + }, + task::{Context, Wake, Waker}, + time::{Duration, Instant}, +}; +use tokio::sync::oneshot; + +#[derive(Debug)] +struct Invalid; + +#[derive(Default)] +struct ImmediatelyInvalid(OnceLock); + +impl ManageConnection for ImmediatelyInvalid { + type Connection = (); + type Error = Invalid; + + async fn connect(&self) -> Result<(), Invalid> { + Ok(()) + } + + async fn is_valid(&self, _: &mut ()) -> Result<(), Invalid> { + // Eventually succeed so a checkout that stops yielding fails instead of hanging. + if self.0.get_or_init(Instant::now).elapsed() < Duration::from_millis(250) { + Err(Invalid) + } else { + Ok(()) + } + } + + fn has_broken(&self, _: &mut ()) -> bool { + false + } +} + +#[tokio::test] +async fn immediate_validation_failures_obey_get_timeout() { + let pool = Pool::builder() + .max_size(1) + .test_on_check_out(true) + .connection_timeout(Duration::from_millis(50)) + .build_unchecked(ImmediatelyInvalid::default()); + pool.add(()).unwrap(); + let mut get = Box::pin(pool.get()); + let waker = futures_util::task::noop_waker(); + assert!(get + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending()); + assert_eq!(pool.state().statistics.connections_closed_invalid, 1); + assert!(matches!(get.await, Err(RunError::TimedOut))); +} + +struct SlowValidator { + pause_next: Arc, + release: Mutex>>, +} + +impl ManageConnection for SlowValidator { + type Connection = (); + type Error = Invalid; + + async fn connect(&self) -> Result<(), Invalid> { + Ok(()) + } + + async fn is_valid(&self, _: &mut ()) -> Result<(), Invalid> { + if self.pause_next.swap(false, Ordering::SeqCst) { + let release = self.release.lock().unwrap().take().unwrap(); + release.await.unwrap(); + } + Ok(()) + } + + fn has_broken(&self, _: &mut ()) -> bool { + false + } +} + +#[derive(Default)] +struct Wakes(AtomicUsize); + +impl Wake for Wakes { + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } +} + +#[tokio::test] +async fn slow_validation_does_not_consume_capacity_waiter_notification() { + let pause_next = Arc::new(AtomicBool::new(false)); + let (release, receiver) = oneshot::channel(); + let pool = Pool::builder() + .max_size(2) + .min_idle(Some(2)) + .test_on_check_out(true) + .connection_timeout(Duration::from_secs(1)) + .build(SlowValidator { + pause_next: pause_next.clone(), + release: Mutex::new(Some(receiver)), + }) + .await + .unwrap(); + let held = pool.get().await.unwrap(); + + pause_next.store(true, Ordering::SeqCst); + let mut slow = Box::pin(pool.get()); + let slow_waker = Waker::from(Arc::new(Wakes::default())); + assert!(slow + .as_mut() + .poll(&mut Context::from_waker(&slow_waker)) + .is_pending()); + assert!(!pause_next.load(Ordering::SeqCst)); + assert_eq!(pool.state().idle_connections, 0); + + let wakes = Arc::new(Wakes::default()); + let waiter_waker = Waker::from(wakes.clone()); + let mut waiter = Box::pin(pool.get()); + assert!(waiter + .as_mut() + .poll(&mut Context::from_waker(&waiter_waker)) + .is_pending()); + let before = wakes.0.load(Ordering::SeqCst); + + drop(held); + assert_eq!(pool.state().idle_connections, 1); + assert!(wakes.0.load(Ordering::SeqCst) > before); + + let _acquired = waiter.await.unwrap(); + release.send(()).unwrap(); + slow.await.unwrap(); +} diff --git a/bb8/tests/demand.rs b/bb8/tests/demand.rs new file mode 100644 index 0000000..f97d79e --- /dev/null +++ b/bb8/tests/demand.rs @@ -0,0 +1,468 @@ +use bb8::{ErrorSink, ManageConnection, Pool, PooledConnection, RunError}; +use std::{ + cell::Cell, + future::{poll_fn, Future}, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, + time::Duration, +}; +use tokio::{ + sync::{mpsc, oneshot}, + time::timeout, +}; + +// Deliberately non-Clone and non-Sync, but Send: bb8's existing Error bound. +#[derive(Debug)] +struct Failure { + id: usize, + fatal: bool, + _not_sync: Cell, +} +impl Failure { + fn new(id: usize, fatal: bool) -> Self { + Self { + id, + fatal, + _not_sync: Cell::new(0), + } + } +} +#[derive(Debug)] +struct Connection { + id: usize, + broken: bool, +} +#[derive(Debug)] +struct Attempt { + id: usize, + finish: oneshot::Sender>, +} +impl Attempt { + fn fail(self, fatal: bool) { + self.finish + .send(Err(Failure::new(self.id, fatal))) + .expect("live connect future"); + } + fn succeed(self) { + self.finish + .send(Ok(Connection { + id: self.id, + broken: false, + })) + .expect("live connect future"); + } +} +struct Manager { + attempts: mpsc::UnboundedSender, + count: AtomicUsize, +} +impl ManageConnection for Manager { + type Connection = Connection; + type Error = Failure; + async fn connect(&self) -> Result { + let id = self.count.fetch_add(1, Ordering::SeqCst) + 1; + let (finish, result) = oneshot::channel(); + self.attempts.send(Attempt { id, finish }).unwrap(); + result.await.expect("controller must finish attempt") + } + async fn is_valid(&self, _: &mut Connection) -> Result<(), Failure> { + Ok(()) + } + fn has_broken(&self, conn: &mut Connection) -> bool { + conn.broken + } + fn error_is_fatal(&self, e: &Failure) -> bool { + e.fatal + } +} +#[derive(Clone, Debug)] +struct Sink { + errors: Arc>>, +} +impl ErrorSink for Sink { + fn sink(&self, e: Failure) { + self.errors.lock().unwrap().push(e.id); + } + fn boxed_clone(&self) -> Box> { + Box::new(self.clone()) + } +} +fn setup( + retry: bool, + max: u32, + ms: u64, + idle: Option, +) -> (Pool, mpsc::UnboundedReceiver, Sink) { + let (attempts, rx) = mpsc::unbounded_channel(); + let manager = Manager { + attempts, + count: AtomicUsize::new(0), + }; + let sink = Sink { + errors: Arc::new(Mutex::new(Vec::new())), + }; + let pool = Pool::builder() + .retry_connection(retry) + .max_size(max) + .min_idle(idle) + .connection_timeout(Duration::from_millis(ms)) + .test_on_check_out(false) + .error_sink(Box::new(sink.clone())) + .build_unchecked(manager); + (pool, rx, sink) +} +async fn attempt(rx: &mut mpsc::UnboundedReceiver) -> Attempt { + timeout(Duration::from_millis(900), rx.recv()) + .await + .expect("connection attempt should start") + .unwrap() +} +async fn start( + pool: &Pool, +) -> tokio::task::JoinHandle, RunError>> { + let pool = pool.clone(); + let (polled, ready) = oneshot::channel(); + let task = tokio::spawn(async move { + let mut get = Box::pin(pool.get_owned()); + let mut polled = Some(polled); + poll_fn(|cx| { + let result = get.as_mut().poll(cx); + if let Some(polled) = polled.take() { + polled.send(()).unwrap(); + } + result + }) + .await + }); + ready.await.unwrap(); + task +} +async fn user_error( + task: tokio::task::JoinHandle, RunError>>, + id: usize, +) { + match timeout(Duration::from_millis(100), task) + .await + .expect("must fail promptly, before pool deadline") + .unwrap() + { + Err(RunError::User(e)) => assert_eq!(e.id, id), + other => panic!("expected owned manager error {id}, got {other:?}"), + } +} + +#[tokio::test] +async fn fatal_error_stops_enabled_retries_promptly() { + let (pool, mut rx, sink) = setup(true, 1, 2000, None); + let task = start(&pool).await; + attempt(&mut rx).await.fail(true); + user_error(task, 1).await; + assert!(rx.try_recv().is_err()); + assert!(sink.errors.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn transient_error_retries_then_recovers_and_is_logged() { + let (pool, mut rx, sink) = setup(true, 1, 2000, None); + let task = start(&pool).await; + attempt(&mut rx).await.fail(false); + let next = attempt(&mut rx).await; + assert_eq!(next.id, 2); + next.succeed(); + assert_eq!(task.await.unwrap().unwrap().id, 2); + assert_eq!(*sink.errors.lock().unwrap(), vec![1]); +} + +#[tokio::test] +async fn get_timeout_cancels_pending_attempt_and_releases_capacity() { + for retry in [false, true] { + let (pool, mut rx, sink) = setup(retry, 1, 150, None); + let task = start(&pool).await; + let mut pending = attempt(&mut rx).await; + if retry { + pending.fail(false); + pending = attempt(&mut rx).await; + } + assert!(matches!(task.await.unwrap(), Err(RunError::TimedOut))); + assert!(pending.finish.is_closed()); + let expected_errors = if retry { vec![1] } else { vec![] }; + assert_eq!(*sink.errors.lock().unwrap(), expected_errors); + + let next = start(&pool).await; + let new = attempt(&mut rx).await; + assert_eq!(new.id, pending.id + 1); + new.succeed(); + assert_eq!(next.await.unwrap().unwrap().id, pending.id + 1); + } +} + +#[tokio::test] +async fn concurrent_errors_go_to_their_own_attempts() { + let (pool, mut rx, sink) = setup(false, 2, 500, None); + let first = start(&pool).await; + let a = attempt(&mut rx).await; + let second = start(&pool).await; + let b = attempt(&mut rx).await; + b.fail(false); + user_error(second, 2).await; + a.succeed(); + assert_eq!(first.await.unwrap().unwrap().id, 1); + assert!(sink.errors.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn no_stale_error_after_recovery_or_in_full_pool() { + let (pool, mut rx, _) = setup(false, 1, 150, None); + let first = start(&pool).await; + attempt(&mut rx).await.fail(false); + user_error(first, 1).await; + let second = start(&pool).await; + attempt(&mut rx).await.succeed(); + let held = second.await.unwrap().unwrap(); + assert!(matches!(pool.get().await, Err(RunError::TimedOut))); + drop(held); + assert_eq!(pool.get().await.unwrap().id, 2); +} + +#[tokio::test] +async fn background_error_is_logged_and_capacity_notifies_waiter() { + let (pool, mut rx, sink) = setup(false, 1, 500, Some(1)); + let background = attempt(&mut rx).await; + let task = start(&pool).await; + background.fail(false); + let demand = attempt(&mut rx).await; + assert_eq!(demand.id, 2); + demand.succeed(); + assert_eq!(task.await.unwrap().unwrap().id, 2); + assert_eq!(*sink.errors.lock().unwrap(), vec![1]); +} + +#[tokio::test] +async fn idle_wins_cancels_owned_attempt_and_preserves_min_idle() { + let (pool, mut rx, _) = setup(false, 2, 1000, Some(1)); + let background = attempt(&mut rx).await; + let first = start(&pool).await; + let second = start(&pool).await; + let demand = attempt(&mut rx).await; + assert_eq!(demand.id, 2); + background.succeed(); + let held = first.await.unwrap().unwrap(); + assert_eq!(held.id, 1); + drop(held); + let second = timeout(Duration::from_millis(100), second) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(second.id, 1); + assert!(demand.finish.is_closed()); + let refill = attempt(&mut rx).await; + assert_eq!(refill.id, 3); + refill.succeed(); + tokio::task::yield_now().await; + assert_eq!(pool.state().idle_connections, 1); + assert_eq!(pool.state().connections, 2); +} + +#[tokio::test] +async fn cancelling_owned_attempt_notifies_an_existing_capacity_waiter() { + let (pool, mut rx, _) = setup(false, 1, 1000, None); + let first = start(&pool).await; + let old = attempt(&mut rx).await; + let waiting = start(&pool).await; + first.abort(); + assert!(first.await.unwrap_err().is_cancelled()); + assert!(old.finish.is_closed()); + attempt(&mut rx).await.succeed(); + assert_eq!(waiting.await.unwrap().unwrap().id, 2); +} + +#[tokio::test] +async fn pending_background_attempt_survives_get_timeout() { + let (pool, mut rx, _) = setup(false, 1, 70, Some(1)); + let background = attempt(&mut rx).await; + assert!(matches!(pool.get().await, Err(RunError::TimedOut))); + assert!(!background.finish.is_closed()); + assert!(matches!(pool.get().await, Err(RunError::TimedOut))); + assert!(rx.try_recv().is_err()); + background.succeed(); + tokio::task::yield_now().await; + assert_eq!(pool.get().await.unwrap().id, 1); +} + +#[tokio::test] +async fn build_keeps_its_existing_connection_error_and_timeout_behavior() { + let (attempts, mut rx) = mpsc::unbounded_channel(); + let manager = Manager { + attempts, + count: AtomicUsize::new(0), + }; + let mut build = tokio::spawn( + Pool::builder() + .max_size(1) + .min_idle(1) + .retry_connection(false) + .connection_timeout(Duration::from_millis(20)) + .build(manager), + ); + let pending = attempt(&mut rx).await; + assert!(timeout(Duration::from_millis(50), &mut build) + .await + .is_err()); + assert!(!pending.finish.is_closed()); + pending.fail(false); + assert_eq!(build.await.unwrap().unwrap_err().id, 1); +} + +#[derive(Debug)] +struct Customizer { + requests: mpsc::UnboundedSender>>, +} +impl bb8::CustomizeConnection for Customizer { + fn on_acquire<'a>( + &'a self, + _: &'a mut Connection, + ) -> std::pin::Pin> + Send + 'a>> { + Box::pin(async move { + let (tx, rx) = oneshot::channel(); + self.requests.send(tx).unwrap(); + rx.await.unwrap() + }) + } +} + +#[tokio::test] +async fn customizer_errors_are_owned_and_pending_customizer_can_be_cancelled() { + let (attempts, mut rx) = mpsc::unbounded_channel(); + let (requests, mut custom) = mpsc::unbounded_channel(); + let manager = Manager { + attempts, + count: AtomicUsize::new(0), + }; + let pool = Pool::builder() + .retry_connection(true) + .max_size(1) + .connection_timeout(Duration::from_millis(300)) + .test_on_check_out(false) + .connection_customizer(Box::new(Customizer { requests })) + .build_unchecked(manager); + let first = start(&pool).await; + attempt(&mut rx).await.succeed(); + custom + .recv() + .await + .unwrap() + .send(Err(Failure::new(123, true))) + .unwrap(); + user_error(first, 123).await; + let pending = start(&pool).await; + attempt(&mut rx).await.succeed(); + let custom_pending = custom.recv().await.unwrap(); + pending.abort(); + assert!(pending.await.unwrap_err().is_cancelled()); + assert!(custom_pending.is_closed()); + let next = start(&pool).await; + attempt(&mut rx).await.succeed(); + custom.recv().await.unwrap().send(Ok(())).unwrap(); + assert_eq!(next.await.unwrap().unwrap().id, 3); +} + +#[tokio::test] +async fn fatal_demand_min_idle_refill_fails_once_without_hot_loop() { + let (pool, mut rx, sink) = setup(true, 1, 1000, Some(1)); + attempt(&mut rx).await.fail(true); + tokio::task::yield_now().await; + let task = start(&pool).await; + attempt(&mut rx).await.fail(true); + user_error(task, 2).await; + let refill = attempt(&mut rx).await; + assert_eq!(refill.id, 3); + refill.fail(true); + tokio::task::yield_now().await; + tokio::task::yield_now().await; + assert!(rx.try_recv().is_err()); + assert_eq!(*sink.errors.lock().unwrap(), vec![1, 3]); +} + +#[test] +fn cancelling_demand_outside_runtime_replenishes_min_idle() { + use std::task::Poll; + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); + let (pool, mut rx, _) = runtime.block_on(async { setup(false, 2, 1000, Some(1)) }); + let mut demand = Box::pin(pool.get()); + let (held, pending) = runtime.block_on(async { + let background = attempt(&mut rx).await; + let first = start(&pool).await; + poll_fn(|cx| { + assert!(demand.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + let pending = attempt(&mut rx).await; + background.succeed(); + (first.await.unwrap().unwrap(), pending) + }); + + drop(demand); + assert!(pending.finish.is_closed()); + runtime.block_on(async { + let refill = attempt(&mut rx).await; + assert_eq!(refill.id, 3); + refill.succeed(); + tokio::task::yield_now().await; + assert_eq!(pool.state().idle_connections, 1); + assert_eq!(pool.state().connections, 2); + }); + drop(held); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_capacity_stress_no_lost_wakeups_or_deadlocks() { + // Exceed Tokio's 32-waker batch size while mixing failures and idle returns. + let (pool, mut rx, _) = setup(false, 4, 2000, None); + let controller = tokio::spawn(async move { + while let Some(a) = rx.recv().await { + let result = if a.id % 5 == 0 { + Err(Failure::new(a.id, false)) + } else { + Ok(Connection { + id: a.id, + broken: false, + }) + }; + // An idle return can legitimately cancel this demand attempt. + let _ = a.finish.send(result); + } + }); + let mut tasks = Vec::new(); + for n in 0..48 { + let pool = pool.clone(); + tasks.push(tokio::spawn(async move { + for i in 0..80 { + match pool.get().await { + Ok(mut conn) => { + tokio::task::yield_now().await; + conn.broken = (n + i) % 11 == 0; + } + Err(RunError::User(_)) => {} + Err(RunError::TimedOut) => panic!("capacity got stranded"), + } + } + })); + } + timeout(Duration::from_secs(5), async { + for t in tasks { + t.await.unwrap(); + } + }) + .await + .unwrap(); + assert!(pool.state().connections <= 4); + drop(pool); + controller.abort(); +} diff --git a/bb8/tests/test.rs b/bb8/tests/test.rs index 3f4f30f..c23720d 100644 --- a/bb8/tests/test.rs +++ b/bb8/tests/test.rs @@ -280,7 +280,7 @@ async fn test_lazy_initialization_failure_no_retry() { .build_unchecked(manager); let res = pool.get().await; - assert_eq!(res.unwrap_err(), RunError::TimedOut); + assert_eq!(res.unwrap_err(), RunError::User(Error)); } #[tokio::test] From f1abeb5b54b58e25c4956a9965037ec8e430e6fd Mon Sep 17 00:00:00 2001 From: kdcokenny Date: Wed, 7 Oct 2026 15:58:36 +0000 Subject: [PATCH 5/6] Move background test before its helpers --- bb8/tests/background.rs | 64 ++++++++++++++++++++--------------------- 1 file changed, 32 insertions(+), 32 deletions(-) diff --git a/bb8/tests/background.rs b/bb8/tests/background.rs index b402bd2..60fdc3d 100644 --- a/bb8/tests/background.rs +++ b/bb8/tests/background.rs @@ -7,6 +7,38 @@ use std::{ }; use tokio::sync::{mpsc, oneshot}; +#[tokio::test] +async fn failed_background_attempt_wakes_capacity_waiter() { + let (attempts, mut rx) = mpsc::unbounded_channel(); + let pool = Pool::builder() + .max_size(1) + .min_idle(1) + .retry_connection(false) + .connection_timeout(Duration::from_secs(1)) + .build_unchecked(Controlled { + attempts, + fatal: false, + }); + let background = rx.recv().await.unwrap(); + let wakes = Arc::new(Wakes::default()); + let waker = Waker::from(wakes.clone()); + let mut get = Box::pin(pool.get()); + assert!(get + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending()); + let before = wakes.0.load(std::sync::atomic::Ordering::SeqCst); + background.send(Err(Failure)).unwrap(); + tokio::task::yield_now().await; + assert!(wakes.0.load(std::sync::atomic::Ordering::SeqCst) > before); + assert!(get + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending()); + rx.recv().await.unwrap().send(Ok(())).unwrap(); + get.await.unwrap(); +} + #[derive(Debug)] struct Failure; @@ -51,38 +83,6 @@ impl Wake for Wakes { } } -#[tokio::test] -async fn failed_background_attempt_wakes_capacity_waiter() { - let (attempts, mut rx) = mpsc::unbounded_channel(); - let pool = Pool::builder() - .max_size(1) - .min_idle(1) - .retry_connection(false) - .connection_timeout(Duration::from_secs(1)) - .build_unchecked(Controlled { - attempts, - fatal: false, - }); - let background = rx.recv().await.unwrap(); - let wakes = Arc::new(Wakes::default()); - let waker = Waker::from(wakes.clone()); - let mut get = Box::pin(pool.get()); - assert!(get - .as_mut() - .poll(&mut Context::from_waker(&waker)) - .is_pending()); - let before = wakes.0.load(std::sync::atomic::Ordering::SeqCst); - background.send(Err(Failure)).unwrap(); - tokio::task::yield_now().await; - assert!(wakes.0.load(std::sync::atomic::Ordering::SeqCst) > before); - assert!(get - .as_mut() - .poll(&mut Context::from_waker(&waker)) - .is_pending()); - rx.recv().await.unwrap().send(Ok(())).unwrap(); - get.await.unwrap(); -} - #[derive(Debug)] struct FailingCustomizer; From 86f8e5b901cf7064b0069703b874e0f4d92b2d3d Mon Sep 17 00:00:00 2001 From: kdcokenny Date: Wed, 7 Oct 2026 19:12:13 +0000 Subject: [PATCH 6/6] Wake all capacity waiters when a checked-out connection is discarded --- bb8/src/inner.rs | 3 ++- bb8/tests/demand.rs | 25 +++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/bb8/src/inner.rs b/bb8/src/inner.rs index 4d1d6ee..4eb4029 100644 --- a/bb8/src/inner.rs +++ b/bb8/src/inner.rs @@ -215,7 +215,8 @@ impl PoolInner { let approvals = locked.dropped(1, &self.inner.statics); self.spawn_replenishing_approvals(approvals); - self.inner.notify.notify_one(); + // A getter with a pending attempt cannot use the freed slot. + self.inner.notify.notify_waiters(); } /// Adds an external connection to the pool if there is capacity for it. diff --git a/bb8/tests/demand.rs b/bb8/tests/demand.rs index f97d79e..ef239e6 100644 --- a/bb8/tests/demand.rs +++ b/bb8/tests/demand.rs @@ -6,6 +6,7 @@ use std::{ atomic::{AtomicUsize, Ordering}, Arc, Mutex, }, + task::Context, time::Duration, }; use tokio::{ @@ -279,6 +280,30 @@ async fn cancelling_owned_attempt_notifies_an_existing_capacity_waiter() { assert_eq!(waiting.await.unwrap().unwrap().id, 2); } +#[tokio::test] +async fn broken_return_notifies_capacity_waiter_behind_pending_attempt() { + let (pool, mut rx, _) = setup(false, 2, 1000, None); + let first = start(&pool).await; + attempt(&mut rx).await.succeed(); + let mut broken = first.await.unwrap().unwrap(); + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + let mut owner = Box::pin(pool.get()); + assert!(owner.as_mut().poll(&mut cx).is_pending()); + let pending = attempt(&mut rx).await; + let waiting = start(&pool).await; + + broken.broken = true; + drop(broken); + // The owner consumes the freed-slot wakeup while its attempt is still pending. + assert!(owner.as_mut().poll(&mut cx).is_pending()); + let next = attempt(&mut rx).await; + pending.succeed(); + next.succeed(); + let _owned = owner.await.unwrap(); + waiting.await.unwrap().unwrap(); +} + #[tokio::test] async fn pending_background_attempt_survives_get_timeout() { let (pool, mut rx, _) = setup(false, 1, 70, Some(1));