Skip to main content

machi_obs/
prometheus.rs

1//! Prometheus text exposition for captured metrics (zero external deps).
2
3use std::collections::{BTreeMap, BTreeSet};
4use std::fmt::Write as _;
5use std::sync::Mutex;
6
7use crate::metrics::MetricsSink;
8
9type LabelSet = Vec<(String, String)>;
10type SeriesKey = (String, LabelSet);
11
12/// Metrics sink that accumulates samples and can render Prometheus text format 0.0.4.
13#[derive(Debug, Default)]
14pub struct PrometheusRecorder {
15    counters: Mutex<BTreeMap<SeriesKey, u64>>,
16    histograms: Mutex<BTreeMap<SeriesKey, Vec<f64>>>,
17    gauges: Mutex<BTreeMap<SeriesKey, f64>>,
18}
19
20impl PrometheusRecorder {
21    /// Empty recorder.
22    #[must_use]
23    pub fn new() -> Self {
24        Self::default()
25    }
26
27    /// Render Prometheus exposition text (sorted, one `# TYPE` per metric name).
28    #[must_use]
29    pub fn render(&self) -> String {
30        let mut out = String::new();
31        let mut typed = BTreeSet::new();
32
33        render_counters(&self.counters, &mut out, &mut typed);
34        render_gauges(&self.gauges, &mut out, &mut typed);
35        render_histograms(&self.histograms, &mut out, &mut typed);
36        out
37    }
38
39    /// Series names that appear in the current capture (counters/histograms/gauges).
40    #[must_use]
41    pub fn series_names(&self) -> BTreeSet<String> {
42        let mut names = BTreeSet::new();
43        for (n, _) in self
44            .counters
45            .lock()
46            .unwrap_or_else(std::sync::PoisonError::into_inner)
47            .keys()
48        {
49            names.insert(n.clone());
50        }
51        for (n, _) in self
52            .histograms
53            .lock()
54            .unwrap_or_else(std::sync::PoisonError::into_inner)
55            .keys()
56        {
57            names.insert(n.clone());
58        }
59        for (n, _) in self
60            .gauges
61            .lock()
62            .unwrap_or_else(std::sync::PoisonError::into_inner)
63            .keys()
64        {
65            names.insert(n.clone());
66        }
67        names
68    }
69}
70
71fn render_counters(
72    counters: &Mutex<BTreeMap<SeriesKey, u64>>,
73    out: &mut String,
74    typed: &mut BTreeSet<String>,
75) {
76    let guard = counters
77        .lock()
78        .unwrap_or_else(std::sync::PoisonError::into_inner);
79    for ((name, labels), value) in guard.iter() {
80        emit_type_once(out, typed, name, "counter", "Machi kernel counter");
81        let _ = writeln!(out, "{name}{} {value}", format_labels(labels));
82    }
83}
84
85fn render_gauges(
86    gauges: &Mutex<BTreeMap<SeriesKey, f64>>,
87    out: &mut String,
88    typed: &mut BTreeSet<String>,
89) {
90    let guard = gauges
91        .lock()
92        .unwrap_or_else(std::sync::PoisonError::into_inner);
93    for ((name, labels), value) in guard.iter() {
94        emit_type_once(out, typed, name, "gauge", "Machi kernel gauge");
95        let _ = writeln!(out, "{name}{} {value}", format_labels(labels));
96    }
97}
98
99fn render_histograms(
100    histograms: &Mutex<BTreeMap<SeriesKey, Vec<f64>>>,
101    out: &mut String,
102    typed: &mut BTreeSet<String>,
103) {
104    let guard = histograms
105        .lock()
106        .unwrap_or_else(std::sync::PoisonError::into_inner);
107    for ((name, labels), samples) in guard.iter() {
108        emit_type_once(out, typed, name, "summary", "Machi kernel summary");
109        let count = samples.len();
110        let sum: f64 = samples.iter().sum();
111        let labs = format_labels(labels);
112        let _ = writeln!(out, "{name}_count{labs} {count}");
113        let _ = writeln!(out, "{name}_sum{labs} {sum}");
114    }
115}
116
117fn emit_type_once(
118    out: &mut String,
119    typed: &mut BTreeSet<String>,
120    name: &str,
121    type_name: &str,
122    help: &str,
123) {
124    if typed.insert(name.to_owned()) {
125        let _ = writeln!(out, "# HELP {name} {help}");
126        let _ = writeln!(out, "# TYPE {name} {type_name}");
127    }
128}
129
130impl MetricsSink for PrometheusRecorder {
131    fn counter(&self, name: &str, value: u64, labels: &[(&str, &str)]) {
132        let key = (name.to_owned(), owned_labels(labels));
133        *self
134            .counters
135            .lock()
136            .unwrap_or_else(std::sync::PoisonError::into_inner)
137            .entry(key)
138            .or_insert(0) += value;
139    }
140
141    fn histogram(&self, name: &str, value: f64, labels: &[(&str, &str)]) {
142        let key = (name.to_owned(), owned_labels(labels));
143        self.histograms
144            .lock()
145            .unwrap_or_else(std::sync::PoisonError::into_inner)
146            .entry(key)
147            .or_default()
148            .push(value);
149    }
150
151    fn gauge(&self, name: &str, value: f64, labels: &[(&str, &str)]) {
152        let key = (name.to_owned(), owned_labels(labels));
153        self.gauges
154            .lock()
155            .unwrap_or_else(std::sync::PoisonError::into_inner)
156            .insert(key, value);
157    }
158}
159
160fn owned_labels(labels: &[(&str, &str)]) -> Vec<(String, String)> {
161    labels
162        .iter()
163        .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
164        .collect()
165}
166
167fn format_labels(labels: &[(String, String)]) -> String {
168    if labels.is_empty() {
169        return String::new();
170    }
171    let parts: Vec<String> = labels
172        .iter()
173        .map(|(k, v)| format!("{k}=\"{}\"", escape_label(v)))
174        .collect();
175    format!("{{{}}}", parts.join(","))
176}
177
178fn escape_label(s: &str) -> String {
179    s.replace('\\', "\\\\")
180        .replace('"', "\\\"")
181        .replace('\n', "\\n")
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187    use crate::metrics::{
188        METRIC_SAMPLE_DURATION_MS, METRIC_TURNS_TOTAL, emit_catalogue_smoke, record_turn,
189        required_metric_names,
190    };
191
192    #[test]
193    fn renders_counter() {
194        let p = PrometheusRecorder::new();
195        record_turn(&p, "ok", 1, 2.0);
196        let text = p.render();
197        assert!(text.contains(METRIC_TURNS_TOTAL), "{text}");
198        assert!(text.contains("status=\"ok\""), "{text}");
199        assert!(text.contains("# TYPE"), "{text}");
200        assert!(text.contains("# HELP"), "{text}");
201        let type_lines = text
202            .lines()
203            .filter(|l| *l == format!("# TYPE {METRIC_TURNS_TOTAL} counter"))
204            .count();
205        assert_eq!(type_lines, 1, "{text}");
206    }
207
208    #[test]
209    fn golden_export_covers_catalogue() {
210        let p = PrometheusRecorder::new();
211        emit_catalogue_smoke(&p);
212        let text = p.render();
213        let names = p.series_names();
214        for required in required_metric_names() {
215            assert!(
216                names.contains(*required),
217                "export missing {required}\n{text}"
218            );
219        }
220        assert!(
221            text.contains(&format!("{METRIC_SAMPLE_DURATION_MS}_count")),
222            "{text}"
223        );
224        assert!(text.contains("# HELP"));
225        assert!(text.contains("# TYPE"));
226        assert!(text.contains("status=\"ok\"") || text.contains("outcome=\"completed\""));
227    }
228
229    #[test]
230    fn escapes_label_values() {
231        let p = PrometheusRecorder::new();
232        p.counter("machi_test", 1, &[("path", r#"a"b\c"#)]);
233        let text = p.render();
234        assert!(text.contains(r#"path="a\"b\\c""#), "{text}");
235    }
236}