stratum-apps 0.3.2

//! Prometheus metrics definitions for SV2 monitoring

use prometheus::{Gauge, GaugeVec, Opts, Registry};

/// Prometheus metrics for the monitoring server.
/// Metrics are optional - only registered when the corresponding monitoring type is enabled.
#[derive(Clone)]
pub struct PrometheusMetrics {
    pub registry: Registry,
    // System metrics
    pub sv2_uptime_seconds: Gauge,
    // Server metrics (upstream connection)
    pub sv2_server_channels: Option<GaugeVec>,
    pub sv2_server_hashrate_total: Option<Gauge>,
    pub sv2_server_channel_hashrate: Option<GaugeVec>,
    pub sv2_server_shares_accepted_total: Option<GaugeVec>,
    pub sv2_server_blocks_found_total: Option<Gauge>,
    // Clients metrics (downstream connections)
    pub sv2_clients_total: Option<Gauge>,
    pub sv2_client_channels: Option<GaugeVec>,
    pub sv2_client_hashrate_total: Option<Gauge>,
    pub sv2_client_channel_hashrate: Option<GaugeVec>,
    pub sv2_client_shares_accepted_total: Option<GaugeVec>,
    pub sv2_client_blocks_found_total: Option<Gauge>,
    // SV1 metrics
    pub sv1_clients_total: Option<Gauge>,
    pub sv1_hashrate_total: Option<Gauge>,
}

impl PrometheusMetrics {
    pub fn new(
        enable_server_metrics: bool,
        enable_clients_metrics: bool,
        enable_sv1_metrics: bool,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let registry = Registry::new();

        // System metrics (always enabled)
        let sv2_uptime_seconds = Gauge::new("sv2_uptime_seconds", "Server uptime in seconds")?;
        registry.register(Box::new(sv2_uptime_seconds.clone()))?;

        // Server metrics (upstream connection)
        let (
            sv2_server_channels,
            sv2_server_hashrate_total,
            sv2_server_channel_hashrate,
            sv2_server_shares_accepted_total,
            sv2_server_blocks_found_total,
        ) = if enable_server_metrics {
            let channels = GaugeVec::new(
                Opts::new("sv2_server_channels", "Number of server channels by type"),
                &["channel_type"],
            )?;
            registry.register(Box::new(channels.clone()))?;

            let hashrate = Gauge::new(
                "sv2_server_hashrate_total",
                "Total hashrate for channels opened with the server",
            )?;
            registry.register(Box::new(hashrate.clone()))?;

            let channel_hashrate = GaugeVec::new(
                Opts::new(
                    "sv2_server_channel_hashrate",
                    "Hashrate for individual server channels",
                ),
                &["channel_id", "user_identity"],
            )?;
            registry.register(Box::new(channel_hashrate.clone()))?;

            let shares_accepted = GaugeVec::new(
                Opts::new(
                    "sv2_server_shares_accepted_total",
                    "Total shares accepted per server channel",
                ),
                &["channel_id", "user_identity"],
            )?;
            registry.register(Box::new(shares_accepted.clone()))?;

            let blocks_found = Gauge::new(
                "sv2_server_blocks_found_total",
                "Total blocks found across all current server channels",
            )?;
            registry.register(Box::new(blocks_found.clone()))?;

            (
                Some(channels),
                Some(hashrate),
                Some(channel_hashrate),
                Some(shares_accepted),
                Some(blocks_found),
            )
        } else {
            (None, None, None, None, None)
        };

        // Clients metrics (downstream connections)
        let (
            sv2_clients_total,
            sv2_client_channels,
            sv2_client_hashrate_total,
            sv2_client_channel_hashrate,
            sv2_client_shares_accepted_total,
            sv2_client_blocks_found_total,
        ) = if enable_clients_metrics {
            let clients_total =
                Gauge::new("sv2_clients_total", "Total number of connected clients")?;
            registry.register(Box::new(clients_total.clone()))?;

            let channels = GaugeVec::new(
                Opts::new("sv2_client_channels", "Number of client channels by type"),
                &["channel_type"],
            )?;
            registry.register(Box::new(channels.clone()))?;

            let hashrate = Gauge::new(
                "sv2_client_hashrate_total",
                "Total hashrate for channels opened with clients",
            )?;
            registry.register(Box::new(hashrate.clone()))?;

            let channel_hashrate = GaugeVec::new(
                Opts::new(
                    "sv2_client_channel_hashrate",
                    "Hashrate for individual client channels",
                ),
                &["client_id", "channel_id", "user_identity"],
            )?;
            registry.register(Box::new(channel_hashrate.clone()))?;

            let shares_accepted = GaugeVec::new(
                Opts::new(
                    "sv2_client_shares_accepted_total",
                    "Total shares accepted per client channel",
                ),
                &["client_id", "channel_id", "user_identity"],
            )?;
            registry.register(Box::new(shares_accepted.clone()))?;

            let blocks_found = Gauge::new(
                "sv2_client_blocks_found_total",
                "Total blocks found across all current client channels",
            )?;
            registry.register(Box::new(blocks_found.clone()))?;

            (
                Some(clients_total),
                Some(channels),
                Some(hashrate),
                Some(channel_hashrate),
                Some(shares_accepted),
                Some(blocks_found),
            )
        } else {
            (None, None, None, None, None, None)
        };

        // SV1 metrics
        let (sv1_clients_total, sv1_hashrate_total) = if enable_sv1_metrics {
            let clients = Gauge::new("sv1_clients_total", "Total number of SV1 clients")?;
            registry.register(Box::new(clients.clone()))?;

            let hashrate = Gauge::new("sv1_hashrate_total", "Total hashrate from SV1 clients")?;
            registry.register(Box::new(hashrate.clone()))?;

            (Some(clients), Some(hashrate))
        } else {
            (None, None)
        };

        Ok(Self {
            registry,
            sv2_uptime_seconds,
            sv2_server_channels,
            sv2_server_hashrate_total,
            sv2_server_channel_hashrate,
            sv2_server_shares_accepted_total,
            sv2_server_blocks_found_total,
            sv2_clients_total,
            sv2_client_channels,
            sv2_client_hashrate_total,
            sv2_client_channel_hashrate,
            sv2_client_shares_accepted_total,
            sv2_client_blocks_found_total,
            sv1_clients_total,
            sv1_hashrate_total,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use prometheus::Encoder;

    #[test]
    fn uptime_metric_always_registered() {
        let m = PrometheusMetrics::new(false, false, false).unwrap();
        m.sv2_uptime_seconds.set(42.0);

        let families = m.registry.gather();
        let names: Vec<&str> = families.iter().map(|f| f.get_name()).collect();
        assert!(names.contains(&"sv2_uptime_seconds"));
    }

    #[test]
    fn server_metrics_enabled() {
        let m = PrometheusMetrics::new(true, false, false).unwrap();
        assert!(m.sv2_server_channels.is_some());
        assert!(m.sv2_server_hashrate_total.is_some());
        assert!(m.sv2_server_channel_hashrate.is_some());
        assert!(m.sv2_server_shares_accepted_total.is_some());
        // clients and sv1 should be None
        assert!(m.sv2_clients_total.is_none());
        assert!(m.sv1_clients_total.is_none());
    }

    #[test]
    fn clients_metrics_enabled() {
        let m = PrometheusMetrics::new(false, true, false).unwrap();
        assert!(m.sv2_clients_total.is_some());
        assert!(m.sv2_client_channels.is_some());
        assert!(m.sv2_client_hashrate_total.is_some());
        assert!(m.sv2_client_channel_hashrate.is_some());
        assert!(m.sv2_client_shares_accepted_total.is_some());
        // server and sv1 should be None
        assert!(m.sv2_server_channels.is_none());
        assert!(m.sv1_clients_total.is_none());
    }

    #[test]
    fn sv1_metrics_enabled() {
        let m = PrometheusMetrics::new(false, false, true).unwrap();
        assert!(m.sv1_clients_total.is_some());
        assert!(m.sv1_hashrate_total.is_some());
        // server and clients should be None
        assert!(m.sv2_server_channels.is_none());
        assert!(m.sv2_clients_total.is_none());
    }

    #[test]
    fn all_metrics_disabled_only_uptime() {
        let m = PrometheusMetrics::new(false, false, false).unwrap();
        assert!(m.sv2_server_channels.is_none());
        assert!(m.sv2_clients_total.is_none());
        assert!(m.sv1_clients_total.is_none());

        let families = m.registry.gather();
        // Only uptime should be registered (with default value 0)
        assert_eq!(families.len(), 1);
        assert_eq!(families[0].get_name(), "sv2_uptime_seconds");
    }

    #[test]
    fn metrics_encode_to_prometheus_text_format() {
        let m = PrometheusMetrics::new(true, true, true).unwrap();
        m.sv2_uptime_seconds.set(123.0);
        m.sv2_server_hashrate_total.as_ref().unwrap().set(1000.0);
        m.sv2_clients_total.as_ref().unwrap().set(5.0);
        m.sv1_clients_total.as_ref().unwrap().set(3.0);

        let encoder = prometheus::TextEncoder::new();
        let families = m.registry.gather();
        let mut buf = Vec::new();
        encoder.encode(&families, &mut buf).unwrap();
        let output = String::from_utf8(buf).unwrap();

        assert!(output.contains("sv2_uptime_seconds 123"));
        assert!(output.contains("sv2_server_hashrate_total 1000"));
        assert!(output.contains("sv2_clients_total 5"));
        assert!(output.contains("sv1_clients_total 3"));
    }
}