wire_framework/vlog/
prometheus.rs1use 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#[derive(Debug)]
23pub struct PrometheusExporterConfig {
24 transport: PrometheusTransport,
25}
26
27impl PrometheusExporterConfig {
28 pub const fn pull(port: u16) -> Self {
30 Self {
31 transport: PrometheusTransport::Pull { port },
32 }
33 }
34
35 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 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}