Skip to main content

wire_framework/vlog/
prometheus.rs

1//! Prometheus-related functionality, such as [`PrometheusExporterConfig`].
2
3use std::{net::Ipv4Addr, time::Duration};
4
5use eyre::Context as _;
6use tokio::sync::watch;
7use vise::MetricsCollection;
8use vise_exporter::MetricsExporter;
9
10#[derive(Debug)]
11enum PrometheusTransport {
12    Pull {
13        port: u16,
14    },
15    Push {
16        gateway_uri: String,
17        interval: Duration,
18    },
19}
20
21/// Configuration of a Prometheus exporter.
22#[derive(Debug)]
23pub struct PrometheusExporterConfig {
24    transport: PrometheusTransport,
25}
26
27impl PrometheusExporterConfig {
28    /// Creates an exporter that will run an HTTP server on the specified `port`.
29    pub const fn pull(port: u16) -> Self {
30        Self {
31            transport: PrometheusTransport::Pull { port },
32        }
33    }
34
35    /// Creates an exporter that will push metrics to the specified Prometheus gateway endpoint.
36    pub const fn push(gateway_uri: String, interval: Duration) -> Self {
37        Self {
38            transport: PrometheusTransport::Push {
39                gateway_uri,
40                interval,
41            },
42        }
43    }
44
45    /// Runs the exporter. This future should be spawned in a separate Tokio task.
46    pub async fn run(self, mut stop_receiver: watch::Receiver<bool>) -> eyre::Result<()> {
47        let registry = MetricsCollection::lazy().collect();
48        let metrics_exporter =
49            MetricsExporter::new(registry.into()).with_graceful_shutdown(async move {
50                stop_receiver.changed().await.ok();
51            });
52
53        match self.transport {
54            PrometheusTransport::Pull { port } => {
55                let prom_bind_address = (Ipv4Addr::UNSPECIFIED, port).into();
56                metrics_exporter
57                    .start(prom_bind_address)
58                    .await
59                    .context("Failed starting metrics server")?;
60            }
61            PrometheusTransport::Push {
62                gateway_uri,
63                interval,
64            } => {
65                let endpoint = gateway_uri
66                    .parse()
67                    .context("Failed parsing Prometheus push gateway endpoint")?;
68                metrics_exporter.push_to_gateway(endpoint, interval).await;
69            }
70        }
71        Ok(())
72    }
73}