Polish for fixes from #223 and #224

.github/workflows/main.yml

@@ -22,64 +22,37 @@ jobs:
     runs-on: ${{ matrix.os }}
 
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions-rs/toolchain@v1
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@master
         with:
-          profile: minimal
           toolchain: ${{ matrix.rust }}
-          override: true
-      - uses: actions-rs/cargo@v1
-        with:
-          command: build
-          args: --workspace --all-targets
-      - uses: actions-rs/cargo@v1
-        with:
-          command: test
-          args: --workspace
+      - run: cargo check --all-targets
+      - run: cargo test --all-features -- --nocapture
 
   msrv:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions-rs/toolchain@v1
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@master
         with:
-          profile: minimal
           toolchain: 1.70.0
-          override: true
-      - uses: actions-rs/cargo@v1
-        with:
-          command: build
-          args: -p bb8 --all-targets
-      - uses: actions-rs/cargo@v1
-        with:
-          command: test
-          args: -p bb8
+      - run: cargo test -p bb8
 
   lint:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions-rs/toolchain@v1
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@stable
         with:
-          profile: minimal
-          toolchain: stable
-          override: true
           components: rustfmt, clippy
-      - uses: actions-rs/cargo@v1
-        with:
-          command: fmt
-          args: --all -- --check
-      - uses: actions-rs/cargo@v1
-        if: always()
-        with:
-          command: clippy
-          args: --workspace --all-targets -- -D warnings
+      - run: cargo fmt -- --check
+      - run: cargo clippy --all-features --all-targets -- -D warnings
 
   audit:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v4
     - uses: EmbarkStudios/cargo-deny-action@v1
 
   coverage:
@@ -88,11 +61,7 @@ jobs:
       CARGO_TERM_COLOR: always
     steps:
       - uses: actions/checkout@v3
-      - uses: actions-rs/toolchain@v1
-        with:
-          profile: minimal
-          toolchain: stable
-          override: true
+      - uses: dtolnay/rust-toolchain@stable
       - name: Install cargo-llvm-cov
         uses: taiki-e/install-action@cargo-llvm-cov
       - name: Generate code coverage

bb8/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "bb8"
-version = "0.8.5"
+version = "0.8.6"
 edition = "2021"
 rust-version = "1.70"
 description = "Full-featured async (tokio-based) connection pool (like r2d2)"

bb8/src/inner.rs

@@ -91,8 +91,9 @@
         let mut wait_time_start = None;
 
         let future = async {
+            let getting = self.inner.start_get();
             loop {
-                let (conn, approvals) = self.inner.pop();
+                let (conn, approvals) = getting.get();
                 self.spawn_replenishing_approvals(approvals);
 
                 // Cancellation safety: make sure to wrap the connection in a `PooledConnection`
@@ -149,18 +150,25 @@
             "handled in caller"
         );
 
+        let is_broken = self.inner.manager.has_broken(&mut conn.conn);
+        let is_expired = match self.inner.statics.max_lifetime {
+            Some(lt) => conn.is_expired(Instant::now(), lt),
+            None => false,
+        };
+
         let mut locked = self.inner.internals.lock();
-        match (state, self.inner.manager.has_broken(&mut conn.conn)) {
-            (ConnectionState::Present, false) => locked.put(conn, None, self.inner.clone()),
-            (_, is_broken) => {
-                if is_broken {
-                    self.inner.statistics.record(StatsKind::ClosedBroken);
-                }
-                let approvals = locked.dropped(1, &self.inner.statics);
-                self.spawn_replenishing_approvals(approvals);
-                self.inner.notify.notify_one();
-            }
+        if let (ConnectionState::Present, false) = (state, is_broken || is_expired) {
+            locked.put(conn, None, self.inner.clone());
+            return;
+        } else if is_broken {
+            self.inner.statistics.record(StatsKind::ClosedBroken);
+        } else if is_expired {
+            self.inner.statistics.record_connections_reaped(0, 1);
         }
+
+        let approvals = locked.dropped(1, &self.inner.statics);
+        self.spawn_replenishing_approvals(approvals);
+        self.inner.notify.notify_one();
     }
 
     /// Adds an external connection to the pool if there is capacity for it.

bb8/src/internals.rs

@@ -36,17 +36,6 @@
         }
     }
 
-    pub(crate) fn pop(&self) -> (Option<Conn<M::Connection>>, ApprovalIter) {
-        let mut locked = self.internals.lock();
-        let conn = locked.conns.pop_front().map(|idle| idle.conn);
-        let approvals = match &conn {
-            Some(_) => locked.wanted(&self.statics),
-            None => locked.approvals(&self.statics, 1),
-        };
-
-        (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);
@@ -67,6 +56,10 @@
         iter
     }
 
+    pub(crate) fn start_get(self: &Arc<Self>) -> Getting<M> {
+        Getting::new(self.clone())
+    }
+
     pub(crate) fn forward_error(&self, err: M::Error) {
         self.statics.error_sink.sink(err);
     }
@@ -81,6 +74,7 @@
     conns: VecDeque<IdleConn<M::Connection>>,
     num_conns: u32,
     pending_conns: u32,
+    in_flight: u32,
 }
 
 impl<M> PoolInternals<M>
@@ -169,7 +163,7 @@
                 }
             }
             if let Some(lifetime) = config.max_lifetime {
-                if now - conn.conn.birth >= lifetime {
+                if conn.conn.is_expired(now, lifetime) {
                     closed_max_lifetime += 1;
                     keep &= false;
                 }
@@ -202,6 +196,7 @@
             conns: VecDeque::new(),
             num_conns: 0,
             pending_conns: 0,
+            in_flight: 0,
         }
     }
 }
@@ -236,6 +231,43 @@
     _priv: (),
 }
 
+pub(crate) struct Getting<M: ManageConnection + Send> {
+    inner: Arc<SharedPool<M>>,
+}
+
+impl<M: ManageConnection + Send> Getting<M> {
+    pub(crate) fn get(&self) -> (Option<Conn<M::Connection>>, ApprovalIter) {
+        let mut locked = self.inner.internals.lock();
+        if let Some(IdleConn { conn, .. }) = locked.conns.pop_front() {
+            return (Some(conn), locked.wanted(&self.inner.statics));
+        }
+
+        let approvals = match locked.in_flight > locked.pending_conns {
+            true => 1,
+            false => 0,
+        };
+
+        (None, locked.approvals(&self.inner.statics, approvals))
+    }
+}
+
+impl<M: ManageConnection + Send> Getting<M> {
+    fn new(inner: Arc<SharedPool<M>>) -> Self {
+        {
+            let mut locked = inner.internals.lock();
+            locked.in_flight += 1;
+        }
+        Getting { inner }
+    }
+}
+
+impl<M: ManageConnection + Send> Drop for Getting<M> {
+    fn drop(&mut self) {
+        let mut locked = self.inner.internals.lock();
+        locked.in_flight -= 1;
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct Conn<C>
 where
@@ -252,6 +284,11 @@
             birth: Instant::now(),
         }
     }
+
+    pub(crate) fn is_expired(&self, now: Instant, max: Duration) -> bool {
+        // The current age of the connection is longer than the maximum allowed lifetime
+        now - self.birth >= max
+    }
 }
 
 impl<C: Send> From<IdleConn<C>> for Conn<C> {
@@ -357,3 +394,31 @@
     ClosedBroken,
     ClosedInvalid,
 }
+
+#[cfg(test)]
+mod tests {
+    use std::time::{Duration, Instant};
+
+    use crate::internals::Conn;
+
+    #[test]
+    fn test_conn_is_expired() {
+        let conn = Conn {
+            conn: (),
+            birth: Instant::now(),
+        };
+
+        assert!(
+            !conn.is_expired(conn.birth, Duration::from_nanos(1)),
+            "at birth, the connection has not expired"
+        );
+        assert!(
+            !conn.is_expired(Instant::now(), Duration::from_secs(5)),
+            "from setup to now is shorter than 5s, so the connection has not expired"
+        );
+        assert!(
+            conn.is_expired(Instant::now(), Duration::from_nanos(1)),
+            "from setup to now is longer than 1ns, so the connection has expired"
+        );
+    }
+}

bb8/tests/test.rs

@@ -12,7 +12,8 @@ use std::{error, fmt};
 use async_trait::async_trait;
 use futures_util::future::{err, lazy, ok, pending, ready, try_join_all, FutureExt};
 use futures_util::stream::{FuturesUnordered, TryStreamExt};
-use tokio::{sync::oneshot, time::timeout};
+use tokio::sync::oneshot;
+use tokio::time::{sleep, timeout};
 
 #[derive(Debug, PartialEq, Eq)]
 pub struct Error;
@@ -585,6 +586,40 @@ async fn test_max_lifetime() {
     assert_eq!(pool.state().statistics.connections_closed_max_lifetime, 5);
 }
 
+#[tokio::test]
+async fn test_max_lifetime_reap_on_drop() {
+    static DROPPED: AtomicUsize = AtomicUsize::new(0);
+
+    #[derive(Default)]
+    struct Connection;
+
+    impl Drop for Connection {
+        fn drop(&mut self) {
+            DROPPED.fetch_add(1, Ordering::SeqCst);
+        }
+    }
+
+    let manager = OkManager::<Connection>::new();
+    let pool = Pool::builder()
+        .max_lifetime(Some(Duration::from_secs(1)))
+        .connection_timeout(Duration::from_secs(1))
+        .reaper_rate(Duration::from_secs(999))
+        .build(manager)
+        .await
+        .unwrap();
+
+    let conn = pool.get().await;
+
+    // And wait.
+    sleep(Duration::from_secs(2)).await;
+    assert_eq!(DROPPED.load(Ordering::SeqCst), 0);
+
+    // Connection is reaped on drop.
+    drop(conn);
+    assert_eq!(DROPPED.load(Ordering::SeqCst), 1);
+    assert_eq!(pool.state().statistics.connections_closed_max_lifetime, 1);
+}
+
 #[tokio::test]
 async fn test_min_idle() {
     let pool = Pool::builder()
@@ -1068,3 +1103,37 @@ async fn test_add_checks_broken_connections() {
     let res = pool.add(conn);
     assert!(matches!(res, Err(AddError::Broken(_))));
 }
+
+#[tokio::test]
+async fn test_reuse_on_drop() {
+    let pool = Pool::builder()
+        .min_idle(0)
+        .max_size(100)
+        .queue_strategy(QueueStrategy::Lifo)
+        .build(OkManager::<FakeConnection>::new())
+        .await
+        .unwrap();
+
+    // The first get should
+    // 1) see nothing in the pool,
+    // 2) spawn a single replenishing approval,
+    // 3) get notified of the new connection and grab it from the pool
+    let conn_0 = pool.get().await.expect("should connect");
+    // Dropping the connection queues up a notify
+    drop(conn_0);
+
+    // The second get should
+    // 1) see the first connection in the pool and grab it
+    let _conn_1 = pool.get().await.expect("should connect");
+
+    // The third get will
+    // 1) see nothing in the pool,
+    // 2) spawn a single replenishing approval,
+    // 3) get notified of the new connection,
+    // 4) see nothing in the pool,
+    // 5) _not_ spawn a single replenishing approval,
+    // 6) get notified of the new connection and grab it from the pool
+    let _conn_2 = pool.get().await.expect("should connect");
+
+    assert_eq!(pool.state().connections, 2);
+}