Add Pool::add

bb8/src/api.rs

@@ -75,6 +75,14 @@ impl<M: ManageConnection> Pool<M> {
     pub fn state(&self) -> State {
         self.inner.state()
     }
+
+    /// Adds a connection to the pool.
+    ///
+    /// If the connection is broken, or the pool is at capacity, the
+    /// connection is not added and instead returned to the caller in Err.
+    pub fn add(&self, conn: M::Connection) -> Result<(), AddError<M::Connection>> {
+        self.inner.try_put(conn)
+    }
 }
 
 /// Information about the state of a `Pool`.
@@ -526,6 +534,33 @@ where
     }
 }
 
+/// Error type returned by `Pool::add(conn)`
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum AddError<C> {
+    /// The connection was broken before it could be added.
+    Broken(C),
+    /// Unable to add the connection to the pool due to insufficient capacity.
+    NoCapacity(C),
+}
+
+impl<E: error::Error + 'static> fmt::Display for AddError<E> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match *self {
+            AddError::Broken(_) => write!(f, "The connection was broken before it could be added"),
+            AddError::NoCapacity(_) => write!(
+                f,
+                "Unable to add the connection to the pool due to insufficient capacity"
+            ),
+        }
+    }
+}
+
+impl<E: error::Error + 'static> error::Error for AddError<E> {
+    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
+        None
+    }
+}
+
 /// A trait to receive errors generated by connection management that aren't
 /// tied to any particular caller.
 pub trait ErrorSink<E>: fmt::Debug + Send + Sync + 'static {

bb8/src/inner.rs

@@ -9,7 +9,9 @@ use futures_util::TryFutureExt;
 use tokio::spawn;
 use tokio::time::{interval_at, sleep, timeout, Interval};
 
-use crate::api::{Builder, ConnectionState, ManageConnection, PooledConnection, RunError, State};
+use crate::api::{
+    AddError, Builder, ConnectionState, ManageConnection, PooledConnection, RunError, State,
+};
 use crate::internals::{Approval, ApprovalIter, Conn, SharedPool, StatsGetKind, StatsKind};
 
 pub(crate) struct PoolInner<M>
@@ -161,6 +163,15 @@ where
         }
     }
 
+    /// Adds an external connection to the pool if there is capacity for it.
+    pub(crate) fn try_put(&self, mut conn: M::Connection) -> Result<(), AddError<M::Connection>> {
+        if self.inner.manager.has_broken(&mut conn) {
+            Err(AddError::Broken(conn))
+        } else {
+            self.inner.try_put(conn).map_err(AddError::NoCapacity)
+        }
+    }
+
     /// Returns information about the current state of the pool.
     pub(crate) fn state(&self) -> State {
         self.inner

bb8/src/internals.rs

@@ -47,6 +47,17 @@ where
         (conn, approvals)
     }
 
+    pub(crate) fn try_put(self: &Arc<Self>, conn: M::Connection) -> Result<(), M::Connection> {
+        let mut locked = self.internals.lock();
+        let mut approvals = locked.approvals(&self.statics, 1);
+        let Some(approval) = approvals.next() else {
+            return Err(conn);
+        };
+        let conn = Conn::new(conn);
+        locked.put(conn, Some(approval), self.clone());
+        Ok(())
+    }
+
     pub(crate) fn reap(&self) -> ApprovalIter {
         let mut locked = self.internals.lock();
         let (iter, closed_idle_timeout, closed_max_lifetime) = locked.reap(&self.statics);

bb8/src/lib.rs

@@ -35,7 +35,7 @@
 
 mod api;
 pub use api::{
-    Builder, CustomizeConnection, ErrorSink, ManageConnection, NopErrorSink, Pool,
+    AddError, Builder, CustomizeConnection, ErrorSink, ManageConnection, NopErrorSink, Pool,
     PooledConnection, QueueStrategy, RunError, State, Statistics,
 };
 

bb8/tests/test.rs

@@ -207,6 +207,30 @@ fn test_is_send_sync() {
     is_send_sync::<Pool<OkManager<FakeConnection>>>();
 }
 
+// A connection manager that always returns `true` for `has_broken()`
+#[derive(Default)]
+struct BrokenConnectionManager<C> {
+    _c: PhantomData<C>,
+}
+
+#[async_trait]
+impl<C: Default + Send + Sync + 'static> ManageConnection for BrokenConnectionManager<C> {
+    type Connection = C;
+    type Error = Error;
+
+    async fn connect(&self) -> Result<Self::Connection, Self::Error> {
+        Ok(C::default())
+    }
+
+    async fn is_valid(&self, _conn: &mut Self::Connection) -> Result<(), Self::Error> {
+        Ok(())
+    }
+
+    fn has_broken(&self, _: &mut Self::Connection) -> bool {
+        true
+    }
+}
+
 #[tokio::test]
 async fn test_drop_on_broken() {
     static DROPPED: AtomicBool = AtomicBool::new(false);
@@ -220,27 +244,11 @@ async fn test_drop_on_broken() {
         }
     }
 
-    struct Handler;
-
-    #[async_trait]
-    impl ManageConnection for Handler {
-        type Connection = Connection;
-        type Error = Error;
-
-        async fn connect(&self) -> Result<Self::Connection, Self::Error> {
-            Ok(Default::default())
-        }
-
-        async fn is_valid(&self, _conn: &mut Self::Connection) -> Result<(), Self::Error> {
-            Ok(())
-        }
-
-        fn has_broken(&self, _: &mut Self::Connection) -> bool {
-            true
-        }
-    }
+    let pool = Pool::builder()
+        .build(BrokenConnectionManager::<Connection>::default())
+        .await
+        .unwrap();
 
-    let pool = Pool::builder().build(Handler).await.unwrap();
     {
         let _ = pool.get().await.unwrap();
     }
@@ -1012,3 +1020,51 @@ async fn test_statistics_connections_created() {
 
     assert_eq!(pool.state().statistics.connections_created, 1);
 }
+
+#[tokio::test]
+async fn test_can_use_added_connections() {
+    let pool = Pool::builder()
+        .connection_timeout(Duration::from_millis(1))
+        .build_unchecked(NthConnectionFailManager::<FakeConnection>::new(0));
+
+    // Assert pool can't replenish connections on its own
+    let res = pool.get().await;
+    assert_eq!(res.unwrap_err(), RunError::TimedOut);
+
+    pool.add(FakeConnection).unwrap();
+    let res = pool.get().await;
+    assert!(res.is_ok());
+}
+
+#[tokio::test]
+async fn test_add_ok_until_max_size() {
+    let pool = Pool::builder()
+        .min_idle(1)
+        .max_size(3)
+        .build(OkManager::<FakeConnection>::new())
+        .await
+        .unwrap();
+
+    for _ in 0..2 {
+        let conn = pool.dedicated_connection().await.unwrap();
+        pool.add(conn).unwrap();
+    }
+
+    let conn = pool.dedicated_connection().await.unwrap();
+    let res = pool.add(conn);
+    assert!(matches!(res, Err(AddError::NoCapacity(_))));
+}
+
+#[tokio::test]
+async fn test_add_checks_broken_connections() {
+    let pool = Pool::builder()
+        .min_idle(1)
+        .max_size(3)
+        .build(BrokenConnectionManager::<FakeConnection>::default())
+        .await
+        .unwrap();
+
+    let conn = pool.dedicated_connection().await.unwrap();
+    let res = pool.add(conn);
+    assert!(matches!(res, Err(AddError::Broken(_))));
+}