Skip to main content

zenkey_fleet/model/
tree.rs

1//! The live key tree (issue #15): an immutable snapshot of everything the
2//! monitor has seen, grouped by key chunks, with per-node statistics.
3//!
4//! Snapshots are rebuilt on the monitor's stats tick and published through
5//! an `ArcSwap` — render loops *pull* the latest snapshot at their own pace
6//! and never contend with the per-sample hot path (a hot bus cannot melt a
7//! zengui redraw).
8
9use std::collections::BTreeMap;
10use std::sync::Arc;
11use std::time::Instant;
12
13use crate::model::stats::StatsTable;
14
15/// One key's contribution to the fold: everything the tree reads out of a
16/// [`KeyStats`](crate::model::stats::KeyStats), and nothing else.
17///
18/// All `Copy` but the key, which is the table's own `Arc<str>` — so copying
19/// the whole table's rows is a walk plus a refcount bump per key, which is
20/// what makes the ingest lock's critical section O(keys) rather than
21/// O(keys × chunks) (#330).
22#[derive(Debug, Clone)]
23pub struct TreeRow {
24    pub key: Arc<str>,
25    pub count: u64,
26    pub bytes: u64,
27    pub rate_hz: f64,
28    pub last_seen: Instant,
29}
30
31/// A whole table's [`TreeRow`]s and its O6 counters, as of one read —
32/// [`StatsTable::rows`](crate::model::stats::StatsTable::rows) produces it
33/// under the lock, [`KeyTreeSnapshot::fold`] consumes it outside.
34#[derive(Debug, Clone, Default)]
35pub struct TreeRows {
36    pub rows: Vec<TreeRow>,
37    /// Distinct keys the table held — [`KeyTreeSnapshot::keys`].
38    pub keys: usize,
39    /// Keys retired to stay within the table's bound.
40    pub evicted: u64,
41    /// Keys retired because their watch was released.
42    pub unwatched: u64,
43}
44
45/// Fold one key's row into every node on its path (the node itself included).
46fn accumulate(node: &mut TreeNode, s: &TreeRow) {
47    node.subtree_count += s.count;
48    node.subtree_bytes += s.bytes;
49    node.subtree_rate_hz += s.rate_hz;
50    node.subtree_keys += 1;
51    node.subtree_last_seen = node.subtree_last_seen.max(Some(s.last_seen));
52}
53
54/// One node of the snapshot: a key chunk, its subtree, and — when a sample
55/// has landed exactly here — its stats.
56///
57/// The `subtree_*` fields are what a **collapsed** node shows: a UI that can
58/// only report the traffic of keys it happens to have expanded is reporting a
59/// number the user will misread as the total.
60#[derive(Debug, Clone, Default)]
61pub struct TreeNode {
62    pub children: BTreeMap<String, TreeNode>,
63    /// Samples observed at exactly this key (leaf traffic).
64    pub count: u64,
65    pub bytes: u64,
66    pub rate_hz: f64,
67    /// When a sample last landed exactly here.
68    pub last_seen: Option<Instant>,
69    /// Aggregates over the whole subtree (this node included).
70    pub subtree_count: u64,
71    pub subtree_bytes: u64,
72    pub subtree_rate_hz: f64,
73    /// Most recent sample anywhere in the subtree.
74    pub subtree_last_seen: Option<Instant>,
75    /// Distinct keys that have carried traffic in this subtree.
76    pub subtree_keys: usize,
77}
78
79/// An immutable point-in-time view of the observed keyspace.
80///
81/// Carries the table's O6 counters too, so a render loop can report what the
82/// bound cost **without taking the ingest lock**: `root.subtree_count/bytes/
83/// rate_hz` are already the fold `StatsTable::totals` performs, and `keys` is
84/// its `len`. A consumer that pulled this `Arc` and then locked the table
85/// anyway was walking 50k entries a second time, four times a second, on the
86/// same mutex 100k samples/s need (`docs/zero-copy.md`).
87#[derive(Debug, Clone, Default)]
88pub struct KeyTreeSnapshot {
89    pub root: TreeNode,
90    pub keys: usize,
91    /// Keys retired to stay within the table's bound, as of this snapshot.
92    pub evicted: u64,
93    /// Keys retired because their watch was released.
94    pub unwatched: u64,
95}
96
97impl KeyTreeSnapshot {
98    /// Build from the stats table: copy the rows, then fold them.
99    ///
100    /// The convenience form, for callers that hold the table exclusively
101    /// (tests, offline projections). The monitor's tick deliberately spells
102    /// the two halves out — [`StatsTable::rows`] under the ingest lock,
103    /// [`fold`](Self::fold) after releasing it — because only the first half
104    /// may run while a network callback thread is waiting (#330).
105    pub fn build(stats: &StatsTable) -> KeyTreeSnapshot {
106        KeyTreeSnapshot::fold(stats.rows())
107    }
108
109    /// Fold copied rows into the snapshot. O(keys × chunks), and never to be
110    /// run under the ingest lock (#330).
111    pub fn fold(rows: TreeRows) -> KeyTreeSnapshot {
112        let TreeRows {
113            rows,
114            keys,
115            evicted,
116            unwatched,
117        } = rows;
118        let mut root = TreeNode::default();
119        for s in &rows {
120            let mut node = &mut root;
121            accumulate(node, s);
122            for chunk in s.key.split('/') {
123                // `entry` would need an owned key, so `chunk.to_string()`
124                // would run — and be dropped — on every *hit*, which is
125                // almost every chunk of almost every key. At 50k keys of 6
126                // chunks that is 300 000 wasted allocations per tick, four
127                // times a second. `BTreeMap<String, _>` looks up by `&str`
128                // through `Borrow`, so the owned key is built only when the
129                // node is genuinely new (`docs/zero-copy.md`).
130                if !node.children.contains_key(chunk) {
131                    node.children.insert(chunk.to_string(), TreeNode::default());
132                }
133                node = node
134                    .children
135                    .get_mut(chunk)
136                    .expect("just inserted if it was missing");
137                accumulate(node, s);
138            }
139            node.count = s.count;
140            node.bytes = s.bytes;
141            node.rate_hz = s.rate_hz;
142            node.last_seen = Some(s.last_seen);
143        }
144        KeyTreeSnapshot {
145            root,
146            keys,
147            evicted,
148            unwatched,
149        }
150    }
151
152    /// Walk to a node by its chunk path.
153    pub fn node(&self, path: &[&str]) -> Option<&TreeNode> {
154        let mut node = &self.root;
155        for chunk in path {
156            node = node.children.get(*chunk)?;
157        }
158        Some(node)
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    use std::time::Instant;
166
167    #[test]
168    fn builds_grouped_counts() {
169        let mut stats = StatsTable::new();
170        let now = Instant::now();
171        stats.record("zs/v1/h-a/telemetry/x/m1", 4, None, now, None, None);
172        stats.record("zs/v1/h-a/telemetry/x/m1", 4, None, now, None, None);
173        stats.record("zs/v1/h-a/telemetry/x/m2", 4, None, now, None, None);
174        stats.record("zs/v1/h-b/state/x/health", 4, None, now, None, None);
175
176        let snap = KeyTreeSnapshot::build(&stats);
177        assert_eq!(snap.keys, 3);
178        assert_eq!(snap.root.subtree_count, 4);
179        let telemetry = snap.node(&["zs", "v1", "h-a", "telemetry", "x"]).unwrap();
180        assert_eq!(telemetry.subtree_count, 3);
181        let m1 = snap
182            .node(&["zs", "v1", "h-a", "telemetry", "x", "m1"])
183            .unwrap();
184        assert_eq!(m1.count, 2);
185        assert_eq!(m1.bytes, 8);
186        assert!(snap.node(&["zs", "v1", "h-c"]).is_none());
187    }
188
189    /// A collapsed node must be able to report its subtree's traffic — bytes
190    /// and distinct keys, not only the sample count.
191    #[test]
192    fn collapsed_nodes_aggregate_their_subtree() {
193        let mut stats = StatsTable::new();
194        let now = Instant::now();
195        stats.record("zs/v1/h-a/telemetry/x/m1", 4, None, now, None, None);
196        stats.record("zs/v1/h-a/telemetry/x/m1", 4, None, now, None, None);
197        stats.record("zs/v1/h-a/telemetry/x/m2", 10, None, now, None, None);
198        stats.record("zs/v1/h-b/state/x/health", 7, None, now, None, None);
199
200        let snap = KeyTreeSnapshot::build(&stats);
201        let root = &snap.root;
202        assert_eq!(root.subtree_count, 4);
203        assert_eq!(root.subtree_bytes, 4 + 4 + 10 + 7);
204        assert_eq!(root.subtree_keys, 3, "three distinct keys carried traffic");
205        assert!(root.subtree_last_seen.is_some());
206        // The root itself is not a leaf: no sample landed exactly there.
207        assert_eq!(root.count, 0);
208        assert_eq!(root.last_seen, None);
209
210        let x = snap.node(&["zs", "v1", "h-a", "telemetry", "x"]).unwrap();
211        assert_eq!(x.subtree_count, 3);
212        assert_eq!(x.subtree_bytes, 18);
213        assert_eq!(x.subtree_keys, 2);
214
215        let m1 = snap
216            .node(&["zs", "v1", "h-a", "telemetry", "x", "m1"])
217            .unwrap();
218        assert_eq!(m1.last_seen, Some(now));
219        assert_eq!(m1.subtree_keys, 1, "a leaf counts only itself");
220    }
221}