use std::collections::VecDeque;
use std::pin::Pin;
use std::sync::{Arc, Condvar, Mutex};
use std::time::{Duration, Instant};
use deadpool::managed::{self, Manager, Metrics, RecycleError, RecycleResult, Timeouts};
use deadpool::Runtime;
use tokio::sync::Mutex as AsyncMutex;
use crate::async_connection::AsyncConnection;
use crate::connection::Connection;
use crate::error::{Error, Result};
use crate::CreateMode;
pub type HookFuture<'a> = Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + 'a>>;
pub type AfterConnectHook = Arc<dyn Fn(&AsyncConnection) -> HookFuture<'_> + Send + Sync + 'static>;
pub type BeforeAcquireHook =
Arc<dyn Fn(&AsyncConnection) -> HookFuture<'_> + Send + Sync + 'static>;
pub type RecycleCheck = Arc<dyn Fn(&AsyncConnection) -> HookFuture<'_> + Send + Sync + 'static>;
#[derive(Clone, Default)]
pub enum RecycleStrategy {
#[default]
SelectOne,
Ping,
None,
Custom(RecycleCheck),
}
impl std::fmt::Debug for RecycleStrategy {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::SelectOne => f.write_str("SelectOne"),
Self::Ping => f.write_str("Ping"),
Self::None => f.write_str("None"),
Self::Custom(_) => f.write_str("Custom(<fn>)"),
}
}
}
#[derive(Clone)]
pub struct PoolConfig {
pub endpoint: String,
pub database: String,
pub create_mode: CreateMode,
pub user: Option<String>,
pub password: Option<String>,
pub max_size: usize,
pub health_check: bool,
pub recycle: RecycleStrategy,
pub wait_timeout: Option<Duration>,
pub create_timeout: Option<Duration>,
pub recycle_timeout: Option<Duration>,
pub max_lifetime: Option<Duration>,
pub idle_timeout: Option<Duration>,
pub min_idle: Option<u32>,
pub after_connect: Option<AfterConnectHook>,
pub before_acquire: Option<BeforeAcquireHook>,
}
impl std::fmt::Debug for PoolConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PoolConfig")
.field("endpoint", &self.endpoint)
.field("database", &self.database)
.field("create_mode", &self.create_mode)
.field("user", &self.user)
.field("password", &self.password.as_ref().map(|_| "<redacted>"))
.field("max_size", &self.max_size)
.field("health_check", &self.health_check)
.field("recycle", &self.recycle)
.field("wait_timeout", &self.wait_timeout)
.field("create_timeout", &self.create_timeout)
.field("recycle_timeout", &self.recycle_timeout)
.field("max_lifetime", &self.max_lifetime)
.field("idle_timeout", &self.idle_timeout)
.field("min_idle", &self.min_idle)
.field(
"after_connect",
&self.after_connect.as_ref().map(|_| "<fn>"),
)
.field(
"before_acquire",
&self.before_acquire.as_ref().map(|_| "<fn>"),
)
.finish()
}
}
impl PoolConfig {
pub fn new(endpoint: impl Into<String>, database: impl Into<String>) -> Self {
Self {
endpoint: endpoint.into(),
database: database.into(),
create_mode: CreateMode::DoNotCreate,
user: None,
password: None,
max_size: 16,
health_check: true,
recycle: RecycleStrategy::SelectOne,
wait_timeout: None,
create_timeout: None,
recycle_timeout: None,
max_lifetime: None,
idle_timeout: None,
min_idle: None,
after_connect: None,
before_acquire: None,
}
}
#[must_use]
pub fn create_mode(mut self, mode: CreateMode) -> Self {
self.create_mode = mode;
self
}
#[must_use]
pub fn auth(mut self, user: impl Into<String>, password: impl Into<String>) -> Self {
self.user = Some(user.into());
self.password = Some(password.into());
self
}
#[must_use]
pub fn max_size(mut self, size: usize) -> Self {
self.max_size = size;
self
}
#[must_use]
pub fn health_check(mut self, enabled: bool) -> Self {
self.health_check = enabled;
self.recycle = if enabled {
RecycleStrategy::SelectOne
} else {
RecycleStrategy::None
};
self
}
#[must_use]
pub fn recycle(mut self, strategy: RecycleStrategy) -> Self {
self.health_check = !matches!(strategy, RecycleStrategy::None);
self.recycle = strategy;
self
}
#[must_use]
pub fn wait_timeout(mut self, timeout: Option<Duration>) -> Self {
self.wait_timeout = timeout;
self
}
#[must_use]
pub fn create_timeout(mut self, timeout: Option<Duration>) -> Self {
self.create_timeout = timeout;
self
}
#[must_use]
pub fn recycle_timeout(mut self, timeout: Option<Duration>) -> Self {
self.recycle_timeout = timeout;
self
}
#[must_use]
pub fn max_lifetime(mut self, lifetime: Option<Duration>) -> Self {
self.max_lifetime = lifetime;
self
}
#[must_use]
pub fn idle_timeout(mut self, timeout: Option<Duration>) -> Self {
self.idle_timeout = timeout;
self
}
#[must_use]
pub fn min_idle(mut self, min_idle: Option<u32>) -> Self {
self.min_idle = min_idle;
self
}
#[must_use]
pub fn after_connect<F>(mut self, hook: F) -> Self
where
F: Fn(&AsyncConnection) -> HookFuture<'_> + Send + Sync + 'static,
{
self.after_connect = Some(Arc::new(hook));
self
}
#[must_use]
pub fn before_acquire<F>(mut self, hook: F) -> Self
where
F: Fn(&AsyncConnection) -> HookFuture<'_> + Send + Sync + 'static,
{
self.before_acquire = Some(Arc::new(hook));
self
}
fn has_timeout(&self) -> bool {
self.wait_timeout.is_some()
|| self.create_timeout.is_some()
|| self.recycle_timeout.is_some()
}
}
#[derive(Debug)]
pub struct ConnectionManager {
config: Arc<PoolConfig>,
init_lock: Arc<AsyncMutex<bool>>,
}
impl ConnectionManager {
#[must_use]
pub fn new(config: PoolConfig) -> Self {
Self {
config: Arc::new(config),
init_lock: Arc::new(AsyncMutex::new(false)),
}
}
async fn open(&self, mode: CreateMode) -> Result<AsyncConnection> {
if let (Some(user), Some(password)) = (&self.config.user, &self.config.password) {
AsyncConnection::connect_with_auth(
&self.config.endpoint,
&self.config.database,
mode,
user,
password,
)
.await
} else {
AsyncConnection::connect(&self.config.endpoint, &self.config.database, mode).await
}
}
}
impl Manager for ConnectionManager {
type Type = AsyncConnection;
type Error = Error;
async fn create(&self) -> Result<AsyncConnection> {
let conn = {
let initialized = self.init_lock.lock().await;
if *initialized {
drop(initialized);
self.open(CreateMode::DoNotCreate).await?
} else {
drop(initialized);
let mut initialized = self.init_lock.lock().await;
if *initialized {
drop(initialized);
self.open(CreateMode::DoNotCreate).await?
} else {
let result = self.open(self.config.create_mode).await;
if result.is_ok() {
*initialized = true;
}
result?
}
}
};
if let Some(hook) = self.config.after_connect.as_ref() {
hook(&conn).await?;
}
Ok(conn)
}
async fn recycle(
&self,
conn: &mut AsyncConnection,
metrics: &Metrics,
) -> RecycleResult<Self::Error> {
if let Some(max_lifetime) = self.config.max_lifetime {
if metrics.age() >= max_lifetime {
return Err(RecycleError::message("connection exceeded max_lifetime"));
}
}
if let Some(idle_timeout) = self.config.idle_timeout {
if metrics.last_used() >= idle_timeout {
return Err(RecycleError::message("connection exceeded idle_timeout"));
}
}
match &self.config.recycle {
RecycleStrategy::SelectOne => {
conn.execute_command("SELECT 1")
.await
.map_err(RecycleError::Backend)?;
}
RecycleStrategy::Ping => {
conn.ping().await.map_err(RecycleError::Backend)?;
}
RecycleStrategy::None => {}
RecycleStrategy::Custom(check) => {
check(conn).await.map_err(RecycleError::Backend)?;
}
}
if let Some(hook) = self.config.before_acquire.as_ref() {
hook(conn).await.map_err(RecycleError::Backend)?;
}
Ok(())
}
}
pub type Pool = managed::Pool<ConnectionManager>;
pub type PooledConnection = managed::Object<ConnectionManager>;
pub fn create_pool(config: PoolConfig) -> Result<Pool> {
let max_size = config.max_size;
let timeouts = Timeouts {
wait: config.wait_timeout,
create: config.create_timeout,
recycle: config.recycle_timeout,
};
let needs_runtime = config.has_timeout();
let manager = ConnectionManager::new(config);
let mut builder = Pool::builder(manager).max_size(max_size).timeouts(timeouts);
if needs_runtime {
builder = builder.runtime(Runtime::Tokio1);
}
builder
.build()
.map_err(|e| Error::config(format!("Failed to create pool: {e}")))
}
pub type SyncRecycleCheck = Arc<dyn Fn(&Connection) -> Result<()> + Send + Sync + 'static>;
#[derive(Clone, Default)]
pub enum SyncRecycleStrategy {
#[default]
SelectOne,
Ping,
None,
Custom(SyncRecycleCheck),
}
impl std::fmt::Debug for SyncRecycleStrategy {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::SelectOne => f.write_str("SelectOne"),
Self::Ping => f.write_str("Ping"),
Self::None => f.write_str("None"),
Self::Custom(_) => f.write_str("Custom(<fn>)"),
}
}
}
#[derive(Clone)]
pub struct SyncPoolConfig {
pub endpoint: String,
pub database: String,
pub create_mode: CreateMode,
pub user: Option<String>,
pub password: Option<String>,
pub max_size: usize,
pub recycle: SyncRecycleStrategy,
pub wait_timeout: Option<Duration>,
pub max_lifetime: Option<Duration>,
pub idle_timeout: Option<Duration>,
pub min_idle: Option<u32>,
}
impl std::fmt::Debug for SyncPoolConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SyncPoolConfig")
.field("endpoint", &self.endpoint)
.field("database", &self.database)
.field("create_mode", &self.create_mode)
.field("user", &self.user)
.field("password", &self.password.as_ref().map(|_| "<redacted>"))
.field("max_size", &self.max_size)
.field("recycle", &self.recycle)
.field("wait_timeout", &self.wait_timeout)
.field("max_lifetime", &self.max_lifetime)
.field("idle_timeout", &self.idle_timeout)
.field("min_idle", &self.min_idle)
.finish()
}
}
impl SyncPoolConfig {
pub fn new(endpoint: impl Into<String>, database: impl Into<String>) -> Self {
Self {
endpoint: endpoint.into(),
database: database.into(),
create_mode: CreateMode::DoNotCreate,
user: None,
password: None,
max_size: 16,
recycle: SyncRecycleStrategy::SelectOne,
wait_timeout: None,
max_lifetime: None,
idle_timeout: None,
min_idle: None,
}
}
#[must_use]
pub fn create_mode(mut self, mode: CreateMode) -> Self {
self.create_mode = mode;
self
}
#[must_use]
pub fn auth(mut self, user: impl Into<String>, password: impl Into<String>) -> Self {
self.user = Some(user.into());
self.password = Some(password.into());
self
}
#[must_use]
pub fn max_size(mut self, size: usize) -> Self {
self.max_size = size;
self
}
#[must_use]
pub fn recycle(mut self, strategy: SyncRecycleStrategy) -> Self {
self.recycle = strategy;
self
}
#[must_use]
pub fn wait_timeout(mut self, timeout: Option<Duration>) -> Self {
self.wait_timeout = timeout;
self
}
#[must_use]
pub fn max_lifetime(mut self, lifetime: Option<Duration>) -> Self {
self.max_lifetime = lifetime;
self
}
#[must_use]
pub fn idle_timeout(mut self, timeout: Option<Duration>) -> Self {
self.idle_timeout = timeout;
self
}
#[must_use]
pub fn min_idle(mut self, min_idle: Option<u32>) -> Self {
self.min_idle = min_idle;
self
}
#[must_use]
pub fn build(self) -> ConnectionPool {
ConnectionPool {
inner: Arc::new(SyncPoolInner {
config: self,
state: Mutex::new(SyncPoolState {
idle: VecDeque::new(),
size: 0,
initialized: false,
init_in_progress: false,
}),
available: Condvar::new(),
}),
}
}
}
struct IdleConn {
conn: Connection,
created: Instant,
last_used: Instant,
}
struct SyncPoolState {
idle: VecDeque<IdleConn>,
size: usize,
initialized: bool,
init_in_progress: bool,
}
struct SyncPoolInner {
config: SyncPoolConfig,
state: Mutex<SyncPoolState>,
available: Condvar,
}
impl SyncPoolInner {
fn open(&self, first: bool) -> Result<Connection> {
let mode = if first {
self.config.create_mode
} else {
CreateMode::DoNotCreate
};
if let (Some(user), Some(password)) = (&self.config.user, &self.config.password) {
Connection::connect_with_auth(
&self.config.endpoint,
&self.config.database,
mode,
user,
password,
)
} else {
Connection::connect(&self.config.endpoint, &self.config.database, mode)
}
}
fn should_evict(&self, idle: &IdleConn, live_size: usize) -> bool {
if let Some(max_lifetime) = self.config.max_lifetime {
if idle.created.elapsed() >= max_lifetime {
return true;
}
}
if let Some(idle_timeout) = self.config.idle_timeout {
let min_idle = self.config.min_idle.unwrap_or(0) as usize;
if idle.last_used.elapsed() >= idle_timeout && live_size > min_idle {
return true;
}
}
false
}
fn recycle(&self, conn: &Connection) -> Result<()> {
if !conn.is_alive() {
return Err(Error::connection("pooled connection is no longer alive"));
}
match &self.config.recycle {
SyncRecycleStrategy::SelectOne => {
conn.execute_command("SELECT 1")?;
}
SyncRecycleStrategy::Ping => {
conn.ping()?;
}
SyncRecycleStrategy::None => {}
SyncRecycleStrategy::Custom(check) => {
check(conn)?;
}
}
Ok(())
}
fn checkin(&self, conn: Connection, created: Instant) {
{
let mut state = self.state.lock().expect("pool mutex poisoned");
state.idle.push_back(IdleConn {
conn,
created,
last_used: Instant::now(),
});
}
self.available.notify_one();
}
fn discard(&self) {
{
let mut state = self.state.lock().expect("pool mutex poisoned");
state.size = state.size.saturating_sub(1);
}
self.available.notify_one();
}
}
#[derive(Clone)]
pub struct ConnectionPool {
inner: Arc<SyncPoolInner>,
}
impl std::fmt::Debug for ConnectionPool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ConnectionPool")
.field("config", &self.inner.config)
.field("status", &self.status())
.finish()
}
}
impl ConnectionPool {
pub fn get(&self) -> Result<SyncPooledConnection> {
self.get_timeout(self.inner.config.wait_timeout)
}
pub fn get_timeout(&self, timeout: Option<Duration>) -> Result<SyncPooledConnection> {
let deadline = timeout.map(|t| Instant::now() + t);
loop {
enum Action {
Reuse(IdleConn),
Create { first: bool },
Wait,
}
let action = {
let mut state = self.inner.state.lock().expect("pool mutex poisoned");
if let Some(idle) = state.idle.pop_back() {
Action::Reuse(idle)
} else if !state.initialized {
if state.init_in_progress {
Action::Wait
} else {
state.init_in_progress = true;
state.size += 1;
Action::Create { first: true }
}
} else if state.size < self.inner.config.max_size {
state.size += 1;
Action::Create { first: false }
} else {
Action::Wait
}
};
match action {
Action::Reuse(idle) => {
let live_size = {
let state = self.inner.state.lock().expect("pool mutex poisoned");
state.size
};
if self.inner.should_evict(&idle, live_size) {
self.inner.discard();
continue;
}
if self.inner.recycle(&idle.conn).is_ok() {
return Ok(SyncPooledConnection {
pool: Arc::clone(&self.inner),
conn: Some(idle.conn),
created: idle.created,
});
}
self.inner.discard();
}
Action::Create { first } => match self.inner.open(first) {
Ok(conn) => {
if first {
let mut state = self.inner.state.lock().expect("pool mutex poisoned");
state.initialized = true;
state.init_in_progress = false;
drop(state);
self.inner.available.notify_all();
}
return Ok(SyncPooledConnection {
pool: Arc::clone(&self.inner),
conn: Some(conn),
created: Instant::now(),
});
}
Err(e) => {
{
let mut state = self.inner.state.lock().expect("pool mutex poisoned");
state.size = state.size.saturating_sub(1);
if first {
state.init_in_progress = false;
}
}
self.inner.available.notify_all();
return Err(e);
}
},
Action::Wait => {
let state = self.inner.state.lock().expect("pool mutex poisoned");
if !state.idle.is_empty()
|| (state.initialized && state.size < self.inner.config.max_size)
|| !state.initialized && !state.init_in_progress
{
continue;
}
match deadline {
Some(dl) => {
let now = Instant::now();
if now >= dl {
return Err(Error::timeout(
"timed out waiting for an available pool connection",
));
}
let (_guard, res) = self
.inner
.available
.wait_timeout(state, dl - now)
.expect("pool mutex poisoned");
if res.timed_out() {
return Err(Error::timeout(
"timed out waiting for an available pool connection",
));
}
}
None => {
let _guard = self
.inner
.available
.wait(state)
.expect("pool mutex poisoned");
}
}
}
}
}
}
#[must_use]
pub fn status(&self) -> PoolStatus {
let state = self.inner.state.lock().expect("pool mutex poisoned");
PoolStatus {
idle: state.idle.len(),
size: state.size,
max_size: self.inner.config.max_size,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PoolStatus {
pub idle: usize,
pub size: usize,
pub max_size: usize,
}
pub struct SyncPooledConnection {
pool: Arc<SyncPoolInner>,
conn: Option<Connection>,
created: Instant,
}
impl SyncPooledConnection {
#[must_use]
pub fn take(mut self) -> Connection {
let conn = self.conn.take().expect("connection already taken");
self.pool.discard();
conn
}
}
impl std::ops::Deref for SyncPooledConnection {
type Target = Connection;
fn deref(&self) -> &Self::Target {
self.conn.as_ref().expect("connection already taken")
}
}
impl std::fmt::Debug for SyncPooledConnection {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SyncPooledConnection")
.field("checked_out", &self.conn.is_some())
.finish_non_exhaustive()
}
}
impl Drop for SyncPooledConnection {
fn drop(&mut self) {
if let Some(conn) = self.conn.take() {
if conn.is_alive() {
self.pool.checkin(conn, self.created);
} else {
drop(conn);
self.pool.discard();
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pool_config_builder() {
let config = PoolConfig::new("localhost:7483", "test.hyper")
.create_mode(CreateMode::CreateIfNotExists)
.auth("user", "pass")
.max_size(32);
assert_eq!(config.endpoint, "localhost:7483");
assert_eq!(config.database, "test.hyper");
assert_eq!(config.create_mode, CreateMode::CreateIfNotExists);
assert_eq!(config.user, Some("user".to_string()));
assert_eq!(config.password, Some("pass".to_string()));
assert_eq!(config.max_size, 32);
}
#[test]
fn test_pool_config_defaults_are_additive() {
let config = PoolConfig::new("localhost:7483", "test.hyper");
assert!(config.health_check);
assert!(matches!(config.recycle, RecycleStrategy::SelectOne));
assert_eq!(config.wait_timeout, None);
assert_eq!(config.create_timeout, None);
assert_eq!(config.recycle_timeout, None);
assert_eq!(config.max_lifetime, None);
assert_eq!(config.idle_timeout, None);
assert_eq!(config.min_idle, None);
assert!(!config.has_timeout());
}
#[test]
fn test_health_check_and_recycle_stay_in_sync() {
let off = PoolConfig::new("e", "d").health_check(false);
assert!(!off.health_check);
assert!(matches!(off.recycle, RecycleStrategy::None));
let on = PoolConfig::new("e", "d").health_check(true);
assert!(on.health_check);
assert!(matches!(on.recycle, RecycleStrategy::SelectOne));
let via_recycle = PoolConfig::new("e", "d").recycle(RecycleStrategy::None);
assert!(!via_recycle.health_check);
let via_ping = PoolConfig::new("e", "d").recycle(RecycleStrategy::Ping);
assert!(via_ping.health_check);
assert!(matches!(via_ping.recycle, RecycleStrategy::Ping));
}
#[test]
fn test_pool_config_timeout_builders() {
let config = PoolConfig::new("e", "d")
.wait_timeout(Some(Duration::from_secs(1)))
.create_timeout(Some(Duration::from_secs(2)))
.recycle_timeout(Some(Duration::from_secs(3)))
.max_lifetime(Some(Duration::from_secs(60)))
.idle_timeout(Some(Duration::from_secs(30)))
.min_idle(Some(2));
assert_eq!(config.wait_timeout, Some(Duration::from_secs(1)));
assert_eq!(config.create_timeout, Some(Duration::from_secs(2)));
assert_eq!(config.recycle_timeout, Some(Duration::from_secs(3)));
assert_eq!(config.max_lifetime, Some(Duration::from_secs(60)));
assert_eq!(config.idle_timeout, Some(Duration::from_secs(30)));
assert_eq!(config.min_idle, Some(2));
assert!(config.has_timeout());
}
#[test]
fn test_sync_pool_config_builder_and_defaults() {
let config = SyncPoolConfig::new("localhost:7483", "test.hyper");
assert_eq!(config.max_size, 16);
assert!(matches!(config.recycle, SyncRecycleStrategy::SelectOne));
assert_eq!(config.wait_timeout, None);
assert_eq!(config.max_lifetime, None);
assert_eq!(config.idle_timeout, None);
assert_eq!(config.min_idle, None);
let tuned = SyncPoolConfig::new("e", "d")
.create_mode(CreateMode::CreateIfNotExists)
.auth("u", "p")
.max_size(4)
.recycle(SyncRecycleStrategy::Ping)
.wait_timeout(Some(Duration::from_millis(500)))
.max_lifetime(Some(Duration::from_secs(10)))
.idle_timeout(Some(Duration::from_secs(5)))
.min_idle(Some(1));
assert_eq!(tuned.max_size, 4);
assert!(matches!(tuned.recycle, SyncRecycleStrategy::Ping));
assert_eq!(tuned.user, Some("u".to_string()));
assert_eq!(tuned.wait_timeout, Some(Duration::from_millis(500)));
}
#[test]
fn test_debug_redacts_password() {
let dbg = format!("{:?}", PoolConfig::new("e", "d").auth("u", "secret"));
assert!(dbg.contains("<redacted>"));
assert!(!dbg.contains("secret"));
let sync_dbg = format!("{:?}", SyncPoolConfig::new("e", "d").auth("u", "secret"));
assert!(sync_dbg.contains("<redacted>"));
assert!(!sync_dbg.contains("secret"));
}
}