Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds feature for users to customize connections created by the pool #89

Merged
merged 8 commits into from
Dec 14, 2020
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 47 additions & 0 deletions bb8/src/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 and cleaning-up connections
pub(crate) connection_customizer: Box<dyn CustomizeConnection<M::Connection, M::Error>>,
_p: PhantomData<M>,
}

Expand All @@ -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,
}
}
Expand Down Expand Up @@ -204,6 +207,18 @@ impl<M: ManageConnection> Builder<M> {
self
}

/// Set the connection customizer which will be used to initialize and cleanup
/// connections created and destroyed 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!(
Expand Down Expand Up @@ -253,6 +268,38 @@ pub trait ManageConnection: Sized + Send + Sync + 'static {
fn has_broken(&self, conn: &mut Self::Connection) -> bool;
}

/// A trait which provides functionality to initialize and cleanup 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(())
}

/// Called with connections when they are removed from the pool.
///
/// The connections may be broken (as reported by `ManageConnection::is_valid` or
/// `ManageConnection::has_broken`), or have simply timed out.
///
/// The default implementation does nothing.
#[allow(unused_variables)]
fn on_release(&self, connection: &mut C) {}
agersant marked this conversation as resolved.
Show resolved Hide resolved
}

#[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
Expand Down
37 changes: 34 additions & 3 deletions bb8/src/inner.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,15 @@
use std::cmp::{max, min};
use std::fmt;
use std::future::Future;
use std::ops::DerefMut;
use std::sync::{Arc, Weak};
use std::time::{Duration, Instant};

use futures_channel::oneshot;
use futures_util::stream::{FuturesUnordered, StreamExt};
use futures_util::{
stream::{FuturesUnordered, StreamExt},
TryFutureExt,
};
use parking_lot::Mutex;
use tokio::spawn;
use tokio::time::{interval_at, sleep, timeout, Interval};
Expand Down Expand Up @@ -106,6 +110,7 @@ where
match self.inner.manager.is_valid(&mut conn).await {
Ok(()) => return Ok(conn),
Err(_) => {
self.on_release_connection(conn.deref_mut());
conn.drop_invalid();
continue;
}
Expand All @@ -126,7 +131,11 @@ where
}

pub(crate) async fn connect(&self) -> Result<M::Connection, M::Error> {
self.inner.manager.connect().await
self.inner
.manager
.connect()
.and_then(|conn| self.on_acquire_connection(conn))
.await
}

/// Return connection back in to the pool
Expand All @@ -135,6 +144,7 @@ where
if !self.inner.manager.has_broken(&mut conn.conn) {
Some(conn)
} else {
self.on_release_connection(&mut conn.conn);
None
}
});
Expand Down Expand Up @@ -174,7 +184,12 @@ 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(|conn| self.on_acquire_connection(conn))
.await;
match conn {
Ok(conn) => {
let conn = Conn::new(conn);
shared.internals.lock().put(conn, Some(approval));
Expand All @@ -194,6 +209,22 @@ where
}
}
}

async fn on_acquire_connection(
&self,
mut conn: M::Connection,
agersant marked this conversation as resolved.
Show resolved Hide resolved
) -> Result<M::Connection, M::Error> {
self.inner
.statics
.connection_customizer
.on_acquire(&mut conn)
.await
.map(|_| conn)
}

fn on_release_connection(&self, conn: &mut M::Connection) {
self.inner.statics.connection_customizer.on_release(conn);
}
}

impl<M> Clone for PoolInner<M>
Expand Down
3 changes: 2 additions & 1 deletion bb8/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
157 changes: 156 additions & 1 deletion bb8/tests/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use std::iter::FromIterator;
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Mutex;
use std::sync::{Arc, Mutex};
use std::task::Poll;
use std::time::Duration;
use std::{error, fmt, mem};
Expand Down Expand Up @@ -746,3 +746,158 @@ 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);
}

#[tokio::test]
async fn test_customize_connection_release() {
#[derive(Debug)]
struct CountingCustomizer {
num_conn_released: Arc<AtomicUsize>,
}

impl CountingCustomizer {
fn new(num_conn_released: Arc<AtomicUsize>) -> Self {
Self { num_conn_released }
}
}

impl<E: 'static> CustomizeConnection<FakeConnection, E> for CountingCustomizer {
fn on_release(&self, _connection: &mut FakeConnection) {
self.num_conn_released.fetch_add(1, Ordering::SeqCst);
}
}

#[derive(Debug)]
struct BreakableManager<C> {
_c: PhantomData<C>,
valid: Arc<AtomicBool>,
broken: Arc<AtomicBool>,
};

impl<C> BreakableManager<C> {
fn new(valid: Arc<AtomicBool>, broken: Arc<AtomicBool>) -> Self {
Self {
valid,
broken,
_c: PhantomData,
}
}
}

#[async_trait]
impl<C> ManageConnection for BreakableManager<C>
where
C: Default + Send + Sync + 'static,
{
type Connection = C;
type Error = Error;

async fn connect(&self) -> Result<Self::Connection, Self::Error> {
Ok(Default::default())
}

async fn is_valid(
&self,
_conn: &mut PooledConnection<'_, Self>,
) -> Result<(), Self::Error> {
if self.valid.load(Ordering::SeqCst) {
Ok(())
} else {
Err(Error)
}
}

fn has_broken(&self, _: &mut Self::Connection) -> bool {
self.broken.load(Ordering::SeqCst)
}
}

let valid = Arc::new(AtomicBool::new(true));
let broken = Arc::new(AtomicBool::new(false));
let manager = BreakableManager::<FakeConnection>::new(valid.clone(), broken.clone());

let num_conn_released = Arc::new(AtomicUsize::new(0));
let customizer = CountingCustomizer::new(num_conn_released.clone());

let pool = Pool::builder()
.max_size(2)
.connection_customizer(Box::new(customizer))
.build(manager)
.await
.unwrap();

// Connections go in and out of the pool without being released
{
{
let _connection_1 = pool.get().await.unwrap();
let _connection_2 = pool.get().await.unwrap();
assert_eq!(num_conn_released.load(Ordering::SeqCst), 0);
}
{
let _connection_1 = pool.get().await.unwrap();
let _connection_2 = pool.get().await.unwrap();
assert_eq!(num_conn_released.load(Ordering::SeqCst), 0);
}
}

// Invalid connections get released
{
valid.store(false, Ordering::SeqCst);
let _connection_1 = pool.get().await.unwrap();
assert_eq!(num_conn_released.load(Ordering::SeqCst), 2);
let _connection_2 = pool.get().await.unwrap();
assert_eq!(num_conn_released.load(Ordering::SeqCst), 2);
valid.store(true, Ordering::SeqCst);
}

// Broken connections get released
{
num_conn_released.store(0, Ordering::SeqCst);
broken.store(true, Ordering::SeqCst);
{
let _connection_1 = pool.get().await.unwrap();
let _connection_2 = pool.get().await.unwrap();
assert_eq!(num_conn_released.load(Ordering::SeqCst), 0);
}
assert_eq!(num_conn_released.load(Ordering::SeqCst), 2);
}
}