promformat/
lib.rs

1use std::fmt::Write;
2
3#[derive(Default)]
4pub struct Metrics {
5    buffer: String,
6}
7
8impl Metrics {
9    pub fn new() -> Self {
10        Self::default()
11    }
12
13    pub fn gauge<'a>(&'a mut self, name: &str, help: &str) -> MetricGroup<'a> {
14        self.metric_group("gauge", name, help)
15    }
16
17    pub fn counter<'a>(&'a mut self, name: &str, help: &str) -> MetricGroup<'a> {
18        self.metric_group("counter", name, help)
19    }
20
21    fn metric_group<'a>(
22        &'a mut self,
23        type_: &'static str,
24        name: &str,
25        help: &str,
26    ) -> MetricGroup<'a> {
27        MetricGroup::new(self, type_, name, help)
28    }
29
30    pub fn render(&self) -> &str {
31        &self.buffer
32    }
33
34    pub fn into_rendered(self) -> String {
35        self.buffer
36    }
37}
38
39pub struct MetricGroup<'a> {
40    metrics: &'a mut Metrics,
41    name: String,
42}
43
44impl<'a> MetricGroup<'a> {
45    fn new(metrics: &'a mut Metrics, metric_type: &str, name: &str, help: &str) -> Self {
46        writeln!(&mut metrics.buffer, "# HELP {name} {help}").unwrap();
47        writeln!(&mut metrics.buffer, "# TYPE {name} {metric_type}").unwrap();
48        Self {
49            name: name.to_owned(),
50            metrics,
51        }
52    }
53
54    pub fn label(&mut self, label: impl AsRef<str>, value: impl AsRef<str>) -> SingleMetric {
55        SingleMetric {
56            name: &self.name,
57            metrics: self.metrics,
58            labels: String::default(),
59        }
60        .label(label, value)
61    }
62
63    pub fn set(self, value: impl std::fmt::Display) {
64        SingleMetric {
65            name: &self.name,
66            metrics: self.metrics,
67            labels: String::from("{"),
68        }
69        .set(value)
70    }
71}
72
73pub struct SingleMetric<'a, 'b> {
74    labels: String,
75    name: &'a String,
76    metrics: &'b mut Metrics,
77}
78
79impl<'a, 'b> SingleMetric<'a, 'b> {
80    pub fn label(mut self, label: impl AsRef<str>, value: impl AsRef<str>) -> Self {
81        let label = label.as_ref();
82        let value = value.as_ref();
83        if self.labels.is_empty() {
84            self.labels.push('{')
85        } else {
86            self.labels.push(',')
87        }
88
89        // TODO: escaping
90        write!(&mut self.labels, "{label}=\"{value}\"").unwrap();
91
92        self
93    }
94
95    pub fn set(mut self, value: impl std::fmt::Display) {
96        self.labels.push('}');
97        writeln!(self.metrics.buffer, "{}{} {value}", self.name, self.labels).unwrap();
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn test_gauge() {
107        let mut metrics = Metrics::default();
108
109        let mut gauge = metrics.gauge("testme", "help here");
110        gauge.label("x", "y").set(2);
111        gauge.label("a", "b").label("c", "d").set(20);
112
113        assert_eq!(
114            metrics.render(),
115            r#"# HELP testme help here
116# TYPE testme gauge
117testme{x="y"} 2
118testme{a="b",c="d"} 20
119"#
120        );
121    }
122
123    #[test]
124    fn test_gauge_no_labels() {
125        let mut metrics = Metrics::default();
126
127        metrics.gauge("testme", "help here").set(20);
128
129        assert_eq!(
130            metrics.render(),
131            r#"# HELP testme help here
132# TYPE testme gauge
133testme{} 20
134"#
135        );
136    }
137}
138
139#[cfg(doctest)]
140mod test_readme {
141    macro_rules! external_doc_test {
142        ($x:expr) => {
143            #[doc = $x]
144            extern "C" {}
145        };
146    }
147
148    external_doc_test!(include_str!("../README.md"));
149}