use crate::error::{Error, Result};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::TcpStream;
use tokio::sync::Mutex;
const DEFAULT_MAX_CONNECTIONS_PER_HOST: usize = 32;
const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(90);
const DEFAULT_CONNECTION_TIMEOUT: Duration = Duration::from_secs(10);
struct PooledConnection {
stream: TcpStream,
last_used: Instant,
}
impl PooledConnection {
fn new(stream: TcpStream) -> Self {
Self {
stream,
last_used: Instant::now(),
}
}
fn is_expired(&self, idle_timeout: Duration) -> bool {
self.last_used.elapsed() > idle_timeout
}
fn update_last_used(&mut self) {
self.last_used = Instant::now();
}
}
#[derive(Clone, Debug)]
pub struct PoolConfig {
pub max_connections_per_host: usize,
pub idle_timeout: Duration,
pub connection_timeout: Duration,
pub keep_alive: bool,
}
impl Default for PoolConfig {
fn default() -> Self {
Self {
max_connections_per_host: DEFAULT_MAX_CONNECTIONS_PER_HOST,
idle_timeout: DEFAULT_IDLE_TIMEOUT,
connection_timeout: DEFAULT_CONNECTION_TIMEOUT,
keep_alive: true,
}
}
}
pub struct ConnectionPool {
config: PoolConfig,
idle_connections: Arc<Mutex<HashMap<String, Vec<PooledConnection>>>>,
active_counts: Arc<Mutex<HashMap<String, usize>>>,
}
impl ConnectionPool {
pub fn new() -> Self {
Self::with_config(PoolConfig::default())
}
pub fn with_config(config: PoolConfig) -> Self {
Self {
config,
idle_connections: Arc::new(Mutex::new(HashMap::new())),
active_counts: Arc::new(Mutex::new(HashMap::new())),
}
}
pub async fn get_connection(&self, host: &str, port: u16) -> Result<TcpStream> {
let key = format!("{}:{}", host, port);
{
let mut idle = self.idle_connections.lock().await;
if let Some(connections) = idle.get_mut(&key) {
connections.retain(|conn| !conn.is_expired(self.config.idle_timeout));
while let Some(mut conn) = connections.pop() {
conn.update_last_used();
let mut active = self.active_counts.lock().await;
*active.entry(key.clone()).or_insert(0) += 1;
return Ok(conn.stream);
}
}
}
{
let active = self.active_counts.lock().await;
let count = active.get(&key).copied().unwrap_or(0);
if count >= self.config.max_connections_per_host {
return Err(Error::TooManyConnections {
host: key,
max: self.config.max_connections_per_host,
});
}
}
let addr = format!("{}:{}", host, port);
let stream = tokio::time::timeout(
self.config.connection_timeout,
TcpStream::connect(&addr),
)
.await
.map_err(|_| Error::ConnectionTimeout { host: addr.clone() })?
.map_err(|e| Error::ConnectionFailed {
host: addr,
source: e,
})?;
if self.config.keep_alive {
#[cfg(unix)]
{
use std::os::unix::io::AsRawFd;
let fd = stream.as_raw_fd();
unsafe {
let optval: libc::c_int = 1;
libc::setsockopt(
fd,
libc::SOL_SOCKET,
libc::SO_KEEPALIVE,
&optval as *const _ as *const libc::c_void,
std::mem::size_of_val(&optval) as libc::socklen_t,
);
}
}
#[cfg(windows)]
{
use std::os::windows::io::AsRawSocket;
let socket = stream.as_raw_socket();
unsafe {
let optval: u32 = 1;
windows_sys::Win32::Networking::WinSock::setsockopt(
socket as usize,
windows_sys::Win32::Networking::WinSock::SOL_SOCKET,
windows_sys::Win32::Networking::WinSock::SO_KEEPALIVE,
&optval as *const _ as *const u8,
std::mem::size_of_val(&optval) as i32,
);
}
}
}
let mut active = self.active_counts.lock().await;
*active.entry(key).or_insert(0) += 1;
Ok(stream)
}
pub async fn return_connection(&self, host: &str, port: u16, stream: TcpStream) {
let key = format!("{}:{}", host, port);
{
let mut active = self.active_counts.lock().await;
if let Some(count) = active.get_mut(&key) {
*count = count.saturating_sub(1);
}
}
let mut idle = self.idle_connections.lock().await;
let connections = idle.entry(key.clone()).or_insert_with(Vec::new);
if connections.len() < self.config.max_connections_per_host {
connections.push(PooledConnection::new(stream));
}
}
pub async fn cleanup_expired(&self) {
let mut idle = self.idle_connections.lock().await;
for connections in idle.values_mut() {
connections.retain(|conn| !conn.is_expired(self.config.idle_timeout));
}
idle.retain(|_, v| !v.is_empty());
}
pub async fn stats(&self) -> PoolStats {
let idle = self.idle_connections.lock().await;
let active = self.active_counts.lock().await;
let total_idle: usize = idle.values().map(|v| v.len()).sum();
let total_active: usize = active.values().sum();
PoolStats {
idle_connections: total_idle,
active_connections: total_active,
hosts: idle.len(),
}
}
pub async fn close_all(&self) {
let mut idle = self.idle_connections.lock().await;
idle.clear();
let mut active = self.active_counts.lock().await;
active.clear();
}
}
impl Default for ConnectionPool {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct PoolStats {
pub idle_connections: usize,
pub active_connections: usize,
pub hosts: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_pool_creation() {
let pool = ConnectionPool::new();
let stats = pool.stats().await;
assert_eq!(stats.idle_connections, 0);
assert_eq!(stats.active_connections, 0);
}
#[tokio::test]
async fn test_pool_config() {
let config = PoolConfig {
max_connections_per_host: 50,
idle_timeout: Duration::from_secs(120),
connection_timeout: Duration::from_secs(5),
keep_alive: true,
};
let pool = ConnectionPool::with_config(config);
assert_eq!(pool.config.max_connections_per_host, 50);
}
#[tokio::test]
async fn test_cleanup_expired() {
let pool = ConnectionPool::new();
pool.cleanup_expired().await;
let stats = pool.stats().await;
assert_eq!(stats.idle_connections, 0);
}
}