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

rws_proxy/proxy/
pool.rs
