use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Mutex;
use tracing::{debug, info};
use crate::error::Result;
use crate::ftp::connection::{FtpClient, FtpMode};
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ConnectionKey {
pub host: String,
pub port: u16,
pub username: String,
pub password: String,
}
impl ConnectionKey {
pub fn new(host: &str, port: u16, username: &str, password: &str) -> Self {
Self {
host: host.to_string(),
port,
username: username.to_string(),
password: password.to_string(),
}
}
}
pub struct PooledConnection {
pub client: FtpClient,
pub key: ConnectionKey,
pub created_at: Instant,
pub last_used: Instant,
pub reuse_count: u64,
pub mode: FtpMode,
}
impl std::fmt::Debug for PooledConnection {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PooledConnection")
.field("key", &self.key)
.field("created_at", &self.created_at)
.field("last_used", &self.last_used)
.field("reuse_count", &self.reuse_count)
.field("mode", &self.mode)
.field("age", &self.age())
.field("idle_time", &self.idle_time())
.finish_non_exhaustive()
}
}
impl PooledConnection {
pub fn new(client: FtpClient, key: ConnectionKey, mode: FtpMode) -> Self {
let now = Instant::now();
Self {
client,
key,
created_at: now,
last_used: now,
reuse_count: 0,
mode,
}
}
pub fn mark_used(&mut self) {
self.last_used = Instant::now();
self.reuse_count += 1;
}
pub fn is_healthy(&self, max_idle_time: Duration) -> bool {
self.last_used.elapsed() < max_idle_time
}
pub fn age(&self) -> Duration {
self.created_at.elapsed()
}
pub fn idle_time(&self) -> Duration {
self.last_used.elapsed()
}
}
#[derive(Debug, Clone)]
struct LruEntry {
key: ConnectionKey,
last_access: Instant,
}
#[derive(Debug, Clone)]
pub struct PoolConfig {
pub max_connections: usize,
pub max_idle_time: Duration,
pub max_connection_age: Duration,
pub connect_timeout: Duration,
pub read_timeout: Duration,
}
impl Default for PoolConfig {
fn default() -> Self {
Self {
max_connections: crate::constants::FTP_POOL_DEFAULT_MAX_CONNECTIONS,
max_idle_time: Duration::from_secs(
crate::constants::FTP_POOL_DEFAULT_MAX_IDLE_TIME_SECS,
),
max_connection_age: Duration::from_secs(
crate::constants::FTP_POOL_DEFAULT_MAX_CONNECTION_AGE_SECS,
),
connect_timeout: Duration::from_secs(
crate::constants::FTP_POOL_DEFAULT_CONNECT_TIMEOUT_SECS,
),
read_timeout: Duration::from_secs(crate::constants::FTP_POOL_DEFAULT_READ_TIMEOUT_SECS),
}
}
}
pub struct FtpConnectionPool {
connections: Arc<Mutex<HashMap<ConnectionKey, PooledConnection>>>,
lru_order: Arc<Mutex<Vec<LruEntry>>>,
config: PoolConfig,
stats: Arc<Mutex<PoolStats>>,
}
#[derive(Debug, Clone, Default)]
pub struct PoolStats {
pub connections_created: u64,
pub connections_reused: u64,
pub connections_evicted: u64,
pub connection_failures: u64,
pub current_size: usize,
pub peak_size: usize,
}
impl FtpConnectionPool {
pub fn new(max_connections: usize) -> Self {
let config = PoolConfig {
max_connections,
..Default::default()
};
Self::with_config(config)
}
pub fn with_config(config: PoolConfig) -> Self {
Self {
connections: Arc::new(Mutex::new(HashMap::new())),
lru_order: Arc::new(Mutex::new(Vec::new())),
config,
stats: Arc::new(Mutex::new(PoolStats::default())),
}
}
pub async fn get_connection(
&self,
host: &str,
port: u16,
username: &str,
password: &str,
mode: FtpMode,
) -> Result<PooledConnection> {
let key = ConnectionKey::new(host, port, username, password);
{
let mut connections = self.connections.lock().await;
if let Some(conn) = connections.get_mut(&key) {
if conn.is_healthy(self.config.max_idle_time) {
conn.mark_used();
self.update_lru_access(&key).await;
let mut stats = self.stats.lock().await;
stats.connections_reused += 1;
debug!(
"Reusing FTP connection to {}:{} (reuse #{})",
host, port, conn.reuse_count
);
let conn = connections.remove(&key).unwrap();
return Ok(conn);
} else {
debug!("Removing stale FTP connection to {}:{}", host, port);
connections.remove(&key);
self.remove_from_lru(&key).await;
let mut stats = self.stats.lock().await;
stats.connections_evicted += 1;
}
}
}
self.evict_if_needed().await?;
debug!("Creating new FTP connection to {}:{}", host, port);
let client = FtpClient::connect(host, port, mode).await?;
{
let mut client = client;
client.login(username, password).await?;
client.set_binary_mode(true).await?;
let pooled = PooledConnection::new(client, key.clone(), mode);
let mut connections = self.connections.lock().await;
connections.insert(key.clone(), pooled);
self.add_to_lru(key.clone()).await;
let mut stats = self.stats.lock().await;
stats.connections_created += 1;
stats.current_size = connections.len();
if connections.len() > stats.peak_size {
stats.peak_size = connections.len();
}
info!(
"FTP connection pool: created new connection to {}:{}",
host, port
);
Ok(connections.remove(&key).unwrap())
}
}
pub async fn return_connection(&self, mut conn: PooledConnection) {
if !conn.is_healthy(self.config.max_idle_time) {
debug!(
"Not returning unhealthy connection to {}:{}",
conn.key.host, conn.key.port
);
let mut stats = self.stats.lock().await;
stats.connections_evicted += 1;
return;
}
if conn.age() > self.config.max_connection_age {
debug!(
"Not returning expired connection to {}:{} (age: {:?})",
conn.key.host,
conn.key.port,
conn.age()
);
let mut stats = self.stats.lock().await;
stats.connections_evicted += 1;
return;
}
conn.mark_used();
let mut connections = self.connections.lock().await;
let key = conn.key.clone();
connections.insert(key.clone(), conn);
self.update_lru_access(&key).await;
let mut stats = self.stats.lock().await;
stats.current_size = connections.len();
debug!("Returned FTP connection to pool: {}:{}", key.host, key.port);
}
async fn evict_if_needed(&self) -> Result<()> {
let mut connections = self.connections.lock().await;
while connections.len() >= self.config.max_connections {
let lru_key = self.find_lru_key().await;
if let Some(key) = lru_key {
debug!(
"Evicting LRU connection to {}:{} (pool full)",
key.host, key.port
);
connections.remove(&key);
self.remove_from_lru(&key).await;
let mut stats = self.stats.lock().await;
stats.connections_evicted += 1;
} else {
break;
}
}
Ok(())
}
async fn add_to_lru(&self, key: ConnectionKey) {
let mut lru = self.lru_order.lock().await;
lru.push(LruEntry {
key,
last_access: Instant::now(),
});
}
async fn update_lru_access(&self, key: &ConnectionKey) {
let mut lru = self.lru_order.lock().await;
if let Some(entry) = lru.iter_mut().find(|e| &e.key == key) {
entry.last_access = Instant::now();
}
}
async fn remove_from_lru(&self, key: &ConnectionKey) {
let mut lru = self.lru_order.lock().await;
lru.retain(|e| &e.key != key);
}
async fn find_lru_key(&self) -> Option<ConnectionKey> {
let lru = self.lru_order.lock().await;
lru.iter()
.min_by_key(|e| e.last_access)
.map(|e| e.key.clone())
}
pub async fn cleanup_stale(&self) {
let mut connections = self.connections.lock().await;
let mut to_remove = Vec::new();
for (key, conn) in connections.iter() {
if !conn.is_healthy(self.config.max_idle_time)
|| conn.age() > self.config.max_connection_age
{
to_remove.push(key.clone());
}
}
for key in to_remove {
connections.remove(&key);
self.remove_from_lru(&key).await;
let mut stats = self.stats.lock().await;
stats.connections_evicted += 1;
}
let mut stats = self.stats.lock().await;
stats.current_size = connections.len();
debug!(
"FTP connection pool cleanup: {} connections remaining",
connections.len()
);
}
pub async fn stats(&self) -> PoolStats {
self.stats.lock().await.clone()
}
pub async fn size(&self) -> usize {
self.connections.lock().await.len()
}
pub async fn clear(&self) {
let mut connections = self.connections.lock().await;
let count = connections.len();
connections.clear();
let mut lru = self.lru_order.lock().await;
lru.clear();
let mut stats = self.stats.lock().await;
stats.connections_evicted += count as u64;
stats.current_size = 0;
info!("FTP connection pool cleared: {} connections removed", count);
}
pub async fn has_connection(&self, host: &str, port: u16, username: &str) -> bool {
let connections = self.connections.lock().await;
connections
.keys()
.any(|k| k.host == host && k.port == port && k.username == username)
}
}
pub fn create_pool(max_connections: usize) -> Arc<FtpConnectionPool> {
Arc::new(FtpConnectionPool::new(max_connections))
}
pub fn create_custom_pool(config: PoolConfig) -> Arc<FtpConnectionPool> {
Arc::new(FtpConnectionPool::with_config(config))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_connection_key_equality() {
let key1 = ConnectionKey::new("example.com", 21, "user", "pass");
let key2 = ConnectionKey::new("example.com", 21, "user", "pass");
let key3 = ConnectionKey::new("example.com", 21, "user2", "pass");
assert_eq!(key1, key2);
assert_ne!(key1, key3);
}
#[test]
fn test_pool_config_default() {
let config = PoolConfig::default();
assert_eq!(config.max_connections, 16);
assert_eq!(config.max_idle_time, Duration::from_secs(300));
assert_eq!(config.max_connection_age, Duration::from_secs(1800));
}
#[tokio::test]
async fn test_pool_creation() {
let pool = FtpConnectionPool::new(10);
assert_eq!(pool.size().await, 0);
}
#[tokio::test]
async fn test_pool_stats_initial() {
let pool = FtpConnectionPool::new(10);
let stats = pool.stats().await;
assert_eq!(stats.connections_created, 0);
assert_eq!(stats.connections_reused, 0);
assert_eq!(stats.connections_evicted, 0);
assert_eq!(stats.current_size, 0);
}
#[tokio::test]
async fn test_pool_clear() {
let pool = FtpConnectionPool::new(10);
pool.clear().await;
assert_eq!(pool.size().await, 0);
}
#[test]
fn test_pooled_connection_health() {
let max_idle_time = Duration::from_secs(300);
let idle_time = Duration::from_secs(10);
assert!(idle_time < max_idle_time);
let idle_time_long = Duration::from_secs(400);
assert!(idle_time_long >= max_idle_time);
}
#[test]
fn test_lru_entry_creation() {
let key = ConnectionKey::new("example.com", 21, "user", "pass");
let entry = LruEntry {
key: key.clone(),
last_access: Instant::now(),
};
assert_eq!(entry.key, key);
assert!(entry.last_access.elapsed() < Duration::from_secs(1));
}
#[tokio::test]
async fn test_create_pool_returns_shared_arc() {
let pool = create_pool(10);
let pool2 = pool.clone();
assert!(Arc::ptr_eq(&pool, &pool2));
}
#[tokio::test]
async fn test_custom_pool_is_different() {
let pool1 = create_pool(10);
let pool2 = create_custom_pool(PoolConfig::default());
assert!(!Arc::ptr_eq(&pool1, &pool2));
}
#[test]
fn test_pool_stats_default() {
let stats = PoolStats::default();
assert_eq!(stats.connections_created, 0);
assert_eq!(stats.connections_reused, 0);
assert_eq!(stats.connections_evicted, 0);
assert_eq!(stats.connection_failures, 0);
assert_eq!(stats.current_size, 0);
assert_eq!(stats.peak_size, 0);
}
#[test]
fn test_connection_key_hash() {
use std::collections::HashSet;
let mut set = HashSet::new();
let key1 = ConnectionKey::new("example.com", 21, "user", "pass");
let key2 = ConnectionKey::new("example.com", 21, "user", "pass");
let key3 = ConnectionKey::new("other.com", 21, "user", "pass");
set.insert(key1.clone());
assert!(set.contains(&key2)); assert!(!set.contains(&key3)); }
}