oxirs_vec/real_time_embedding_pipeline/
monitoring.rs

1//! Performance monitoring and alerting for the real-time embedding pipeline
2
3use parking_lot::Mutex;
4use std::collections::{HashMap, VecDeque};
5use std::fmt;
6use std::sync::Arc;
7use std::time::{Duration, Instant, SystemTime};
8use tokio::sync::RwLock;
9use uuid::Uuid;
10
11use super::config::MonitoringConfig;
12use super::traits::{
13    Alert, AlertCategory, AlertConfig, AlertHandler, AlertSeverity, AlertThrottling, HealthStatus,
14    MetricPoint, MetricsStorage,
15};
16use super::types::PerformanceMetrics;
17use super::PipelineError;
18
19// Display impls needed for formatting in throttle key
20impl fmt::Display for AlertCategory {
21    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
22        match self {
23            AlertCategory::Performance => write!(f, "Performance"),
24            AlertCategory::Quality => write!(f, "Quality"),
25            AlertCategory::Health => write!(f, "Health"),
26            AlertCategory::Security => write!(f, "Security"),
27            AlertCategory::Configuration => write!(f, "Configuration"),
28        }
29    }
30}
31
32impl fmt::Display for AlertSeverity {
33    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34        match self {
35            AlertSeverity::Info => write!(f, "Info"),
36            AlertSeverity::Warning => write!(f, "Warning"),
37            AlertSeverity::Error => write!(f, "Error"),
38            AlertSeverity::Critical => write!(f, "Critical"),
39        }
40    }
41}
42
43/// In-memory metrics storage implementation
44pub struct InMemoryMetricsStorage {
45    metrics: Vec<(String, MetricPoint)>,
46    max_metrics: usize,
47}
48
49impl InMemoryMetricsStorage {
50    pub fn new(max_metrics: usize) -> Self {
51        Self {
52            metrics: Vec::new(),
53            max_metrics,
54        }
55    }
56}
57
58impl MetricsStorage for InMemoryMetricsStorage {
59    fn store_metric(
60        &mut self,
61        name: &str,
62        value: f64,
63        timestamp: SystemTime,
64        tags: HashMap<String, String>,
65    ) -> anyhow::Result<()> {
66        self.metrics.push((
67            name.to_string(),
68            MetricPoint {
69                value,
70                timestamp,
71                tags,
72            },
73        ));
74        while self.metrics.len() > self.max_metrics {
75            self.metrics.remove(0);
76        }
77        Ok(())
78    }
79
80    fn get_metrics(
81        &self,
82        name: &str,
83        start: SystemTime,
84        end: SystemTime,
85    ) -> anyhow::Result<Vec<MetricPoint>> {
86        let filtered = self
87            .metrics
88            .iter()
89            .filter(|(n, m)| n == name && m.timestamp >= start && m.timestamp <= end)
90            .map(|(_, m)| m.clone())
91            .collect();
92        Ok(filtered)
93    }
94
95    fn get_metric_names(&self) -> anyhow::Result<Vec<String>> {
96        let mut names: Vec<String> = self.metrics.iter().map(|(n, _)| n.clone()).collect();
97        names.dedup();
98        Ok(names)
99    }
100
101    fn cleanup_old_metrics(&mut self, cutoff: SystemTime) -> anyhow::Result<usize> {
102        let before = self.metrics.len();
103        self.metrics.retain(|(_, m)| m.timestamp >= cutoff);
104        Ok(before - self.metrics.len())
105    }
106}
107
108/// Console-based alert handler implementation
109pub struct ConsoleAlertHandler;
110
111impl AlertHandler for ConsoleAlertHandler {
112    fn handle_alert(&self, alert: &Alert) -> anyhow::Result<()> {
113        println!(
114            "[ALERT] {} - {} - {} - {}",
115            alert.severity, alert.category, alert.source, alert.message
116        );
117        Ok(())
118    }
119
120    fn get_config(&self) -> AlertConfig {
121        AlertConfig {
122            min_severity: AlertSeverity::Warning,
123            throttling: AlertThrottling {
124                enabled: false,
125                window_duration: Duration::from_secs(60),
126                max_alerts_per_window: 100,
127            },
128            enable_notifications: true,
129        }
130    }
131
132    fn is_enabled(&self) -> bool {
133        true
134    }
135}
136
137/// Alert manager for handling pipeline alerts
138pub struct AlertManager {
139    alert_handler: Arc<dyn AlertHandler>,
140    active_alerts: Arc<RwLock<HashMap<Uuid, Alert>>>,
141    alert_history: Arc<tokio::sync::Mutex<VecDeque<Alert>>>,
142    throttling_state: Arc<RwLock<HashMap<String, Instant>>>,
143}
144
145impl AlertManager {
146    pub fn new(alert_handler: Arc<dyn AlertHandler>) -> Self {
147        Self {
148            alert_handler,
149            active_alerts: Arc::new(RwLock::new(HashMap::new())),
150            alert_history: Arc::new(tokio::sync::Mutex::new(VecDeque::new())),
151            throttling_state: Arc::new(RwLock::new(HashMap::new())),
152        }
153    }
154
155    pub async fn register_alert(&self, alert: Alert) -> Result<(), PipelineError> {
156        if self.should_throttle_alert(&alert).await {
157            return Ok(());
158        }
159
160        {
161            let mut active_alerts = self.active_alerts.write().await;
162            active_alerts.insert(alert.id, alert.clone());
163        }
164
165        {
166            let mut history = self.alert_history.lock().await;
167            history.push_back(alert.clone());
168            while history.len() > 1000 {
169                history.pop_front();
170            }
171        }
172
173        self.alert_handler
174            .handle_alert(&alert)
175            .map_err(|e| PipelineError::MonitoringError {
176                message: format!("Alert handling failed: {}", e),
177            })
178    }
179
180    pub async fn process_alerts(&self) -> Result<(), PipelineError> {
181        let alerts = {
182            self.active_alerts
183                .read()
184                .await
185                .values()
186                .cloned()
187                .collect::<Vec<_>>()
188        };
189        for alert in alerts {
190            self.check_alert_status(alert).await?;
191        }
192        Ok(())
193    }
194
195    pub async fn get_active_alerts(&self) -> Vec<Alert> {
196        self.active_alerts.read().await.values().cloned().collect()
197    }
198
199    async fn should_throttle_alert(&self, alert: &Alert) -> bool {
200        let throttle_key = format!("{}:{}", alert.category, alert.source);
201        let throttling_state = self.throttling_state.read().await;
202
203        if let Some(last_sent) = throttling_state.get(&throttle_key) {
204            let throttle_duration = match alert.severity {
205                AlertSeverity::Critical => Duration::from_secs(60),
206                AlertSeverity::Error => Duration::from_secs(300),
207                AlertSeverity::Warning => Duration::from_secs(600),
208                AlertSeverity::Info => Duration::from_secs(1200),
209            };
210            last_sent.elapsed() < throttle_duration
211        } else {
212            false
213        }
214    }
215
216    async fn check_alert_status(&self, _alert: Alert) -> Result<(), PipelineError> {
217        Ok(())
218    }
219}
220
221/// Health checker for monitoring system health
222pub struct HealthChecker {
223    component_health: Arc<RwLock<HashMap<String, HealthStatus>>>,
224}
225
226impl HealthChecker {
227    pub fn new() -> Self {
228        Self {
229            component_health: Arc::new(RwLock::new(HashMap::new())),
230        }
231    }
232
233    pub async fn check_health(&self) -> Result<HealthStatus, PipelineError> {
234        let component_health = self.component_health.read().await;
235
236        if component_health.is_empty() {
237            return Ok(HealthStatus::Healthy);
238        }
239
240        for status in component_health.values() {
241            if matches!(status, HealthStatus::Unhealthy { .. }) {
242                return Ok(HealthStatus::Unhealthy {
243                    message: "Component unhealthy".to_string(),
244                });
245            }
246        }
247
248        for status in component_health.values() {
249            if matches!(status, HealthStatus::Warning { .. }) {
250                return Ok(HealthStatus::Warning {
251                    message: "Component warning".to_string(),
252                });
253            }
254        }
255
256        Ok(HealthStatus::Healthy)
257    }
258
259    pub async fn get_overall_health(&self) -> Result<HealthStatus, PipelineError> {
260        self.check_health().await
261    }
262
263    pub async fn update_component_health(&self, component: String, status: HealthStatus) {
264        let mut component_health = self.component_health.write().await;
265        component_health.insert(component, status);
266    }
267}
268
269impl Default for HealthChecker {
270    fn default() -> Self {
271        Self::new()
272    }
273}
274
275/// Metrics collector for gathering performance data
276pub struct MetricsCollector {
277    current_metrics: Arc<RwLock<PerformanceMetrics>>,
278}
279
280impl MetricsCollector {
281    pub fn new() -> Self {
282        Self {
283            current_metrics: Arc::new(RwLock::new(PerformanceMetrics::default())),
284        }
285    }
286
287    /// Returns `(name, value, tags)` tuples — MetricPoint has no name field,
288    /// so callers are responsible for associating the name when storing.
289    pub async fn collect_metrics(
290        &self,
291    ) -> Result<Vec<(String, f64, HashMap<String, String>)>, PipelineError> {
292        let metrics = vec![
293            ("cpu_usage".to_string(), 50.0_f64, HashMap::new()),
294            (
295                "memory_usage".to_string(),
296                (1024.0 * 1024.0 * 512.0_f64),
297                HashMap::new(),
298            ),
299            (
300                "embedding_throughput".to_string(),
301                100.0_f64,
302                HashMap::new(),
303            ),
304        ];
305        Ok(metrics)
306    }
307
308    pub async fn get_current_metrics(&self) -> Result<PerformanceMetrics, PipelineError> {
309        let metrics = self.current_metrics.read().await;
310        Ok(metrics.clone())
311    }
312}
313
314impl Default for MetricsCollector {
315    fn default() -> Self {
316        Self::new()
317    }
318}
319
320/// Performance monitoring manager for the pipeline
321pub struct PipelinePerformanceMonitor {
322    config: MonitoringConfig,
323    metrics_storage: Arc<Mutex<InMemoryMetricsStorage>>,
324    alert_manager: Arc<AlertManager>,
325    health_checker: Arc<HealthChecker>,
326    metrics_collector: Arc<MetricsCollector>,
327    is_running: Arc<tokio::sync::Mutex<bool>>,
328}
329
330impl PipelinePerformanceMonitor {
331    /// Create a new performance monitor
332    pub fn new(config: MonitoringConfig, alert_handler: Arc<dyn AlertHandler>) -> Self {
333        let alert_manager = Arc::new(AlertManager::new(alert_handler));
334        let health_checker = Arc::new(HealthChecker::new());
335        let metrics_collector = Arc::new(MetricsCollector::new());
336        let metrics_storage = Arc::new(Mutex::new(InMemoryMetricsStorage::new(10_000)));
337
338        Self {
339            config,
340            metrics_storage,
341            alert_manager,
342            health_checker,
343            metrics_collector,
344            is_running: Arc::new(tokio::sync::Mutex::new(false)),
345        }
346    }
347
348    /// Start the monitoring system
349    pub async fn start(&self) -> Result<(), PipelineError> {
350        let mut running = self.is_running.lock().await;
351        if *running {
352            return Err(PipelineError::AlreadyRunning);
353        }
354        *running = true;
355        drop(running);
356
357        self.start_metrics_collection().await;
358        self.start_health_monitoring().await;
359        self.start_alert_processing().await;
360
361        Ok(())
362    }
363
364    /// Stop the monitoring system
365    pub async fn stop(&self) -> Result<(), PipelineError> {
366        let mut running = self.is_running.lock().await;
367        if !*running {
368            return Err(PipelineError::NotRunning);
369        }
370        *running = false;
371        Ok(())
372    }
373
374    /// Record a performance metric by name
375    pub fn record_metric(
376        &self,
377        name: &str,
378        value: f64,
379        tags: HashMap<String, String>,
380    ) -> Result<(), PipelineError> {
381        self.metrics_storage
382            .lock()
383            .store_metric(name, value, SystemTime::now(), tags)
384            .map_err(|e| PipelineError::MonitoringError {
385                message: format!("Failed to store metric: {}", e),
386            })
387    }
388
389    /// Get current performance metrics
390    pub async fn get_current_metrics(&self) -> Result<PerformanceMetrics, PipelineError> {
391        self.metrics_collector.get_current_metrics().await
392    }
393
394    /// Get health status
395    pub async fn get_health_status(&self) -> Result<HealthStatus, PipelineError> {
396        self.health_checker.get_overall_health().await
397    }
398
399    /// Register an alert
400    pub async fn register_alert(&self, alert: Alert) -> Result<(), PipelineError> {
401        self.alert_manager.register_alert(alert).await
402    }
403
404    async fn start_metrics_collection(&self) {
405        let metrics_collector = Arc::clone(&self.metrics_collector);
406        let metrics_storage = Arc::clone(&self.metrics_storage);
407        let is_running = Arc::clone(&self.is_running);
408        let collection_interval = Duration::from_millis(self.config.metrics_interval_ms);
409
410        tokio::spawn(async move {
411            loop {
412                {
413                    let running = is_running.lock().await;
414                    if !*running {
415                        break;
416                    }
417                }
418                if let Ok(metrics) = metrics_collector.collect_metrics().await {
419                    let timestamp = SystemTime::now();
420                    let mut storage = metrics_storage.lock();
421                    for (name, value, tags) in metrics {
422                        let _ = storage.store_metric(&name, value, timestamp, tags);
423                    }
424                }
425                tokio::time::sleep(collection_interval).await;
426            }
427        });
428    }
429
430    async fn start_health_monitoring(&self) {
431        let health_checker = Arc::clone(&self.health_checker);
432        let alert_manager = Arc::clone(&self.alert_manager);
433        let is_running = Arc::clone(&self.is_running);
434
435        tokio::spawn(async move {
436            loop {
437                {
438                    let running = is_running.lock().await;
439                    if !*running {
440                        break;
441                    }
442                }
443                if let Ok(health_status) = health_checker.check_health().await {
444                    if !matches!(health_status, HealthStatus::Healthy) {
445                        let alert = Alert {
446                            id: Uuid::new_v4(),
447                            category: AlertCategory::Health,
448                            severity: AlertSeverity::Warning,
449                            message: format!("Health check failed: {:?}", health_status),
450                            timestamp: SystemTime::now(),
451                            source: "health_monitor".to_string(),
452                            details: HashMap::new(),
453                        };
454                        let _ = alert_manager.register_alert(alert).await;
455                    }
456                }
457                tokio::time::sleep(Duration::from_secs(30)).await;
458            }
459        });
460    }
461
462    async fn start_alert_processing(&self) {
463        let alert_manager = Arc::clone(&self.alert_manager);
464        let is_running = Arc::clone(&self.is_running);
465
466        tokio::spawn(async move {
467            loop {
468                {
469                    let running = is_running.lock().await;
470                    if !*running {
471                        break;
472                    }
473                }
474                let _ = alert_manager.process_alerts().await;
475                tokio::time::sleep(Duration::from_secs(1)).await;
476            }
477        });
478    }
479}
480
481#[cfg(test)]
482mod tests {
483    use super::*;
484    use crate::real_time_embedding_pipeline::config::MonitoringConfig;
485
486    #[tokio::test]
487    async fn test_metrics_collector_creation() {
488        let collector = MetricsCollector::new();
489        let metrics = collector.collect_metrics().await.expect("should collect");
490        assert!(!metrics.is_empty());
491    }
492
493    #[tokio::test]
494    async fn test_health_checker() {
495        let checker = HealthChecker::new();
496        let health = checker.get_overall_health().await.expect("should check");
497        assert!(matches!(health, HealthStatus::Healthy));
498    }
499
500    #[tokio::test]
501    async fn test_alert_manager() {
502        let handler = Arc::new(ConsoleAlertHandler);
503        let manager = AlertManager::new(handler);
504
505        let alert = Alert {
506            id: Uuid::new_v4(),
507            category: AlertCategory::Performance,
508            severity: AlertSeverity::Warning,
509            message: "Test alert".to_string(),
510            timestamp: SystemTime::now(),
511            source: "test".to_string(),
512            details: HashMap::new(),
513        };
514
515        let result = manager.register_alert(alert).await;
516        assert!(result.is_ok());
517    }
518
519    #[test]
520    fn test_in_memory_metrics_storage() {
521        let mut storage = InMemoryMetricsStorage::new(100);
522        let result = storage.store_metric("test_metric", 42.0, SystemTime::now(), HashMap::new());
523        assert!(result.is_ok());
524    }
525
526    #[tokio::test]
527    async fn test_monitor_creation() {
528        let config = MonitoringConfig::default();
529        let handler = Arc::new(ConsoleAlertHandler);
530        let monitor = PipelinePerformanceMonitor::new(config, handler);
531        let health = monitor
532            .get_health_status()
533            .await
534            .expect("should check health");
535        assert!(matches!(health, HealthStatus::Healthy));
536    }
537
538    #[test]
539    fn test_console_alert_handler_config() {
540        let handler = ConsoleAlertHandler;
541        assert!(handler.is_enabled());
542        let config = handler.get_config();
543        assert!(config.enable_notifications);
544    }
545
546    #[test]
547    fn test_alert_severity_display() {
548        assert_eq!(format!("{}", AlertSeverity::Critical), "Critical");
549        assert_eq!(format!("{}", AlertSeverity::Warning), "Warning");
550    }
551
552    #[test]
553    fn test_alert_category_display() {
554        assert_eq!(format!("{}", AlertCategory::Performance), "Performance");
555        assert_eq!(format!("{}", AlertCategory::Health), "Health");
556    }
557
558    #[tokio::test]
559    async fn test_health_checker_with_component() {
560        let checker = HealthChecker::new();
561        checker
562            .update_component_health(
563                "test_component".to_string(),
564                HealthStatus::Warning {
565                    message: "test warning".to_string(),
566                },
567            )
568            .await;
569        let health = checker.check_health().await.expect("should check");
570        assert!(matches!(health, HealthStatus::Warning { .. }));
571    }
572
573    #[test]
574    fn test_in_memory_metrics_storage_cleanup() {
575        let mut storage = InMemoryMetricsStorage::new(100);
576        let past = SystemTime::now()
577            .checked_sub(Duration::from_secs(3600))
578            .unwrap_or(SystemTime::UNIX_EPOCH);
579        let _ = storage.store_metric("old_metric", 1.0, past, HashMap::new());
580        let _ = storage.store_metric("new_metric", 2.0, SystemTime::now(), HashMap::new());
581
582        let cutoff = SystemTime::now()
583            .checked_sub(Duration::from_secs(1800))
584            .unwrap_or(SystemTime::UNIX_EPOCH);
585        let removed = storage.cleanup_old_metrics(cutoff).expect("cleanup ok");
586        assert_eq!(removed, 1);
587    }
588
589    #[test]
590    fn test_in_memory_metrics_storage_get() {
591        let mut storage = InMemoryMetricsStorage::new(100);
592        let now = SystemTime::now();
593        let _ = storage.store_metric("cpu", 75.0, now, HashMap::new());
594
595        let start = now.checked_sub(Duration::from_secs(1)).unwrap_or(now);
596        let end = now
597            .checked_add(Duration::from_secs(1))
598            .unwrap_or_else(SystemTime::now);
599        let results = storage.get_metrics("cpu", start, end).expect("should get");
600        assert_eq!(results.len(), 1);
601        assert!((results[0].value - 75.0).abs() < f64::EPSILON);
602    }
603}
oxirs_vec/real_time_embedding_pipeline/monitoring.rs

oxirs_vec/real_time_embedding_pipeline/
monitoring.rs
