rpki/src/analysis/timing.rs

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

use serde::{Deserialize, Serialize};

#[derive(Clone)]
pub struct TimingHandle {
    inner: Arc<Mutex<TimingCollector>>,
}

impl TimingHandle {
    pub fn new(meta: TimingMeta) -> Self {
        Self {
            inner: Arc::new(Mutex::new(TimingCollector::new(meta))),
        }
    }

    pub fn span_phase(&self, phase: &'static str) -> TimingSpanGuard<'_> {
        TimingSpanGuard {
            handle: self.clone(),
            kind: TimingSpanKind::Phase(phase),
            start: Instant::now(),
        }
    }

    pub fn span_rrdp_repo<'a>(&self, repo_uri: &'a str) -> TimingSpanGuard<'a> {
        TimingSpanGuard {
            handle: self.clone(),
            kind: TimingSpanKind::RrdpRepo(repo_uri),
            start: Instant::now(),
        }
    }

    pub fn span_rrdp_repo_step<'a>(
        &self,
        repo_uri: &'a str,
        step: &'static str,
    ) -> TimingSpanGuard<'a> {
        TimingSpanGuard {
            handle: self.clone(),
            kind: TimingSpanKind::RrdpRepoStep { repo_uri, step },
            start: Instant::now(),
        }
    }

    pub fn span_publication_point<'a>(&self, manifest_rsync_uri: &'a str) -> TimingSpanGuard<'a> {
        TimingSpanGuard {
            handle: self.clone(),
            kind: TimingSpanKind::PublicationPoint(manifest_rsync_uri),
            start: Instant::now(),
        }
    }

    pub fn set_meta(&self, update: TimingMetaUpdate<'_>) {
        let mut g = self.inner.lock().expect("timing lock");
        if let Some(v) = update.tal_url {
            g.meta.tal_url = Some(v.to_string());
        }
        if let Some(v) = update.db_path {
            g.meta.db_path = Some(v.to_string());
        }
    }

    pub fn record_count(&self, key: &'static str, inc: u64) {
        let mut g = self.inner.lock().expect("timing lock");
        g.counts
            .entry(key)
            .and_modify(|v| *v = v.saturating_add(inc))
            .or_insert(inc);
    }

    pub fn counts_snapshot(&self) -> HashMap<String, u64> {
        let g = self.inner.lock().expect("timing lock");
        g.counts
            .iter()
            .map(|(key, value)| ((*key).to_string(), *value))
            .collect()
    }

    pub fn report_snapshot(&self, top_n: usize) -> TimingReportV1 {
        let g = self.inner.lock().expect("timing lock");
        g.to_report(top_n)
    }

    /// Record a phase duration directly in nanoseconds.
    ///
    /// This is useful when aggregating sub-phase timings locally (to reduce lock contention)
    /// and then emitting a single record per publication point.
    pub fn record_phase_nanos(&self, phase: &'static str, nanos: u64) {
        let mut g = self.inner.lock().expect("timing lock");
        g.phases.record(phase, nanos);
    }

    pub fn record_publication_point_nanos(&self, manifest_rsync_uri: &str, nanos: u64) {
        let mut g = self.inner.lock().expect("timing lock");
        g.publication_points.record(manifest_rsync_uri, nanos);
    }

    pub fn record_publication_point_step_nanos(
        &self,
        manifest_rsync_uri: &str,
        step: &'static str,
        nanos: u64,
    ) {
        let mut g = self.inner.lock().expect("timing lock");
        g.publication_point_steps
            .record(&format!("{manifest_rsync_uri}::{step}"), nanos);
    }

    pub fn write_json(&self, path: &Path, top_n: usize) -> Result<(), String> {
        let report = {
            let g = self.inner.lock().expect("timing lock");
            g.to_report(top_n)
        };

        let f = std::fs::File::create(path)
            .map_err(|e| format!("create timing json failed: {}: {e}", path.display()))?;
        serde_json::to_writer_pretty(f, &report)
            .map_err(|e| format!("write timing json failed: {e}"))?;
        Ok(())
    }

    fn record_duration(&self, kind: TimingSpanKind<'_>, duration: Duration) {
        let nanos_u64 = duration.as_nanos().min(u128::from(u64::MAX)) as u64;
        let mut g = self.inner.lock().expect("timing lock");
        match kind {
            TimingSpanKind::Phase(name) => g.phases.record(name, nanos_u64),
            TimingSpanKind::RrdpRepo(uri) => g.rrdp_repos.record(uri, nanos_u64),
            TimingSpanKind::RrdpRepoStep { repo_uri, step } => g
                .rrdp_repo_steps
                .record(&format!("{repo_uri}::{step}"), nanos_u64),
            TimingSpanKind::PublicationPoint(uri) => g.publication_points.record(uri, nanos_u64),
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TimingMeta {
    pub recorded_at_utc_rfc3339: String,
    pub validation_time_utc_rfc3339: String,
    pub tal_url: Option<String>,
    pub db_path: Option<String>,
}

#[derive(Clone, Debug, Default)]
pub struct TimingMetaUpdate<'a> {
    pub tal_url: Option<&'a str>,
    pub db_path: Option<&'a str>,
}

pub struct TimingSpanGuard<'a> {
    handle: TimingHandle,
    kind: TimingSpanKind<'a>,
    start: Instant,
}

impl Drop for TimingSpanGuard<'_> {
    fn drop(&mut self) {
        self.handle
            .record_duration(self.kind.clone(), self.start.elapsed());
    }
}

#[derive(Clone, Debug)]
enum TimingSpanKind<'a> {
    Phase(&'static str),
    RrdpRepo(&'a str),
    RrdpRepoStep {
        repo_uri: &'a str,
        step: &'static str,
    },
    PublicationPoint(&'a str),
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct DurationStats {
    pub count: u64,
    pub total_nanos: u64,
}

impl DurationStats {
    fn record(&mut self, nanos: u64) {
        self.count = self.count.saturating_add(1);
        self.total_nanos = self.total_nanos.saturating_add(nanos);
    }
}

#[derive(Clone, Debug, Default)]
struct DurationStatsMap {
    map: HashMap<String, DurationStats>,
}

impl DurationStatsMap {
    fn record(&mut self, key: &str, nanos: u64) {
        self.map.entry(key.to_string()).or_default().record(nanos);
    }

    fn top(&self, n: usize) -> Vec<TopDurationEntry> {
        let mut v = self
            .map
            .iter()
            .map(|(k, s)| TopDurationEntry {
                key: k.clone(),
                count: s.count,
                total_nanos: s.total_nanos,
            })
            .collect::<Vec<_>>();
        v.sort_by(|a, b| b.total_nanos.cmp(&a.total_nanos));
        v.truncate(n);
        v
    }
}

struct TimingCollector {
    meta: TimingMeta,
    counts: HashMap<&'static str, u64>,
    phases: DurationStatsMap,
    rrdp_repos: DurationStatsMap,
    rrdp_repo_steps: DurationStatsMap,
    publication_points: DurationStatsMap,
    publication_point_steps: DurationStatsMap,
}

impl TimingCollector {
    fn new(meta: TimingMeta) -> Self {
        Self {
            meta,
            counts: HashMap::new(),
            phases: DurationStatsMap::default(),
            rrdp_repos: DurationStatsMap::default(),
            rrdp_repo_steps: DurationStatsMap::default(),
            publication_points: DurationStatsMap::default(),
            publication_point_steps: DurationStatsMap::default(),
        }
    }

    fn to_report(&self, top_n: usize) -> TimingReportV1 {
        TimingReportV1 {
            format_version: 1,
            meta: self.meta.clone(),
            counts: self
                .counts
                .iter()
                .map(|(k, v)| ((*k).to_string(), *v))
                .collect(),
            phases: self
                .phases
                .map
                .iter()
                .map(|(k, s)| (k.clone(), s.clone()))
                .collect(),
            top_rrdp_repos: self.rrdp_repos.top(top_n),
            top_rrdp_repo_steps: self.rrdp_repo_steps.top(top_n),
            top_publication_points: self.publication_points.top(top_n),
            top_publication_point_steps: self.publication_point_steps.top(top_n),
        }
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct TimingReportV1 {
    pub format_version: u64,
    pub meta: TimingMeta,
    pub counts: HashMap<String, u64>,
    pub phases: HashMap<String, DurationStats>,
    pub top_rrdp_repos: Vec<TopDurationEntry>,
    pub top_rrdp_repo_steps: Vec<TopDurationEntry>,
    pub top_publication_points: Vec<TopDurationEntry>,
    pub top_publication_point_steps: Vec<TopDurationEntry>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TopDurationEntry {
    pub key: String,
    pub count: u64,
    pub total_nanos: u64,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn timing_handle_writes_json_with_phases_and_tops() {
        let meta = TimingMeta {
            recorded_at_utc_rfc3339: "2026-02-28T00:00:00Z".to_string(),
            validation_time_utc_rfc3339: "2026-02-28T00:00:00Z".to_string(),
            tal_url: Some("https://example.test/x.tal".to_string()),
            db_path: Some("db".to_string()),
        };
        let h = TimingHandle::new(meta);

        {
            let _p = h.span_phase("tal_bootstrap");
        }
        {
            let _r = h.span_rrdp_repo("https://rrdp.example.test/notification.xml");
        }
        {
            let _s = h.span_rrdp_repo_step(
                "https://rrdp.example.test/notification.xml",
                "fetch_notification",
            );
        }
        {
            let _pp = h.span_publication_point("rsync://example.test/repo/manifest.mft");
        }
        h.record_count("vrps", 42);
        h.record_publication_point_nanos("rsync://example.test/repo/manifest.mft", 1_000_000);
        h.record_publication_point_step_nanos(
            "rsync://example.test/repo/manifest.mft",
            "fresh_snapshot_prepare",
            1_000_000,
        );

        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("timing.json");
        h.write_json(&path, 10).expect("write_json");

        let rep: TimingReportV1 =
            serde_json::from_slice(&std::fs::read(&path).expect("read timing.json"))
                .expect("parse timing.json");
        assert_eq!(rep.format_version, 1);
        assert!(rep.phases.contains_key("tal_bootstrap"));
        assert_eq!(rep.counts.get("vrps").copied(), Some(42));
        assert!(
            rep.top_rrdp_repos
                .iter()
                .any(|e| e.key.contains("rrdp.example.test")),
            "expected repo in top list"
        );
        assert!(
            rep.top_rrdp_repo_steps
                .iter()
                .any(|e| e.key.contains("fetch_notification")),
            "expected repo step in top list"
        );
        assert!(
            rep.top_publication_points
                .iter()
                .any(|e| e.key.contains("manifest.mft")),
            "expected PP in top list"
        );
        assert!(
            rep.top_publication_point_steps
                .iter()
                .any(|e| e.key.contains("fresh_snapshot_prepare")),
            "expected PP step in top list"
        );
    }
}