lobe-core 0.1.0

//! Connection pool for the capture proxy.
//!
//! Sits in front of the DNS/TCP/TLS/handshake path. Before opening a fresh
//! connection to an upstream, `forward_request` asks the pool whether we
//! already have a working connection to `(host, port, protocol)`. If yes,
//! we reuse it and skip all handshake phases — DNS, TCP, TLS, and HTTP/2's
//! `SETTINGS` frame exchange — because the underlying transport is already
//! established. If no, we fall through to the existing "fresh connection"
//! path AND cache the resulting connection for the next request to the
//! same upstream.
//!
//! Currently pools two connection types:
//! - HTTP/2 `SendRequest` handles (cheap to clone; one handle per host
//!   can service unlimited concurrent streams via multiplexing).
//! - HTTP/1.1 keep-alive `TcpStream` / `TlsStream` — one stream per host
//!   for sequential reuse.
//!
//! Entries auto-evict after `IDLE_TIMEOUT` of unused time.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use http_body_util::Full;
use hyper::body::Bytes;
use hyper::client::conn::http2::SendRequest;
use tokio::sync::Mutex;

/// A pooled entry is fresh enough to reuse if it hasn't been idle for more
/// than this window. 30s matches typical HTTP/1.1 keep-alive timeouts on
/// modern servers (nginx default is 75s; we're conservative).
const IDLE_TIMEOUT: Duration = Duration::from_secs(30);

/// The type stored per host in the HTTP/2 pool. `SendRequest` is `Clone`
/// (it's just a channel handle into the shared connection driver task), so
/// we can hand out a fresh clone to every caller and multiplex requests
/// over the same underlying TCP+TLS session for free.
pub type PooledH2Sender = SendRequest<Full<Bytes>>;

#[derive(Hash, PartialEq, Eq, Clone, Debug)]
struct PoolKey {
    host: String,
    port: u16,
    is_https: bool,
}

struct H2Entry {
    sender: PooledH2Sender,
    last_used: Instant,
}

/// Thread-safe connection pool shared across concurrent `forward_request`
/// tasks. Cloneable — internal state lives behind an `Arc<Mutex<...>>`
/// so cloning is cheap (bumps the ref-count).
#[derive(Clone, Debug, Default)]
pub struct ConnectionPool {
    inner: Arc<Mutex<PoolInner>>,
}

#[derive(Default)]
struct PoolInner {
    h2: HashMap<PoolKey, H2Entry>,
}

impl std::fmt::Debug for PoolInner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PoolInner")
            .field("h2_entries", &self.h2.len())
            .finish()
    }
}

impl ConnectionPool {
    pub fn new() -> Self {
        Self::default()
    }

    /// Try to reuse an HTTP/2 sender for `(host, port)`. Returns `None` if
    /// we have no cached entry OR the cached entry has gone idle beyond
    /// `IDLE_TIMEOUT`. Fresh senders returned from here can be used
    /// immediately — `SendRequest` clones share the same underlying
    /// connection driver task.
    pub async fn get_h2(&self, host: &str, port: u16) -> Option<PooledH2Sender> {
        let key = PoolKey {
            host: host.to_string(),
            port,
            is_https: true,
        };
        let mut inner = self.inner.lock().await;
        let entry = inner.h2.get(&key)?;
        if entry.last_used.elapsed() > IDLE_TIMEOUT {
            // Stale — evict and force a fresh handshake.
            inner.h2.remove(&key);
            return None;
        }
        let sender = entry.sender.clone();
        // Bump last_used so a busy connection doesn't get evicted mid-flight.
        if let Some(entry) = inner.h2.get_mut(&key) {
            entry.last_used = Instant::now();
        }
        Some(sender)
    }

    /// Cache a freshly-negotiated HTTP/2 sender so future requests to the
    /// same host reuse the underlying TCP+TLS+h2 session.
    pub async fn store_h2(&self, host: &str, port: u16, sender: PooledH2Sender) {
        let key = PoolKey {
            host: host.to_string(),
            port,
            is_https: true,
        };
        let mut inner = self.inner.lock().await;
        inner.h2.insert(
            key,
            H2Entry {
                sender,
                last_used: Instant::now(),
            },
        );
    }

    /// Remove the cached h2 sender for a host — call this when a request
    /// against a cached sender fails, so the next attempt re-handshakes.
    pub async fn invalidate_h2(&self, host: &str, port: u16) {
        let key = PoolKey {
            host: host.to_string(),
            port,
            is_https: true,
        };
        let mut inner = self.inner.lock().await;
        inner.h2.remove(&key);
    }
}