rust_web_server/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/// use std::sync::Arc;
31/// use std::time::Duration;
32/// use rust_web_server::proxy::ConnPool;
33/// use rust_web_server::proxy::ReverseProxy;
34///
35/// let pool = Arc::new(ConnPool::new(16, Duration::from_secs(30)));
36/// let _proxy = ReverseProxy::new(["http://backend:8080"])
37///     .with_pool(Arc::clone(&pool));
38/// ```
39pub struct Pool<S> {
40    inner: Mutex<HashMap<String, VecDeque<PoolEntry<S>>>>,
41    max_idle: usize,
42    idle_timeout: Duration,
43}
44
45struct PoolEntry<S> {
46    stream: S,
47    added: Instant,
48}
49
50impl<S> Pool<S> {
51    /// Create a pool with the given per-backend idle limit and idle timeout.
52    pub fn new(max_idle: usize, idle_timeout: Duration) -> Self {
53        Pool {
54            inner: Mutex::new(HashMap::new()),
55            max_idle,
56            idle_timeout,
57        }
58    }
59
60    /// Create a pool with defaults: 8 idle connections per backend, 60-second timeout.
61    pub fn new_default() -> Self {
62        Self::new(DEFAULT_MAX_IDLE, DEFAULT_IDLE_TIMEOUT)
63    }
64
65    /// Try to acquire an idle connection for `key = "host:port"`.
66    ///
67    /// Stale entries (older than `idle_timeout`) are discarded automatically.
68    /// Returns `None` if no usable connection is available.
69    pub fn acquire(&self, key: &str) -> Option<S> {
70        let mut map = self.inner.lock().unwrap_or_else(|e| e.into_inner());
71        let queue = map.get_mut(key)?;
72        let now = Instant::now();
73        while let Some(entry) = queue.pop_front() {
74            if now.duration_since(entry.added) < self.idle_timeout {
75                return Some(entry.stream);
76            }
77            // stale — drop, which closes the TCP connection
78        }
79        None
80    }
81
82    /// Return a keep-alive connection to the pool.
83    ///
84    /// If the backend slot is already at `max_idle`, the stream is dropped
85    /// (the TCP connection closes) rather than exceeding the limit.
86    pub fn release(&self, key: &str, stream: S) {
87        let mut map = self.inner.lock().unwrap_or_else(|e| e.into_inner());
88        // The key is copied only the first time a backend is seen.
89        if !map.contains_key(key) {
90            map.insert(key.to_string(), Default::default());
91        }
92        let Some(queue) = map.get_mut(key) else { return };
93        if queue.len() < self.max_idle {
94            queue.push_back(PoolEntry { stream, added: Instant::now() });
95        }
96        // over limit — stream dropped here, closing the connection
97    }
98
99    /// Total idle connections across all backends (useful for testing/metrics).
100    pub fn idle_count(&self) -> usize {
101        let map = self.inner.lock().unwrap_or_else(|e| e.into_inner());
102        map.values().map(|q| q.len()).sum()
103    }
104}
rust_web_server/proxy/pool.rs

rust_web_server/proxy/
pool.rs
