prometheus_json_encoder/
lib.rs

1// Adapted from : https://docs.rs/prometheus/0.13.0/src/prometheus/encoder/text.rs.html#21
2    // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
3
4
5use std::io::{self, Write};
6
7use prometheus::Result;
8use prometheus::proto::{self, MetricFamily, MetricType};
9
10use prometheus::{Encoder};
11use serde_json::{Map, Value, json};
12
13/// The text format of metric family.
14pub const TEXT_FORMAT: &str = "text/plain; version=0.0.4";
15
16const POSITIVE_INF: &str = "+Inf";
17
18/// An implementation of an [`Encoder`] that converts a [`MetricFamily`] proto message
19/// into json format.
20
21#[derive(Debug, Default)]
22pub struct JsonEncoder;
23
24impl JsonEncoder {
25    /// Create a new json encoder.
26    pub fn new() -> JsonEncoder {
27        JsonEncoder
28    }
29    /// Appends metrics to a given `String` buffer.
30    ///
31    /// This is a convenience wrapper around `<JsonEncoder as Encoder>::encode`.
32    pub fn encode_utf8(&self, metric_families: &[MetricFamily], buf: &mut String) -> Result<()> {
33        // Note: it's important to *not* re-validate UTF8-validity for the
34        // entirety of `buf`. Otherwise, repeatedly appending metrics to the
35        // same `buf` will lead to quadratic behavior. That's why we use
36        // `WriteUtf8` abstraction to skip the validation.
37        self.encode_impl(metric_families, &mut StringBuf(buf))?;
38        Ok(())
39    }
40    /// Converts metrics to `String`.
41    ///
42    /// This is a convenience wrapper around `<JsonEncoder as Encoder>::encode`.
43    pub fn encode_to_string(&self, metric_families: &[MetricFamily]) -> Result<String> {
44        let mut buf = String::new();
45        self.encode_utf8(metric_families, &mut buf)?;
46        Ok(buf)
47    }
48
49    fn encode_impl(
50        &self,
51        metric_families: &[MetricFamily],
52        writer: &mut dyn WriteUtf8,
53    ) -> Result<()> {
54
55        let mut map = Map::new();
56
57        for mf in metric_families {
58            
59            // Add entry for the metric
60            let name : &str = mf.get_name();
61            
62            let mut mf_map = Map::new();
63
64            // Add Help
65            let help : &str = mf.get_help();
66            mf_map.insert("help".to_string(), json!{help}); 
67
68            // Write `# TYPE` header.
69            let metric_type : MetricType = mf.get_field_type();
70            let lowercase_type = json!(format!("{:?}", metric_type).to_lowercase());
71            mf_map.insert("type".to_string(), lowercase_type); 
72
73            let mut debug_counter = 0;
74
75            for m in mf.get_metric() {
76                println!("{}", debug_counter);
77                debug_counter += 1;         // TODO : Remove
78                // Metric
79                match metric_type {
80                    MetricType::COUNTER => {
81                        mf_map.insert("counter".to_string(), json!(m.get_counter().get_value()));
82                        extra_info(&mut mf_map, m);
83                        // f64
84                    }
85                    MetricType::GAUGE => {
86                        mf_map.insert("gauge".to_string(), json!(m.get_gauge().get_value()));
87                        extra_info(&mut mf_map, m);
88                        // f64
89                    }
90                    MetricType::HISTOGRAM => {
91                        let h = m.get_histogram();
92                        let mut upper_bounds : Vec<Value> = vec![];
93                        let mut cumulative_counts : Vec<Value> = vec![];
94                        let mut inf_seen = false;
95                        
96                        for b in h.get_bucket() {
97                            let upper_bound = b.get_upper_bound();           // f64
98                            let cumulative_count = b.get_cumulative_count(); // f64
99
100                            upper_bounds.push(json!(upper_bound));
101                            cumulative_counts.push(json!(cumulative_count));
102
103                            if upper_bound.is_sign_positive() && upper_bound.is_infinite() {
104                                inf_seen = true;
105                            }
106                        }
107                        if !inf_seen {
108                            upper_bounds.push(json!(POSITIVE_INF));
109                            cumulative_counts.push(json!(h.get_sample_count()));
110                        }
111                        let names = [
112                            "cumulative_counts".to_string(),
113                            "upper_bounds".to_string(),
114                            "sum".to_string(),
115                            "counts".to_string()
116                        ];
117
118                        let values = [
119                            json!(cumulative_counts),
120                            json!(upper_bounds),
121                            json!(h.get_sample_sum()),
122                            json!(h.get_sample_count())
123                        ];
124                        for (key, value) in names.into_iter().zip(values.into_iter()) {
125                            mf_map.insert(key, value);
126                        }
127                        extra_info(&mut mf_map, m);
128                    }
129
130                    MetricType::SUMMARY => {
131                        let s = m.get_summary();
132                        let mut quantiles = vec![];
133                        let mut values = vec![];
134
135                        for q in s.get_quantile() {
136                            quantiles.push(json!(q.get_quantile()));
137                            values.push(q.get_value());
138                        }
139
140                        let names = [
141                            "sum".to_string(),
142                            "quantiles".to_string()
143                        ];
144
145                        let values = [
146                            json!(s.get_sample_sum()),
147                            json!(s.get_sample_count())
148                        ];
149                        for (key, value) in names.into_iter().zip(values.into_iter()) {
150                            mf_map.insert(key, value);
151                        }
152                        extra_info(&mut mf_map, m);
153                    }
154                    MetricType::UNTYPED => {
155                        unimplemented!();
156                    }
157                }
158            }
159            map.insert(name.to_string(), json!(mf_map));
160        }
161
162        let x = serde_json::to_vec(&map).unwrap();
163        //println!{"{}", &x};
164        writer.write_all(&x)?; // String 
165        Ok(())
166    }
167    
168}
169
170impl Encoder for JsonEncoder {
171    fn encode<W: Write>(&self, metric_families: &[MetricFamily], writer: &mut W) -> Result<()> {
172        self.encode_impl(metric_families, &mut *writer)
173    }
174
175    fn format_type(&self) -> &str {
176        TEXT_FORMAT
177    }
178}
179
180// Adds into a map m.timestamp and m.LabelPair.
181// names and values must be of the same length
182fn extra_info(
183    map : &mut Map<String, Value>,
184    mc: &proto::Metric
185) -> () {
186
187    let timestamp = mc.get_timestamp_ms();
188    if timestamp != 0 {
189        map.insert("timestamp".to_string(), json!(timestamp));
190    }
191
192    for lp in mc.get_label() {
193        map.insert(lp.get_name().to_string(), json!(lp.get_value()));
194    }
195}
196
197trait WriteUtf8 {
198    fn write_all(&mut self, text: &[u8]) -> io::Result<()>;
199}
200
201impl<W: Write> WriteUtf8 for W {
202    fn write_all(&mut self, text: &[u8]) -> io::Result<()> {
203        Write::write_all(self, text)
204    }
205}
206
207/// Coherence forbids to impl `WriteUtf8` directly on `String`, need this
208/// wrapper as a work-around.
209struct StringBuf<'a>(&'a mut String);
210
211impl WriteUtf8 for StringBuf<'_> {
212    fn write_all(&mut self, text: &[u8]) -> io::Result<()> {
213        self.0.push_str(std::str::from_utf8(text).unwrap());
214        Ok(())
215    }
216}
217
218
219#[cfg(test)]
220mod tests {
221
222    use super::*;
223    use prometheus::Counter;
224    use prometheus::Gauge;
225    use prometheus::{Histogram, HistogramOpts};
226    use prometheus::Opts;
227    use prometheus::core::Collector;
228
229    #[test]
230    fn test_json_encoder() {
231        let counter_opts = Opts::new("test_counter", "test help")
232            .const_label("a", "1")
233            .const_label("b", "2");
234        let counter = Counter::with_opts(counter_opts).unwrap();
235        counter.inc();
236
237        let mf = counter.collect();
238        let mut writer = Vec::<u8>::new();
239        let encoder = JsonEncoder::new();
240        let txt = encoder.encode(&mf, &mut writer);
241        assert!(txt.is_ok());
242
243        // Object({"test_counter": Object({"a": String("1"), "b": String("2"), "counter": Number(1.0), "help": String("test help"), "type": String("counter")})})
244        let v: Value = serde_json::from_slice(&writer).unwrap();
245        
246
247        assert_eq!(v["test_counter"]["help"], "test help");
248        assert_eq!(v["test_counter"]["a"], "1");
249        assert_eq!(v["test_counter"]["b"], "2");
250        assert_eq!(v["test_counter"]["type"], "counter");
251        assert_eq!(v["test_counter"]["counter"], 1.0);
252        
253
254
255        let gauge_opts = Opts::new("test_gauge", "test help")
256            .const_label("a", "1")
257            .const_label("b", "2");
258        let gauge = Gauge::with_opts(gauge_opts).unwrap();
259        gauge.inc();
260        gauge.set(42.0);
261
262        let mf = gauge.collect();
263        writer.clear();
264        let txt = encoder.encode(&mf, &mut writer);
265        assert!(txt.is_ok());
266        let v: Value = serde_json::from_slice(&writer).unwrap();
267        println!("{:?}", v);
268        assert_eq!(v["test_gauge"]["help"], "test help");
269        assert_eq!(v["test_gauge"]["a"], "1");
270        assert_eq!(v["test_gauge"]["b"], "2");
271        assert_eq!(v["test_gauge"]["type"], "gauge");
272        assert_eq!(v["test_gauge"]["gauge"], json!(42.0));
273
274    }
275
276    
277    #[test]
278    fn test_text_encoder_histogram() {
279        let opts = HistogramOpts::new("test_histogram", "test help").const_label("a", "1");
280        let histogram = Histogram::with_opts(opts).unwrap();
281        histogram.observe(0.25);
282
283        let mf = histogram.collect();
284        let mut writer = Vec::<u8>::new();
285        let encoder = JsonEncoder::new();
286        let res = encoder.encode(&mf, &mut writer);
287        assert!(res.is_ok());
288
289        let v: Value = serde_json::from_slice(&writer).unwrap();
290
291        println!("{:?}", v);
292        assert_eq!(v["test_histogram"]["help"], "test help");
293        assert_eq!(v["test_histogram"]["type"], "histogram");
294        assert_eq!(v["test_histogram"]["a"], "1");
295        assert_eq!(v["test_histogram"]["counts"], json!(1));
296        assert_eq!(v["test_histogram"]["sum"], json!(0.25));
297        assert_eq!(v["test_histogram"]["cumulative_counts"], json!([0,0,0,0,0,1,1,1,1,1,1,1]));
298        let array_test = [json!(0.005), json!(0.01), json!(0.025), json!(0.05), json!(0.1), json!(0.25), json!(0.5), json!(1.0), json!(2.5), json!(5.0), json!(10.0), json!(POSITIVE_INF)];
299        assert_eq!(v["test_histogram"]["upper_bounds"], json!(array_test));
300
301
302    }
303
304    #[test]
305    fn test_text_encoder_summary() {
306        use prometheus::proto::{Metric, Quantile, Summary};
307        use protobuf::RepeatedField;
308
309        let mut metric_family = MetricFamily::default();
310        metric_family.set_name("test_summary".to_string());
311        metric_family.set_help("This is a test summary statistic".to_string());
312        metric_family.set_field_type(MetricType::SUMMARY);
313
314        let mut summary = Summary::default();
315        summary.set_sample_count(5.0 as u64);
316        summary.set_sample_sum(15.0);
317
318        let mut quantile1 = Quantile::default();
319        quantile1.set_quantile(50.0);
320        quantile1.set_value(3.0);
321
322        let mut quantile2 = Quantile::default();
323        quantile2.set_quantile(100.0);
324        quantile2.set_value(5.0);
325
326        summary.set_quantile(RepeatedField::from_vec(vec!(quantile1, quantile2)));
327
328        let mut metric = Metric::default();
329        metric.set_summary(summary);
330        metric_family.set_metric(RepeatedField::from_vec(vec!(metric)));
331
332        let mut writer = Vec::<u8>::new();
333        let encoder = JsonEncoder::new();
334        let res = encoder.encode(&vec![metric_family], &mut writer);
335        assert!(res.is_ok());
336        
337        let v: Value = serde_json::from_slice(&writer).unwrap();
338        println!("{:?}", v);
339        assert_eq!(v["test_summary"]["help"], "This is a test summary statistic");
340        assert_eq!(v["test_summary"]["type"], "summary");
341        assert_eq!(v["test_summary"]["quantiles"], json!(5));
342        assert_eq!(v["test_summary"]["sum"], json!(15));
343
344    }
345 
346    #[test]
347    fn test_text_encoder_to_string() {
348        let counter_opts = Opts::new("test_counter", "test help")
349            .const_label("a", "1")
350            .const_label("b", "2");
351        let counter = Counter::with_opts(counter_opts).unwrap();
352        counter.inc();
353
354        let mf = counter.collect();
355
356        let encoder = JsonEncoder::new();
357        let txt = encoder.encode_to_string(&mf);
358        let txt = txt.unwrap();
359        let v : Value = serde_json::from_str(txt.as_str()).unwrap();
360
361
362        assert_eq!(v["test_counter"]["help"], "test help");
363        assert_eq!(v["test_counter"]["a"], "1");
364        assert_eq!(v["test_counter"]["b"], "2");
365        assert_eq!(v["test_counter"]["counter"], json!(1.0));
366    }
367    
368}