30 lines
818 B
Rust
30 lines
818 B
Rust
|
use std::{collections::HashMap, net::SocketAddr, sync::Arc};
|
||
|
|
|
||
|
|
use tokio::net::TcpStream;
|
||
|
|
use tokio::sync::Mutex;
|
||
|
|
|
||
|
|
#[derive(Clone)]
|
||
|
|
pub struct ConnectionPool {
|
||
|
|
pool: Arc<Mutex<HashMap<SocketAddr, Arc<Mutex<TcpStream>>>>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl ConnectionPool {
|
||
|
|
pub fn new(size: usize) -> Self {
|
||
|
|
Self {
|
||
|
|
pool: Arc::new(Mutex::new(HashMap::with_capacity(size))),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub async fn get_or_init(&self, addr: &SocketAddr) -> std::io::Result<Arc<Mutex<TcpStream>>> {
|
||
|
|
let mut pool = self.pool.lock().await;
|
||
|
|
match pool.get(addr).cloned() {
|
||
|
|
Some(tcp) => Ok(tcp),
|
||
|
|
None => {
|
||
|
|
let tcp = Arc::new(Mutex::new(TcpStream::connect(addr).await?));
|
||
|
|
pool.insert(*addr, tcp.clone());
|
||
|
|
Ok(tcp)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|