v0.25.0

src/traffic.rs

@@ -0,0 +1,343 @@
+//! 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 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;
+        }
+        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, 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, 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,
+    client_ips: TopTable,
+    methods: 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.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);
+        }
+        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.client_ips.merge(&other.client_ips);
+        self.methods.merge(&other.methods);
+        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(),
+            client_ips: self.client_ips.view(),
+            methods: self.methods.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 client_ips: Vec<TopEntry>,
+    pub methods: 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) {
+        if self.tx.try_send(sample).is_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 truncate(value: &str, max_bytes: usize) -> String {
+    if value.len() <= max_bytes {
+        return value.to_owned();
+    }
+    let mut end = max_bytes;
+    while !value.is_char_boundary(end) {
+        end -= 1;
+    }
+    value[..end].to_owned()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn top_table_is_bounded() {
+        let mut table = TopTable::default();
+        for i in 0..1000 {
+            table.add(&format!("value-{i}"), 1);
+        }
+        assert_eq!(table.values.len(), TOP_CAPACITY);
+    }
+
+    #[test]
+    fn labels_are_truncated_at_utf8_boundaries() {
+        let value = "ą".repeat(200);
+        let trimmed = truncate(&value, 255);
+        assert!(trimmed.len() <= 255);
+        assert!(trimmed.is_char_boundary(trimmed.len()));
+    }
+
+    #[test]
+    fn client_ip_is_exposed_in_the_bounded_top_view() {
+        let mut counters = Counters::default();
+        counters.record(Sample {
+            vhost: "site".into(),
+            route: None,
+            method: "GET".into(),
+            path: "/".into(),
+            client_ip: Some("203.0.113.7".into()),
+            status: 200,
+            duration_ms: 2,
+            request_bytes: 10,
+            response_bytes: 20,
+            cache: "MISS",
+            user_agent: Some("test-agent".into()),
+            country: Some("PL".into()),
+        });
+        let view = counters.view();
+        assert_eq!(view.client_ips.len(), 1);
+        assert_eq!(view.client_ips[0].value, "203.0.113.7");
+        assert_eq!(view.client_ips[0].count, 1);
+    }
+}