use crate::network::{NetworkConfig, NetworkError};
use std::collections::BTreeMap;
use std::net::{SocketAddr, TcpStream};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use std::collections::HashMap;
struct ConnectionInfo {
stream: TcpStream,
created_at: Instant,
last_used: Instant,
remote_addr: SocketAddr,
}
#[derive(Debug, Clone)]
pub struct ConnectionPoolStats {
pub total_connections: usize,
pub active_connections: usize,
pub idle_connections: usize,
pub reuse_count: u64,
pub new_connection_count: u64,
pub timeout_closed: u64,
}
pub struct ConnectionPool {
config: NetworkConfig,
pools: Arc<Mutex<HashMap<SocketAddr, Vec<ConnectionInfo>>>>,
stats: Arc<Mutex<ConnectionPoolStats>>,
}
impl ConnectionPool {
pub fn new(config: NetworkConfig) -> Result<Self, NetworkError> {
Ok(Self {
config,
pools: Arc::new(Mutex::new(HashMap::new())),
stats: Arc::new(Mutex::new(ConnectionPoolStats {
total_connections: 0,
active_connections: 0,
idle_connections: 0,
reuse_count: 0,
new_connection_count: 0,
timeout_closed: 0,
})),
})
}
pub fn get_connection(&mut self, addr: &str) -> Result<Option<TcpStream>, NetworkError> {
let socket_addr: SocketAddr = addr.parse()
.map_err(|_| NetworkError::Connection(format!("Invalid address: {}", addr)))?;
let mut pools = self.pools.lock().unwrap();
if let Some(connections) = pools.get_mut(&socket_addr) {
connections.retain(|conn| {
conn.last_used.elapsed() < Duration::from_secs(300) });
if let Some(conn_info) = connections.pop() {
{
let mut stats = self.stats.lock().unwrap();
stats.reuse_count += 1;
stats.active_connections += 1;
stats.idle_connections = connections.len();
stats.timeout_closed = 0; }
let stream: _ = conn_info.stream;
stream.set_nonblocking(true)?;
return Ok(Some(stream));
}
}
let stream: _ = TcpStream::connect_timeout(&socket_addr, Duration::from_secs(10))?;
stream.set_nonblocking(true)?;
{
let mut stats = self.stats.lock().unwrap();
stats.new_connection_count += 1;
stats.active_connections += 1;
}
Ok(Some(stream))
}
pub fn release_connection(&mut self, addr: &str, conn: TcpStream) -> Result<(), NetworkError> {
let socket_addr: SocketAddr = addr.parse()
.map_err(|_| NetworkError::Connection(format!("Invalid address: {}", addr)))?;
let mut pools = self.pools.lock().unwrap();
let pool_size_limit: _ = self.config.pool_size;
let connections: _ = pools.entry(socket_addr).or_insert_with(Vec::new);
if connections.len() >= pool_size_limit {
return Ok(());
}
connections.push(ConnectionInfo {
stream: conn,
created_at: Instant::now(),
last_used: Instant::now(),
remote_addr: socket_addr,
});
{
let mut stats = self.stats.lock().unwrap();
stats.active_connections = stats.active_connections.saturating_sub(1);
}
Ok(())
}
pub fn warm_up(&mut self, addr: &str, count: usize) -> Result<(), NetworkError> {
for _ in 0..count {
if let Some(stream) = self.get_connection(addr)? {
let _: _ = self.release_connection(addr, stream);
}
}
Ok(())
}
pub fn get_stats(&self) -> ConnectionPoolStats {
let stats: _ = self.stats.lock().unwrap();
let pools: _ = self.pools.lock().unwrap();
let total_idle: _ = pools.values().map(|v| v.len()).sum();
ConnectionPoolStats {
total_connections: stats.total_connections,
active_connections: stats.active_connections,
idle_connections: total_idle,
reuse_count: stats.reuse_count,
new_connection_count: stats.new_connection_count,
timeout_closed: stats.timeout_closed,
}
}
pub fn clear(&mut self) {
let mut pools = self.pools.lock().unwrap();
pools.clear();
let mut stats = self.stats.lock().unwrap();
stats.active_connections = 0;
stats.idle_connections = 0;
}
}
impl Default for ConnectionPool {
fn default() -> Self {
ConnectionPool::new(NetworkConfig::default()).unwrap()
}
}