zenkey-fleet 0.11.1

//! The live key tree (issue #15): an immutable snapshot of everything the
//! monitor has seen, grouped by key chunks, with per-node statistics.
//!
//! Snapshots are rebuilt on the monitor's stats tick and published through
//! an `ArcSwap` — render loops *pull* the latest snapshot at their own pace
//! and never contend with the per-sample hot path (a hot bus cannot melt a
//! zengui redraw).

use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::Instant;

use crate::model::stats::StatsTable;

/// One key's contribution to the fold: everything the tree reads out of a
/// [`KeyStats`](crate::model::stats::KeyStats), and nothing else.
///
/// All `Copy` but the key, which is the table's own `Arc<str>` — so copying
/// the whole table's rows is a walk plus a refcount bump per key, which is
/// what makes the ingest lock's critical section O(keys) rather than
/// O(keys × chunks) (#330).
#[derive(Debug, Clone)]
pub struct TreeRow {
    pub key: Arc<str>,
    pub count: u64,
    pub bytes: u64,
    pub rate_hz: f64,
    pub last_seen: Instant,
}

/// A whole table's [`TreeRow`]s and its O6 counters, as of one read —
/// [`StatsTable::rows`](crate::model::stats::StatsTable::rows) produces it
/// under the lock, [`KeyTreeSnapshot::fold`] consumes it outside.
#[derive(Debug, Clone, Default)]
pub struct TreeRows {
    pub rows: Vec<TreeRow>,
    /// Distinct keys the table held — [`KeyTreeSnapshot::keys`].
    pub keys: usize,
    /// Keys retired to stay within the table's bound.
    pub evicted: u64,
    /// Keys retired because their watch was released.
    pub unwatched: u64,
}

/// Fold one key's row into every node on its path (the node itself included).
fn accumulate(node: &mut TreeNode, s: &TreeRow) {
    node.subtree_count += s.count;
    node.subtree_bytes += s.bytes;
    node.subtree_rate_hz += s.rate_hz;
    node.subtree_keys += 1;
    node.subtree_last_seen = node.subtree_last_seen.max(Some(s.last_seen));
}

/// One node of the snapshot: a key chunk, its subtree, and — when a sample
/// has landed exactly here — its stats.
///
/// The `subtree_*` fields are what a **collapsed** node shows: a UI that can
/// only report the traffic of keys it happens to have expanded is reporting a
/// number the user will misread as the total.
#[derive(Debug, Clone, Default)]
pub struct TreeNode {
    pub children: BTreeMap<String, TreeNode>,
    /// Samples observed at exactly this key (leaf traffic).
    pub count: u64,
    pub bytes: u64,
    pub rate_hz: f64,
    /// When a sample last landed exactly here.
    pub last_seen: Option<Instant>,
    /// Aggregates over the whole subtree (this node included).
    pub subtree_count: u64,
    pub subtree_bytes: u64,
    pub subtree_rate_hz: f64,
    /// Most recent sample anywhere in the subtree.
    pub subtree_last_seen: Option<Instant>,
    /// Distinct keys that have carried traffic in this subtree.
    pub subtree_keys: usize,
}

/// An immutable point-in-time view of the observed keyspace.
///
/// Carries the table's O6 counters too, so a render loop can report what the
/// bound cost **without taking the ingest lock**: `root.subtree_count/bytes/
/// rate_hz` are already the fold `StatsTable::totals` performs, and `keys` is
/// its `len`. A consumer that pulled this `Arc` and then locked the table
/// anyway was walking 50k entries a second time, four times a second, on the
/// same mutex 100k samples/s need (`docs/zero-copy.md`).
#[derive(Debug, Clone, Default)]
pub struct KeyTreeSnapshot {
    pub root: TreeNode,
    pub keys: usize,
    /// Keys retired to stay within the table's bound, as of this snapshot.
    pub evicted: u64,
    /// Keys retired because their watch was released.
    pub unwatched: u64,
}

impl KeyTreeSnapshot {
    /// Build from the stats table: copy the rows, then fold them.
    ///
    /// The convenience form, for callers that hold the table exclusively
    /// (tests, offline projections). The monitor's tick deliberately spells
    /// the two halves out — [`StatsTable::rows`] under the ingest lock,
    /// [`fold`](Self::fold) after releasing it — because only the first half
    /// may run while a network callback thread is waiting (#330).
    pub fn build(stats: &StatsTable) -> KeyTreeSnapshot {
        KeyTreeSnapshot::fold(stats.rows())
    }

    /// Fold copied rows into the snapshot. O(keys × chunks), and never to be
    /// run under the ingest lock (#330).
    pub fn fold(rows: TreeRows) -> KeyTreeSnapshot {
        let TreeRows {
            rows,
            keys,
            evicted,
            unwatched,
        } = rows;
        let mut root = TreeNode::default();
        for s in &rows {
            let mut node = &mut root;
            accumulate(node, s);
            for chunk in s.key.split('/') {
                // `entry` would need an owned key, so `chunk.to_string()`
                // would run — and be dropped — on every *hit*, which is
                // almost every chunk of almost every key. At 50k keys of 6
                // chunks that is 300 000 wasted allocations per tick, four
                // times a second. `BTreeMap<String, _>` looks up by `&str`
                // through `Borrow`, so the owned key is built only when the
                // node is genuinely new (`docs/zero-copy.md`).
                if !node.children.contains_key(chunk) {
                    node.children.insert(chunk.to_string(), TreeNode::default());
                }
                node = node
                    .children
                    .get_mut(chunk)
                    .expect("just inserted if it was missing");
                accumulate(node, s);
            }
            node.count = s.count;
            node.bytes = s.bytes;
            node.rate_hz = s.rate_hz;
            node.last_seen = Some(s.last_seen);
        }
        KeyTreeSnapshot {
            root,
            keys,
            evicted,
            unwatched,
        }
    }

    /// Walk to a node by its chunk path.
    pub fn node(&self, path: &[&str]) -> Option<&TreeNode> {
        let mut node = &self.root;
        for chunk in path {
            node = node.children.get(*chunk)?;
        }
        Some(node)
    }
}

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

    #[test]
    fn builds_grouped_counts() {
        let mut stats = StatsTable::new();
        let now = Instant::now();
        stats.record("zs/v1/h-a/telemetry/x/m1", 4, None, now, None, None);
        stats.record("zs/v1/h-a/telemetry/x/m1", 4, None, now, None, None);
        stats.record("zs/v1/h-a/telemetry/x/m2", 4, None, now, None, None);
        stats.record("zs/v1/h-b/state/x/health", 4, None, now, None, None);

        let snap = KeyTreeSnapshot::build(&stats);
        assert_eq!(snap.keys, 3);
        assert_eq!(snap.root.subtree_count, 4);
        let telemetry = snap.node(&["zs", "v1", "h-a", "telemetry", "x"]).unwrap();
        assert_eq!(telemetry.subtree_count, 3);
        let m1 = snap
            .node(&["zs", "v1", "h-a", "telemetry", "x", "m1"])
            .unwrap();
        assert_eq!(m1.count, 2);
        assert_eq!(m1.bytes, 8);
        assert!(snap.node(&["zs", "v1", "h-c"]).is_none());
    }

    /// A collapsed node must be able to report its subtree's traffic — bytes
    /// and distinct keys, not only the sample count.
    #[test]
    fn collapsed_nodes_aggregate_their_subtree() {
        let mut stats = StatsTable::new();
        let now = Instant::now();
        stats.record("zs/v1/h-a/telemetry/x/m1", 4, None, now, None, None);
        stats.record("zs/v1/h-a/telemetry/x/m1", 4, None, now, None, None);
        stats.record("zs/v1/h-a/telemetry/x/m2", 10, None, now, None, None);
        stats.record("zs/v1/h-b/state/x/health", 7, None, now, None, None);

        let snap = KeyTreeSnapshot::build(&stats);
        let root = &snap.root;
        assert_eq!(root.subtree_count, 4);
        assert_eq!(root.subtree_bytes, 4 + 4 + 10 + 7);
        assert_eq!(root.subtree_keys, 3, "three distinct keys carried traffic");
        assert!(root.subtree_last_seen.is_some());
        // The root itself is not a leaf: no sample landed exactly there.
        assert_eq!(root.count, 0);
        assert_eq!(root.last_seen, None);

        let x = snap.node(&["zs", "v1", "h-a", "telemetry", "x"]).unwrap();
        assert_eq!(x.subtree_count, 3);
        assert_eq!(x.subtree_bytes, 18);
        assert_eq!(x.subtree_keys, 2);

        let m1 = snap
            .node(&["zs", "v1", "h-a", "telemetry", "x", "m1"])
            .unwrap();
        assert_eq!(m1.last_seen, Some(now));
        assert_eq!(m1.subtree_keys, 1, "a leaf counts only itself");
    }
}