zombienet_orchestrator/network_helper/
metrics.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use reqwest::Url;
5
6#[async_trait]
7pub trait MetricsHelper {
8    async fn metric(&self, metric_name: &str) -> Result<f64, anyhow::Error>;
9    async fn metric_with_url(
10        metric: impl AsRef<str> + Send,
11        endpoint: impl Into<Url> + Send,
12    ) -> Result<f64, anyhow::Error>;
13}
14
15pub struct Metrics {
16    endpoint: Url,
17}
18
19impl Metrics {
20    fn new(endpoint: impl Into<Url>) -> Self {
21        Self {
22            endpoint: endpoint.into(),
23        }
24    }
25
26    async fn fetch_metrics(
27        endpoint: impl AsRef<str>,
28    ) -> Result<HashMap<String, f64>, anyhow::Error> {
29        let response = reqwest::get(endpoint.as_ref()).await?;
30        Ok(prom_metrics_parser::parse(&response.text().await?)?)
31    }
32
33    fn get_metric(
34        metrics_map: HashMap<String, f64>,
35        metric_name: &str,
36    ) -> Result<f64, anyhow::Error> {
37        let treat_not_found_as_zero = true;
38        if let Some(val) = metrics_map.get(metric_name) {
39            Ok(*val)
40        } else if treat_not_found_as_zero {
41            Ok(0_f64)
42        } else {
43            Err(anyhow::anyhow!("MetricNotFound: {metric_name}"))
44        }
45    }
46}
47
48#[async_trait]
49impl MetricsHelper for Metrics {
50    async fn metric(&self, metric_name: &str) -> Result<f64, anyhow::Error> {
51        let metrics_map = Metrics::fetch_metrics(self.endpoint.as_str()).await?;
52        Metrics::get_metric(metrics_map, metric_name)
53    }
54
55    async fn metric_with_url(
56        metric_name: impl AsRef<str> + Send,
57        endpoint: impl Into<Url> + Send,
58    ) -> Result<f64, anyhow::Error> {
59        let metrics_map = Metrics::fetch_metrics(endpoint.into()).await?;
60        Metrics::get_metric(metrics_map, metric_name.as_ref())
61    }
62}