use async_trait::async_trait;
use std::collections::VecDeque;
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::{Mutex, Notify};
use crate::error::PoolError;
pub type QueryRows = Vec<std::collections::HashMap<String, crate::value::Value>>;
pub trait Connection: Send + Sync {
fn execute<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<u64, crate::DbError>> + Send + 'a>>;
fn query<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<QueryRows, crate::DbError>> + Send + 'a>>;
fn begin_transaction<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
fn commit<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
fn rollback<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>>;
fn is_connected(&self) -> bool;
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<(), crate::DbError>> + Send + 'a>>;
}
pub struct PooledConnection {
conn: Box<dyn Connection>,
created_at: Instant,
last_used_at: Instant,
}
impl PooledConnection {
fn new(conn: Box<dyn Connection>) -> Self {
let now = Instant::now();
Self {
conn,
created_at: now,
last_used_at: now,
}
}
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(self) -> Box<dyn Connection> {
self.conn
}
}
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()
}
}
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,
}
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),
}
}
}
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,
}
}
}
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,
}
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)
.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 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>, crate::DbError>;
}
pub struct Pool {
config: PoolConfig,
factory: Arc<dyn ConnectionFactory>,
idle: Arc<Mutex<VecDeque<PooledConnection>>>,
total_count: Arc<AtomicU32>,
closed: Arc<AtomicBool>,
notify: Notify,
}
impl Pool {
pub fn new(config: PoolConfig, factory: Arc<dyn ConnectionFactory>) -> Result<Self, PoolError> {
config.validate()?;
Ok(Self {
config,
factory,
idle: Arc::new(Mutex::new(VecDeque::new())),
total_count: Arc::new(AtomicU32::new(0)),
closed: Arc::new(AtomicBool::new(false)),
notify: Notify::new(),
})
}
pub fn config(&self) -> &PoolConfig {
&self.config
}
#[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);
}
let deadline = Instant::now() + self.config.acquire_timeout;
loop {
let mut to_close: Vec<PooledConnection> = Vec::new();
let acquired: Option<PooledConnection> = {
let mut idle = self.idle.lock().await;
let mut found: Option<PooledConnection> = None;
while let Some(pooled) = idle.pop_front() {
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(pooled) = acquired {
return Ok(pooled);
}
let created = loop {
let current = self.total_count.load(Ordering::Acquire);
if current >= self.config.max_size {
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)) => return Ok(PooledConnection::new(conn)),
Ok(Err(e)) => {
self.total_count.fetch_sub(1, Ordering::SeqCst);
return Err(PoolError::ConnectionFailed(e.to_string()));
}
Err(_) => {
self.total_count.fetch_sub(1, Ordering::SeqCst);
return Err(PoolError::Timeout);
}
}
}
let now = Instant::now();
if now >= deadline {
return Err(PoolError::Timeout);
}
let _ = tokio::time::timeout(deadline - now, self.notify.notified()).await;
}
}
#[tracing::instrument(skip(self, pooled))]
pub async fn release(&self, mut pooled: PooledConnection) {
if self.closed.load(Ordering::Acquire) {
let _ = pooled.conn.close().await;
self.total_count.fetch_sub(1, Ordering::SeqCst);
return;
}
if !pooled.conn.is_connected() {
let _ = pooled.conn.close().await;
self.total_count.fetch_sub(1, Ordering::SeqCst);
return;
}
pooled.last_used_at = Instant::now();
{
let mut idle = self.idle.lock().await;
idle.push_back(pooled);
}
self.notify.notify_one();
}
pub async fn status(&self) -> PoolStatus {
let idle_count = self.idle.lock().await.len() as u32;
let active = self.total_count.load(Ordering::Acquire);
PoolStatus {
idle: idle_count,
active,
max: self.config.max_size,
min: self.config.min_idle,
}
}
#[tracing::instrument(skip(self))]
pub async fn reap_idle(&self) {
let mut idle = self.idle.lock().await;
let mut to_close = Vec::new();
let mut remaining = VecDeque::new();
while let Some(pooled) = idle.pop_front() {
if pooled.is_idle_too_long(self.config.idle_timeout)
|| pooled.is_expired(self.config.max_lifetime)
{
to_close.push(pooled);
} else {
remaining.push_back(pooled);
}
}
*idle = remaining;
drop(idle);
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 to_close: Vec<PooledConnection> = {
let mut idle = self.idle.lock().await;
let mut collected = Vec::with_capacity(idle.len());
while let Some(pooled) = idle.pop_front() {
collected.push(pooled);
}
collected
};
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> = {
let mut idle = self.idle.lock().await;
let mut collected = Vec::with_capacity(idle.len());
while let Some(pooled) = idle.pop_front() {
collected.push(pooled);
}
collected
};
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;
{
let mut idle = self.idle.lock().await;
for pooled in alive {
idle.push_back(pooled);
}
}
if removed > 0 {
self.total_count.fetch_sub(removed, Ordering::SeqCst);
}
if alive_count > 0 {
self.notify.notify_one();
}
removed
}
}
#[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, crate::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, crate::value::Value>>,
crate::DbError,
>,
> + Send
+ 'a,
>,
> {
Box::pin(async move { Ok(vec![]) })
}
fn begin_transaction<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
Box::pin(async move { Ok(()) })
}
fn commit<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), crate::DbError>> + Send + 'a>> {
Box::pin(async move { Ok(()) })
}
fn rollback<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), crate::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<(), crate::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>, crate::DbError> {
Ok(Box::new(MockConnection::new()))
}
}
#[tokio::test]
async fn test_pool_config_builder() {
let config = PoolConfigBuilder::new()
.max_size(50)
.min_idle(10)
.build()
.unwrap();
assert_eq!(config.max_size, 50);
assert_eq!(config.min_idle, 10);
}
#[test]
fn test_pool_status_display() {
let status = PoolStatus {
idle: 5,
active: 10,
max: 100,
min: 5,
};
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() {
let builder = PoolConfigBuilder::new();
let config = builder.build().unwrap();
assert_eq!(config.max_size, 100);
}
#[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() {
let config = PoolConfigBuilder::new()
.max_size(5)
.min_idle(1)
.build()
.unwrap();
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory).unwrap();
let conn = pool.acquire().await.unwrap();
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.unwrap();
let status = pool.status().await;
assert_eq!(status.idle, 0);
}
#[tokio::test]
async fn test_pool_status() {
let config = PoolConfigBuilder::new()
.max_size(10)
.min_idle(2)
.build()
.unwrap();
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory).unwrap();
let status = pool.status().await;
assert_eq!(status.max, 10);
assert_eq!(status.min, 2);
assert_eq!(status.active, 0);
}
#[tokio::test]
async fn test_pool_close_all() {
let config = PoolConfigBuilder::new().max_size(5).build().unwrap();
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory).unwrap();
let conn1 = pool.acquire().await.unwrap();
let conn2 = pool.acquire().await.unwrap();
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);
}
#[tokio::test]
async fn test_pool_reap_idle() {
let config = PoolConfigBuilder::new()
.max_size(5)
.idle_timeout(0) .build()
.unwrap();
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory).unwrap();
let conn = pool.acquire().await.unwrap();
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);
}
#[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() {
let config = PoolConfigBuilder::new()
.max_size(1)
.acquire_timeout(5) .build()
.unwrap();
assert_eq!(config.acquire_timeout, Duration::from_secs(5));
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory).unwrap();
let _conn1 = pool.acquire().await.unwrap();
let fast_config = PoolConfigBuilder::new()
.max_size(1)
.acquire_timeout(0) .build()
.unwrap();
let fast_pool = Pool::new(fast_config, Arc::new(MockConnectionFactory)).unwrap();
let _fast_conn = fast_pool.acquire().await.unwrap(); let result = fast_pool.acquire().await;
assert!(
matches!(result, Err(PoolError::Timeout)),
"H-7: 应返回 Timeout"
);
}
#[tokio::test]
async fn test_m7_health_check_removes_nothing_when_all_healthy() {
let config = PoolConfigBuilder::new().max_size(5).build().unwrap();
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory).unwrap();
let conn1 = pool.acquire().await.unwrap();
let conn2 = pool.acquire().await.unwrap();
let conn3 = pool.acquire().await.unwrap();
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);
}
#[tokio::test]
async fn test_m7_health_check_returns_zero_for_empty_pool() {
let config = PoolConfigBuilder::new().max_size(5).build().unwrap();
let factory = Arc::new(MockConnectionFactory);
let pool = Pool::new(config, factory).unwrap();
let removed = pool.health_check().await;
assert_eq!(removed, 0);
}
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>, crate::DbError> {
self.count.fetch_add(1, Ordering::SeqCst);
Ok(Box::new(MockConnection::new()))
}
}
#[tokio::test]
async fn test_production_bug_max_lifetime_never_expires() {
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),
};
let factory = Arc::new(CountingFactory::new());
let pool = Pool::new(config, factory.clone()).unwrap();
let conn = pool.acquire().await.unwrap();
assert_eq!(factory.created_count(), 1, "应创建 1 个连接");
pool.release(conn).await;
tokio::time::sleep(Duration::from_millis(150)).await;
let conn2 = pool.acquire().await.unwrap();
assert_eq!(
factory.created_count(),
2,
"超过 max_lifetime 后应创建新连接(旧连接应被回收)"
);
pool.release(conn2).await;
}
}