diff --git a/bb8/src/inner.rs b/bb8/src/inner.rs index 2dbf6af..7da32a4 100644 --- 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 where @@ -86,7 +86,7 @@ where } pub(crate) async fn get(&self) -> Result, RunError> { - 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 { diff --git a/bb8/src/internals.rs b/bb8/src/internals.rs index 1cd9039..ac8bcff 100644 --- 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,6 +249,12 @@ impl From> for IdleConn { } } +pub(crate) enum GetKind { + Direct, + Waited, + TimedOut, +} + #[derive(Default)] pub(crate) struct AtomicStatistics { pub(crate) get_direct: AtomicU64, @@ -256,6 +262,16 @@ pub(crate) struct AtomicStatistics { 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]