rws_proxy/proxy/
pool.rs

1use std::collections::{HashMap, VecDeque};
2use std::net::TcpStream;
3use std::sync::Mutex;
4use std::time::{Duration, Instant};
5
6const DEFAULT_MAX_IDLE: usize = 8;
7/// How long the proxy's own pools keep an idle backend connection: shorter
8/// than the 30 s after which an rws backend closes an idle keep-alive
9/// connection, so a pooled one is normally still open. Backends that close
10/// sooner are caught by the liveness check before reuse (`UpstreamIo::peer_closed`).
11pub(crate) const IDLE_TIMEOUT: Duration = Duration::from_secs(25);
12
13/// Per-backend HTTP/1.1 connection pool of plain TCP connections — the pool
14/// `ReverseProxy::with_pool` takes. (HTTPS backends use an internal
15/// `Pool` of TLS streams.)
16pub type ConnPool = Pool<TcpStream>;
17
18/// Per-backend HTTP/1.1 connection pool, generic over the connection type.
19///
20/// Holds idle connections keyed by `"host:port"`.  When a
21/// backend responds with `Connection: keep-alive` (or the HTTP/1.1 default),
22/// the stream is returned here and reused for the next request to the same
23/// backend, eliminating the TCP-handshake cost and reducing ephemeral-port
24/// exhaustion under load.
25///
26/// # Thread safety
27///
28/// All methods take `&self` and are safe to call from multiple threads.
29/// The inner map is protected by a `Mutex`.
30///
31/// # Example
32///
33/// ```rust,ignore
34/// # extern crate rws_proxy as rust_web_server;
35/// use std::sync::Arc;
36/// use std::time::Duration;
37/// use rust_web_server::proxy::ConnPool;
38/// use rust_web_server::proxy::ReverseProxy;
39///
40/// let pool = Arc::new(ConnPool::new(16, Duration::from_secs(30)));
41/// let _proxy = ReverseProxy::new(["http://backend:8080"])
42///     .with_pool(Arc::clone(&pool));
43/// ```
44pub struct Pool<S> {
45    inner: Mutex<HashMap<String, VecDeque<PoolEntry<S>>>>,
46    max_idle: usize,
47    idle_timeout: Duration,
48}
49
50struct PoolEntry<S> {
51    stream: S,
52    added: Instant,
53}
54
55impl<S> Pool<S> {
56    /// Create a pool with the given per-backend idle limit and idle timeout.
57    pub fn new(max_idle: usize, idle_timeout: Duration) -> Self {
58        Pool {
59            inner: Mutex::new(HashMap::new()),
60            max_idle,
61            idle_timeout,
62        }
63    }
64
65    /// Create a pool with defaults: 8 idle connections per backend, 25-second timeout.
66    pub fn new_default() -> Self {
67        Self::new(DEFAULT_MAX_IDLE, IDLE_TIMEOUT)
68    }
69
70    /// Try to acquire an idle connection for `key = "host:port"`.
71    ///
72    /// Stale entries (older than `idle_timeout`) are discarded automatically.
73    /// Returns `None` if no usable connection is available.
74    pub fn acquire(&self, key: &str) -> Option<S> {
75        let mut map = self.inner.lock().unwrap_or_else(|e| e.into_inner());
76        let queue = map.get_mut(key)?;
77        let now = Instant::now();
78        while let Some(entry) = queue.pop_front() {
79            if now.duration_since(entry.added) < self.idle_timeout {
80                return Some(entry.stream);
81            }
82            // stale — drop, which closes the TCP connection
83        }
84        None
85    }
86
87    /// Return a keep-alive connection to the pool.
88    ///
89    /// If the backend slot is already at `max_idle`, the stream is dropped
90    /// (the TCP connection closes) rather than exceeding the limit.
91    pub fn release(&self, key: &str, stream: S) {
92        let mut map = self.inner.lock().unwrap_or_else(|e| e.into_inner());
93        // The key is copied only the first time a backend is seen.
94        if !map.contains_key(key) {
95            map.insert(key.to_string(), Default::default());
96        }
97        let Some(queue) = map.get_mut(key) else { return };
98        if queue.len() < self.max_idle {
99            queue.push_back(PoolEntry { stream, added: Instant::now() });
100        }
101        // over limit — stream dropped here, closing the connection
102    }
103
104    /// Total idle connections across all backends (useful for testing/metrics).
105    pub fn idle_count(&self) -> usize {
106        let map = self.inner.lock().unwrap_or_else(|e| e.into_inner());
107        map.values().map(|q| q.len()).sum()
108    }
109}
rws_proxy/proxy/pool.rs

rws_proxy/proxy/
pool.rs
