Move some logic about recording

djc · May 30, 2024 · bc9fba4 · bc9fba4
1 parent 95565a5
commit bc9fba4
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 24 deletions.
diff --git a/bb8/src/inner.rs b/bb8/src/inner.rs
@@ -11,7 +11,7 @@ use tokio::spawn;
 use tokio::time::{interval_at, sleep, timeout, Interval};
 
 use crate::api::{Builder, ConnectionState, ManageConnection, PooledConnection, RunError, State};
-use crate::internals::{Approval, ApprovalIter, AtomicStatistics, Conn, SharedPool};
+use crate::internals::{Approval, ApprovalIter, AtomicStatistics, Conn, GetKind, SharedPool};
 
 pub(crate) struct PoolInner<M>
 where
@@ -86,7 +86,7 @@ where
     }
 
     pub(crate) async fn get(&self) -> Result<PooledConnection<'_, M>, RunError<M::Error>> {
-        let mut get_direct = true;
+        let mut get_kind = GetKind::Direct;
 
         let future = async {
             loop {
@@ -99,7 +99,7 @@ where
                 let mut conn = match conn {
                     Some(conn) => PooledConnection::new(self, conn),
                     None => {
-                        get_direct = false;
+                        get_kind = GetKind::Waited;
                         self.inner.notify.notified().await;
                         continue;
                     }
@@ -120,29 +120,17 @@ where
             }
         };
 
-        match timeout(self.inner.statics.connection_timeout, future).await {
-            Ok(result) => {
-                if get_direct {
-                    self.inner
-                        .statistics
-                        .get_direct
-                        .fetch_add(1, Ordering::SeqCst);
-                } else {
-                    self.inner
-                        .statistics
-                        .get_waited
-                        .fetch_add(1, Ordering::SeqCst);
-                }
-                result
-            }
+        let result = match timeout(self.inner.statics.connection_timeout, future).await {
+            Ok(result) => result,
             _ => {
-                self.inner
-                    .statistics
-                    .get_timed_out
-                    .fetch_add(1, Ordering::SeqCst);
+                get_kind = GetKind::TimedOut;
                 Err(RunError::TimedOut)
             }
-        }
+        };
+
+        self.inner.statistics.record_get(get_kind);
+
+        result
     }
 
     pub(crate) async fn connect(&self) -> Result<M::Connection, M::Error> {

diff --git a/bb8/src/internals.rs b/bb8/src/internals.rs
@@ -1,5 +1,5 @@
 use std::cmp::min;
-use std::sync::atomic::AtomicU64;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 use std::time::Instant;
 
@@ -249,13 +249,29 @@ impl<C: Send> From<Conn<C>> for IdleConn<C> {
     }
 }
 
+pub(crate) enum GetKind {
+    Direct,
+    Waited,
+    TimedOut,
+}
+
 #[derive(Default)]
 pub(crate) struct AtomicStatistics {
     pub(crate) get_direct: AtomicU64,
     pub(crate) get_waited: AtomicU64,
     pub(crate) get_timed_out: AtomicU64,
 }
 
+impl AtomicStatistics {
+    pub(crate) fn record_get(&self, get_kind: GetKind) {
+        match get_kind {
+            GetKind::Direct => self.get_direct.fetch_add(1, Ordering::SeqCst),
+            GetKind::Waited => self.get_waited.fetch_add(1, Ordering::SeqCst),
+            GetKind::TimedOut => self.get_timed_out.fetch_add(1, Ordering::SeqCst),
+        };
+    }
+}
+
 /// Information about the state of a `Pool`.
 #[derive(Debug)]
 #[non_exhaustive]