use async_trait::async_trait;
use crossbeam_queue::ArrayQueue;
#[cfg(feature = "circuit-breaker")]
use parking_lot::Mutex as PlMutex;
#[cfg(feature = "rate-limit")]
use parking_lot::RwLock as PlRwLock;
use std::future::Future;
use std::ops::{Deref, DerefMut};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Notify;
#[cfg(feature = "circuit-breaker")]
use crate::circuit_breaker::{CircuitBreaker, CircuitState, DefaultCircuitBreaker};
#[cfg(feature = "rate-limit")]
use crate::rate_limiter::RateLimiter;
use sz_orm_model::PoolError;
pub use sz_orm_model::QueryRows;
pub type QueryStreamItem =
Result<std::collections::HashMap<String, sz_orm_model::Value>, sz_orm_model::DbError>;
pub trait Connection: Send + Sync {
fn execute<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>>;
fn query<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<QueryRows, sz_orm_model::DbError>> + Send + 'a>>;
fn begin_transaction<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>>;
fn commit<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>>;
fn rollback<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>>;
fn is_connected(&self) -> bool;
fn in_transaction(&self) -> bool {
false
}
fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>>;
fn close<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>>;
fn execute_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [sz_orm_model::Value],
) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>> {
let _ = (sql, params);
Box::pin(async move {
Err(sz_orm_model::DbError::Internal(
"execute_with_params not implemented for this adapter".to_string(),
))
})
}
fn query_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [sz_orm_model::Value],
) -> Pin<Box<dyn Future<Output = Result<QueryRows, sz_orm_model::DbError>> + Send + 'a>> {
let _ = (sql, params);
Box::pin(async move {
Err(sz_orm_model::DbError::Internal(
"query_with_params not implemented for this adapter".to_string(),
))
})
}
fn query_values<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<
Box<
dyn Future<Output = Result<sz_orm_model::QueryValues, sz_orm_model::DbError>>
+ Send
+ 'a,
>,
> {
let _ = sql;
Box::pin(async move {
Err(sz_orm_model::DbError::Internal(
"query_values not implemented for this adapter".to_string(),
))
})
}
fn query_values_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [sz_orm_model::Value],
) -> Pin<
Box<
dyn Future<Output = Result<sz_orm_model::QueryValues, sz_orm_model::DbError>>
+ Send
+ 'a,
>,
> {
let _ = (sql, params);
Box::pin(async move {
Err(sz_orm_model::DbError::Internal(
"query_values_with_params not implemented for this adapter".to_string(),
))
})
}
fn query_stream<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn futures::Stream<Item = QueryStreamItem> + Send + 'a>> {
let _ = sql;
Box::pin(futures::stream::empty())
}
fn execute_batch<'a>(
&'a mut self,
sqls: &'a [String],
) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async move {
let mut total = 0u64;
for sql in sqls {
total += self.execute(sql).await?;
}
Ok(total)
})
}
fn execute_batch_params<'a>(
&'a mut self,
sql: &'a str,
params_batch: &'a [Vec<sz_orm_model::Value>],
) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async move {
let mut total = 0u64;
for params in params_batch {
total += self.execute_with_params(sql, params).await?;
}
Ok(total)
})
}
}
pub struct PooledConnection {
conn: Box<dyn Connection>,
created_at: Instant,
last_used_at: Instant,
pool: Option<Pool>,
}
impl PooledConnection {
fn new(conn: Box<dyn Connection>, pool: Pool) -> Self {
let now = Instant::now();
Self {
conn,
created_at: now,
last_used_at: now,
pool: Some(pool),
}
}
fn is_expired(&self, max_lifetime: Duration) -> bool {
self.created_at.elapsed() >= max_lifetime
}
fn is_idle_too_long(&self, idle_timeout: Duration) -> bool {
self.last_used_at.elapsed() >= idle_timeout
}
pub fn created_at(&self) -> Instant {
self.created_at
}
pub fn into_inner(mut self) -> Box<dyn Connection> {
self.pool = None; std::mem::replace(&mut self.conn, Box::new(ClosedConnection))
}
}
impl Drop for PooledConnection {
fn drop(&mut self) {
if let Some(pool) = self.pool.take() {
let conn = std::mem::replace(&mut self.conn, Box::new(ClosedConnection));
let pooled = PooledConnection {
conn,
created_at: self.created_at,
last_used_at: self.last_used_at,
pool: None,
};
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
pool.release(pooled).await;
});
} else {
drop(pooled);
pool.total_count.fetch_sub(1, Ordering::SeqCst);
}
}
}
}
struct ClosedConnection;
impl Connection for ClosedConnection {
fn execute<'a>(
&'a mut self,
_sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async {
Err(sz_orm_model::DbError::ConnectionError(
"connection already returned to pool".to_string(),
))
})
}
fn query<'a>(
&'a mut self,
_sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<QueryRows, sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async {
Err(sz_orm_model::DbError::ConnectionError(
"connection already returned to pool".to_string(),
))
})
}
fn begin_transaction<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async {
Err(sz_orm_model::DbError::ConnectionError(
"connection already returned to pool".to_string(),
))
})
}
fn commit<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async { Ok(()) })
}
fn rollback<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async { Ok(()) })
}
fn is_connected(&self) -> bool {
false
}
fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
Box::pin(async { false })
}
fn close<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async { Ok(()) })
}
}
impl Deref for PooledConnection {
type Target = dyn Connection;
fn deref(&self) -> &Self::Target {
self.conn.as_ref()
}
}
impl DerefMut for PooledConnection {
fn deref_mut(&mut self) -> &mut Self::Target {
self.conn.as_mut()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TlsVersion {
#[default]
Tls12,
Tls13,
}
#[derive(Debug, Clone, Default)]
pub struct TlsConfig {
pub enabled: bool,
pub ca_cert_path: Option<String>,
pub client_cert_path: Option<String>,
pub client_key_path: Option<String>,
pub min_version: TlsVersion,
}
#[derive(Debug, Clone)]
pub enum PoolEvent {
ConnectionCreated,
ConnectionClosed,
ConnectionAcquired,
ConnectionReleased,
AcquireTimeout,
}
pub type PoolEventCallback = Arc<dyn Fn(PoolEvent) + Send + Sync>;
pub struct PoolConfig {
pub max_size: u32,
pub min_idle: u32,
pub acquire_timeout: Duration,
pub idle_timeout: Duration,
pub max_lifetime: Duration,
pub connection_timeout: Duration,
pub tls: Option<TlsConfig>,
pub query_timeout: Option<Duration>,
pub max_rows: Option<usize>,
pub memory_limit: Option<usize>,
pub on_event: Option<PoolEventCallback>,
}
impl Default for PoolConfig {
fn default() -> Self {
Self {
max_size: 100,
min_idle: 0,
acquire_timeout: Duration::from_secs(30),
idle_timeout: Duration::from_secs(600),
max_lifetime: Duration::from_secs(1800),
connection_timeout: Duration::from_secs(10),
tls: None,
query_timeout: Some(Duration::from_secs(30)),
max_rows: None,
memory_limit: None,
on_event: None,
}
}
}
impl Clone for PoolConfig {
fn clone(&self) -> Self {
Self {
max_size: self.max_size,
min_idle: self.min_idle,
acquire_timeout: self.acquire_timeout,
idle_timeout: self.idle_timeout,
max_lifetime: self.max_lifetime,
connection_timeout: self.connection_timeout,
tls: self.tls.clone(),
query_timeout: self.query_timeout,
max_rows: self.max_rows,
memory_limit: self.memory_limit,
on_event: self.on_event.clone(),
}
}
}
impl PoolConfig {
pub fn validate(&self) -> Result<(), PoolError> {
if self.max_size == 0 {
return Err(PoolError::InvalidConfig("max_size cannot be 0".to_string()));
}
if self.min_idle > self.max_size {
return Err(PoolError::InvalidConfig(
"min_idle cannot exceed max_size".to_string(),
));
}
Ok(())
}
}
pub struct PoolStatus {
pub idle: u32,
pub active: u32,
pub max: u32,
pub min: u32,
pub waiters: u32,
}
impl std::fmt::Debug for PoolStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PoolStatus")
.field("idle", &self.idle)
.field("active", &self.active)
.field("max", &self.max)
.field("min", &self.min)
.field("waiters", &self.waiters)
.finish()
}
}
pub struct PoolConfigBuilder {
config: PoolConfig,
}
impl PoolConfigBuilder {
pub fn new() -> Self {
Self {
config: PoolConfig::default(),
}
}
pub fn max_size(mut self, size: u32) -> Self {
self.config.max_size = size;
self
}
pub fn min_idle(mut self, count: u32) -> Self {
self.config.min_idle = count;
self
}
pub fn acquire_timeout(mut self, timeout_secs: u64) -> Self {
self.config.acquire_timeout = Duration::from_secs(timeout_secs);
self
}
pub fn idle_timeout(mut self, timeout_secs: u64) -> Self {
self.config.idle_timeout = Duration::from_secs(timeout_secs);
self
}
pub fn max_lifetime(mut self, lifetime_secs: u64) -> Self {
self.config.max_lifetime = Duration::from_secs(lifetime_secs);
self
}
pub fn tls(mut self, tls: TlsConfig) -> Self {
self.config.tls = Some(tls);
self
}
pub fn query_timeout(mut self, timeout: Duration) -> Self {
self.config.query_timeout = Some(timeout);
self
}
pub fn max_rows(mut self, max_rows: usize) -> Self {
self.config.max_rows = Some(max_rows);
self
}
pub fn memory_limit(mut self, memory_limit: usize) -> Self {
self.config.memory_limit = Some(memory_limit);
self
}
pub fn on_event(mut self, callback: PoolEventCallback) -> Self {
self.config.on_event = Some(callback);
self
}
pub fn build(self) -> Result<PoolConfig, PoolError> {
self.config.validate()?;
Ok(self.config)
}
}
impl Default for PoolConfigBuilder {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
pub trait ConnectionFactory: Send + Sync {
async fn create(&self) -> Result<Box<dyn Connection>, sz_orm_model::DbError>;
}
pub struct Pool {
config: PoolConfig,
factory: Arc<dyn ConnectionFactory>,
idle: Arc<ArrayQueue<PooledConnection>>,
total_count: Arc<AtomicU32>,
closed: Arc<AtomicBool>,
notify: Arc<Notify>,
waiters_count: Arc<AtomicU32>,
dynamic_max_size: Arc<AtomicU32>,
#[cfg(feature = "circuit-breaker")]
circuit_breaker: Arc<PlMutex<DefaultCircuitBreaker>>,
#[cfg(feature = "rate-limit")]
rate_limiter: Arc<PlRwLock<Option<Arc<dyn RateLimiter>>>>,
#[cfg(feature = "rate-limit")]
rate_limit_key: String,
}
impl Clone for Pool {
fn clone(&self) -> Self {
Self {
config: self.config.clone(),
factory: self.factory.clone(),
idle: self.idle.clone(),
total_count: self.total_count.clone(),
closed: self.closed.clone(),
notify: Arc::clone(&self.notify),
waiters_count: self.waiters_count.clone(),
dynamic_max_size: self.dynamic_max_size.clone(),
#[cfg(feature = "circuit-breaker")]
circuit_breaker: Arc::clone(&self.circuit_breaker),
#[cfg(feature = "rate-limit")]
rate_limiter: Arc::clone(&self.rate_limiter),
#[cfg(feature = "rate-limit")]
rate_limit_key: self.rate_limit_key.clone(),
}
}
}
impl Pool {
pub fn new(config: PoolConfig, factory: Arc<dyn ConnectionFactory>) -> Result<Self, PoolError> {
config.validate()?;
let max_size = config.max_size as usize;
let dynamic_max = config.max_size;
Ok(Self {
config,
factory,
idle: Arc::new(ArrayQueue::new(max_size)),
total_count: Arc::new(AtomicU32::new(0)),
closed: Arc::new(AtomicBool::new(false)),
notify: Arc::new(Notify::new()),
waiters_count: Arc::new(AtomicU32::new(0)),
dynamic_max_size: Arc::new(AtomicU32::new(dynamic_max)),
#[cfg(feature = "circuit-breaker")]
circuit_breaker: Arc::new(PlMutex::new(DefaultCircuitBreaker::new(
5,
std::time::Duration::from_secs(30),
))),
#[cfg(feature = "rate-limit")]
rate_limiter: Arc::new(PlRwLock::new(None)),
#[cfg(feature = "rate-limit")]
rate_limit_key: "pool".to_string(),
})
}
pub fn config(&self) -> &PoolConfig {
&self.config
}
#[cfg(feature = "circuit-breaker")]
pub fn configure_circuit_breaker(
&self,
failure_threshold: usize,
reset_timeout: std::time::Duration,
) {
let new_cb = DefaultCircuitBreaker::new(failure_threshold, reset_timeout);
let mut guard = self.circuit_breaker.lock();
*guard = new_cb;
}
#[cfg(feature = "circuit-breaker")]
pub fn reset_circuit_breaker(&self) -> bool {
let mut guard = self.circuit_breaker.lock();
guard.reset()
}
#[cfg(feature = "circuit-breaker")]
pub fn circuit_state(&self) -> CircuitState {
let guard = self.circuit_breaker.lock();
guard.state()
}
#[cfg(feature = "rate-limit")]
pub fn set_rate_limiter(&self, limiter: Option<Arc<dyn RateLimiter>>) {
let mut guard = self.rate_limiter.write();
*guard = limiter;
}
#[cfg(feature = "rate-limit")]
pub fn with_rate_limit_key(mut self, key: impl Into<String>) -> Self {
self.rate_limit_key = key.into();
self
}
fn emit_event(&self, event: PoolEvent) {
if let Some(ref callback) = self.config.on_event {
callback(event);
}
}
#[tracing::instrument(skip(self), fields(max_size = self.config.max_size, acquire_timeout = ?self.config.acquire_timeout))]
pub async fn acquire(&self) -> Result<PooledConnection, PoolError> {
if self.closed.load(Ordering::Acquire) {
return Err(PoolError::Closed);
}
#[cfg(feature = "circuit-breaker")]
{
let mut guard = self.circuit_breaker.lock();
if !guard.can_execute() {
return Err(PoolError::CircuitOpen);
}
}
#[cfg(feature = "rate-limit")]
{
let guard = self.rate_limiter.read();
if let Some(ref limiter) = *guard {
match limiter.try_acquire(&self.rate_limit_key) {
Ok(result) if !result.allowed => {
return Err(PoolError::RateLimited {
remaining: result.remaining,
reset_at: result.reset_at,
});
}
Ok(_) => {} Err(_) => {
}
}
}
}
let deadline = Instant::now() + self.config.acquire_timeout;
let mut backoff = Duration::from_millis(1);
const MAX_BACKOFF: Duration = Duration::from_millis(100);
loop {
let mut to_close: Vec<PooledConnection> = Vec::new();
let acquired: Option<PooledConnection> = {
let mut found: Option<PooledConnection> = None;
while let Some(pooled) = self.idle.pop() {
if pooled.is_expired(self.config.max_lifetime) {
to_close.push(pooled);
continue;
}
if pooled.is_idle_too_long(self.config.idle_timeout) {
to_close.push(pooled);
continue;
}
if !pooled.conn.is_connected() {
to_close.push(pooled);
continue;
}
found = Some(pooled);
break;
}
found
};
for mut pooled in to_close {
let _ = pooled.conn.close().await;
self.total_count.fetch_sub(1, Ordering::SeqCst);
}
if let Some(mut pooled) = acquired {
pooled.pool = Some(self.clone());
return Ok(pooled);
}
let current_max = self.dynamic_max_size.load(Ordering::Acquire);
let created = loop {
let current = self.total_count.load(Ordering::Acquire);
if current >= current_max {
break None; }
match self.total_count.compare_exchange(
current,
current + 1,
Ordering::SeqCst,
Ordering::Acquire,
) {
Ok(_) => break Some(()), Err(_) => continue, }
};
if created.is_some() {
match tokio::time::timeout(self.config.connection_timeout, self.factory.create())
.await
{
Ok(Ok(conn)) => {
#[cfg(feature = "circuit-breaker")]
{
self.circuit_breaker.lock().record_success();
}
self.emit_event(PoolEvent::ConnectionCreated);
self.emit_event(PoolEvent::ConnectionAcquired);
return Ok(PooledConnection::new(conn, self.clone()));
}
Ok(Err(e)) => {
self.total_count.fetch_sub(1, Ordering::SeqCst);
#[cfg(feature = "circuit-breaker")]
{
self.circuit_breaker.lock().record_failure();
}
return Err(PoolError::ConnectionFailed(e.to_string()));
}
Err(_) => {
self.total_count.fetch_sub(1, Ordering::SeqCst);
#[cfg(feature = "circuit-breaker")]
{
self.circuit_breaker.lock().record_failure();
}
return Err(PoolError::Timeout);
}
}
}
let now = Instant::now();
if now >= deadline {
self.emit_event(PoolEvent::AcquireTimeout);
return Err(PoolError::Timeout);
}
self.waiters_count.fetch_add(1, Ordering::SeqCst);
let wait = std::cmp::min(backoff, deadline - now);
match tokio::time::timeout(wait, self.notify.notified()).await {
Ok(()) => {
backoff = Duration::from_millis(1);
}
Err(_) => {
backoff = std::cmp::min(backoff * 2, MAX_BACKOFF);
}
}
self.waiters_count.fetch_sub(1, Ordering::SeqCst);
}
}
#[tracing::instrument(skip(self, pooled))]
pub async fn release(&self, mut pooled: PooledConnection) {
pooled.pool = None;
if self.closed.load(Ordering::Acquire) {
let _ = pooled.conn.close().await;
self.total_count.fetch_sub(1, Ordering::SeqCst);
self.emit_event(PoolEvent::ConnectionClosed);
return;
}
if !pooled.conn.is_connected() {
let _ = pooled.conn.close().await;
self.total_count.fetch_sub(1, Ordering::SeqCst);
self.emit_event(PoolEvent::ConnectionClosed);
return;
}
if pooled.conn.in_transaction() {
let _ = pooled.conn.rollback().await;
}
pooled.last_used_at = Instant::now();
if let Err(mut rejected) = self.idle.push(pooled) {
let _ = rejected.conn.close().await;
self.total_count.fetch_sub(1, Ordering::SeqCst);
self.emit_event(PoolEvent::ConnectionClosed);
} else {
self.emit_event(PoolEvent::ConnectionReleased);
}
self.notify.notify_one();
}
pub async fn status(&self) -> PoolStatus {
let idle_count = self.idle.len() as u32;
let active = self.total_count.load(Ordering::Acquire);
let waiters = self.waiters_count.load(Ordering::Acquire);
PoolStatus {
idle: idle_count,
active,
max: self.dynamic_max_size.load(Ordering::Acquire),
min: self.config.min_idle,
waiters,
}
}
#[tracing::instrument(skip(self))]
pub async fn reap_idle(&self) {
let mut all: Vec<PooledConnection> = Vec::new();
while let Some(pooled) = self.idle.pop() {
all.push(pooled);
}
let mut to_close = Vec::new();
for pooled in all {
if pooled.is_idle_too_long(self.config.idle_timeout)
|| pooled.is_expired(self.config.max_lifetime)
{
to_close.push(pooled);
} else {
if let Err(mut rejected) = self.idle.push(pooled) {
let _ = rejected.conn.close().await;
self.total_count.fetch_sub(1, Ordering::SeqCst);
}
}
}
for mut pooled in to_close {
let _ = pooled.conn.close().await;
self.total_count.fetch_sub(1, Ordering::SeqCst);
}
}
pub async fn close_all(&self) {
self.closed.store(true, Ordering::Release);
let mut to_close: Vec<PooledConnection> = Vec::new();
while let Some(pooled) = self.idle.pop() {
to_close.push(pooled);
}
let closed_count: u32 = to_close.len() as u32;
for mut pooled in to_close {
let _ = pooled.conn.close().await;
}
self.total_count.fetch_sub(closed_count, Ordering::SeqCst);
}
pub async fn health_check(&self) -> u32 {
let mut to_check: Vec<PooledConnection> = Vec::new();
while let Some(pooled) = self.idle.pop() {
to_check.push(pooled);
}
let mut removed: u32 = 0;
let mut alive: Vec<PooledConnection> = Vec::with_capacity(to_check.len());
for mut pooled in to_check.drain(..) {
if !pooled.conn.is_connected() {
let _ = pooled.conn.close().await;
removed += 1;
continue;
}
let ping_timeout = self.config.connection_timeout / 2;
match tokio::time::timeout(ping_timeout, pooled.conn.ping()).await {
Ok(true) => alive.push(pooled),
Ok(false) => {
let _ = pooled.conn.close().await;
removed += 1;
}
Err(_) => {
let _ = pooled.conn.close().await;
removed += 1;
}
}
}
let alive_count: u32 = alive.len() as u32;
for pooled in alive {
if let Err(mut rejected) = self.idle.push(pooled) {
let _ = rejected.conn.close().await;
removed += 1;
}
}
if removed > 0 {
self.total_count.fetch_sub(removed, Ordering::SeqCst);
}
if alive_count > 0 {
self.notify.notify_one();
}
removed
}
pub async fn shutdown(&self) {
self.closed.store(true, Ordering::SeqCst);
self.notify.notify_waiters();
self.close_all().await;
let deadline = Instant::now() + Duration::from_secs(30);
while self.total_count.load(Ordering::SeqCst) > 0 {
if Instant::now() >= deadline {
break;
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
pub fn resize(&self, new_max: usize) {
self.set_max_size(new_max as u32);
}
pub fn set_max_size(&self, new_max: u32) {
self.dynamic_max_size.store(new_max, Ordering::SeqCst);
}
pub fn max_size(&self) -> u32 {
self.dynamic_max_size.load(Ordering::Acquire)
}
pub async fn warmup(&self, min_idle: usize) -> Result<(), PoolError> {
for _ in 0..min_idle {
let current_max = self.dynamic_max_size.load(Ordering::Acquire);
let current = self.total_count.load(Ordering::Acquire);
if current >= current_max {
break;
}
match self.total_count.compare_exchange(
current,
current + 1,
Ordering::SeqCst,
Ordering::Acquire,
) {
Ok(_) => {}
Err(_) => continue, }
match self.factory.create().await {
Ok(conn) => {
let now = Instant::now();
let pooled = PooledConnection {
conn,
created_at: now,
last_used_at: now,
pool: None,
};
if let Err(mut rejected) = self.idle.push(pooled) {
let _ = rejected.conn.close().await;
self.total_count.fetch_sub(1, Ordering::SeqCst);
}
self.emit_event(PoolEvent::ConnectionCreated);
}
Err(_) => {
self.total_count.fetch_sub(1, Ordering::SeqCst);
break;
}
}
}
Ok(())
}
pub async fn query_with_timeout(&self, sql: &str) -> Result<QueryRows, sz_orm_model::DbError> {
let timeout = self.config.query_timeout.unwrap_or(Duration::from_secs(30));
let mut conn = self
.acquire()
.await
.map_err(sz_orm_model::DbError::PoolError)?;
tokio::time::timeout(timeout, conn.query(sql))
.await
.map_err(|_| {
sz_orm_model::DbError::QueryError(format!("Query timeout after {:?}", timeout))
})?
}
}
#[cfg(test)]
mod tests {
use super::*;
struct MockConnection {
connected: bool,
}
impl MockConnection {
fn new() -> Self {
Self { connected: true }
}
}
impl Connection for MockConnection {
fn execute<'a>(
&'a mut self,
_sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<u64, sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async move { Ok(1) })
}
fn query<'a>(
&'a mut self,
_sql: &'a str,
) -> Pin<
Box<
dyn Future<
Output = Result<
Vec<std::collections::HashMap<String, sz_orm_model::Value>>,
sz_orm_model::DbError,
>,
> + Send
+ 'a,
>,
> {
Box::pin(async move { Ok(vec![]) })
}
fn begin_transaction<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async move { Ok(()) })
}
fn commit<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async move { Ok(()) })
}
fn rollback<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async move { Ok(()) })
}
fn is_connected(&self) -> bool {
self.connected
}
fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
Box::pin(async move { true })
}
fn close<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), sz_orm_model::DbError>> + Send + 'a>> {
Box::pin(async move {
self.connected = false;
Ok(())
})
}
}
struct MockConnectionFactory;
#[async_trait]
impl ConnectionFactory for MockConnectionFactory {
async fn create(&self) -> Result<Box<dyn Connection>, sz_orm_model::DbError> {
Ok(Box::new(MockConnection::new()))
}
}
#[tokio::test]
async fn test_pool_config_builder() -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new().max_size(50).min_idle(10).build()?;
assert_eq!(config.max_size, 50);
assert_eq!(config.min_idle, 10);
Ok(())
}
#[test]
fn test_pool_status_display() {
let status = PoolStatus {
idle: 5,
active: 10,
max: 100,
min: 5,
waiters: 0,
};
let display = format!("{:?}", status);
assert!(display.contains("idle"));
assert!(display.contains("active"));
}
#[test]
fn test_default_pool_config() {
let config = PoolConfig::default();
assert_eq!(config.max_size, 100);
assert_eq!(config.min_idle, 0);
assert_eq!(config.acquire_timeout.as_secs(), 30);
assert_eq!(config.idle_timeout.as_secs(), 600);
assert_eq!(config.max_lifetime.as_secs(), 1800);
}
#[tokio::test]
async fn test_pool_config_clone() {
let config = PoolConfig::default();
let cloned = config.clone();
assert_eq!(cloned.max_size, config.max_size);
assert_eq!(cloned.min_idle, config.min_idle);
}
#[test]
fn test_pool_config_builder_default() -> Result<(), Box<dyn std::error::Error>> {
let builder = PoolConfigBuilder::new();
let config = builder.build()?;
assert_eq!(config.max_size, 100);
Ok(())
}
#[test]
fn test_pool_config_validate() {
let result = PoolConfigBuilder::new().max_size(0).build();
assert!(result.is_err());
let result = PoolConfigBuilder::new().max_size(10).min_idle(20).build();
assert!(result.is_err());
}
#[tokio::test]
async fn test_pool_acquire_and_release() -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new().max_size(5).min_idle(1).build()?;
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory)?;
let conn = pool.acquire().await?;
let status = pool.status().await;
assert_eq!(status.active, 1);
assert_eq!(status.idle, 0);
pool.release(conn).await;
let status = pool.status().await;
assert_eq!(status.idle, 1);
let _conn2 = pool.acquire().await?;
let status = pool.status().await;
assert_eq!(status.idle, 0);
Ok(())
}
#[tokio::test]
async fn test_pool_status() -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new().max_size(10).min_idle(2).build()?;
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory)?;
let status = pool.status().await;
assert_eq!(status.max, 10);
assert_eq!(status.min, 2);
assert_eq!(status.active, 0);
Ok(())
}
#[tokio::test]
async fn test_pool_close_all() -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new().max_size(5).build()?;
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory)?;
let conn1 = pool.acquire().await?;
let conn2 = pool.acquire().await?;
pool.release(conn1).await;
pool.release(conn2).await;
pool.close_all().await;
let status = pool.status().await;
assert_eq!(status.idle, 0);
assert_eq!(status.active, 0);
Ok(())
}
#[tokio::test]
async fn test_pool_reap_idle() -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new()
.max_size(5)
.idle_timeout(0) .build()?;
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory)?;
let conn = pool.acquire().await?;
pool.release(conn).await;
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
pool.reap_idle().await;
let status = pool.status().await;
assert_eq!(status.idle, 0);
Ok(())
}
#[tokio::test]
async fn test_h7_acquire_timeout_default_30s() {
let config = PoolConfig::default();
assert_eq!(
config.acquire_timeout,
Duration::from_secs(30),
"H-7: acquire_timeout 默认应为 30s"
);
}
#[tokio::test]
async fn test_h7_acquire_timeout_configurable() -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new()
.max_size(1)
.acquire_timeout(5) .build()?;
assert_eq!(config.acquire_timeout, Duration::from_secs(5));
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory)?;
let _conn1 = pool.acquire().await?;
let fast_config = PoolConfigBuilder::new()
.max_size(1)
.acquire_timeout(0) .build()?;
let fast_pool = Pool::new(fast_config, Arc::new(MockConnectionFactory))?;
let _fast_conn = fast_pool.acquire().await?; let result = fast_pool.acquire().await;
assert!(
matches!(result, Err(PoolError::Timeout)),
"H-7: 应返回 Timeout"
);
Ok(())
}
#[tokio::test]
async fn test_m7_health_check_removes_nothing_when_all_healthy(
) -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new().max_size(5).build()?;
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory)?;
let conn1 = pool.acquire().await?;
let conn2 = pool.acquire().await?;
let conn3 = pool.acquire().await?;
pool.release(conn1).await;
pool.release(conn2).await;
pool.release(conn3).await;
let removed = pool.health_check().await;
assert_eq!(removed, 0, "Healthy connections should not be removed");
let status = pool.status().await;
assert_eq!(status.idle, 3);
assert_eq!(status.active, 3);
Ok(())
}
#[tokio::test]
async fn test_m7_health_check_returns_zero_for_empty_pool(
) -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new().max_size(5).build()?;
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory)?;
let removed = pool.health_check().await;
assert_eq!(removed, 0);
Ok(())
}
struct CountingFactory {
count: AtomicU32,
}
impl CountingFactory {
fn new() -> Self {
Self {
count: AtomicU32::new(0),
}
}
fn created_count(&self) -> u32 {
self.count.load(Ordering::SeqCst)
}
}
#[async_trait]
impl ConnectionFactory for CountingFactory {
async fn create(&self) -> Result<Box<dyn Connection>, sz_orm_model::DbError> {
self.count.fetch_add(1, Ordering::SeqCst);
Ok(Box::new(MockConnection::new()))
}
}
#[tokio::test]
async fn test_production_bug_max_lifetime_never_expires(
) -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfig {
max_size: 5,
min_idle: 0,
acquire_timeout: Duration::from_secs(30),
idle_timeout: Duration::from_secs(600),
max_lifetime: Duration::from_millis(100), connection_timeout: Duration::from_secs(10),
tls: None,
query_timeout: None,
max_rows: None,
memory_limit: None,
on_event: None,
};
let factory = Arc::new(CountingFactory::new());
let pool = Pool::new(config, factory.clone())?;
let conn = pool.acquire().await?;
assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
pool.release(conn).await;
tokio::time::sleep(Duration::from_millis(150)).await;
let conn2 = pool.acquire().await?;
assert_eq!(
factory.created_count(),
2,
"超过 max_lifetime 后应创建新连接(旧连接应被回收)"
);
pool.release(conn2).await;
Ok(())
}
#[tokio::test]
async fn test_drop_auto_release_connection() -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new().max_size(2).build()?;
let factory = Arc::new(CountingFactory::new());
let pool = Pool::new(config, factory.clone())?;
{
let _conn = pool.acquire().await?;
assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
let status = pool.status().await;
assert_eq!(status.active, 1, "active 应为 1");
assert_eq!(status.idle, 0, "idle 应为 0");
}
tokio::time::sleep(Duration::from_millis(50)).await;
let status = pool.status().await;
assert_eq!(status.idle, 1, "Drop 后连接应自动归还,idle 应为 1");
assert_eq!(status.active, 1, "total_count 应为 1");
assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
Ok(())
}
#[tokio::test]
async fn test_drop_auto_release_then_reuse() -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new().max_size(1).build()?;
let factory = Arc::new(CountingFactory::new());
let pool = Pool::new(config, factory.clone())?;
{
let _conn = pool.acquire().await?;
}
tokio::time::sleep(Duration::from_millis(50)).await;
let conn = pool.acquire().await?;
assert_eq!(factory.created_count(), 1, "应复用归还的连接,不创建新连接");
pool.release(conn).await;
Ok(())
}
#[tokio::test]
async fn test_into_inner_does_not_return_to_pool() -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new().max_size(2).build()?;
let factory = Arc::new(CountingFactory::new());
let pool = Pool::new(config, factory.clone())?;
let conn = pool.acquire().await?;
assert_eq!(factory.created_count(), 1);
let _raw_conn = conn.into_inner();
tokio::time::sleep(Duration::from_millis(50)).await;
let status = pool.status().await;
assert_eq!(status.idle, 0, "into_inner 后连接不应归还");
assert_eq!(status.active, 1, "total_count 仍为 1(连接被外部持有)");
Ok(())
}
#[tokio::test]
async fn test_explicit_release_no_double_return() -> Result<(), Box<dyn std::error::Error>> {
let config = PoolConfigBuilder::new().max_size(2).build()?;
let factory = Arc::new(CountingFactory::new());
let pool = Pool::new(config, factory.clone())?;
let conn = pool.acquire().await?;
pool.release(conn).await;
let status = pool.status().await;
assert_eq!(status.idle, 1, "release 后 idle 应为 1");
let conn = pool.acquire().await?;
pool.release(conn).await;
let status = pool.status().await;
assert_eq!(status.idle, 1, "再次 release 后 idle 仍应为 1(不重复)");
assert_eq!(status.active, 1, "total_count 应为 1");
Ok(())
}
}