1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use std::io::Write;

use protobuf::Message;

use crate::errors::Result;
use crate::proto::MetricFamily;

use super::{check_metric_family, Encoder};

/// The protocol buffer format of metric family.
pub const PROTOBUF_FORMAT: &str = "application/vnd.google.protobuf; \
                                   proto=io.prometheus.client.MetricFamily; \
                                   encoding=delimited";

/// An implementation of an [`Encoder`] that converts a [`MetricFamily`] proto
/// message into the binary wire format of protobuf.
#[derive(Debug, Default)]
pub struct ProtobufEncoder;

impl ProtobufEncoder {
    /// Create a new protobuf encoder.
    pub fn new() -> ProtobufEncoder {
        ProtobufEncoder
    }
}

impl Encoder for ProtobufEncoder {
    fn encode<W: Write>(&self, metric_families: &[MetricFamily], writer: &mut W) -> Result<()> {
        for mf in metric_families {
            // Fail-fast checks.
            check_metric_family(mf)?;
            mf.write_length_delimited_to_writer(writer)?;
        }
        Ok(())
    }

    fn format_type(&self) -> &str {
        PROTOBUF_FORMAT
    }
}

#[cfg(test)]
mod tests {
    use crate::counter::CounterVec;
    use crate::encoder::Encoder;
    use crate::metrics::Opts;
    use crate::registry;

    // TODO: add more tests.
    #[rustfmt::skip]
    #[test]
    fn test_protobuf_encoder() {
        let cv = CounterVec::new(Opts::new("test_counter_vec", "help information"),
                                 &["labelname"])
            .unwrap();
        let reg = registry::Registry::new();
        reg.register(Box::new(cv.clone())).unwrap();

        cv.get_metric_with_label_values(&["2230"]).unwrap().inc();
        let mf = reg.gather();
        let mut writer = Vec::<u8>::new();
        let encoder = super::ProtobufEncoder::new();
        let res = encoder.encode(&mf, &mut writer);
        assert!(res.is_ok());

        // Generated by a golang demo,
        // see more: https://gist.github.com/overvenus/bd39bde014b0cba87c9bde20dbea6ce0
        let ans = vec![70, 10, 16, 116, 101, 115, 116, 95, 99, 111, 117, 110, 116, 101, 114, 95,
                       118, 101, 99, 18, 16, 104, 101, 108, 112, 32, 105, 110, 102, 111, 114, 109,
                       97, 116, 105, 111, 110, 24, 0, 34, 30, 10, 17, 10, 9, 108, 97, 98, 101,
                       108, 110, 97, 109, 101, 18, 4, 50, 50, 51, 48, 26, 9, 9, 0, 0, 0, 0, 0, 0,
                       240, 63];
        assert_eq!(ans, writer);
    }
}