1use 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#[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 #[must_use]
23 pub fn new() -> Self {
24 Self::default()
25 }
26
27 #[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 #[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}