Fix: stale metrics in distributed deployment of orchestrator (#425)

* fix stale metrics serving when orchestrator is deployed in distributed env

Author: JannikSt

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

@@ -98,17 +98,12 @@ async fn heartbeat(
             .map(|metric| (&metric.key.task_id, &metric.key.label))
             .collect();
 
-        // Clean up stale metrics
+        // Clean up stale metrics from Redis only
+        // The sync service will handle all Prometheus updates
         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
                     if let Err(e) = app_state
                         .store_context
@@ -122,7 +117,7 @@ async fn heartbeat(
             }
         }
 
-        // Store new metrics and update Prometheus
+        // Store new metrics in Redis only
         if let Err(e) = app_state
             .store_context
             .metrics_store
@@ -131,15 +126,6 @@ async fn heartbeat(
         {
             error!("Error storing metrics: {}", e);
         }
-
-        for metric in metrics {
-            app_state.metrics.record_compute_task_gauge(
-                &node_address.to_string(),
-                &metric.key.task_id,
-                &metric.key.label,
-                metric.value,
-            );
-        }
     }
 
     let current_task = app_state.scheduler.get_task_for_node(node_address).await;
@@ -171,7 +157,9 @@ mod tests {
 
     use super::*;
     use crate::api::tests::helper::create_test_app_state;
+    use crate::metrics::sync_service::MetricsSyncService;
     use crate::models::node::OrchestratorNode;
+    use crate::ServerMode;
 
     use actix_web::http::StatusCode;
     use actix_web::test;
@@ -254,10 +242,38 @@ 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"));
+        // Verify metrics are stored in Redis (heartbeat API responsibility)
+        let redis_metrics = app_state
+            .store_context
+            .metrics_store
+            .get_metrics_for_node(node_address)
+            .await
+            .unwrap();
+        assert!(redis_metrics.contains_key("long-task-1234"));
+        assert!(redis_metrics["long-task-1234"].contains_key("performance/batch_avg_seq_length"));
+        assert!(redis_metrics["long-task-1234"].contains_key("performance/batch_min_seq_length"));
+
+        // Test metrics sync service: Redis -> Prometheus
+        // Verify Prometheus registry is initially empty (no sync service has run)
+        let prometheus_metrics_before = app_state.metrics.export_metrics().unwrap();
+        assert!(!prometheus_metrics_before.contains("performance/batch_avg_seq_length"));
+
+        // Create and run sync service manually to test the sync functionality
+        let sync_service = MetricsSyncService::new(
+            app_state.store_context.clone(),
+            app_state.metrics.clone(),
+            ServerMode::Full, // Test app uses Full mode
+            10,
+        );
+
+        // Manually trigger a sync operation
+        sync_service.sync_metrics_from_redis().await.unwrap();
+
+        // Verify Prometheus registry now contains the metrics from Redis
+        let prometheus_metrics_after = app_state.metrics.export_metrics().unwrap();
+        assert!(prometheus_metrics_after.contains("performance/batch_avg_seq_length"));
+        assert!(prometheus_metrics_after.contains("performance/batch_min_seq_length"));
+        assert!(prometheus_metrics_after.contains("long-task-1234"));
 
         let heartbeat_two = json!({"address": address, "metrics": [
             {"key": {"task_id": "long-task-1235", "label": "performance/batch_len"}, "value": 10.0},
@@ -272,11 +288,17 @@ mod tests {
         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"));
+        // Verify new metrics in Redis and old metrics cleaned up
+        let redis_metrics = app_state
+            .store_context
+            .metrics_store
+            .get_metrics_for_node(node_address)
+            .await
+            .unwrap();
+        assert!(redis_metrics.contains_key("long-task-1235"));
+        assert!(redis_metrics["long-task-1235"].contains_key("performance/batch_len"));
+        assert!(redis_metrics["long-task-1235"].contains_key("performance/batch_min_len"));
+        assert!(!redis_metrics.contains_key("long-task-1234")); // Stale metrics cleaned up
         let aggregated_metrics = app_state
             .store_context
             .metrics_store
@@ -293,6 +315,12 @@ mod tests {
             aggregated_metrics.get("performance/batch_avg_seq_length"),
             None
         );
+        sync_service.sync_metrics_from_redis().await.unwrap();
+        let prometheus_metrics_after_two = app_state.metrics.export_metrics().unwrap();
+        assert!(prometheus_metrics_after_two.contains("performance/batch_len"));
+        assert!(prometheus_metrics_after_two.contains("performance/batch_min_len"));
+        assert!(prometheus_metrics_after_two.contains("long-task-1235"));
+        assert!(!prometheus_metrics_after_two.contains("long-task-1234"));
 
         let heartbeat_three = json!({"address": address, "metrics": [
         ]});
@@ -305,15 +333,22 @@ mod tests {
         let resp = test::call_service(&app, req).await;
         assert_eq!(resp.status(), StatusCode::OK);
 
-        let metrics = app_state.metrics.export_metrics().unwrap();
+        // Verify all metrics cleaned up from Redis
+        let redis_metrics = app_state
+            .store_context
+            .metrics_store
+            .get_metrics_for_node(node_address)
+            .await
+            .unwrap();
+        assert!(redis_metrics.is_empty()); // All metrics for this node should be gone
+
         let aggregated_metrics = app_state
             .store_context
             .metrics_store
             .get_aggregate_metrics_for_all_tasks()
             .await
             .unwrap();
         assert_eq!(aggregated_metrics, HashMap::new());
-        assert_eq!(metrics, "");
     }
 
     #[actix_web::test]

crates/orchestrator/src/main.rs

@@ -25,6 +25,7 @@ use log::debug;
 use log::error;
 use log::info;
 use log::LevelFilter;
+use metrics::sync_service::MetricsSyncService;
 use metrics::webhook_sender::MetricsWebhookSender;
 use metrics::MetricsContext;
 use plugins::node_groups::NodeGroupConfiguration;
@@ -268,6 +269,19 @@ async fn main() -> Result<()> {
 
     // Only spawn processor tasks if in ProcessorOnly or Full mode
     if matches!(server_mode, ServerMode::ProcessorOnly | ServerMode::Full) {
+        // Start metrics sync service to centralize metrics from Redis to Prometheus
+        let metrics_sync_store_context = store_context.clone();
+        let metrics_sync_context = metrics_context.clone();
+        tasks.spawn(async move {
+            let sync_service = MetricsSyncService::new(
+                metrics_sync_store_context,
+                metrics_sync_context,
+                server_mode,
+                10,
+            );
+            sync_service.run().await
+        });
+
         if let Some(group_plugin) = node_groups_plugin.clone() {
             tasks.spawn(async move {
                 group_plugin

crates/orchestrator/src/metrics/mod.rs

@@ -1,4 +1,5 @@
 use prometheus::{GaugeVec, Opts, Registry, TextEncoder};
+pub mod sync_service;
 pub mod webhook_sender;
 
 pub struct MetricsContext {
@@ -53,4 +54,11 @@ impl MetricsContext {
         let metric_families = self.registry.gather();
         encoder.encode_to_string(&metric_families)
     }
+
+    /// Clear all metrics from the registry
+    pub fn clear_compute_task_metrics(&self) {
+        // Clear all time series from the compute_task_gauges metric family
+        // This removes all existing metrics so we can rebuild from Redis
+        self.compute_task_gauges.reset();
+    }
 }

crates/orchestrator/src/metrics/sync_service.rs

@@ -0,0 +1,92 @@
+use crate::metrics::MetricsContext;
+use crate::store::core::StoreContext;
+use crate::ServerMode;
+use log::{debug, error, info};
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::time::interval;
+
+pub struct MetricsSyncService {
+    store_context: Arc<StoreContext>,
+    metrics_context: Arc<MetricsContext>,
+    server_mode: ServerMode,
+    sync_interval: Duration,
+}
+
+impl MetricsSyncService {
+    pub fn new(
+        store_context: Arc<StoreContext>,
+        metrics_context: Arc<MetricsContext>,
+        server_mode: ServerMode,
+        sync_interval_seconds: u64,
+    ) -> Self {
+        Self {
+            store_context,
+            metrics_context,
+            server_mode,
+            sync_interval: Duration::from_secs(sync_interval_seconds),
+        }
+    }
+
+    pub async fn run(&self) -> anyhow::Result<()> {
+        // Only run the sync service on ProcessorOnly or Full mode instances
+        if !matches!(
+            self.server_mode,
+            ServerMode::ProcessorOnly | ServerMode::Full
+        ) {
+            debug!("Metrics sync service disabled for ApiOnly mode");
+            return Ok(());
+        }
+
+        info!(
+            "Starting metrics sync service (interval: {:?})",
+            self.sync_interval
+        );
+        let mut interval = interval(self.sync_interval);
+
+        loop {
+            interval.tick().await;
+            if let Err(e) = self.sync_metrics_from_redis().await {
+                error!("Error syncing metrics from Redis: {}", e);
+            }
+        }
+    }
+
+    pub async fn sync_metrics_from_redis(&self) -> anyhow::Result<()> {
+        debug!("Syncing metrics from Redis to Prometheus");
+
+        // Get all metrics from Redis
+        let redis_metrics = match self.store_context.metrics_store.get_all_metrics().await {
+            Ok(metrics) => metrics,
+            Err(e) => {
+                error!("Failed to get metrics from Redis: {}", e);
+                return Err(e);
+            }
+        };
+
+        // Clear existing Prometheus metrics
+        self.metrics_context.clear_compute_task_metrics();
+
+        // Rebuild metrics from Redis data
+        let mut total_metrics = 0;
+        for (task_id, task_metrics) in redis_metrics {
+            for (label, node_metrics) in task_metrics {
+                for (node_address, value) in node_metrics {
+                    self.metrics_context.record_compute_task_gauge(
+                        &node_address,
+                        &task_id,
+                        &label,
+                        value,
+                    );
+                    total_metrics += 1;
+                }
+            }
+        }
+
+        debug!(
+            "Synced {} metric entries from Redis to Prometheus",
+            total_metrics
+        );
+        Ok(())
+    }
+}