tracker-redis-exporter/src/metrics/project_metrics.rs

use prometheus::{GaugeVec, IntGaugeVec, Opts, Registry};
use prometheus::core::{Atomic, GenericGaugeVec};
use redis::{Commands, FromRedisValue};
use crate::client_manager::{ClientManager, RedisConnection};
use crate::helpers::get_last_requests_time;
use crate::metrics::MetricsProvider;
use crate::models::TrackerInfo;

fn update_gauge_from_redis<A: Atomic<T=T>, T: FromRedisValue>(
    con: &mut RedisConnection,
    project: &str,
    redis_key: &str,
    metric: &GenericGaugeVec<A>,
) -> color_eyre::Result<()> {
    let key = format!("{project}:{redis_key}");
    let val: Option<T> = con.get(&key)?;
    if let Some(v) = val {
        let m = metric.get_metric_with_label_values(&[project])?;
        m.set(v);
    }
    Ok(())
}

pub struct ProjectMetrics {
    registry: Registry,
    item_request_serve_rate: GaugeVec,
    reclaim_serve_rate: GaugeVec,
    reclaim_rate: GaugeVec,
    requests_processed: IntGaugeVec,
    requests_granted: IntGaugeVec,
    item_fail_rate: GaugeVec,
    rtt: GaugeVec,
    rtt_real: GaugeVec,
    rtt_count: IntGaugeVec,
}

impl ProjectMetrics {
    pub fn new() -> color_eyre::Result<Self> {
        let registry = Registry::new_custom(Some("at_tracker".to_string()), None)?;

        let s = Self {
            registry,
            item_request_serve_rate: GaugeVec::new(Opts::new("item_request_serve_rate", "item_request_serve_rate"), &["project"])?,
            reclaim_serve_rate: GaugeVec::new(Opts::new("reclaim_serve_rate", "reclaim_serve_rate"), &["project"])?,
            reclaim_rate: GaugeVec::new(Opts::new("reclaim_rate", "reclaim_rate"), &["project"])?,
            requests_processed: IntGaugeVec::new(Opts::new("requests_processed", "requests_processed"), &["project"])?,
            requests_granted: IntGaugeVec::new(Opts::new("requests_granted", "requests_granted"), &["project"])?,
            item_fail_rate: GaugeVec::new(Opts::new("item_fail_rate", "item_fail_rate"), &["project"])?,
            rtt: GaugeVec::new(Opts::new("rtt", "rtt"), &["project"])?,
            rtt_real: GaugeVec::new(Opts::new("rtt_real", "rtt_real"), &["project"])?,
            rtt_count: IntGaugeVec::new(Opts::new("rtt_count", "rtt_count"), &["project"])?,
        };

        s.registry.register(Box::new(s.item_request_serve_rate.clone()))?;
        s.registry.register(Box::new(s.reclaim_serve_rate.clone()))?;
        s.registry.register(Box::new(s.reclaim_rate.clone()))?;
        s.registry.register(Box::new(s.requests_processed.clone()))?;
        s.registry.register(Box::new(s.requests_granted.clone()))?;
        s.registry.register(Box::new(s.item_fail_rate.clone()))?;

        Ok(s)
    }
}

impl MetricsProvider for ProjectMetrics {
    fn get_registry(&self) -> color_eyre::Result<&Registry> {
        Ok(&self.registry)
    }

    fn update(&self, client_mgr: &ClientManager, project: &str, info: &TrackerInfo) -> color_eyre::Result<()> {
        let mut con = client_mgr.get_connection_from_tracker(info)?;

        update_gauge_from_redis(&mut con, project, "item_request_serve_rate", &self.item_request_serve_rate)?;
        update_gauge_from_redis(&mut con, project, "reclaim_serve_rate", &self.reclaim_serve_rate)?;
        update_gauge_from_redis(&mut con, project, "reclaim_rate", &self.reclaim_rate)?;
        update_gauge_from_redis(&mut con, project, "item_fail_rate", &self.item_fail_rate)?;
        update_gauge_from_redis(&mut con, project, "rtt", &self.rtt)?;
        update_gauge_from_redis(&mut con, project, "rtt_real", &self.rtt_real)?;
        update_gauge_from_redis(&mut con, project, "rtt_count", &self.rtt_count)?;
        let t = get_last_requests_time()?;
        let rp_key = format!("requests_processed:{t}");
        update_gauge_from_redis(&mut con, project, &rp_key, &self.requests_processed)?;
        let rg_key = format!("requests_granted:{t}");
        update_gauge_from_redis(&mut con, project, &rg_key, &self.requests_granted)?;

        Ok(())
    }

    fn get_metric_type_name(&self) -> &str {
        "project"
    }
}