1use std::fmt;
2use std::io;
3use std::iter::once;
4
5#[cfg(test)]
6mod tests;
7
8struct FormattedValue(f64);
9
10impl fmt::Display for FormattedValue {
11 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
12 let value = self.0;
18 if value.is_nan() {
19 write!(f, "NaN")
20 } else if value == f64::INFINITY {
21 write!(f, "+Inf")
22 } else if value == f64::NEG_INFINITY {
23 write!(f, "-Inf")
24 } else {
25 write!(f, "{}", value)
26 }
27 }
28}
29
30pub struct LabeledMetricsBuilder<'a, W>
34where
35 W: io::Write,
36{
37 encoder: &'a mut MetricsEncoder<W>,
38 name: &'a str,
39}
40
41impl<W: io::Write> LabeledMetricsBuilder<'_, W> {
42 pub fn value(self, labels: &[(&str, &str)], value: f64) -> io::Result<Self> {
50 self.encoder
51 .encode_value_with_labels(self.name, labels, value)?;
52 Ok(self)
53 }
54}
55
56pub struct LabeledHistogramBuilder<'a, W>
60where
61 W: io::Write,
62{
63 encoder: &'a mut MetricsEncoder<W>,
64 name: &'a str,
65}
66
67impl<W: io::Write> LabeledHistogramBuilder<'_, W> {
68 pub fn histogram(
76 self,
77 labels: &[(&str, &str)],
78 buckets: impl Iterator<Item = (f64, f64)>,
79 sum: f64,
80 ) -> io::Result<Self> {
81 for (label, _) in labels.iter() {
82 validate_prometheus_name(label);
83 }
84
85 let mut total: f64 = 0.0;
86 let mut saw_infinity = false;
87 for (bucket, v) in buckets {
88 total += v;
89 if bucket == std::f64::INFINITY {
90 saw_infinity = true;
91 writeln!(
92 self.encoder.writer,
93 "{}_bucket{{{}}} {} {}",
94 self.name,
95 MetricsEncoder::<W>::encode_labels(labels.iter().chain(once(&("le", "+Inf")))),
96 total,
97 self.encoder.now_millis
98 )?;
99 } else {
100 let bucket_str = bucket.to_string();
101 writeln!(
102 self.encoder.writer,
103 "{}_bucket{{{}}} {} {}",
104 self.name,
105 MetricsEncoder::<W>::encode_labels(
106 labels.iter().chain(once(&("le", bucket_str.as_str())))
107 ),
108 total,
109 self.encoder.now_millis
110 )?;
111 }
112 }
113 if !saw_infinity {
114 writeln!(
115 self.encoder.writer,
116 "{}_bucket{{{}}} {} {}",
117 self.name,
118 MetricsEncoder::<W>::encode_labels(labels.iter().chain(once(&("le", "+Inf")))),
119 total,
120 self.encoder.now_millis
121 )?;
122 }
123
124 if labels.is_empty() {
125 writeln!(
126 self.encoder.writer,
127 "{}_sum {} {}",
128 self.name,
129 FormattedValue(sum),
130 self.encoder.now_millis
131 )?;
132 writeln!(
133 self.encoder.writer,
134 "{}_count {} {}",
135 self.name,
136 FormattedValue(total),
137 self.encoder.now_millis
138 )?;
139 } else {
140 writeln!(
141 self.encoder.writer,
142 "{}_sum{{{}}} {} {}",
143 self.name,
144 MetricsEncoder::<W>::encode_labels(labels.iter()),
145 FormattedValue(sum),
146 self.encoder.now_millis
147 )?;
148 writeln!(
149 self.encoder.writer,
150 "{}_count{{{}}} {} {}",
151 self.name,
152 MetricsEncoder::<W>::encode_labels(labels.iter()),
153 FormattedValue(total),
154 self.encoder.now_millis
155 )?;
156 }
157
158 Ok(self)
159 }
160}
161pub struct MetricsEncoder<W: io::Write> {
172 writer: W,
173 now_millis: i64,
174}
175
176impl<W: io::Write> MetricsEncoder<W> {
177 pub fn new(writer: W, now_millis: i64) -> Self {
180 Self { writer, now_millis }
181 }
182
183 pub fn into_inner(self) -> W {
186 self.writer
187 }
188
189 fn encode_header(&mut self, name: &str, help: &str, typ: &str) -> io::Result<()> {
190 writeln!(self.writer, "# HELP {} {}", name, help)?;
191 writeln!(self.writer, "# TYPE {} {}", name, typ)
192 }
193
194 pub fn encode_histogram(
203 &mut self,
204 name: &str,
205 buckets: impl Iterator<Item = (f64, f64)>,
206 sum: f64,
207 help: &str,
208 ) -> io::Result<()> {
209 self.histogram_vec(name, help)?
210 .histogram(&[], buckets, sum)?;
211 Ok(())
212 }
213
214 pub fn histogram_vec<'a>(
215 &'a mut self,
216 name: &'a str,
217 help: &'a str,
218 ) -> io::Result<LabeledHistogramBuilder<'a, W>> {
219 validate_prometheus_name(name);
220 self.encode_header(name, help, "histogram")?;
221 Ok(LabeledHistogramBuilder {
222 encoder: self,
223 name,
224 })
225 }
226
227 pub fn encode_single_value(
228 &mut self,
229 typ: &str,
230 name: &str,
231 value: f64,
232 help: &str,
233 ) -> io::Result<()> {
234 validate_prometheus_name(name);
235 self.encode_header(name, help, typ)?;
236 writeln!(
237 self.writer,
238 "{} {} {}",
239 name,
240 FormattedValue(value),
241 self.now_millis
242 )
243 }
244
245 pub fn encode_counter(&mut self, name: &str, value: f64, help: &str) -> io::Result<()> {
251 self.encode_single_value("counter", name, value, help)
252 }
253
254 pub fn encode_gauge(&mut self, name: &str, value: f64, help: &str) -> io::Result<()> {
260 self.encode_single_value("gauge", name, value, help)
261 }
262
263 pub fn counter_vec<'a>(
270 &'a mut self,
271 name: &'a str,
272 help: &'a str,
273 ) -> io::Result<LabeledMetricsBuilder<'a, W>> {
274 validate_prometheus_name(name);
275 self.encode_header(name, help, "counter")?;
276 Ok(LabeledMetricsBuilder {
277 encoder: self,
278 name,
279 })
280 }
281
282 pub fn gauge_vec<'a>(
289 &'a mut self,
290 name: &'a str,
291 help: &'a str,
292 ) -> io::Result<LabeledMetricsBuilder<'a, W>> {
293 validate_prometheus_name(name);
294 self.encode_header(name, help, "gauge")?;
295 Ok(LabeledMetricsBuilder {
296 encoder: self,
297 name,
298 })
299 }
300
301 fn encode_labels<'a>(labels: impl Iterator<Item = &'a (&'a str, &'a str)>) -> String {
302 let mut buf = String::new();
303 for (i, (k, v)) in labels.enumerate() {
304 validate_prometheus_name(k);
305 if i > 0 {
306 buf.push(',')
307 }
308 buf.push_str(k);
309 buf.push('=');
310 buf.push('"');
311 for c in v.chars() {
312 match c {
313 '\\' => {
314 buf.push('\\');
315 buf.push('\\');
316 }
317 '\n' => {
318 buf.push('\\');
319 buf.push('n');
320 }
321 '"' => {
322 buf.push('\\');
323 buf.push('"');
324 }
325 _ => buf.push(c),
326 }
327 }
328 buf.push('"');
329 }
330 buf
331 }
332
333 fn encode_value_with_labels(
334 &mut self,
335 name: &str,
336 label_values: &[(&str, &str)],
337 value: f64,
338 ) -> io::Result<()> {
339 writeln!(
340 self.writer,
341 "{}{{{}}} {} {}",
342 name,
343 Self::encode_labels(label_values.iter()),
344 FormattedValue(value),
345 self.now_millis
346 )
347 }
348}
349
350fn validate_prometheus_name(name: &str) {
353 if name.is_empty() {
354 panic!("Empty names are not allowed");
355 }
356 let bytes = name.as_bytes();
357 if (!bytes[0].is_ascii_alphabetic() && bytes[0] != b'_')
358 || !bytes[1..]
359 .iter()
360 .all(|c| c.is_ascii_alphanumeric() || *c == b'_')
361 {
362 panic!(
363 "Name '{}' does not match pattern [a-zA-Z_][a-zA-Z0-9_]",
364 name
365 );
366 }
367}