RustProxy/src/traffic.rs

//! Bounded, asynchronous per-vhost traffic analytics.
//!
//! The request path only uses `try_send`; aggregation runs on one background
//! worker. High-cardinality dimensions use bounded Space-Saving tables so a
//! hostile path/User-Agent stream cannot grow process memory without bound.
use serde::Serialize;
use std::{
    collections::{BTreeMap, HashMap},
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc, Mutex,
    },
};
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

const QUEUE_CAPACITY: usize = 2048;
const TOP_CAPACITY: usize = 32;
const MAX_VHOSTS: usize = 64;
const MAX_LABEL_BYTES: usize = 256;

#[derive(Debug, Clone)]
pub struct Sample {
    pub vhost: String,
    pub route: Option<String>,
    pub method: String,
    pub path: String,
    pub client_ip: Option<String>,
    pub status: u16,
    pub duration_ms: u64,
    pub request_bytes: u64,
    pub response_bytes: u64,
    pub cache: &'static str,
    pub source: &'static str,
    pub user_agent: Option<String>,
    pub country: Option<String>,
}

#[derive(Debug, Clone, Serialize)]
pub struct TopEntry {
    pub value: String,
    pub count: u64,
}

#[derive(Debug, Clone, Default)]
struct TopTable {
    values: HashMap<String, u64>,
}
impl TopTable {
    fn add(&mut self, raw: &str, weight: u64) {
        if raw.is_empty() || weight == 0 {
            return;
        }
        // Probe with the borrowed label. Under production traffic most methods,
        // statuses, routes and user agents repeat; allocating a fresh String
        // for every hit creates avoidable allocator churn.
        let value = truncate(raw, MAX_LABEL_BYTES);
        if let Some(count) = self.values.get_mut(value) {
            *count = count.saturating_add(weight);
            return;
        }
        if self.values.len() < TOP_CAPACITY {
            self.values.insert(value.to_owned(), weight);
            return;
        }
        let Some((smallest, minimum)) = self
            .values
            .iter()
            .min_by_key(|(_, count)| **count)
            .map(|(value, count)| (value.clone(), *count))
        else {
            return;
        };
        self.values.remove(&smallest);
        self.values.insert(value.to_owned(), minimum.saturating_add(weight));
    }
    fn merge(&mut self, other: &Self) {
        for (value, count) in &other.values {
            self.add(value, *count);
        }
    }
    fn view(&self) -> Vec<TopEntry> {
        let mut entries = self
            .values
            .iter()
            .map(|(value, count)| TopEntry {
                value: value.clone(),
                count: *count,
            })
            .collect::<Vec<_>>();
        entries.sort_by(|a, b| b.count.cmp(&a.count).then_with(|| a.value.cmp(&b.value)));
        entries
    }
}

#[derive(Debug, Clone, Default)]
struct Counters {
    requests: u64,
    errors: u64,
    duration_ms: u64,
    request_bytes: u64,
    response_bytes: u64,
    countries: TopTable,
    user_agents: TopTable,
    user_agent_families: TopTable,
    client_ips: TopTable,
    methods: TopTable,
    method_families: TopTable,
    sources: TopTable,
    statuses: TopTable,
    routes: TopTable,
    paths: TopTable,
    cache: TopTable,
}
impl Counters {
    fn record(&mut self, sample: Sample) {
        self.requests = self.requests.saturating_add(1);
        self.errors = self.errors.saturating_add((sample.status >= 500) as u64);
        self.duration_ms = self.duration_ms.saturating_add(sample.duration_ms);
        self.request_bytes = self.request_bytes.saturating_add(sample.request_bytes);
        self.response_bytes = self.response_bytes.saturating_add(sample.response_bytes);
        self.methods.add(&sample.method, 1);
        self.method_families.add(method_family(&sample.method), 1);
        self.sources.add(sample.source, 1);
        self.statuses.add(&format!("{}xx", sample.status / 100), 1);
        self.routes.add(sample.route.as_deref().unwrap_or("(main)"), 1);
        self.paths.add(&sample.path, 1);
        self.cache.add(sample.cache, 1);
        if let Some(country) = sample.country.as_deref() {
            self.countries.add(country, 1);
        }
        if let Some(user_agent) = sample.user_agent.as_deref() {
            self.user_agents.add(user_agent, 1);
            self.user_agent_families.add(user_agent_family(user_agent), 1);
        }
        if let Some(client_ip) = sample.client_ip.as_deref() {
            self.client_ips.add(client_ip, 1);
        }
    }
    fn merge(&mut self, other: &Self) {
        self.requests = self.requests.saturating_add(other.requests);
        self.errors = self.errors.saturating_add(other.errors);
        self.duration_ms = self.duration_ms.saturating_add(other.duration_ms);
        self.request_bytes = self.request_bytes.saturating_add(other.request_bytes);
        self.response_bytes = self.response_bytes.saturating_add(other.response_bytes);
        self.countries.merge(&other.countries);
        self.user_agents.merge(&other.user_agents);
        self.user_agent_families.merge(&other.user_agent_families);
        self.client_ips.merge(&other.client_ips);
        self.methods.merge(&other.methods);
        self.method_families.merge(&other.method_families);
        self.sources.merge(&other.sources);
        self.statuses.merge(&other.statuses);
        self.routes.merge(&other.routes);
        self.paths.merge(&other.paths);
        self.cache.merge(&other.cache);
    }
    fn view(&self) -> CountersView {
        CountersView {
            requests: self.requests,
            errors: self.errors,
            duration_ms: self.duration_ms,
            request_bytes: self.request_bytes,
            response_bytes: self.response_bytes,
            mean_duration_ms: (self.requests > 0).then(|| self.duration_ms as f64 / self.requests as f64),
            countries: self.countries.view(),
            user_agents: self.user_agents.view(),
            user_agent_families: self.user_agent_families.view(),
            client_ips: self.client_ips.view(),
            methods: self.methods.view(),
            method_families: self.method_families.view(),
            sources: self.sources.view(),
            statuses: self.statuses.view(),
            routes: self.routes.view(),
            paths: self.paths.view(),
            cache: self.cache.view(),
        }
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct CountersView {
    pub requests: u64,
    pub errors: u64,
    pub duration_ms: u64,
    pub request_bytes: u64,
    pub response_bytes: u64,
    pub mean_duration_ms: Option<f64>,
    pub countries: Vec<TopEntry>,
    pub user_agents: Vec<TopEntry>,
    pub user_agent_families: Vec<TopEntry>,
    pub client_ips: Vec<TopEntry>,
    pub methods: Vec<TopEntry>,
    pub method_families: Vec<TopEntry>,
    pub sources: Vec<TopEntry>,
    pub statuses: Vec<TopEntry>,
    pub routes: Vec<TopEntry>,
    pub paths: Vec<TopEntry>,
    pub cache: Vec<TopEntry>,
}

#[derive(Debug, Clone, Serialize)]
pub struct View {
    pub scope: String,
    pub approximate_top_lists: bool,
    pub queue_capacity: usize,
    pub queue_dropped: u64,
    pub data: CountersView,
}

#[derive(Default)]
struct State {
    vhosts: BTreeMap<String, Counters>,
}

pub struct Analytics {
    state: Arc<Mutex<State>>,
    tx: mpsc::Sender<Sample>,
    dropped: Arc<AtomicU64>,
}
impl Analytics {
    pub fn start(shutdown: CancellationToken) -> Arc<Self> {
        let state = Arc::new(Mutex::new(State::default()));
        let dropped = Arc::new(AtomicU64::new(0));
        let (tx, mut rx) = mpsc::channel::<Sample>(QUEUE_CAPACITY);
        let worker_state = state.clone();
        tokio::spawn(async move {
            loop {
                tokio::select! {
                    _ = shutdown.cancelled() => break,
                    sample = rx.recv() => {
                        let Some(sample) = sample else { break; };
                        let mut state = worker_state.lock().unwrap_or_else(|e| e.into_inner());
                        if !state.vhosts.contains_key(&sample.vhost) && state.vhosts.len() >= MAX_VHOSTS {
                            continue;
                        }
                        state.vhosts.entry(sample.vhost.clone()).or_default().record(sample);
                    }
                }
            }
        });
        Arc::new(Self { state, tx, dropped })
    }

    pub fn push(&self, sample: Sample) {
        self.push_lazy(|| sample);
    }

    /// Acquire queue capacity before constructing owned analytics fields.
    /// Overflow accounting is identical to `push`, and the queue remains
    /// bounded. In saturated traffic the rejected sample allocates nothing.
    pub fn push_lazy(&self, make_sample: impl FnOnce() -> Sample) {
        match self.tx.try_reserve() {
            Ok(permit) => permit.send(make_sample()),
            Err(_) => {
                self.dropped.fetch_add(1, Ordering::Relaxed);
            }
        }
    }

    pub fn view(&self, scope: Option<&str>) -> View {
        let state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        let selected = scope.filter(|value| !value.is_empty()).unwrap_or("*");
        let data = if selected == "*" {
            let mut total = Counters::default();
            for counters in state.vhosts.values() {
                total.merge(counters);
            }
            total.view()
        } else {
            state.vhosts.get(selected).cloned().unwrap_or_default().view()
        };
        View {
            scope: selected.to_owned(),
            approximate_top_lists: true,
            queue_capacity: QUEUE_CAPACITY,
            queue_dropped: self.dropped.load(Ordering::Relaxed),
            data,
        }
    }

    pub fn summaries(&self) -> BTreeMap<String, CountersView> {
        self.state
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .vhosts
            .iter()
            .map(|(id, counters)| (id.clone(), counters.view()))
            .collect()
    }

    pub fn reset(&self, scope: Option<&str>) -> usize {
        let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        match scope.filter(|value| !value.is_empty() && *value != "*") {
            Some(scope) => {
                if state.vhosts.remove(scope).is_some() { 1 } else { 0 }
            },
            None => {
                let count = state.vhosts.len();
                state.vhosts.clear();
                count
            }
        }
    }
}

fn method_family(method: &str) -> &'static str {
    match method {
        "GET" => "GET",
        "HEAD" => "HEAD",
        "POST" => "POST",
        "PUT" => "PUT",
        "PATCH" => "PATCH",
        "DELETE" => "DELETE",
        "OPTIONS" => "OPTIONS",
        "CONNECT" => "CONNECT",
        "TRACE" => "TRACE",
        _ => "OTHER",
    }
}

fn user_agent_family(value: &str) -> &'static str {
    let lower = value.to_ascii_lowercase();
    if lower.contains("bot") || lower.contains("spider") || lower.contains("crawler") || lower.contains("slurp") {
        "bot"
    } else if lower.contains("curl/") {
        "curl"
    } else if lower.contains("wget/") {
        "wget"
    } else if lower.contains("python-requests") || lower.contains("python/") || lower.contains("aiohttp/") {
        "python"
    } else if lower.contains("go-http-client") {
        "go"
    } else if lower.contains("java/") || lower.contains("okhttp/") {
        "java"
    } else if lower.contains("postmanruntime/") {
        "postman"
    } else if lower.contains("edg/") || lower.contains("edge/") {
        "edge"
    } else if lower.contains("opr/") || lower.contains("opera") {
        "opera"
    } else if lower.contains("chrome/") || lower.contains("crios/") {
        "chrome"
    } else if lower.contains("firefox/") || lower.contains("fxios/") {
        "firefox"
    } else if lower.contains("safari/") {
        "safari"
    } else {
        "other"
    }
}

fn truncate(value: &str, max_bytes: usize) -> &str {
    if value.len() <= max_bytes {
        return value;
    }
    let mut end = max_bytes;
    while !value.is_char_boundary(end) {
        end -= 1;
    }
    &value[..end]
}

#[cfg(test)]
include!("traffic/tests.rs");