Skip to main content

prometheus_exposition_format_rs/
lib.rs

1use crate::comment::{comment_parser, CommentType};
2use crate::common::empty_line_parser;
3use crate::samples::{parse_sample, SampleEntry};
4use crate::types::{Err, Metric, MetricType, Sample};
5use nom::branch::alt;
6use nom::combinator::map;
7use nom::IResult;
8use std::collections::HashMap;
9
10// Restrict this to internal visibility only
11pub(crate) mod comment;
12pub(crate) mod common;
13pub(crate) mod samples;
14pub mod types;
15
16#[derive(Debug)]
17enum LineType<'a> {
18    Empty,
19    Sample(SampleEntry<'a>),
20    Comment(CommentType<'a>),
21}
22
23fn parse_line(input: &str) -> IResult<&str, LineType> {
24    alt((
25        map(comment_parser, |l| LineType::Comment(l)),
26        map(parse_sample, |l| LineType::Sample(l)),
27        map(empty_line_parser, |_| LineType::Empty),
28    ))(input)
29}
30
31struct InputIter<'a>(&'a str);
32
33impl<'a> Iterator for InputIter<'a> {
34    type Item = Result<LineType<'a>, Err>;
35
36    fn next(&mut self) -> Option<Self::Item> {
37        if self.0.len() == 0 {
38            None
39        } else {
40            match parse_line(self.0) {
41                Ok(res) => {
42                    self.0 = res.0;
43                    Some(Ok(res.1))
44                }
45                Result::Err(err) => Some(Result::Err(Err::from(err))),
46            }
47        }
48    }
49}
50
51impl<'a> Into<Metric> for SampleEntry<'a> {
52    fn into(self) -> Metric {
53        Metric {
54            name: self.name.to_string(),
55            data_type: MetricType::Untyped,
56            samples: vec![self.into()],
57        }
58    }
59}
60
61impl<'a> Into<Sample> for SampleEntry<'a> {
62    fn into(self) -> Sample {
63        Sample {
64            labels: self
65                .labels
66                .iter()
67                .map(|(&k, v)| (k.to_string(), v.to_string()))
68                .collect(),
69            value: self.value,
70            timestamp: self.timestamp_ms,
71        }
72    }
73}
74
75impl Metric {
76    fn append_sample_entry(&mut self, s: SampleEntry) {
77        assert_eq!(
78            s.name, self.name,
79            "Names should be equal when calling update on a metric"
80        );
81        self.push_sample(s.into());
82    }
83    fn append_type_def(&mut self, s: &str, t: MetricType) {
84        assert_eq!(
85            s,
86            &self.name[..],
87            "Names should be equal when calling update on a metric"
88        );
89        self.data_type = t;
90    }
91}
92
93fn add_comment<'a, 'b>(map: &mut HashMap<&'a str, Metric>, c: CommentType<'a>) {
94    if let CommentType::Type(s, t) = c {
95        if let Some(x) = map.get_mut(s) {
96            x.append_type_def(s, t);
97        } else {
98            map.insert(s, Metric::new(s, t));
99        }
100    }
101}
102
103fn add_sample<'a, 'b>(map: &'b mut HashMap<&'a str, Metric>, s: SampleEntry<'a>) {
104    if let Some(x) = map.get_mut(s.name) {
105        x.append_sample_entry(s);
106    } else {
107        map.insert(s.name, s.into());
108    };
109}
110
111/// Parse a string and return a vector of metrics extracted from it.
112pub fn parse_complete<'a>(input: &'a str) -> Result<Vec<Metric>, Err> {
113    let mut acc: HashMap<&'a str, Metric> = HashMap::new();
114    for l in InputIter(input) {
115        match l? {
116            LineType::Comment(c) => add_comment(&mut acc, c),
117            LineType::Sample(s) => add_sample(&mut acc, s),
118            LineType::Empty => {}
119        };
120    }
121    let mut res: Vec<Metric> = acc.drain().map(|(_, v)| v).collect();
122    // Make the order constant
123    res.sort_unstable_by(|a, b| a.name.cmp(&b.name));
124    Ok(res)
125}
126
127#[cfg(test)]
128fn assert_metric(m: &Metric, name: &str, tpe: MetricType, samples: Vec<Sample>) {
129    assert_eq!(m.name, name, "name {:?}", m);
130    assert_eq!(m.data_type, tpe, "type {:?}", m);
131    assert_eq!(m.samples, samples);
132}
133
134#[test]
135fn test_parse_summary() {
136    let res = parse_complete(
137        r#"
138# TYPE chain_account_commits summary
139chain_account_commits {quantile="0.5"} 0
140
141# TYPE chain_account_commits summary
142chain_account_commits {quantile="0.75"} 123
143
144# TYPE chain_account_commits summary
145chain_account_commits {quantile="0.95"} 50
146"#,
147    )
148    .unwrap();
149    assert_eq!(res.len(), 1);
150    assert_metric(
151        &res[0],
152        "chain_account_commits",
153        MetricType::Summary,
154        vec![
155            Sample::new(0f64, Option::None, vec!["quantile", "0.5"]),
156            Sample::new(123f64, Option::None, vec!["quantile", "0.75"]),
157            Sample::new(50f64, Option::None, vec!["quantile", "0.95"]),
158        ],
159    );
160}
161
162#[test]
163fn test_parse_complete() {
164    let res = parse_complete(
165        r#"
166# HELP http_requests_total The total number of HTTP requests.
167# TYPE http_requests_total counter
168http_requests_total{method="post",code="200"} 1027 1395066363000
169http_requests_total{method="post",code="400"} 1028 1395066363000
170
171rpc_duration_seconds_count 2693
172"#,
173    )
174    .unwrap();
175    assert_eq!(res.len(), 2);
176    assert_metric(
177        &res[0],
178        "http_requests_total",
179        MetricType::Counter,
180        vec![
181            Sample::new(
182                1027f64,
183                Option::Some(1395066363000),
184                vec!["method", "post", "code", "200"],
185            ),
186            Sample::new(
187                1028f64,
188                Option::Some(1395066363000),
189                vec!["method", "post", "code", "400"],
190            ),
191        ],
192    );
193    assert_metric(
194        &res[1],
195        "rpc_duration_seconds_count",
196        MetricType::Untyped,
197        vec![Sample::new(2693f64, None, vec![])],
198    );
199}