ic_metrics_encoder/
lib.rs

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        // If the value is not numeric, we need to represent it using Go library conventions.
13        //
14        // > `value` is a float represented as required by Go's ParseFloat() function.
15        // > In addition to standard numerical values, NaN, +Inf, and -Inf are
16        // > valid values representing not a number, positive infinity, and negative infinity, respectively.
17        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
30/// A helper for encoding metrics that use
31/// [labels](https://prometheus.io/docs/practices/naming/#labels).
32/// See [MetricsEncoder::counter_vec] and [MetricsEncoder::gauge_vec].
33pub 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    /// Encodes the metrics value observed for the specified values of labels.
43    ///
44    /// # Panics
45    ///
46    /// This function panics if one of the labels does not match pattern
47    /// [a-zA-Z_][a-zA-Z0-9_]. See
48    /// https://prometheus.io/docs/concepts/data_model/#metric-names-and-labels.
49    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
56/// A helper for encoding histograms that use
57/// [labels](https://prometheus.io/docs/practices/naming/#labels).
58/// See [MetricsEncoder::histogram_vec].
59pub 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    /// Encodes the metrics histogram observed for the given values of labels.
69    ///
70    /// # Panics
71    ///
72    /// This function panics if one of the labels does not match pattern
73    /// [a-zA-Z_][a-zA-Z0-9_]. See
74    /// https://prometheus.io/docs/concepts/data_model/#metric-names-and-labels.
75    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}
161/// `MetricsEncoder` provides methods to encode metrics in a text format
162/// that can be understood by Prometheus.
163///
164/// Metrics are encoded with the block time included, to allow Prometheus
165/// to discard out-of-order samples collected from replicas that are behind.
166///
167/// See [Exposition Formats][1] for an informal specification of the text
168/// format.
169///
170/// [1]: https://github.com/prometheus/docs/blob/master/content/docs/instrumenting/exposition_formats.md
171pub struct MetricsEncoder<W: io::Write> {
172    writer: W,
173    now_millis: i64,
174}
175
176impl<W: io::Write> MetricsEncoder<W> {
177    /// Constructs a new encoder dumping metrics with the given timestamp into
178    /// the specified writer.
179    pub fn new(writer: W, now_millis: i64) -> Self {
180        Self { writer, now_millis }
181    }
182
183    /// Returns the internal buffer that was used to record the
184    /// metrics.
185    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    /// Encodes the metadata and the value of a histogram.
195    ///
196    /// SUM is the sum of all observed values, before they were put
197    /// into buckets.
198    ///
199    /// BUCKETS is a list (key, value) pairs, where KEY is the bucket
200    /// and VALUE is the number of items *in* this bucket (i.e., it's
201    /// not a cumulative value).
202    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    /// Encodes the metadata and the value of a counter.
246    ///
247    /// # Panics
248    ///
249    /// This function panics if the `name` argument does not match pattern [a-zA-Z_][a-zA-Z0-9_].
250    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    /// Encodes the metadata and the value of a gauge.
255    ///
256    /// # Panics
257    ///
258    /// This function panics if the `name` argument does not match pattern [a-zA-Z_][a-zA-Z0-9_].
259    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    /// Starts encoding of a counter that uses
264    /// [labels](https://prometheus.io/docs/practices/naming/#labels).
265    ///
266    /// # Panics
267    ///
268    /// This function panics if the `name` argument does not match pattern [a-zA-Z_][a-zA-Z0-9_].
269    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    /// Starts encoding of a gauge that uses
283    /// [labels](https://prometheus.io/docs/practices/naming/#labels).
284    ///
285    /// # Panics
286    ///
287    /// This function panics if the `name` argument does not match pattern [a-zA-Z_][a-zA-Z0-9_].
288    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
350/// Panics if the specified string is not a valid Prometheus metric/label name.
351/// See https://prometheus.io/docs/concepts/data_model/#metric-names-and-labels.
352fn 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}