diff --git a/bb8/src/api.rs b/bb8/src/api.rs index 77bf219..92967f3 100644 --- a/bb8/src/api.rs +++ b/bb8/src/api.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::error; use std::fmt; use std::marker::PhantomData; @@ -54,6 +55,14 @@ impl Pool { self.inner.get().await } + /// Retrieves an owned connection from the pool + /// + /// 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. + pub async fn get_owned(&self) -> Result, RunError> { + self.inner.get_owned().await + } + /// Get a new dedicated connection that will not be managed by the pool. /// An application may want a persistent connection (e.g. to do a /// postgres LISTEN) that will not be closed or repurposed by the pool. @@ -285,7 +294,7 @@ pub struct PooledConnection<'a, M> where M: ManageConnection, { - pool: &'a PoolInner, + pool: Cow<'a, PoolInner>, conn: Option>, } @@ -295,7 +304,7 @@ where { pub(crate) fn new(pool: &'a PoolInner, conn: Conn) -> Self { Self { - pool, + pool: Cow::Borrowed(pool), conn: Some(conn), } } @@ -305,6 +314,18 @@ where } } +impl PooledConnection<'static, M> +where + M: ManageConnection, +{ + pub(crate) fn new_owned(pool: PoolInner, conn: Conn) -> Self { + Self { + pool: Cow::Owned(pool), + conn: Some(conn), + } + } +} + impl<'a, M> Deref for PooledConnection<'a, M> where M: ManageConnection, @@ -340,7 +361,7 @@ where M: ManageConnection, { fn drop(&mut self) { - self.pool.put_back(self.conn.take()); + self.pool.as_ref().put_back(self.conn.take()) } } diff --git a/bb8/src/inner.rs b/bb8/src/inner.rs index 4c20e42..662bb86 100644 --- a/bb8/src/inner.rs +++ b/bb8/src/inner.rs @@ -83,13 +83,36 @@ where } pub(crate) async fn get(&self) -> Result, RunError> { + self.make_pooled(|this, conn| PooledConnection::new(this, conn)) + .await + } + + pub(crate) async fn get_owned( + &self, + ) -> Result, RunError> { + self.make_pooled(|this, conn| { + let pool = PoolInner { + inner: Arc::clone(&this.inner), + }; + PooledConnection::new_owned(pool, conn) + }) + .await + } + + pub(crate) async fn make_pooled<'a, 'b, F>( + &'a self, + make_pool: F, + ) -> Result, RunError> + where + F: Fn(&'a Self, Conn) -> PooledConnection<'b, M>, + { loop { let mut conn = { let mut locked = self.inner.internals.lock(); match locked.pop(&self.inner.statics) { Some((conn, approvals)) => { self.spawn_replenishing_approvals(approvals); - PooledConnection::new(self, conn) + make_pool(self, conn) } None => break, } @@ -117,7 +140,7 @@ where }; match timeout(self.inner.statics.connection_timeout, rx).await { - Ok(Ok(mut guard)) => Ok(PooledConnection::new(self, guard.extract())), + Ok(Ok(mut guard)) => Ok(make_pool(self, guard.extract())), _ => Err(RunError::TimedOut), } }