diff --git a/bb8/Cargo.toml b/bb8/Cargo.toml index aab5878..77bb8d3 100644 --- a/bb8/Cargo.toml +++ b/bb8/Cargo.toml @@ -12,13 +12,13 @@ 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" [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 6a80e70..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(), @@ -369,6 +379,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] @@ -456,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, @@ -463,6 +476,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 +493,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 d7ad1d5..4eb4029 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; } } @@ -164,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. @@ -186,16 +238,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(true).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, background: bool) -> 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,20 +269,17 @@ 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(()); + if !background { + self.inner.statistics.record(StatsKind::Created); + } + return Ok(conn); } Err(e) => { if !self.inner.statics.retry_connection - || Instant::now() - start > self.inner.statics.connection_timeout + || self.inner.manager.error_is_fatal(&e) + || (background + && Instant::now() - start > self.inner.statics.connection_timeout) { - let mut locked = shared.internals.lock(); - locked.connect_failed(approval); return Err(e); } else { self.inner.forward_error(e); 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/background.rs b/bb8/tests/background.rs new file mode 100644 index 0000000..60fdc3d --- /dev/null +++ b/bb8/tests/background.rs @@ -0,0 +1,121 @@ +use bb8::{ManageConnection, Pool}; +use std::{ + future::Future, + sync::Arc, + task::{Context, Wake, Waker}, + time::Duration, +}; +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; + +struct Controlled { + attempts: mpsc::UnboundedSender>>, + fatal: bool, +} + +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 + } + + fn error_is_fatal(&self, _: &Failure) -> bool { + self.fatal + } +} + +#[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); + } +} + +#[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()); + } +} 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..ef239e6 --- /dev/null +++ b/bb8/tests/demand.rs @@ -0,0 +1,493 @@ +use bb8::{ErrorSink, ManageConnection, Pool, PooledConnection, RunError}; +use std::{ + cell::Cell, + future::{poll_fn, Future}, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, + task::Context, + 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 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)); + 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]