prometheus/encoder/
text.rs

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

prometheus/encoder/
text.rs
