RustProxy/src/cache/stats.rs

use arc_swap::ArcSwap;

const DOMAIN_STATS_LIMIT: usize = 4096;
const FAST_REASONS: [&str; 4] = [
    "cache_disabled",
    "cache_hit",
    "cache_miss",
    "request_cookie",
];

type ReasonIndex = BTreeMap<&'static str, Arc<AtomicU64>>;
type DomainIndex = BTreeMap<String, BTreeMap<String, Arc<DomainStats>>>;

#[derive(Clone, Default)]
struct PublishedDomains {
    index: DomainIndex,
    count: usize,
}

struct DomainStats {
    hits: AtomicU64,
    misses: AtomicU64,
    bypasses: AtomicU64,
    stored: AtomicU64,
    saved_body_bytes: AtomicU64,
    fast_reasons: [AtomicU64; FAST_REASONS.len()],
    reasons: ArcSwap<ReasonIndex>,
    register_reason: Mutex<()>,
}

impl Default for DomainStats {
    fn default() -> Self {
        Self {
            hits: AtomicU64::new(0),
            misses: AtomicU64::new(0),
            bypasses: AtomicU64::new(0),
            stored: AtomicU64::new(0),
            saved_body_bytes: AtomicU64::new(0),
            fast_reasons: std::array::from_fn(|_| AtomicU64::new(0)),
            reasons: ArcSwap::from_pointee(ReasonIndex::new()),
            register_reason: Mutex::new(()),
        }
    }
}

impl DomainStats {
    fn record(&self, outcome: &str, reason: &'static str) {
        match outcome {
            "HIT" => &self.hits,
            "MISS" => &self.misses,
            _ => &self.bypasses,
        }
        .fetch_add(1, Ordering::Relaxed);
        self.record_reason(reason);
    }

    fn record_reason(&self, reason: &'static str) {
        if reason == "cache_disabled" {
            self.fast_reasons[0].fetch_add(1, Ordering::Relaxed);
            return;
        }
        if let Some(index) = FAST_REASONS.iter().position(|label| *label == reason) {
            self.fast_reasons[index].fetch_add(1, Ordering::Relaxed);
            return;
        }
        let published = self.reasons.load();
        if let Some(counter) = published.get(reason) {
            counter.fetch_add(1, Ordering::Relaxed);
            return;
        }
        drop(published);
        let _guard = self
            .register_reason
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        let published = self.reasons.load_full();
        if let Some(counter) = published.get(reason) {
            counter.fetch_add(1, Ordering::Relaxed);
            return;
        }
        let mut next = (*published).clone();
        next.insert(reason, Arc::new(AtomicU64::new(1)));
        self.reasons.store(Arc::new(next));
    }

    fn snapshot(&self) -> DomainCounters {
        let load = |value: &AtomicU64| value.load(Ordering::Relaxed);
        let mut reasons = BTreeMap::new();
        for (label, counter) in FAST_REASONS.iter().zip(&self.fast_reasons) {
            let value = load(counter);
            if value > 0 {
                reasons.insert((*label).to_owned(), value);
            }
        }
        for (label, counter) in self.reasons.load().iter() {
            reasons.insert((*label).to_owned(), load(counter));
        }
        DomainCounters {
            hits: load(&self.hits),
            misses: load(&self.misses),
            bypasses: load(&self.bypasses),
            stored: load(&self.stored),
            saved_body_bytes: load(&self.saved_body_bytes),
            reasons,
        }
    }
}

struct DomainStatsStore {
    domains: ArcSwap<PublishedDomains>,
    register: Mutex<()>,
}

impl Default for DomainStatsStore {
    fn default() -> Self {
        Self {
            domains: ArcSwap::from_pointee(PublishedDomains::default()),
            register: Mutex::new(()),
        }
    }
}

impl DomainStatsStore {
    fn with(&self, scope: &str, domain: &str, action: impl FnOnce(&DomainStats)) {
        let published = self.domains.load();
        if let Some(counters) = published
            .index
            .get(scope)
            .and_then(|hosts| hosts.get(domain))
        {
            action(counters);
            return;
        }
        drop(published);
        let _guard = self.register.lock().unwrap_or_else(|e| e.into_inner());
        let published = self.domains.load_full();
        if let Some(counters) = published
            .index
            .get(scope)
            .and_then(|hosts| hosts.get(domain))
        {
            action(counters);
            return;
        }
        if published.count >= DOMAIN_STATS_LIMIT {
            return;
        }
        let mut next = (*published).clone();
        let counters = Arc::new(DomainStats::default());
        next.index
            .entry(scope.to_owned())
            .or_default()
            .insert(domain.to_owned(), Arc::clone(&counters));
        next.count += 1;
        self.domains.store(Arc::new(next));
        action(&counters);
    }

    fn record(&self, scope: &str, domain: &str, outcome: &str, reason: &'static str) {
        self.with(scope, domain, |counters| counters.record(outcome, reason));
    }

    fn record_reason(&self, scope: &str, domain: &str, reason: &'static str) {
        self.with(scope, domain, |counters| counters.record_reason(reason));
    }

    fn stored(&self, scope: &str, domain: &str) {
        self.with(scope, domain, |counters| {
            counters.stored.fetch_add(1, Ordering::Relaxed);
        });
    }

    fn saved_body_hit(&self, scope: &str, domain: &str, bytes: usize) {
        self.with(scope, domain, |counters| {
            counters
                .saved_body_bytes
                .fetch_add(bytes as u64, Ordering::Relaxed);
        });
    }

    fn clear(&self) {
        let _guard = self.register.lock().unwrap_or_else(|e| e.into_inner());
        self.domains.store(Arc::new(PublishedDomains::default()));
    }

    fn snapshot(&self, scope: Option<&str>) -> BTreeMap<(String, String), DomainCounters> {
        let published = self.domains.load();
        let mut counters = BTreeMap::new();
        for (vhost, domains) in &published.index {
            if !scope.is_none_or(|selected| selected == vhost) {
                continue;
            }
            for (domain, metric) in domains {
                counters.insert((vhost.clone(), domain.clone()), metric.snapshot());
            }
        }
        counters
    }

    fn scope_totals(&self) -> BTreeMap<String, ScopeCacheTotals> {
        let published = self.domains.load();
        let load = |value: &AtomicU64| value.load(Ordering::Relaxed);
        published
            .index
            .iter()
            .map(|(scope, domains)| {
                let mut totals = ScopeCacheTotals::default();
                for stats in domains.values() {
                    totals.hits = totals.hits.saturating_add(load(&stats.hits));
                    totals.misses = totals.misses.saturating_add(load(&stats.misses));
                    totals.bypasses = totals.bypasses.saturating_add(load(&stats.bypasses));
                    totals.saved_body_bytes = totals
                        .saved_body_bytes
                        .saturating_add(load(&stats.saved_body_bytes));
                }
                (scope.clone(), totals)
            })
            .collect()
    }
}