use std::collections::{HashMap, VecDeque};
use std::net::TcpStream;
use std::sync::Mutex;
use std::time::{Duration, Instant};
const DEFAULT_MAX_IDLE: usize = 8;
const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
pub struct ConnPool {
inner: Mutex<HashMap<String, VecDeque<PoolEntry>>>,
max_idle: usize,
idle_timeout: Duration,
}
struct PoolEntry {
stream: TcpStream,
added: Instant,
}
impl ConnPool {
pub fn new(max_idle: usize, idle_timeout: Duration) -> Self {
ConnPool {
inner: Mutex::new(HashMap::new()),
max_idle,
idle_timeout,
}
}
pub fn new_default() -> Self {
Self::new(DEFAULT_MAX_IDLE, DEFAULT_IDLE_TIMEOUT)
}
pub fn acquire(&self, key: &str) -> Option<TcpStream> {
let mut map = self.inner.lock().unwrap_or_else(|e| e.into_inner());
let queue = map.get_mut(key)?;
let now = Instant::now();
while let Some(entry) = queue.pop_front() {
if now.duration_since(entry.added) < self.idle_timeout {
return Some(entry.stream);
}
}
None
}
pub fn release(&self, key: &str, stream: TcpStream) {
let mut map = self.inner.lock().unwrap_or_else(|e| e.into_inner());
let queue = map.entry(key.to_string()).or_default();
if queue.len() < self.max_idle {
queue.push_back(PoolEntry { stream, added: Instant::now() });
}
}
pub fn idle_count(&self) -> usize {
let map = self.inner.lock().unwrap_or_else(|e| e.into_inner());
map.values().map(|q| q.len()).sum()
}
}