rustic_rs/metrics/
prometheus.rs

1use anyhow::{Context, Result, bail};
2use log::debug;
3use prometheus::register_gauge;
4use reqwest::Url;
5use std::collections::BTreeMap;
6
7use crate::metrics::MetricValue::*;
8
9use super::{Metric, MetricsExporter};
10
11pub struct PrometheusExporter {
12    pub endpoint: Url,
13    pub job_name: String,
14    pub grouping: BTreeMap<String, String>,
15    pub prometheus_user: Option<String>,
16    pub prometheus_pass: Option<String>,
17}
18
19impl MetricsExporter for PrometheusExporter {
20    fn push_metrics(&self, metrics: &[Metric]) -> Result<()> {
21        use prometheus::{Encoder, ProtobufEncoder};
22        use reqwest::{StatusCode, blocking::Client, header::CONTENT_TYPE};
23
24        for metric in metrics {
25            let gauge = register_gauge!(metric.name, metric.description,)
26                .context("registering prometheus gauge")?;
27
28            gauge.set(match metric.value {
29                Int(i) => i as f64,
30                Float(f) => f,
31            });
32        }
33
34        let (full_url, encoded_metrics) = self.make_url_and_encoded_metrics()?;
35
36        debug!("using url: {full_url}");
37
38        let mut builder = Client::new()
39            .post(full_url)
40            .header(CONTENT_TYPE, ProtobufEncoder::new().format_type())
41            .body(encoded_metrics);
42
43        if let Some(username) = &self.prometheus_user {
44            debug!(
45                "using auth {} {}",
46                username,
47                self.prometheus_pass.as_deref().unwrap_or("[NOT SET]")
48            );
49            builder = builder.basic_auth(username, self.prometheus_pass.as_ref());
50        }
51
52        let response = builder.send()?;
53
54        match response.status() {
55            StatusCode::ACCEPTED | StatusCode::OK => Ok(()),
56            _ => bail!(
57                "unexpected status code {} while pushing to {}",
58                response.status(),
59                self.endpoint
60            ),
61        }
62    }
63}
64
65impl PrometheusExporter {
66    // TODO: This should be actually part of the prometheus crate, see https://github.com/tikv/rust-prometheus/issues/536
67    fn make_url_and_encoded_metrics(&self) -> Result<(Url, Vec<u8>)> {
68        use base64::prelude::*;
69        use prometheus::{Encoder, ProtobufEncoder};
70
71        let mut url_components = vec![
72            "metrics".to_string(),
73            "job@base64".to_string(),
74            BASE64_URL_SAFE_NO_PAD.encode(&self.job_name),
75        ];
76
77        for (ln, lv) in &self.grouping {
78            // See https://github.com/tikv/rust-prometheus/issues/535
79            if !lv.is_empty() {
80                // TODO: check label name
81                let name = ln.to_string() + "@base64";
82                url_components.push(name);
83                url_components.push(BASE64_URL_SAFE_NO_PAD.encode(lv));
84            }
85        }
86        let url = self.endpoint.join(&url_components.join("/"))?;
87
88        let encoder = ProtobufEncoder::new();
89        let mut buf = Vec::new();
90        for mf in prometheus::gather() {
91            // Note: We don't check here for pre-existing grouping labels, as we don't set them
92
93            // Ignore error, `no metrics` and `no name`.
94            let _ = encoder.encode(&[mf], &mut buf);
95        }
96
97        Ok((url, buf))
98    }
99}
100
101#[cfg(feature = "prometheus")]
102#[test]
103fn test_make_url_and_encoded_metrics() -> Result<()> {
104    use std::str::FromStr;
105
106    let grouping = [
107        ("abc", "xyz"),
108        ("path", "/my/path"),
109        ("tags", "a,b,cde"),
110        ("nogroup", ""),
111    ]
112    .into_iter()
113    .map(|(a, b)| (a.to_string(), b.to_string()))
114    .collect();
115
116    let exporter = PrometheusExporter {
117        endpoint: Url::from_str("http://host")?,
118        job_name: "test_job".to_string(),
119        grouping,
120        prometheus_user: None,
121        prometheus_pass: None,
122    };
123
124    let (url, _) = exporter.make_url_and_encoded_metrics()?;
125    assert_eq!(
126        url.to_string(),
127        "http://host/metrics/job@base64/dGVzdF9qb2I/abc@base64/eHl6/path@base64/L215L3BhdGg/tags@base64/YSxiLGNkZQ"
128    );
129    Ok(())
130}
rustic_rs/metrics/prometheus.rs

rustic_rs/metrics/
prometheus.rs
