Merge pull request #117 from seanwevans/codex/modify-health-monitor-for-per-node-tracking

Track health per node

Author: seanwevans

src/health.rs

@@ -6,6 +6,7 @@ use tokio::sync::broadcast;
 use tokio_util::sync::CancellationToken;
 use tracing::{info, error};
 use std::error::Error;
+use std::collections::HashMap;
 
 #[derive(Clone, Debug)]
 pub struct HealthCheck {
@@ -47,27 +48,29 @@ pub async fn monitor_health(
     unhealthy_threshold: u32,
     cancel: CancellationToken,
 ) -> Result<(), Box<dyn Error>> {
-    let mut unhealthy_count = 0;
+    let mut unhealthy_counts: HashMap<String, u32> = HashMap::new();
 
     loop {
         tokio::select! {
             result = rx.recv() => {
                 match result {
                     Ok(health_check) => {
-                        if !health_check.is_healthy {
-                            unhealthy_count += 1;
-                            error!("Node {} is unhealthy. Unhealthy count: {}", health_check.node_id, unhealthy_count);
-                        } else {
-                            unhealthy_count = 0;
-                            info!("Node {} is healthy.", health_check.node_id);
-                        }
+                        let threshold_reached = update_unhealthy_counts(&mut unhealthy_counts, &health_check, unhealthy_threshold);
 
-                        if unhealthy_count >= unhealthy_threshold {
+                        if health_check.is_healthy {
+                            info!("Node {} is healthy.", health_check.node_id);
+                        } else if threshold_reached {
                             error!(
                                 "Node {} has been unhealthy for {} consecutive checks. Taking corrective action.",
                                 health_check.node_id, unhealthy_threshold
                             );
                             // Add corrective actions here, such as restarting the node or notifying other services.
+                        } else {
+                            if let Some(count) = unhealthy_counts.get(&health_check.node_id) {
+                                error!("Node {} is unhealthy. Unhealthy count: {}", health_check.node_id, count);
+                            } else {
+                                error!("Node {} is unhealthy.", health_check.node_id);
+                            }
                         }
                     }
                     Err(broadcast::error::RecvError::Closed) => {
@@ -88,10 +91,29 @@ pub async fn monitor_health(
     }
 }
 
+fn update_unhealthy_counts(
+    unhealthy_counts: &mut HashMap<String, u32>,
+    health_check: &HealthCheck,
+    unhealthy_threshold: u32,
+) -> bool {
+    if health_check.is_healthy {
+        unhealthy_counts.remove(&health_check.node_id);
+        return false;
+    }
+
+    let count = unhealthy_counts
+        .entry(health_check.node_id.clone())
+        .and_modify(|counter| *counter = counter.saturating_add(1))
+        .or_insert(1);
+
+    *count >= unhealthy_threshold
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use tokio::time::Duration;
+    use std::collections::HashMap;
 
     #[tokio::test]
     async fn monitor_health_returns_on_channel_close() {
@@ -129,4 +151,51 @@ mod tests {
             .expect("start_heartbeat did not stop in time")
             .expect("task panicked");
     }
+
+    #[test]
+    fn update_unhealthy_counts_tracks_per_node_thresholds() {
+        let mut counts = HashMap::new();
+        let threshold = 2;
+
+        let node_a = "node-a".to_string();
+        let node_b = "node-b".to_string();
+
+        // First unhealthy report for node A.
+        let alert = update_unhealthy_counts(
+            &mut counts,
+            &HealthCheck { node_id: node_a.clone(), is_healthy: false },
+            threshold,
+        );
+        assert!(!alert);
+        assert_eq!(counts.get(&node_a), Some(&1));
+
+        // First unhealthy report for node B.
+        let alert = update_unhealthy_counts(
+            &mut counts,
+            &HealthCheck { node_id: node_b.clone(), is_healthy: false },
+            threshold,
+        );
+        assert!(!alert);
+        assert_eq!(counts.get(&node_b), Some(&1));
+
+        // Second unhealthy report for node A should trigger the alert for node A only.
+        let alert = update_unhealthy_counts(
+            &mut counts,
+            &HealthCheck { node_id: node_a.clone(), is_healthy: false },
+            threshold,
+        );
+        assert!(alert);
+        assert_eq!(counts.get(&node_a), Some(&2));
+        assert_eq!(counts.get(&node_b), Some(&1));
+
+        // Healthy report for node B should reset its counter without affecting node A.
+        let alert = update_unhealthy_counts(
+            &mut counts,
+            &HealthCheck { node_id: node_b.clone(), is_healthy: true },
+            threshold,
+        );
+        assert!(!alert);
+        assert!(counts.get(&node_b).is_none());
+        assert_eq!(counts.get(&node_a), Some(&2));
+    }
 }