Allow users to customize connections created by the pool (#89, fixes #88)

Author: agersant

bb8/src/api.rs

@@ -84,6 +84,8 @@ pub struct Builder<M: ManageConnection> {
     pub(crate) error_sink: Box<dyn ErrorSink<M::Error>>,
     /// The time interval used to wake up and reap connections.
     pub(crate) reaper_rate: Duration,
+    /// User-supplied trait object responsible for initializing connections
+    pub(crate) connection_customizer: Box<dyn CustomizeConnection<M::Connection, M::Error>>,
     _p: PhantomData<M>,
 }
 
@@ -98,6 +100,7 @@ impl<M: ManageConnection> Default for Builder<M> {
             connection_timeout: Duration::from_secs(30),
             error_sink: Box::new(NopErrorSink),
             reaper_rate: Duration::from_secs(30),
+            connection_customizer: Box::new(NopConnectionCustomizer {}),
             _p: PhantomData,
         }
     }
@@ -204,6 +207,18 @@ impl<M: ManageConnection> Builder<M> {
         self
     }
 
+    /// Set the connection customizer which will be used to initialize
+    /// connections created by the pool.
+    ///
+    /// Defaults to `NopConnectionCustomizer`.
+    pub fn connection_customizer(
+        mut self,
+        connection_customizer: Box<dyn CustomizeConnection<M::Connection, M::Error>>,
+    ) -> Builder<M> {
+        self.connection_customizer = connection_customizer;
+        self
+    }
+
     fn build_inner(self, manager: M) -> Pool<M> {
         if let Some(min_idle) = self.min_idle {
             assert!(
@@ -253,6 +268,29 @@ pub trait ManageConnection: Sized + Send + Sync + 'static {
     fn has_broken(&self, conn: &mut Self::Connection) -> bool;
 }
 
+/// A trait which provides functionality to initialize a connection
+#[async_trait]
+pub trait CustomizeConnection<C: Send + 'static, E: 'static>:
+    std::fmt::Debug + Send + Sync + 'static
+{
+    /// Called with connections immediately after they are returned from
+    /// `ManageConnection::connect`.
+    ///
+    /// The default implementation simply returns `Ok(())`.
+    ///
+    /// # Errors
+    ///
+    /// If this method returns an error, the connection will be discarded.
+    #[allow(unused_variables)]
+    async fn on_acquire(&self, connection: &mut C) -> Result<(), E> {
+        Ok(())
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+struct NopConnectionCustomizer;
+impl<C: Send + 'static, E: 'static> CustomizeConnection<C, E> for NopConnectionCustomizer {}
+
 /// A smart pointer wrapping a connection.
 pub struct PooledConnection<'a, M>
 where

bb8/src/inner.rs

@@ -6,6 +6,7 @@ use std::time::{Duration, Instant};
 
 use futures_channel::oneshot;
 use futures_util::stream::{FuturesUnordered, StreamExt};
+use futures_util::TryFutureExt;
 use parking_lot::Mutex;
 use tokio::spawn;
 use tokio::time::{interval_at, sleep, timeout, Interval};
@@ -126,7 +127,9 @@ where
     }
 
     pub(crate) async fn connect(&self) -> Result<M::Connection, M::Error> {
-        self.inner.manager.connect().await
+        let mut conn = self.inner.manager.connect().await?;
+        self.on_acquire_connection(&mut conn).await?;
+        Ok(conn)
     }
 
     /// Return connection back in to the pool
@@ -174,7 +177,13 @@ where
         let start = Instant::now();
         let mut delay = Duration::from_secs(0);
         loop {
-            match shared.manager.connect().await {
+            let conn = shared
+                .manager
+                .connect()
+                .and_then(|mut c| async { self.on_acquire_connection(&mut c).await.map(|_| c) })
+                .await;
+
+            match conn {
                 Ok(conn) => {
                     let conn = Conn::new(conn);
                     shared.internals.lock().put(conn, Some(approval));
@@ -194,6 +203,14 @@ where
             }
         }
     }
+
+    async fn on_acquire_connection(&self, conn: &mut M::Connection) -> Result<(), M::Error> {
+        self.inner
+            .statics
+            .connection_customizer
+            .on_acquire(conn)
+            .await
+    }
 }
 
 impl<M> Clone for PoolInner<M>

bb8/src/lib.rs

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

bb8/tests/test.rs

@@ -746,3 +746,43 @@ async fn test_guard() {
     tx4.send(()).unwrap();
     tx6.send(()).unwrap();
 }
+
+#[tokio::test]
+async fn test_customize_connection_acquire() {
+    #[derive(Debug, Default)]
+    struct Connection {
+        custom_field: usize,
+    };
+
+    #[derive(Debug, Default)]
+    struct CountingCustomizer {
+        count: std::sync::atomic::AtomicUsize,
+    }
+
+    #[async_trait]
+    impl<E: 'static> CustomizeConnection<Connection, E> for CountingCustomizer {
+        async fn on_acquire(&self, connection: &mut Connection) -> Result<(), E> {
+            connection.custom_field = 1 + self.count.fetch_add(1, Ordering::SeqCst);
+            Ok(())
+        }
+    }
+
+    let pool = Pool::builder()
+        .max_size(2)
+        .connection_customizer(Box::new(CountingCustomizer::default()))
+        .build(OkManager::<Connection>::new())
+        .await
+        .unwrap();
+
+    // Each connection gets customized
+    {
+        let connection_1 = pool.get().await.unwrap();
+        assert_eq!(connection_1.custom_field, 1);
+        let connection_2 = pool.get().await.unwrap();
+        assert_eq!(connection_2.custom_field, 2);
+    }
+
+    // Connections don't get customized again on re-use
+    let connection_1_or_2 = pool.get().await.unwrap();
+    assert!(connection_1_or_2.custom_field == 1 || connection_1_or_2.custom_field == 2);
+}