prometheus/encoder/
text.rs

1// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
2
3use std::borrow::Cow;
4use std::io::{self, Write};
5
6use crate::errors::Result;
7use crate::histogram::BUCKET_LABEL;
8use crate::proto::{self, MetricFamily, MetricType};
9#[cfg(feature = "protobuf")]
10use crate::proto_ext::MessageFieldExt;
11
12use super::{check_metric_family, Encoder};
13
14/// The text format of metric family.
15pub const TEXT_FORMAT: &str = "text/plain; version=0.0.4";
16
17const POSITIVE_INF: &str = "+Inf";
18const QUANTILE: &str = "quantile";
19
20/// An implementation of an [`Encoder`] that converts a [`MetricFamily`] proto message
21/// into text format.
22#[derive(Debug, Default)]
23pub struct TextEncoder;
24
25impl TextEncoder {
26    /// Create a new text encoder.
27    pub fn new() -> TextEncoder {
28        TextEncoder
29    }
30    /// Appends metrics to a given `String` buffer.
31    ///
32    /// This is a convenience wrapper around `<TextEncoder as Encoder>::encode`.
33    pub fn encode_utf8(&self, metric_families: &[MetricFamily], buf: &mut String) -> Result<()> {
34        // Note: it's important to *not* re-validate UTF8-validity for the
35        // entirety of `buf`. Otherwise, repeatedly appending metrics to the
36        // same `buf` will lead to quadratic behavior. That's why we use
37        // `WriteUtf8` abstraction to skip the validation.
38        self.encode_impl(metric_families, &mut StringBuf(buf))?;
39        Ok(())
40    }
41    /// Converts metrics to `String`.
42    ///
43    /// This is a convenience wrapper around `<TextEncoder as Encoder>::encode`.
44    pub fn encode_to_string(&self, metric_families: &[MetricFamily]) -> Result<String> {
45        let mut buf = String::new();
46        self.encode_utf8(metric_families, &mut buf)?;
47        Ok(buf)
48    }
49
50    fn encode_impl(
51        &self,
52        metric_families: &[MetricFamily],
53        writer: &mut dyn WriteUtf8,
54    ) -> Result<()> {
55        for mf in metric_families {
56            // Fail-fast checks.
57            check_metric_family(mf)?;
58
59            // Write `# HELP` header.
60            let name = mf.name();
61            let help = mf.help();
62            if !help.is_empty() {
63                writer.write_all("# HELP ")?;
64                writer.write_all(name)?;
65                writer.write_all(" ")?;
66                writer.write_all(&escape_string(help, false))?;
67                writer.write_all("\n")?;
68            }
69
70            // Write `# TYPE` header.
71            let metric_type = mf.get_field_type();
72            let lowercase_type = format!("{:?}", metric_type).to_lowercase();
73            writer.write_all("# TYPE ")?;
74            writer.write_all(name)?;
75            writer.write_all(" ")?;
76            writer.write_all(&lowercase_type)?;
77            writer.write_all("\n")?;
78
79            for m in mf.get_metric() {
80                match metric_type {
81                    MetricType::COUNTER => {
82                        write_sample(writer, name, None, m, None, m.get_counter().get_value())?;
83                    }
84                    MetricType::GAUGE => {
85                        write_sample(writer, name, None, m, None, m.get_gauge().get_value())?;
86                    }
87                    MetricType::HISTOGRAM => {
88                        let h = m.get_histogram();
89
90                        let mut inf_seen = false;
91                        for b in h.get_bucket() {
92                            let upper_bound = b.upper_bound();
93                            write_sample(
94                                writer,
95                                name,
96                                Some("_bucket"),
97                                m,
98                                Some((BUCKET_LABEL, &upper_bound.to_string())),
99                                b.cumulative_count() as f64,
100                            )?;
101                            if upper_bound.is_sign_positive() && upper_bound.is_infinite() {
102                                inf_seen = true;
103                            }
104                        }
105                        if !inf_seen {
106                            write_sample(
107                                writer,
108                                name,
109                                Some("_bucket"),
110                                m,
111                                Some((BUCKET_LABEL, POSITIVE_INF)),
112                                h.get_sample_count() as f64,
113                            )?;
114                        }
115
116                        write_sample(writer, name, Some("_sum"), m, None, h.get_sample_sum())?;
117
118                        write_sample(
119                            writer,
120                            name,
121                            Some("_count"),
122                            m,
123                            None,
124                            h.get_sample_count() as f64,
125                        )?;
126                    }
127                    MetricType::SUMMARY => {
128                        let s = m.get_summary();
129
130                        for q in s.get_quantile().iter() {
131                            write_sample(
132                                writer,
133                                name,
134                                None,
135                                m,
136                                Some((QUANTILE, &q.quantile().to_string())),
137                                q.value(),
138                            )?;
139                        }
140
141                        write_sample(writer, name, Some("_sum"), m, None, s.sample_sum())?;
142
143                        write_sample(
144                            writer,
145                            name,
146                            Some("_count"),
147                            m,
148                            None,
149                            s.sample_count() as f64,
150                        )?;
151                    }
152                    MetricType::UNTYPED => {
153                        unimplemented!();
154                    }
155                }
156            }
157        }
158
159        Ok(())
160    }
161}
162
163impl Encoder for TextEncoder {
164    fn encode<W: Write>(&self, metric_families: &[MetricFamily], writer: &mut W) -> Result<()> {
165        self.encode_impl(metric_families, &mut *writer)
166    }
167
168    fn format_type(&self) -> &str {
169        TEXT_FORMAT
170    }
171}
172
173/// `write_sample` writes a single sample in text format to `writer`, given the
174/// metric name, an optional metric name postfix, the metric proto message
175/// itself, optionally an additional label name and value (use empty strings if
176/// not required), and the value. The function returns the number of bytes
177/// written and any error encountered.
178fn write_sample(
179    writer: &mut dyn WriteUtf8,
180    name: &str,
181    name_postfix: Option<&str>,
182    mc: &proto::Metric,
183    additional_label: Option<(&str, &str)>,
184    value: f64,
185) -> Result<()> {
186    writer.write_all(name)?;
187    if let Some(postfix) = name_postfix {
188        writer.write_all(postfix)?;
189    }
190
191    label_pairs_to_text(mc.get_label(), additional_label, writer)?;
192
193    writer.write_all(" ")?;
194    writer.write_all(&value.to_string())?;
195
196    let timestamp = mc.timestamp_ms();
197    if timestamp != 0 {
198        writer.write_all(" ")?;
199        writer.write_all(&timestamp.to_string())?;
200    }
201
202    writer.write_all("\n")?;
203
204    Ok(())
205}
206
207/// `label_pairs_to_text` converts a slice of `LabelPair` proto messages plus
208/// the explicitly given additional label pair into text formatted as required
209/// by the text format and writes it to `writer`. An empty slice in combination
210/// with an empty string `additional_label_name` results in nothing being
211/// written. Otherwise, the label pairs are written, escaped as required by the
212/// text format, and enclosed in '{...}'. The function returns the number of
213/// bytes written and any error encountered.
214fn label_pairs_to_text(
215    pairs: &[proto::LabelPair],
216    additional_label: Option<(&str, &str)>,
217    writer: &mut dyn WriteUtf8,
218) -> Result<()> {
219    if pairs.is_empty() && additional_label.is_none() {
220        return Ok(());
221    }
222
223    let mut separator = "{";
224    for lp in pairs {
225        writer.write_all(separator)?;
226        writer.write_all(lp.name())?;
227        writer.write_all("=\"")?;
228        writer.write_all(&escape_string(lp.value(), true))?;
229        writer.write_all("\"")?;
230
231        separator = ",";
232    }
233
234    if let Some((name, value)) = additional_label {
235        writer.write_all(separator)?;
236        writer.write_all(name)?;
237        writer.write_all("=\"")?;
238        writer.write_all(&escape_string(value, true))?;
239        writer.write_all("\"")?;
240    }
241
242    writer.write_all("}")?;
243
244    Ok(())
245}
246
247fn find_first_occurence(v: &str, include_double_quote: bool) -> Option<usize> {
248    if include_double_quote {
249        memchr::memchr3(b'\\', b'\n', b'\"', v.as_bytes())
250    } else {
251        memchr::memchr2(b'\\', b'\n', v.as_bytes())
252    }
253}
254
255/// `escape_string` replaces `\` by `\\`, new line character by `\n`, and `"` by `\"` if
256/// `include_double_quote` is true.
257///
258/// Implementation adapted from
259/// https://lise-henry.github.io/articles/optimising_strings.html
260fn escape_string(v: &str, include_double_quote: bool) -> Cow<'_, str> {
261    let first_occurence = find_first_occurence(v, include_double_quote);
262
263    if let Some(first) = first_occurence {
264        let mut escaped = String::with_capacity(v.len() * 2);
265        escaped.push_str(&v[0..first]);
266        let remainder = v[first..].chars();
267
268        for c in remainder {
269            match c {
270                '\\' | '\n' => {
271                    escaped.extend(c.escape_default());
272                }
273                '"' if include_double_quote => {
274                    escaped.extend(c.escape_default());
275                }
276                _ => {
277                    escaped.push(c);
278                }
279            }
280        }
281
282        escaped.shrink_to_fit();
283        escaped.into()
284    } else {
285        // The input string does not contain any characters that would need to
286        // be escaped. Return it as it is.
287        v.into()
288    }
289}
290
291trait WriteUtf8 {
292    fn write_all(&mut self, text: &str) -> io::Result<()>;
293}
294
295impl<W: Write> WriteUtf8 for W {
296    fn write_all(&mut self, text: &str) -> io::Result<()> {
297        Write::write_all(self, text.as_bytes())
298    }
299}
300
301/// Coherence forbids to impl `WriteUtf8` directly on `String`, need this
302/// wrapper as a work-around.
303struct StringBuf<'a>(&'a mut String);
304
305impl WriteUtf8 for StringBuf<'_> {
306    fn write_all(&mut self, text: &str) -> io::Result<()> {
307        self.0.push_str(text);
308        Ok(())
309    }
310}
311
312#[cfg(test)]
313mod tests {
314
315    use super::*;
316    use crate::counter::Counter;
317    use crate::gauge::Gauge;
318    use crate::histogram::{Histogram, HistogramOpts};
319    use crate::metrics::{Collector, Opts};
320
321    #[test]
322    fn test_escape_string() {
323        assert_eq!(r"\\", escape_string("\\", false));
324        assert_eq!(r"a\\", escape_string("a\\", false));
325        assert_eq!(r"\n", escape_string("\n", false));
326        assert_eq!(r"a\n", escape_string("a\n", false));
327        assert_eq!(r"\\n", escape_string("\\n", false));
328
329        assert_eq!(r##"\\n\""##, escape_string("\\n\"", true));
330        assert_eq!(r##"\\\n\""##, escape_string("\\\n\"", true));
331        assert_eq!(r##"\\\\n\""##, escape_string("\\\\n\"", true));
332        assert_eq!(r##"\"\\n\""##, escape_string("\"\\n\"", true));
333    }
334
335    #[test]
336    fn test_text_encoder() {
337        let counter_opts = Opts::new("test_counter", "test help")
338            .const_label("a", "1")
339            .const_label("b", "2");
340        let counter = Counter::with_opts(counter_opts).unwrap();
341        counter.inc();
342
343        let mf = counter.collect();
344        let mut writer = Vec::<u8>::new();
345        let encoder = TextEncoder::new();
346        let txt = encoder.encode(&mf, &mut writer);
347        assert!(txt.is_ok());
348
349        let counter_ans = r##"# HELP test_counter test help
350# TYPE test_counter counter
351test_counter{a="1",b="2"} 1
352"##;
353        assert_eq!(counter_ans.as_bytes(), writer.as_slice());
354
355        let gauge_opts = Opts::new("test_gauge", "test help")
356            .const_label("a", "1")
357            .const_label("b", "2");
358        let gauge = Gauge::with_opts(gauge_opts).unwrap();
359        gauge.inc();
360        gauge.set(42.0);
361
362        let mf = gauge.collect();
363        writer.clear();
364        let txt = encoder.encode(&mf, &mut writer);
365        assert!(txt.is_ok());
366
367        let gauge_ans = r##"# HELP test_gauge test help
368# TYPE test_gauge gauge
369test_gauge{a="1",b="2"} 42
370"##;
371        assert_eq!(gauge_ans.as_bytes(), writer.as_slice());
372    }
373
374    #[test]
375    fn test_text_encoder_histogram() {
376        let opts = HistogramOpts::new("test_histogram", "test help").const_label("a", "1");
377        let histogram = Histogram::with_opts(opts).unwrap();
378        histogram.observe(0.25);
379
380        let mf = histogram.collect();
381        let mut writer = Vec::<u8>::new();
382        let encoder = TextEncoder::new();
383        let res = encoder.encode(&mf, &mut writer);
384        assert!(res.is_ok());
385
386        let ans = r##"# HELP test_histogram test help
387# TYPE test_histogram histogram
388test_histogram_bucket{a="1",le="0.005"} 0
389test_histogram_bucket{a="1",le="0.01"} 0
390test_histogram_bucket{a="1",le="0.025"} 0
391test_histogram_bucket{a="1",le="0.05"} 0
392test_histogram_bucket{a="1",le="0.1"} 0
393test_histogram_bucket{a="1",le="0.25"} 1
394test_histogram_bucket{a="1",le="0.5"} 1
395test_histogram_bucket{a="1",le="1"} 1
396test_histogram_bucket{a="1",le="2.5"} 1
397test_histogram_bucket{a="1",le="5"} 1
398test_histogram_bucket{a="1",le="10"} 1
399test_histogram_bucket{a="1",le="+Inf"} 1
400test_histogram_sum{a="1"} 0.25
401test_histogram_count{a="1"} 1
402"##;
403        assert_eq!(ans.as_bytes(), writer.as_slice());
404    }
405
406    #[test]
407    fn test_text_encoder_summary() {
408        use crate::proto::{Metric, Quantile, Summary};
409        use std::str;
410
411        let mut metric_family = MetricFamily::default();
412        metric_family.set_name("test_summary".to_string());
413        metric_family.set_help("This is a test summary statistic".to_string());
414        metric_family.set_field_type(MetricType::SUMMARY);
415
416        let mut summary = Summary::default();
417        summary.set_sample_count(5.0 as u64);
418        summary.set_sample_sum(15.0);
419
420        let mut quantile1 = Quantile::default();
421        quantile1.set_quantile(50.0);
422        quantile1.set_value(3.0);
423
424        let mut quantile2 = Quantile::default();
425        quantile2.set_quantile(100.0);
426        quantile2.set_value(5.0);
427
428        summary.set_quantile(vec![quantile1, quantile2]);
429
430        let mut metric = Metric::default();
431        metric.set_summary(summary);
432        metric_family.set_metric(vec![metric]);
433
434        let mut writer = Vec::<u8>::new();
435        let encoder = TextEncoder::new();
436        let res = encoder.encode(&vec![metric_family], &mut writer);
437        assert!(res.is_ok());
438
439        let ans = r##"# HELP test_summary This is a test summary statistic
440# TYPE test_summary summary
441test_summary{quantile="50"} 3
442test_summary{quantile="100"} 5
443test_summary_sum 15
444test_summary_count 5
445"##;
446        assert_eq!(ans, str::from_utf8(writer.as_slice()).unwrap());
447    }
448
449    #[test]
450    fn test_text_encoder_to_string() {
451        let counter_opts = Opts::new("test_counter", "test help")
452            .const_label("a", "1")
453            .const_label("b", "2");
454        let counter = Counter::with_opts(counter_opts).unwrap();
455        counter.inc();
456
457        let mf = counter.collect();
458
459        let encoder = TextEncoder::new();
460        let txt = encoder.encode_to_string(&mf);
461        let txt = txt.unwrap();
462
463        let counter_ans = r##"# HELP test_counter test help
464# TYPE test_counter counter
465test_counter{a="1",b="2"} 1
466"##;
467        assert_eq!(counter_ans, txt.as_str());
468    }
469}
prometheus/encoder/text.rs

prometheus/encoder/
text.rs
