improve stale metrics handling (#393)

* improve stale metrics handling

* add metrics testing for stale metrics

Author: JannikSt

crates/orchestrator/src/api/routes/heartbeat.rs

@@ -9,6 +9,7 @@ use shared::models::{
     api::ApiResponse,
     heartbeat::{HeartbeatRequest, HeartbeatResponse},
 };
+use std::collections::HashSet;
 use std::str::FromStr;
 
 async fn heartbeat(
@@ -42,6 +43,40 @@ async fn heartbeat(
 
     app_state.store_context.heartbeat_store.beat(&heartbeat);
     if let Some(metrics) = heartbeat.metrics.clone() {
+        // Get all previously reported metrics for this node
+        let previous_metrics = app_state
+            .store_context
+            .metrics_store
+            .get_metrics_for_node(node_address);
+
+        // Create a HashSet of new metrics for efficient lookup
+        let new_metrics_set: HashSet<_> = metrics
+            .iter()
+            .map(|metric| (&metric.key.task_id, &metric.key.label))
+            .collect();
+
+        // Clean up stale metrics
+        for (task_id, task_metrics) in previous_metrics {
+            for (label, _value) in task_metrics {
+                let prev_key = (&task_id, &label);
+                if !new_metrics_set.contains(&prev_key) {
+                    // Remove from Prometheus metrics
+                    app_state.metrics.remove_compute_task_gauge(
+                        &node_address.to_string(),
+                        &task_id,
+                        &label,
+                    );
+                    // Remove from Redis metrics store
+                    app_state.store_context.metrics_store.delete_metric(
+                        &task_id,
+                        &label,
+                        &node_address.to_string(),
+                    );
+                }
+            }
+        }
+
+        // Store new metrics and update Prometheus
         app_state
             .store_context
             .metrics_store
@@ -82,6 +117,8 @@ pub fn heartbeat_routes() -> Scope {
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashMap;
+
     use super::*;
     use crate::api::tests::helper::create_test_app_state;
 
@@ -158,6 +195,59 @@ mod tests {
         let metrics = app_state.metrics.export_metrics().unwrap();
         assert!(metrics.contains("performance/batch_avg_seq_length"));
         assert!(metrics.contains("performance/batch_min_seq_length"));
+        assert!(metrics.contains("long-task-1234"));
+
+        let heartbeat_two = json!({"address": address, "metrics": [
+            {"key": {"task_id": "long-task-1235", "label": "performance/batch_len"}, "value": 10.0},
+            {"key": {"task_id": "long-task-1235", "label": "performance/batch_min_len"}, "value": 50.0}
+        ]});
+
+        let req = test::TestRequest::post()
+            .uri("/heartbeat")
+            .set_json(&heartbeat_two)
+            .to_request();
+
+        let resp = test::call_service(&app, req).await;
+        assert_eq!(resp.status(), StatusCode::OK);
+
+        let metrics = app_state.metrics.export_metrics().unwrap();
+        assert!(metrics.contains("long-task-1235"));
+        assert!(metrics.contains("performance/batch_len"));
+        assert!(metrics.contains("performance/batch_min_len"));
+        assert!(!metrics.contains("long-task-1234"));
+        let aggregated_metrics = app_state
+            .store_context
+            .metrics_store
+            .get_aggregate_metrics_for_all_tasks();
+        assert_eq!(aggregated_metrics.len(), 2);
+        assert_eq!(aggregated_metrics.get("performance/batch_len"), Some(&10.0));
+        assert_eq!(
+            aggregated_metrics.get("performance/batch_min_len"),
+            Some(&50.0)
+        );
+        assert_eq!(
+            aggregated_metrics.get("performance/batch_avg_seq_length"),
+            None
+        );
+
+        let heartbeat_three = json!({"address": address, "metrics": [
+        ]});
+
+        let req = test::TestRequest::post()
+            .uri("/heartbeat")
+            .set_json(&heartbeat_three)
+            .to_request();
+
+        let resp = test::call_service(&app, req).await;
+        assert_eq!(resp.status(), StatusCode::OK);
+
+        let metrics = app_state.metrics.export_metrics().unwrap();
+        let aggregated_metrics = app_state
+            .store_context
+            .metrics_store
+            .get_aggregate_metrics_for_all_tasks();
+        assert_eq!(aggregated_metrics, HashMap::new());
+        assert_eq!(metrics, "");
     }
 
     #[actix_web::test]

crates/orchestrator/src/api/tests/helper.rs

@@ -150,3 +150,54 @@ pub fn setup_contract() -> Contracts {
         .build()
         .unwrap()
 }
+
+#[cfg(test)]
+pub async fn create_test_app_state_with_metrics() -> Data<AppState> {
+    use shared::utils::MockStorageProvider;
+
+    use crate::{
+        metrics::MetricsContext, scheduler::Scheduler, utils::loop_heartbeats::LoopHeartbeats,
+        ServerMode,
+    };
+
+    let store = Arc::new(RedisStore::new_test());
+    let mut con = store
+        .client
+        .get_connection()
+        .expect("Should connect to test Redis instance");
+
+    redis::cmd("PING")
+        .query::<String>(&mut con)
+        .expect("Redis should be responsive");
+    redis::cmd("FLUSHALL")
+        .query::<String>(&mut con)
+        .expect("Redis should be flushed");
+
+    let store_context = Arc::new(StoreContext::new(store.clone()));
+    let mode = ServerMode::Full;
+    let scheduler = Scheduler::new(store_context.clone(), vec![]);
+
+    let mock_storage = MockStorageProvider::new();
+    let storage_provider = Arc::new(mock_storage);
+    let metrics = Arc::new(MetricsContext::new("0".to_string()));
+
+    Data::new(AppState {
+        store_context: store_context.clone(),
+        contracts: None,
+        pool_id: 1,
+        wallet: Arc::new(
+            Wallet::new(
+                "0xdbda1821b80551c9d65939329250298aa3472ba22feea921c0cf5d620ea67b97",
+                Url::parse("http://localhost:8545").unwrap(),
+            )
+            .unwrap(),
+        ),
+        storage_provider,
+        heartbeats: Arc::new(LoopHeartbeats::new(&mode)),
+        hourly_upload_limit: 12,
+        redis_store: store.clone(),
+        scheduler,
+        node_groups_plugin: None,
+        metrics,
+    })
+}

crates/orchestrator/src/main.rs

@@ -315,6 +315,7 @@ async fn main() -> Result<()> {
         let status_update_store_context = store_context.clone();
         let status_update_heartbeats = heartbeats.clone();
         let status_update_contracts = contracts.clone();
+        let status_update_metrics = metrics_context.clone();
 
         tasks.spawn(async move {
             let status_updater = NodeStatusUpdater::new(
@@ -326,6 +327,7 @@ async fn main() -> Result<()> {
                 args.disable_ejection,
                 status_update_heartbeats.clone(),
                 status_update_plugins,
+                status_update_metrics,
             );
             status_updater.run().await
         });

crates/orchestrator/src/metrics/mod.rs

@@ -37,6 +37,17 @@ impl MetricsContext {
             .set(value);
     }
 
+    pub fn remove_compute_task_gauge(&self, node_address: &str, task_id: &str, label: &str) {
+        if let Err(e) = self.compute_task_gauges.remove_label_values(&[
+            node_address,
+            task_id,
+            label,
+            &self.pool_id,
+        ]) {
+            println!("Error removing compute task gauge: {}", e);
+        }
+    }
+
     pub fn export_metrics(&self) -> Result<String, prometheus::Error> {
         let encoder = TextEncoder::new();
         let metric_families = self.registry.gather();