zombienet_orchestrator/network_helper/
metrics.rs1use 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}