foxhunt/services/ml_training_service/tests/load/metrics_collector.rs

//! Real-Time Metrics Collector
//!
//! Collects system and application metrics during load tests.

use anyhow::Result;
use std::sync::Arc;
use std::time::{Duration, Instant};
use serde::{Deserialize, Serialize};
use sysinfo::{ProcessRefreshKind, RefreshKind, System};

/// Metrics sample collected at a point in time
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MetricsSample {
    pub timestamp: Duration,
    pub cpu_usage_percent: f32,
    pub memory_usage_mb: u64,
    pub active_operations: u64,
    pub throughput_rps: f64,
    pub latency_p50_ms: f64,
    pub latency_p95_ms: f64,
}

/// Real-time metrics collector
pub struct MetricsCollector {
    start_time: Instant,
    samples: Arc<tokio::sync::Mutex<Vec<MetricsSample>>>,
    sample_interval: Duration,
}

impl MetricsCollector {
    /// Create a new metrics collector
    pub fn new(sample_interval: Duration) -> Self {
        Self {
            start_time: Instant::now(),
            samples: Arc::new(tokio::sync::Mutex::new(Vec::new())),
            sample_interval,
        }
    }

    /// Start collecting metrics
    pub async fn start(&self) -> tokio::task::JoinHandle<()> {
        let samples = self.samples.clone();
        let interval = self.sample_interval;
        let start = self.start_time;

        tokio::spawn(async move {
            let mut system = System::new_with_specifics(
                RefreshKind::new().with_processes(ProcessRefreshKind::everything())
            );

            loop {
                tokio::time::sleep(interval).await;

                // Refresh system info
                system.refresh_processes();

                let elapsed = start.elapsed();

                // Get current process stats
                let (cpu_usage, memory_usage) = if let Ok(pid) = sysinfo::get_current_pid() {
                    if let Some(process) = system.process(pid) {
                        (process.cpu_usage(), process.memory() / 1_048_576) // Convert to MB
                    } else {
                        (0.0, 0)
                    }
                } else {
                    (0.0, 0)
                };

                let sample = MetricsSample {
                    timestamp: elapsed,
                    cpu_usage_percent: cpu_usage,
                    memory_usage_mb: memory_usage,
                    active_operations: 0, // Updated by caller
                    throughput_rps: 0.0,   // Updated by caller
                    latency_p50_ms: 0.0,   // Updated by caller
                    latency_p95_ms: 0.0,   // Updated by caller
                };

                samples.lock().await.push(sample);
            }
        })
    }

    /// Update application metrics
    pub async fn update_app_metrics(
        &self,
        active_ops: u64,
        throughput: f64,
        latency_p50: Duration,
        latency_p95: Duration,
    ) {
        let mut samples = self.samples.lock().await;
        if let Some(last) = samples.last_mut() {
            last.active_operations = active_ops;
            last.throughput_rps = throughput;
            last.latency_p50_ms = latency_p50.as_secs_f64() * 1000.0;
            last.latency_p95_ms = latency_p95.as_secs_f64() * 1000.0;
        }
    }

    /// Get all collected samples
    pub async fn get_samples(&self) -> Vec<MetricsSample> {
        self.samples.lock().await.clone()
    }

    /// Get summary statistics
    pub async fn get_summary(&self) -> MetricsSummary {
        let samples = self.samples.lock().await;

        if samples.is_empty() {
            return MetricsSummary::default();
        }

        let cpu_values: Vec<f32> = samples.iter().map(|s| s.cpu_usage_percent).collect();
        let memory_values: Vec<u64> = samples.iter().map(|s| s.memory_usage_mb).collect();
        let throughput_values: Vec<f64> = samples.iter().map(|s| s.throughput_rps).collect();

        MetricsSummary {
            avg_cpu_percent: cpu_values.iter().sum::<f32>() / cpu_values.len() as f32,
            max_cpu_percent: cpu_values.iter().cloned().fold(0.0f32, f32::max),
            avg_memory_mb: memory_values.iter().sum::<u64>() / memory_values.len() as u64,
            max_memory_mb: *memory_values.iter().max().unwrap_or(&0),
            avg_throughput_rps: throughput_values.iter().sum::<f64>() / throughput_values.len() as f64,
            peak_throughput_rps: throughput_values.iter().cloned().fold(0.0f64, f64::max),
        }
    }
}

/// Metrics summary statistics
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct MetricsSummary {
    pub avg_cpu_percent: f32,
    pub max_cpu_percent: f32,
    pub avg_memory_mb: u64,
    pub max_memory_mb: u64,
    pub avg_throughput_rps: f64,
    pub peak_throughput_rps: f64,
}

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

    #[tokio::test]
    async fn test_metrics_collector() {
        let collector = MetricsCollector::new(Duration::from_millis(100));
        let handle = collector.start().await;

        // Let it collect a few samples
        tokio::time::sleep(Duration::from_millis(350)).await;

        handle.abort();

        let samples = collector.get_samples().await;
        assert!(!samples.is_empty(), "Expected metrics samples to be collected");

        let summary = collector.get_summary().await;
        assert!(summary.avg_cpu_percent >= 0.0);
    }
}