use async_trait::async_trait;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;
use sz_orm_core::{
Connection, ConnectionFactory, DbError, Pool, PoolConfig, PoolConfigBuilder, PoolError,
};
struct MockConnection {
connected: bool,
close_count: Arc<AtomicU32>,
}
impl MockConnection {
fn new(close_count: Arc<AtomicU32>) -> Self {
Self {
connected: true,
close_count,
}
}
}
impl Connection for MockConnection {
fn execute<'a>(
&'a mut self,
_sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<u64, 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_core::Value>>,
DbError,
>,
> + Send
+ 'a,
>,
> {
Box::pin(async move { Ok(vec![]) })
}
fn begin_transaction<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move { Ok(()) })
}
fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move { Ok(()) })
}
fn rollback<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), 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 { self.connected })
}
fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
self.connected = false;
self.close_count.fetch_add(1, Ordering::SeqCst);
Ok(())
})
}
}
struct MockConnectionFactory {
close_count: Arc<AtomicU32>,
}
impl MockConnectionFactory {
fn new() -> Self {
Self {
close_count: Arc::new(AtomicU32::new(0)),
}
}
}
#[async_trait]
impl ConnectionFactory for MockConnectionFactory {
async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
Ok(Box::new(MockConnection::new(self.close_count.clone())))
}
}
struct DelayedConnectionFactory {
delay: Duration,
}
impl DelayedConnectionFactory {
fn new(delay: Duration) -> Self {
Self { delay }
}
}
#[async_trait]
impl ConnectionFactory for DelayedConnectionFactory {
async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
tokio::time::sleep(self.delay).await;
Ok(Box::new(MockConnection::new(Arc::new(AtomicU32::new(0)))))
}
}
struct FailingConnectionFactory {
fail_count: Arc<AtomicU32>,
max_failures: u32,
}
impl FailingConnectionFactory {
fn new(max_failures: u32) -> Self {
Self {
fail_count: Arc::new(AtomicU32::new(0)),
max_failures,
}
}
}
#[async_trait]
impl ConnectionFactory for FailingConnectionFactory {
async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
let failures = self.fail_count.fetch_add(1, Ordering::SeqCst);
if failures < self.max_failures {
Err(DbError::ConnectionRefused("simulated failure".to_string()))
} else {
Ok(Box::new(MockConnection::new(Arc::new(AtomicU32::new(0)))))
}
}
}
fn make_pool(max_size: u32, acquire_timeout_secs: u64) -> (Pool, Arc<MockConnectionFactory>) {
let factory = Arc::new(MockConnectionFactory::new());
let config = PoolConfigBuilder::new()
.max_size(max_size)
.min_idle(0)
.acquire_timeout(acquire_timeout_secs)
.build()
.unwrap();
let pool = Pool::new(config, factory.clone()).unwrap();
(pool, factory)
}
#[tokio::test]
async fn chaos_pool_starvation_attack() {
let (pool, _factory) = make_pool(2, 0);
let conn1 = pool.acquire().await.expect("第 1 个连接应成功");
let conn2 = pool.acquire().await.expect("第 2 个连接应成功");
let status = pool.status().await;
assert_eq!(status.active, 2, "池满:active 应为 2");
assert_eq!(status.idle, 0, "池满:idle 应为 0");
let err = match pool.acquire().await {
Err(e) => e,
Ok(_) => panic!("第 3 个 acquire 应失败返回 Timeout"),
};
assert!(
matches!(err, PoolError::Timeout),
"饥荒时 acquire 应返回 Timeout,实际: {}",
err
);
pool.release(conn1).await;
pool.release(conn2).await;
let status = pool.status().await;
assert_eq!(status.idle, 2, "释放后 idle 应为 2");
}
#[tokio::test]
async fn chaos_pool_starvation_recovery() {
let (pool, _factory) = make_pool(2, 5); let conn1 = pool.acquire().await.unwrap();
let conn2 = pool.acquire().await.unwrap();
let pool_clone = pool.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(50)).await;
pool_clone.release(conn1).await;
});
let conn3 = pool.acquire().await.expect("释放后应能获取连接");
pool.release(conn2).await;
pool.release(conn3).await;
let status = pool.status().await;
assert_eq!(status.idle, 2, "恢复后 idle 应为 2");
}
#[tokio::test]
async fn chaos_pool_connection_creation_timeout() {
let factory = Arc::new(DelayedConnectionFactory::new(Duration::from_secs(2)));
let config = PoolConfig {
max_size: 5,
min_idle: 0,
acquire_timeout: Duration::from_secs(5),
idle_timeout: Duration::from_secs(600),
max_lifetime: Duration::from_secs(1800),
connection_timeout: Duration::from_millis(100), tls: None,
query_timeout: None,
max_rows: None,
memory_limit: None,
on_event: None,
};
let pool = Pool::new(config, factory).unwrap();
let err = match pool.acquire().await {
Err(e) => e,
Ok(_) => panic!("连接创建超时应返回 PoolError::Timeout"),
};
assert!(
matches!(err, PoolError::Timeout),
"连接创建超时应返回 PoolError::Timeout,实际: {}",
err
);
let status = pool.status().await;
assert_eq!(
status.active, 0,
"失败的创建不应泄漏计数器,active 应为 0,实际: {}",
status.active
);
}
#[tokio::test]
async fn chaos_pool_transaction_leak_detection() {
let (pool, _factory) = make_pool(2, 5);
let conn = pool.acquire().await.unwrap();
let raw_conn = conn.into_inner();
let mut tx = sz_orm_core::Transaction::new(raw_conn, sz_orm_core::TransactOptions::default());
tx.execute("INSERT INTO test VALUES (1)").await.unwrap();
drop(tx);
tokio::time::sleep(Duration::from_millis(100)).await;
let conn2 = pool.acquire().await.expect("池应能提供连接");
assert!(conn2.is_connected(), "新连接应健康");
pool.release(conn2).await;
let status = pool.status().await;
assert_eq!(status.idle, 1, "最终 idle 应为 1");
assert!(
status.active <= 2,
"active 不应超过 max_size: {}",
status.active
);
}
#[tokio::test]
async fn chaos_pool_acquire_after_close() {
let (pool, _factory) = make_pool(5, 10);
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, "close_all 后 idle 应为 0");
assert_eq!(status.active, 0, "close_all 后 active 应为 0");
let err = match pool.acquire().await {
Err(e) => e,
Ok(_) => panic!("关闭后 acquire 应返回 PoolError::Closed"),
};
assert!(
matches!(err, PoolError::Closed),
"关闭后 acquire 应返回 PoolError::Closed,实际: {}",
err
);
}
#[tokio::test]
async fn chaos_pool_release_after_close() {
let (pool, _factory) = make_pool(5, 10);
let conn1 = pool.acquire().await.unwrap();
let conn2 = pool.acquire().await.unwrap();
pool.close_all().await;
pool.release(conn1).await;
pool.release(conn2).await;
let status = pool.status().await;
assert_eq!(
status.active, 0,
"close_all + release 后 active 应为 0,实际: {}",
status.active
);
assert_eq!(status.idle, 0, "close_all 后 idle 应为 0");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn chaos_pool_concurrent_stress() {
let (pool, _factory) = make_pool(50, 10);
let task_count = 20;
let ops_per_task = 100;
let max_active = Arc::new(AtomicU32::new(0));
let total_ops = Arc::new(AtomicU32::new(0));
let mut handles = Vec::new();
for _ in 0..task_count {
let pool_c = pool.clone();
let max_c = max_active.clone();
let ops_c = total_ops.clone();
handles.push(tokio::spawn(async move {
for _ in 0..ops_per_task {
let conn = pool_c.acquire().await.expect("并发 acquire 应成功");
let status = pool_c.status().await;
max_c.fetch_max(status.active, Ordering::Relaxed);
tokio::task::yield_now().await;
pool_c.release(conn).await;
ops_c.fetch_add(1, Ordering::Relaxed);
}
}));
}
for h in handles {
h.await.unwrap();
}
let completed = total_ops.load(Ordering::Relaxed);
assert_eq!(
completed,
(task_count * ops_per_task) as u32,
"所有 {} 个操作应完成,实际: {}",
task_count * ops_per_task,
completed
);
let observed_max = max_active.load(Ordering::Relaxed);
assert!(
observed_max <= 50,
"active 从未超过 max_size: observed_max={}",
observed_max
);
let status = pool.status().await;
assert_eq!(
status.active, status.idle,
"并发结束后 active={} 应等于 idle={}",
status.active, status.idle
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn chaos_pool_concurrent_with_timeout_no_panic() {
let factory = Arc::new(MockConnectionFactory::new());
let config = PoolConfigBuilder::new()
.max_size(2)
.min_idle(0)
.acquire_timeout(1)
.build()
.unwrap();
let pool = Pool::new(config, factory).unwrap();
let mut handles = Vec::new();
for _ in 0..8 {
let pool_c = pool.clone();
handles.push(tokio::spawn(async move {
for _ in 0..50 {
match pool_c.acquire().await {
Ok(conn) => {
tokio::task::yield_now().await;
pool_c.release(conn).await;
}
Err(e) => {
assert!(
matches!(e, PoolError::Timeout),
"只应出现 Timeout,实际: {:?}",
e
);
}
}
}
}));
}
for h in handles {
h.await.unwrap();
}
let status = pool.status().await;
assert!(
status.active <= 2,
"active 不应超过 max_size: {}",
status.active
);
assert_eq!(
status.active, status.idle,
"最终 active={} 应等于 idle={}",
status.active, status.idle
);
}
#[tokio::test]
async fn chaos_pool_health_check_recycles_expired() {
let factory = Arc::new(MockConnectionFactory::new());
let config = PoolConfig {
max_size: 10,
min_idle: 0,
acquire_timeout: Duration::from_secs(5),
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 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 before = pool.status().await;
assert_eq!(before.idle, 3, "初始应有 3 个 idle 连接");
tokio::time::sleep(Duration::from_millis(150)).await;
pool.reap_idle().await;
let after = pool.status().await;
assert_eq!(after.idle, 0, "回收后 idle 应为 0");
assert_eq!(after.active, 0, "回收后 total_count 应为 0");
}
#[tokio::test]
async fn chaos_pool_health_check_then_serve() {
let factory = Arc::new(MockConnectionFactory::new());
let config = PoolConfig {
max_size: 10,
min_idle: 0,
acquire_timeout: Duration::from_secs(5),
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 pool = Pool::new(config, factory).unwrap();
let conn = pool.acquire().await.unwrap();
pool.release(conn).await;
tokio::time::sleep(Duration::from_millis(150)).await;
pool.reap_idle().await;
let conn = pool.acquire().await.expect("回收后应能创建新连接");
assert!(conn.is_connected(), "新连接应健康");
pool.release(conn).await;
}
#[tokio::test]
async fn chaos_pool_health_check_detects_disconnected() {
use std::sync::atomic::AtomicBool;
struct DisconnectableConnection {
alive: Arc<AtomicBool>,
}
impl Connection for DisconnectableConnection {
fn execute<'a>(
&'a mut self,
_sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
if !self.alive.load(Ordering::SeqCst) {
Box::pin(async { Err(DbError::ConnectionError("disconnected".into())) })
} else {
Box::pin(async { Ok(1) })
}
}
fn query<'a>(
&'a mut self,
_sql: &'a str,
) -> Pin<
Box<
dyn Future<
Output = Result<
Vec<std::collections::HashMap<String, sz_orm_core::Value>>,
DbError,
>,
> + Send
+ 'a,
>,
> {
Box::pin(async { Ok(vec![]) })
}
fn begin_transaction<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async { Ok(()) })
}
fn commit<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async { Ok(()) })
}
fn rollback<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async { Ok(()) })
}
fn is_connected(&self) -> bool {
self.alive.load(Ordering::SeqCst)
}
fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
let a = self.alive.clone();
Box::pin(async move { a.load(Ordering::SeqCst) })
}
fn close<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async { Ok(()) })
}
}
struct DisconnectableFactory {
alive: Arc<AtomicBool>,
}
#[async_trait]
impl ConnectionFactory for DisconnectableFactory {
async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
Ok(Box::new(DisconnectableConnection {
alive: self.alive.clone(),
}))
}
}
let alive = Arc::new(AtomicBool::new(true));
let factory = Arc::new(DisconnectableFactory {
alive: alive.clone(),
});
let config = PoolConfig {
max_size: 10,
min_idle: 0,
acquire_timeout: Duration::from_secs(5),
idle_timeout: Duration::from_secs(600),
max_lifetime: Duration::from_secs(3600),
connection_timeout: Duration::from_secs(10),
tls: None,
query_timeout: None,
max_rows: None,
memory_limit: None,
on_event: None,
};
let pool = Pool::new(config, factory).unwrap();
let conn = pool.acquire().await.unwrap();
pool.release(conn).await;
assert_eq!(pool.status().await.idle, 1);
alive.store(false, Ordering::SeqCst);
let removed = pool.health_check().await;
assert_eq!(removed, 1, "health_check 应移除 1 个断开连接");
assert_eq!(pool.status().await.idle, 0);
}
#[tokio::test]
async fn chaos_pool_health_check_no_false_positive() {
let factory = Arc::new(MockConnectionFactory::new());
let config = PoolConfig {
max_size: 10,
min_idle: 0,
acquire_timeout: Duration::from_secs(5),
idle_timeout: Duration::from_secs(600),
max_lifetime: Duration::from_secs(3600), connection_timeout: Duration::from_secs(10),
tls: None,
query_timeout: None,
max_rows: None,
memory_limit: None,
on_event: None,
};
let pool = Pool::new(config, factory).unwrap();
let conn = pool.acquire().await.unwrap();
pool.release(conn).await;
let removed = pool.health_check().await;
assert_eq!(removed, 0, "未过期的连接不应被回收");
let status = pool.status().await;
assert_eq!(status.idle, 1, "健康连接应保留");
}
#[tokio::test]
async fn chaos_pool_factory_failure_no_counter_leak() {
let factory = Arc::new(FailingConnectionFactory::new(3)); let config = PoolConfig {
max_size: 5,
min_idle: 0,
acquire_timeout: Duration::from_secs(5),
idle_timeout: Duration::from_secs(600),
max_lifetime: Duration::from_secs(3600),
connection_timeout: Duration::from_secs(10),
tls: None,
query_timeout: None,
max_rows: None,
memory_limit: None,
on_event: None,
};
let pool = Pool::new(config, factory).unwrap();
for i in 0..3 {
let result = pool.acquire().await;
assert!(result.is_err(), "第 {} 次 acquire 应失败", i + 1,);
}
let status = pool.status().await;
assert_eq!(
status.active, 0,
"3 次失败后 active 应为 0,实际: {}",
status.active
);
let conn = pool.acquire().await.expect("第 4 次 acquire 应成功");
pool.release(conn).await;
let status = pool.status().await;
assert_eq!(status.idle, 1, "成功创建后 idle 应为 1");
}