connection_pool/
tcp.rs

1use crate::{CleanupConfig, ConnectionManager, ConnectionPool, PooledStream};
2use std::{sync::Arc, time::Duration};
3use tokio::net::{TcpStream, ToSocketAddrs};
4
5/// Example implementation for TcpStream
6pub struct TcpConnectionManager<A: ToSocketAddrs + Send + Sync + Clone + 'static> {
7    pub address: A,
8}
9
10impl<A> ConnectionManager for TcpConnectionManager<A>
11where
12    A: ToSocketAddrs + Send + Sync + Clone + 'static,
13{
14    type Connection = TcpStream;
15    type Error = std::io::Error;
16    type CreateFut = std::pin::Pin<Box<dyn Future<Output = Result<TcpStream, Self::Error>> + Send>>;
17    type ValidFut<'a> = std::pin::Pin<Box<dyn Future<Output = bool> + Send + 'a>>;
18
19    fn create_connection(&self) -> Self::CreateFut {
20        let addr = self.address.clone();
21        Box::pin(async move { TcpStream::connect(addr).await })
22    }
23
24    fn is_valid<'a>(&'a self, stream: &'a Self::Connection) -> Self::ValidFut<'a> {
25        Box::pin(async move {
26            stream
27                .ready(tokio::io::Interest::READABLE | tokio::io::Interest::WRITABLE)
28                .await
29                .is_ok()
30        })
31    }
32}
33
34/// Convenience type aliases for TCP connections
35pub type TcpConnectionPool<A = std::net::SocketAddr> = ConnectionPool<TcpConnectionManager<A>>;
36
37/// Pooled TCP stream
38pub type TcpPooledStream<A = std::net::SocketAddr> = PooledStream<TcpConnectionManager<A>>;
39
40impl<A> TcpConnectionPool<A>
41where
42    A: ToSocketAddrs + Send + Sync + Clone + 'static,
43{
44    /// Create a new TCP connection pool
45    pub fn new_tcp(
46        max_size: Option<usize>,
47        max_idle_time: Option<Duration>,
48        connection_timeout: Option<Duration>,
49        cleanup_config: Option<CleanupConfig>,
50        address: A,
51    ) -> Arc<Self> {
52        log::info!("Creating TCP connection pool");
53        let manager = TcpConnectionManager { address };
54        Self::new(max_size, max_idle_time, connection_timeout, cleanup_config, manager)
55    }
56}