lobe_core/engine/
pool.rs

1//! Connection pool for the capture proxy.
2//!
3//! Sits in front of the DNS/TCP/TLS/handshake path. Before opening a fresh
4//! connection to an upstream, `forward_request` asks the pool whether we
5//! already have a working connection to `(host, port, protocol)`. If yes,
6//! we reuse it and skip all handshake phases — DNS, TCP, TLS, and HTTP/2's
7//! `SETTINGS` frame exchange — because the underlying transport is already
8//! established. If no, we fall through to the existing "fresh connection"
9//! path AND cache the resulting connection for the next request to the
10//! same upstream.
11//!
12//! Currently pools two connection types:
13//! - HTTP/2 `SendRequest` handles (cheap to clone; one handle per host
14//!   can service unlimited concurrent streams via multiplexing).
15//! - HTTP/1.1 keep-alive `TcpStream` / `TlsStream` — one stream per host
16//!   for sequential reuse.
17//!
18//! Entries auto-evict after `IDLE_TIMEOUT` of unused time.
19
20use std::collections::HashMap;
21use std::sync::Arc;
22use std::time::{Duration, Instant};
23
24use http_body_util::Full;
25use hyper::body::Bytes;
26use hyper::client::conn::http2::SendRequest;
27use tokio::sync::Mutex;
28
29/// A pooled entry is fresh enough to reuse if it hasn't been idle for more
30/// than this window. 30s matches typical HTTP/1.1 keep-alive timeouts on
31/// modern servers (nginx default is 75s; we're conservative).
32const IDLE_TIMEOUT: Duration = Duration::from_secs(30);
33
34/// The type stored per host in the HTTP/2 pool. `SendRequest` is `Clone`
35/// (it's just a channel handle into the shared connection driver task), so
36/// we can hand out a fresh clone to every caller and multiplex requests
37/// over the same underlying TCP+TLS session for free.
38pub type PooledH2Sender = SendRequest<Full<Bytes>>;
39
40#[derive(Hash, PartialEq, Eq, Clone, Debug)]
41struct PoolKey {
42    host: String,
43    port: u16,
44    is_https: bool,
45}
46
47struct H2Entry {
48    sender: PooledH2Sender,
49    last_used: Instant,
50}
51
52/// Thread-safe connection pool shared across concurrent `forward_request`
53/// tasks. Cloneable — internal state lives behind an `Arc<Mutex<...>>`
54/// so cloning is cheap (bumps the ref-count).
55#[derive(Clone, Debug, Default)]
56pub struct ConnectionPool {
57    inner: Arc<Mutex<PoolInner>>,
58}
59
60#[derive(Default)]
61struct PoolInner {
62    h2: HashMap<PoolKey, H2Entry>,
63}
64
65impl std::fmt::Debug for PoolInner {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        f.debug_struct("PoolInner")
68            .field("h2_entries", &self.h2.len())
69            .finish()
70    }
71}
72
73impl ConnectionPool {
74    pub fn new() -> Self {
75        Self::default()
76    }
77
78    /// Try to reuse an HTTP/2 sender for `(host, port)`. Returns `None` if
79    /// we have no cached entry OR the cached entry has gone idle beyond
80    /// `IDLE_TIMEOUT`. Fresh senders returned from here can be used
81    /// immediately — `SendRequest` clones share the same underlying
82    /// connection driver task.
83    pub async fn get_h2(&self, host: &str, port: u16) -> Option<PooledH2Sender> {
84        let key = PoolKey {
85            host: host.to_string(),
86            port,
87            is_https: true,
88        };
89        let mut inner = self.inner.lock().await;
90        let entry = inner.h2.get(&key)?;
91        if entry.last_used.elapsed() > IDLE_TIMEOUT {
92            // Stale — evict and force a fresh handshake.
93            inner.h2.remove(&key);
94            return None;
95        }
96        let sender = entry.sender.clone();
97        // Bump last_used so a busy connection doesn't get evicted mid-flight.
98        if let Some(entry) = inner.h2.get_mut(&key) {
99            entry.last_used = Instant::now();
100        }
101        Some(sender)
102    }
103
104    /// Cache a freshly-negotiated HTTP/2 sender so future requests to the
105    /// same host reuse the underlying TCP+TLS+h2 session.
106    pub async fn store_h2(&self, host: &str, port: u16, sender: PooledH2Sender) {
107        let key = PoolKey {
108            host: host.to_string(),
109            port,
110            is_https: true,
111        };
112        let mut inner = self.inner.lock().await;
113        inner.h2.insert(
114            key,
115            H2Entry {
116                sender,
117                last_used: Instant::now(),
118            },
119        );
120    }
121
122    /// Remove the cached h2 sender for a host — call this when a request
123    /// against a cached sender fails, so the next attempt re-handshakes.
124    pub async fn invalidate_h2(&self, host: &str, port: u16) {
125        let key = PoolKey {
126            host: host.to_string(),
127            port,
128            is_https: true,
129        };
130        let mut inner = self.inner.lock().await;
131        inner.h2.remove(&key);
132    }
133}
lobe_core/engine/pool.rs

lobe_core/engine/
pool.rs
