use crate::error::{Result, TdbError};
use crate::store::TdbStore;
use parking_lot::{Mutex, RwLock};
use std::collections::VecDeque;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct ConnectionPoolConfig {
pub min_connections: usize,
pub max_connections: usize,
pub acquire_timeout: Duration,
pub max_idle_time: Duration,
pub enable_health_check: bool,
pub health_check_interval: Duration,
}
impl Default for ConnectionPoolConfig {
fn default() -> Self {
Self {
min_connections: 2,
max_connections: 10,
acquire_timeout: Duration::from_secs(30),
max_idle_time: Duration::from_secs(300), enable_health_check: true,
health_check_interval: Duration::from_secs(60),
}
}
}
pub struct PooledConnection {
store: Option<TdbStore>,
id: u64,
last_used: Instant,
pool: Arc<ConnectionPoolInner>,
}
impl PooledConnection {
pub fn store(&self) -> &TdbStore {
self.store.as_ref().expect("Store should be present")
}
pub fn store_mut(&mut self) -> &mut TdbStore {
self.store.as_mut().expect("Store should be present")
}
pub fn id(&self) -> u64 {
self.id
}
pub fn idle_time(&self) -> Duration {
self.last_used.elapsed()
}
fn touch(&mut self) {
self.last_used = Instant::now();
}
}
impl Drop for PooledConnection {
fn drop(&mut self) {
if let Some(store) = self.store.take() {
self.pool.return_connection(store, self.id);
}
}
}
struct ConnectionPoolInner {
db_path: PathBuf,
config: ConnectionPoolConfig,
available: Mutex<VecDeque<(u64, TdbStore)>>,
next_id: AtomicU64,
current_size: AtomicUsize,
stats: ConnectionPoolStats,
}
impl ConnectionPoolInner {
fn create_connection(&self) -> Result<TdbStore> {
TdbStore::open(&self.db_path)
}
fn return_connection(&self, store: TdbStore, id: u64) {
let mut available = self.available.lock();
if available.len() < self.config.max_connections {
available.push_back((id, store));
self.stats
.returned_connections
.fetch_add(1, Ordering::Relaxed);
} else {
drop(store);
self.current_size.fetch_sub(1, Ordering::Relaxed);
self.stats
.closed_connections
.fetch_add(1, Ordering::Relaxed);
}
}
}
pub struct ConnectionPool {
inner: Arc<ConnectionPoolInner>,
}
impl ConnectionPool {
pub fn new<P: AsRef<Path>>(db_path: P, config: ConnectionPoolConfig) -> Result<Self> {
let db_path = db_path.as_ref().to_path_buf();
if config.min_connections > config.max_connections {
return Err(TdbError::Other(
"min_connections cannot exceed max_connections".to_string(),
));
}
let inner = Arc::new(ConnectionPoolInner {
db_path: db_path.clone(),
config: config.clone(),
available: Mutex::new(VecDeque::with_capacity(config.max_connections)),
next_id: AtomicU64::new(1),
current_size: AtomicUsize::new(0),
stats: ConnectionPoolStats::default(),
});
for _ in 0..config.min_connections {
let store = TdbStore::open(&db_path)?;
let id = inner.next_id.fetch_add(1, Ordering::Relaxed);
inner.available.lock().push_back((id, store));
inner.current_size.fetch_add(1, Ordering::Relaxed);
}
Ok(Self { inner })
}
pub fn acquire(&self) -> Result<PooledConnection> {
self.inner
.stats
.acquire_requests
.fetch_add(1, Ordering::Relaxed);
let start = Instant::now();
loop {
{
let mut available = self.inner.available.lock();
if let Some((id, store)) = available.pop_front() {
self.inner
.stats
.successful_acquires
.fetch_add(1, Ordering::Relaxed);
return Ok(PooledConnection {
store: Some(store),
id,
last_used: Instant::now(),
pool: Arc::clone(&self.inner),
});
}
}
let current_size = self.inner.current_size.load(Ordering::Relaxed);
if current_size < self.inner.config.max_connections {
match self.inner.create_connection() {
Ok(store) => {
let id = self.inner.next_id.fetch_add(1, Ordering::Relaxed);
self.inner.current_size.fetch_add(1, Ordering::Relaxed);
self.inner
.stats
.created_connections
.fetch_add(1, Ordering::Relaxed);
self.inner
.stats
.successful_acquires
.fetch_add(1, Ordering::Relaxed);
return Ok(PooledConnection {
store: Some(store),
id,
last_used: Instant::now(),
pool: Arc::clone(&self.inner),
});
}
Err(e) => {
self.inner
.stats
.failed_acquires
.fetch_add(1, Ordering::Relaxed);
return Err(e);
}
}
}
if start.elapsed() >= self.inner.config.acquire_timeout {
self.inner
.stats
.timeout_acquires
.fetch_add(1, Ordering::Relaxed);
return Err(TdbError::Other(format!(
"Connection acquire timeout after {:?}",
self.inner.config.acquire_timeout
)));
}
std::thread::sleep(Duration::from_millis(10));
}
}
pub fn stats(&self) -> ConnectionPoolStatsSnapshot {
ConnectionPoolStatsSnapshot {
current_size: self.inner.current_size.load(Ordering::Relaxed),
available: self.inner.available.lock().len(),
acquire_requests: self.inner.stats.acquire_requests.load(Ordering::Relaxed),
successful_acquires: self.inner.stats.successful_acquires.load(Ordering::Relaxed),
failed_acquires: self.inner.stats.failed_acquires.load(Ordering::Relaxed),
timeout_acquires: self.inner.stats.timeout_acquires.load(Ordering::Relaxed),
created_connections: self.inner.stats.created_connections.load(Ordering::Relaxed),
returned_connections: self
.inner
.stats
.returned_connections
.load(Ordering::Relaxed),
closed_connections: self.inner.stats.closed_connections.load(Ordering::Relaxed),
}
}
pub fn size(&self) -> usize {
self.inner.current_size.load(Ordering::Relaxed)
}
pub fn available(&self) -> usize {
self.inner.available.lock().len()
}
pub fn close_idle_connections(&self) -> usize {
let mut available = self.inner.available.lock();
let _max_idle = self.inner.config.max_idle_time;
let closed_count = 0;
let _now = Instant::now();
available.retain(|(_, _)| {
true
});
closed_count
}
pub fn resize(&self, new_size: usize) -> Result<()> {
if new_size < self.inner.config.min_connections {
return Err(TdbError::Other(format!(
"New size {} is below minimum {}",
new_size, self.inner.config.min_connections
)));
}
if new_size > self.inner.config.max_connections {
return Err(TdbError::Other(format!(
"New size {} exceeds maximum {}",
new_size, self.inner.config.max_connections
)));
}
let current_size = self.inner.current_size.load(Ordering::Relaxed);
if new_size > current_size {
for _ in current_size..new_size {
let store = self.inner.create_connection()?;
let id = self.inner.next_id.fetch_add(1, Ordering::Relaxed);
self.inner.available.lock().push_back((id, store));
self.inner.current_size.fetch_add(1, Ordering::Relaxed);
self.inner
.stats
.created_connections
.fetch_add(1, Ordering::Relaxed);
}
} else if new_size < current_size {
let to_remove = current_size - new_size;
let mut available = self.inner.available.lock();
for _ in 0..to_remove.min(available.len()) {
if available.pop_back().is_some() {
self.inner.current_size.fetch_sub(1, Ordering::Relaxed);
self.inner
.stats
.closed_connections
.fetch_add(1, Ordering::Relaxed);
}
}
}
Ok(())
}
}
#[derive(Debug, Default)]
struct ConnectionPoolStats {
acquire_requests: AtomicU64,
successful_acquires: AtomicU64,
failed_acquires: AtomicU64,
timeout_acquires: AtomicU64,
created_connections: AtomicU64,
returned_connections: AtomicU64,
closed_connections: AtomicU64,
}
#[derive(Debug, Clone)]
pub struct ConnectionPoolStatsSnapshot {
pub current_size: usize,
pub available: usize,
pub acquire_requests: u64,
pub successful_acquires: u64,
pub failed_acquires: u64,
pub timeout_acquires: u64,
pub created_connections: u64,
pub returned_connections: u64,
pub closed_connections: u64,
}
impl ConnectionPoolStatsSnapshot {
pub fn success_rate(&self) -> f64 {
if self.acquire_requests == 0 {
0.0
} else {
(self.successful_acquires as f64 / self.acquire_requests as f64) * 100.0
}
}
pub fn utilization_rate(&self) -> f64 {
if self.current_size == 0 {
0.0
} else {
let in_use = self.current_size - self.available;
(in_use as f64 / self.current_size as f64) * 100.0
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn create_test_pool() -> (TempDir, ConnectionPool) {
let temp_dir = TempDir::new().unwrap();
let db_path = temp_dir.path().join("test.db");
let config = ConnectionPoolConfig {
min_connections: 2,
max_connections: 5,
acquire_timeout: Duration::from_secs(5),
..Default::default()
};
let pool = ConnectionPool::new(&db_path, config).unwrap();
(temp_dir, pool)
}
#[test]
fn test_connection_pool_creation() {
let (_temp_dir, pool) = create_test_pool();
assert_eq!(pool.size(), 2); assert_eq!(pool.available(), 2);
}
#[test]
fn test_acquire_and_return() {
let (_temp_dir, pool) = create_test_pool();
{
let conn = pool.acquire().unwrap();
assert_eq!(pool.available(), 1);
let _ = conn.store();
}
assert_eq!(pool.available(), 2);
}
#[test]
fn test_multiple_acquires() {
let (_temp_dir, pool) = create_test_pool();
let conn1 = pool.acquire().unwrap();
let conn2 = pool.acquire().unwrap();
let conn3 = pool.acquire().unwrap();
assert_eq!(pool.size(), 3);
assert_eq!(pool.available(), 0);
drop(conn1);
assert_eq!(pool.available(), 1);
drop(conn2);
drop(conn3);
assert_eq!(pool.available(), 3);
}
#[test]
fn test_max_connections_limit() {
let (_temp_dir, pool) = create_test_pool();
let mut connections = Vec::new();
for _ in 0..5 {
connections.push(pool.acquire().unwrap());
}
assert_eq!(pool.size(), 5);
assert_eq!(pool.available(), 0);
}
#[test]
fn test_connection_pool_stats() {
let (_temp_dir, pool) = create_test_pool();
let _conn1 = pool.acquire().unwrap();
let _conn2 = pool.acquire().unwrap();
let stats = pool.stats();
assert_eq!(stats.acquire_requests, 2);
assert_eq!(stats.successful_acquires, 2);
assert!(stats.success_rate() > 99.0);
}
#[test]
fn test_pool_resize_grow() {
let (_temp_dir, pool) = create_test_pool();
assert_eq!(pool.size(), 2);
pool.resize(4).unwrap();
assert_eq!(pool.size(), 4);
assert_eq!(pool.available(), 4);
}
#[test]
fn test_pool_resize_shrink() {
let (_temp_dir, pool) = create_test_pool();
pool.resize(4).unwrap();
assert_eq!(pool.size(), 4);
pool.resize(2).unwrap();
assert_eq!(pool.size(), 2);
}
#[test]
fn test_resize_validation() {
let (_temp_dir, pool) = create_test_pool();
assert!(pool.resize(1).is_err());
assert!(pool.resize(10).is_err());
}
#[test]
fn test_utilization_rate() {
let (_temp_dir, pool) = create_test_pool();
let _conn1 = pool.acquire().unwrap();
let stats = pool.stats();
assert!((stats.utilization_rate() - 50.0).abs() < 1.0);
}
#[test]
fn test_connection_id() {
let (_temp_dir, pool) = create_test_pool();
let conn1 = pool.acquire().unwrap();
let conn2 = pool.acquire().unwrap();
assert_ne!(conn1.id(), conn2.id());
}
#[test]
fn test_pooled_connection_touch() {
let (_temp_dir, pool) = create_test_pool();
let mut conn = pool.acquire().unwrap();
std::thread::sleep(Duration::from_millis(100));
assert!(conn.idle_time() >= Duration::from_millis(100));
conn.touch();
assert!(conn.idle_time() < Duration::from_millis(50));
}
#[test]
fn test_concurrent_acquires() {
use std::thread;
let (_temp_dir, pool) = create_test_pool();
let pool = Arc::new(pool);
let mut handles = vec![];
for _ in 0..3 {
let pool_clone = Arc::clone(&pool);
let handle = thread::spawn(move || {
let _conn = pool_clone.acquire().unwrap();
thread::sleep(Duration::from_millis(50));
});
handles.push(handle);
}
for handle in handles {
handle.join().unwrap();
}
let stats = pool.stats();
assert_eq!(stats.successful_acquires, 3);
}
#[test]
fn test_stats_snapshot_success_rate() {
let stats = ConnectionPoolStatsSnapshot {
current_size: 5,
available: 2,
acquire_requests: 100,
successful_acquires: 95,
failed_acquires: 3,
timeout_acquires: 2,
created_connections: 5,
returned_connections: 90,
closed_connections: 0,
};
assert_eq!(stats.success_rate(), 95.0);
}
}