use std::collections::{BTreeMap, BTreeSet};
use crate::SliceSet;
use crate::judge::common::EXPANSION_CAP;
use crate::model::examples::Examples;
use crate::model::facts::{KeyFacts, KeyShape, OriginKind};
use crate::report::{BudgetCell, BudgetWindow, TopicList};
#[derive(Debug, Clone, Default)]
pub struct BudgetObservation {
families: BTreeMap<(String, String), BTreeMap<String, BTreeSet<String>>>,
}
impl BudgetObservation {
pub fn observe<'a>(
base: &str,
slices: &SliceSet,
keys: impl IntoIterator<Item = &'a str>,
) -> BudgetObservation {
let mut families: BTreeMap<(String, String), BTreeMap<String, BTreeSet<String>>> =
BTreeMap::new();
for key in keys {
let facts = KeyFacts::project(base, key);
let KeyShape::V1(v) = &facts.shape else {
continue;
};
if !v.class_kind.is_data_class() {
continue;
}
let producer = match v.origin_kind {
OriginKind::Host => v.producer.clone(),
OriginKind::Service => slices.by_service_origin(&v.origin).map(|s| s.name.clone()),
};
let Some(producer) = producer else {
continue;
};
let tail: Vec<&str> = v.subject.iter().map(String::as_str).collect();
let Some((decl, _)) = slices.refine(&producer, &v.class, &tail) else {
continue;
};
if !decl.path.contains('{') {
continue; }
families
.entry((producer, decl.path.clone()))
.or_default()
.entry(v.origin.clone())
.or_default()
.insert(key.to_string());
}
BudgetObservation { families }
}
pub fn family(
&self,
producer: &str,
path: &str,
) -> Option<&BTreeMap<String, BTreeSet<String>>> {
self.families.get(&(producer.to_string(), path.to_string()))
}
}
pub fn join_budget(list: &mut TopicList, obs: &BudgetObservation, window: BudgetWindow) {
for row in &mut list.subjects {
if row.deprecated || !row.path.contains('{') {
continue;
}
let empty = BTreeMap::new();
let origins = obs.family(&row.producer, &row.path).unwrap_or(&empty);
let observed: usize = origins.values().map(BTreeSet::len).sum();
let (worst_origin, worst_observed) = origins
.iter()
.max_by_key(|(_, keys)| keys.len())
.map(|(o, keys)| (Some(o.clone()), keys.len()))
.unwrap_or((None, 0));
let examples = worst_origin
.as_ref()
.and_then(|o| origins.get(o))
.map(|keys| {
let mut ex = Examples::new(EXPANSION_CAP);
for key in keys {
ex.push_with(|| key.clone());
}
ex.into_vec()
})
.unwrap_or_default();
let exempt = row
.path
.contains("...")
.then(|| "rest-variable".to_string());
let over = exempt.is_none()
&& row
.cardinality
.is_some_and(|declared| worst_observed as i64 > declared);
row.budget = Some(BudgetCell {
declared: row.cardinality,
observed,
origins: origins.len(),
worst_origin,
worst_observed,
exempt,
over,
examples,
});
}
list.budget = Some(window);
}
#[cfg(test)]
mod tests {
use super::*;
const SLICE: &str = r#"
[registry]
version = "1.0"
app = "t"
convention = 1
[producer]
name = "sysinfo"
[[subject]]
path = "disk/{mount}/used"
class = "telemetry"
type = "Point"
cardinality = 16
[[subject]]
path = "health"
class = "state"
type = "Health"
"#;
#[test]
fn observation_groups_per_origin_and_skips_literals() {
let slices = SliceSet::from_toml_for_tests(SLICE);
let keys = [
"v1/h-aaaaaaaaaaaa/telemetry/sysinfo/disk/root/used",
"v1/h-aaaaaaaaaaaa/telemetry/sysinfo/disk/var/used",
"v1/h-bbbbbbbbbbbb/telemetry/sysinfo/disk/root/used",
"v1/h-aaaaaaaaaaaa/state/sysinfo/health",
"v1/h-aaaaaaaaaaaa/telemetry/sysinfo/not/registered",
];
let obs = BudgetObservation::observe("", &slices, keys);
let fam = obs.family("sysinfo", "disk/{mount}/used").unwrap();
assert_eq!(fam.len(), 2, "two origins expanded the family");
assert_eq!(fam["h-aaaaaaaaaaaa"].len(), 2);
assert_eq!(fam["h-bbbbbbbbbbbb"].len(), 1);
assert!(
obs.family("sysinfo", "health").is_none(),
"literals excluded"
);
}
}