1use std::collections::BTreeMap;
9
10use rustev_contract::redaction::{Action, FieldClass, RedactionPolicy};
11
12pub mod name {
14 pub const SCREENING: &str = "rustev.screening";
15 pub const SCREENING_UNITS: &str = "rustev.screening.units";
16 pub const DECISIONS: &str = "rustev.decisions";
17 pub const UNRESOLVED: &str = "rustev.unresolved";
18 pub const FALLBACKS: &str = "rustev.fallbacks";
19 pub const DECISION_DURATION: &str = "rustev.decision.duration";
20 pub const ATTEMPT_DURATION: &str = "rustev.attempt.duration";
21 pub const DECISION_COST: &str = "rustev.decision.cost";
22 pub const TOKEN_USAGE: &str = "gen_ai.client.token.usage";
23 pub const OPERATION_DURATION: &str = "gen_ai.client.operation.duration";
24 pub const ADMISSION_REJECTED: &str = "rustev.admission.rejected";
25 pub const EVIDENCE_SINK: &str = "rustev.evidence.sink";
26 pub const TELEMETRY_LOST: &str = "rustev.telemetry.lost";
27 pub const SAMPLED_OUT: &str = "rustev.telemetry.sampled_out";
28}
29
30pub const BOUNDARIES: [f64; 15] = [
32 0.0, 5.0, 10.0, 25.0, 50.0, 75.0, 100.0, 250.0, 500.0, 750.0, 1000.0, 2500.0, 5000.0, 7500.0,
33 10000.0,
34];
35
36pub type Attrs = Vec<(String, String)>;
38
39#[derive(Debug, Clone, PartialEq)]
41pub struct Histogram {
42 pub count: u64,
43 pub sum: f64,
44 pub min: f64,
45 pub max: f64,
46 pub buckets: Vec<u64>,
48}
49
50impl Default for Histogram {
51 fn default() -> Self {
52 Histogram {
53 count: 0,
54 sum: 0.0,
55 min: f64::INFINITY,
56 max: f64::NEG_INFINITY,
57 buckets: vec![0; BOUNDARIES.len() + 1],
58 }
59 }
60}
61
62impl Histogram {
63 pub fn record(&mut self, v: f64) {
64 self.count += 1;
65 self.sum += v;
66 self.min = self.min.min(v);
67 self.max = self.max.max(v);
68 let i = BOUNDARIES
69 .iter()
70 .position(|b| v <= *b)
71 .unwrap_or(BOUNDARIES.len());
72 self.buckets[i] += 1;
73 }
74}
75
76#[derive(Debug, Clone, PartialEq)]
78pub enum Points {
79 Counter(BTreeMap<Attrs, u64>),
80 Histogram(BTreeMap<Attrs, Histogram>),
81}
82
83#[derive(Debug, Clone, Default, PartialEq)]
85pub struct MetricsSnapshot {
86 pub instruments: BTreeMap<&'static str, Points>,
87}
88
89impl MetricsSnapshot {
90 pub fn counter(&self, instrument: &str, attrs: &[(&str, &str)]) -> u64 {
92 let key: Attrs = attrs
93 .iter()
94 .map(|(k, v)| (k.to_string(), v.to_string()))
95 .collect();
96 match self.instruments.get(instrument) {
97 Some(Points::Counter(m)) => m.get(&key).copied().unwrap_or(0),
98 _ => 0,
99 }
100 }
101
102 pub fn counter_total(&self, instrument: &str) -> u64 {
104 match self.instruments.get(instrument) {
105 Some(Points::Counter(m)) => m.values().sum(),
106 _ => 0,
107 }
108 }
109}
110
111#[derive(Debug, Default)]
113pub struct Instruments {
114 counters: BTreeMap<&'static str, BTreeMap<Attrs, u64>>,
115 histograms: BTreeMap<&'static str, BTreeMap<Attrs, Histogram>>,
116}
117
118pub struct AttrSet<'p> {
120 policy: &'p RedactionPolicy,
121 attrs: Attrs,
122}
123
124impl<'p> AttrSet<'p> {
125 pub fn new(policy: &'p RedactionPolicy) -> Self {
126 AttrSet {
127 policy,
128 attrs: vec![],
129 }
130 }
131
132 pub fn with(mut self, key: &str, field: &str, class: FieldClass, value: &str) -> Self {
137 let dimension = matches!(
138 class,
139 FieldClass::Enum | FieldClass::PlanMetadata | FieldClass::ModelIdentity
140 );
141 if dimension && self.policy.action(field, class) == Action::Allow {
142 self.attrs.push((key.to_string(), value.to_string()));
143 }
144 self
145 }
146
147 pub fn done(mut self) -> Attrs {
148 self.attrs.sort();
149 self.attrs
150 }
151}
152
153impl Instruments {
154 pub fn add(&mut self, instrument: &'static str, attrs: Attrs, n: u64) {
155 *self
156 .counters
157 .entry(instrument)
158 .or_default()
159 .entry(attrs)
160 .or_default() += n;
161 }
162
163 pub fn record(&mut self, instrument: &'static str, attrs: Attrs, v: f64) {
164 self.histograms
165 .entry(instrument)
166 .or_default()
167 .entry(attrs)
168 .or_default()
169 .record(v);
170 }
171
172 pub fn snapshot(&self) -> MetricsSnapshot {
173 let mut s = MetricsSnapshot::default();
174 for (k, v) in &self.counters {
175 s.instruments.insert(k, Points::Counter(v.clone()));
176 }
177 for (k, v) in &self.histograms {
178 s.instruments.insert(k, Points::Histogram(v.clone()));
179 }
180 s
181 }
182}
183
184#[cfg(test)]
185mod tests {
186 use super::*;
187
188 #[test]
189 fn only_dimension_classes_become_metric_attributes_in_every_build() {
190 let mut p = RedactionPolicy::default();
191 for c in FieldClass::ALL {
192 if !c.is_deny() {
193 p.classes.insert(c, Action::Allow);
194 }
195 }
196 let attrs = AttrSet::new(&p)
198 .with("a", "rustev.step.id", FieldClass::Enum, "x")
199 .with("b", "rustev.step.id", FieldClass::SystemOfRecord, "y")
200 .with(
201 "c",
202 "rustev.step.id",
203 FieldClass::AuthenticatedAppField,
204 "z",
205 )
206 .done();
207 assert_eq!(attrs, vec![("a".to_string(), "x".to_string())]);
208 }
209}