return connection errors from get()

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"]

bb8/src/api.rs

@@ -22,6 +22,15 @@ impl<M: ManageConnection> Pool<M> {
     }
 
     /// 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<PooledConnection<'_, M>, RunError<M::Error>> {
         self.inner.get().await
     }
@@ -30,6 +39,7 @@ impl<M: ManageConnection> Pool<M> {
     ///
     /// 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<PooledConnection<'static, M>, RunError<M::Error>> {
         Ok(PooledConnection {
             conn: self.get().await?.take(),
@@ -369,6 +379,8 @@ impl<M: ManageConnection> Builder<M> {
     /// 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,13 +468,22 @@ pub trait ManageConnection: Sized + Send + Sync + 'static {
 
     /// Attempts to create a new connection.
     fn connect(&self) -> impl Future<Output = Result<Self::Connection, Self::Error>> + Send;
+
     /// Determines if the connection is still connected to the database.
     fn is_valid(
         &self,
         conn: &mut Self::Connection,
     ) -> impl Future<Output = Result<(), Self::Error>> + 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<C: Send + 'static, E: 'static>:
     /// 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,

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<M: ManageConnection + Send> {
     inner: Arc<SharedPool<M>>,
@@ -46,6 +49,10 @@ impl<M: ManageConnection + Send> PoolInner<M> {
         Ok(())
     }
 
+    pub(crate) fn replenish_after_demand(inner: Arc<SharedPool<M>>) {
+        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<M: ManageConnection + Send> PoolInner<M> {
 
         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<M: ManageConnection + Send> PoolInner<M> {
                         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<M: ManageConnection + Send> PoolInner<M> {
 
         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,37 +238,48 @@ impl<M: ManageConnection + Send> PoolInner<M> {
 
     // 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<M::Connection, M::Error> {
         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) })
                 .await;
 
             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);

bb8/src/internals.rs

@@ -225,18 +225,31 @@ pub(crate) struct Getting<M: ManageConnection + Send> {
 }
 
 impl<M: ManageConnection + Send> Getting<M> {
-    pub(crate) fn get(&self) -> (Option<Conn<M::Connection>>, ApprovalIter) {
+    pub(crate) fn get(&self, connecting: bool) -> GettingResult<M> {
         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<M: ManageConnection + Send> Drop for Getting<M> {
     }
 }
 
+pub(crate) struct GettingResult<M: ManageConnection> {
+    pub(crate) conn: Option<Conn<M::Connection>>,
+    pub(crate) approvals: ApprovalIter,
+    pub(crate) demand: Option<Demand<M>>,
+}
+
+pub(crate) struct Demand<M: ManageConnection> {
+    inner: Arc<SharedPool<M>>,
+    approval: Option<Approval>,
+    runtime: tokio::runtime::Handle,
+}
+
+impl<M: ManageConnection> Demand<M> {
+    pub(crate) fn checkout(mut self, conn: M::Connection) -> (Conn<M::Connection>, 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<M: ManageConnection> Drop for Demand<M> {
+    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<C: Send> {
     pub(crate) conn: C,