rustic_rs/metrics/
prometheus.rs

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

rustic_rs/metrics/
prometheus.rs
