use std::collections::BTreeMap;
use crate::stats::StatsTable;
#[derive(Debug, Clone, Default)]
pub struct TreeNode {
pub children: BTreeMap<String, TreeNode>,
pub count: u64,
pub bytes: u64,
pub rate_hz: f64,
pub subtree_count: u64,
}
#[derive(Debug, Clone, Default)]
pub struct KeyTreeSnapshot {
pub root: TreeNode,
pub keys: usize,
}
impl KeyTreeSnapshot {
pub fn build(stats: &StatsTable) -> KeyTreeSnapshot {
let mut root = TreeNode::default();
for (key, s) in stats.iter() {
let mut node = &mut root;
node.subtree_count += s.count;
for chunk in key.split('/') {
node = node.children.entry(chunk.to_string()).or_default();
node.subtree_count += s.count;
}
node.count = s.count;
node.bytes = s.bytes;
node.rate_hz = s.rate_hz;
}
KeyTreeSnapshot {
root,
keys: stats.len(),
}
}
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);
stats.record("zs/v1/h-a/telemetry/x/m1", 4, None, now);
stats.record("zs/v1/h-a/telemetry/x/m2", 4, None, now);
stats.record("zs/v1/h-b/state/x/health", 4, None, now);
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());
}
}