feat(metrics): add NATS metrics to prometheus_metrics_fmt

- Add NATS client metrics collection via callbacks in transports/nats.rs
- Execute callbacks before TextEncoder to include NATS metrics in output
- Update prometheus_metrics_fmt to include dynamically added NATS metrics
- Add comprehensive test coverage for NATS client metrics
- Update tests to filter out dynamic NATS metrics for comparison
- Ensure NATS connection state, bytes, and other client metrics are exposed
- Add ComponentSystemStatusNatsMetrics struct for component service stats
- Implement automatic metrics registration in Namespace::component()
- Centralize NATS metric names in prometheus_names module
- Add filtering support for NATS metrics in tests
- Move ComponentSystemStatusNatsMetrics to service.rs module
- Update metric callbacks to use proper runtime prefix hierarchy

Author: keivenchang

lib/runtime/examples/Cargo.lock

@@ -30,7 +30,10 @@
 //! TODO: Top-level Overview of Endpoints/Functions
 
 use crate::{
-    config::HealthStatus, discovery::Lease, metrics::MetricsRegistry, service::ServiceSet,
+    config::HealthStatus,
+    discovery::Lease,
+    metrics::{prometheus_names, MetricsRegistry},
+    service::ServiceSet,
     transports::etcd::EtcdPath,
 };
 
@@ -247,6 +250,7 @@ impl Component {
         Ok(out)
     }
 
+    /// Scrape user defined stats (embedded in data field of ServiceInfo)
     pub async fn scrape_stats(&self, timeout: Duration) -> Result<ServiceSet> {
         let service_name = self.service_name();
         let service_client = self.drt().service_client();
@@ -255,6 +259,49 @@ impl Component {
             .await
     }
 
+    /// Register Prometheus metrics for this component's service stats
+    pub fn register_metrics(&self) -> Result<crate::service::ComponentSystemStatusNatsMetrics> {
+        let component_metrics =
+            crate::service::ComponentSystemStatusNatsMetrics::from_component(self)?;
+
+        // Create a callback that scrapes stats and updates metrics when called
+        let metrics_clone = component_metrics.clone();
+        let component_clone = self.clone();
+        let prefix = self.prefix();
+        let service_name = self.service_name();
+        let prefix_for_closure = prefix.clone();
+        self.drt().add_metrics_callback(&prefix, move |_runtime| {
+            println!(
+                "[DEBUG]CALLING  metrics callback for component: {}, prefix:{}",
+                service_name, prefix_for_closure
+            );
+            // Use tokio::runtime::Handle to run async code in the callback
+            let handle = tokio::runtime::Handle::try_current();
+            if let Ok(handle) = handle {
+                let metrics_ref = metrics_clone.clone();
+                let comp_ref = component_clone.clone();
+                handle.spawn(async move {
+                    let timeout = std::time::Duration::from_millis(500);
+                    match comp_ref.scrape_stats(timeout).await {
+                        Ok(service_set) => {
+                            metrics_ref.update_from_service_set(&service_set);
+                        }
+                        Err(err) => {
+                            tracing::warn!(
+                                "Failed to scrape stats for component '{}': {}",
+                                comp_ref.service_name(),
+                                err
+                            );
+                        }
+                    }
+                });
+            }
+            Ok("".to_string())
+        });
+
+        Ok(component_metrics)
+    }
+
     /// TODO
     ///
     /// This method will scrape the stats for all available services
@@ -488,11 +535,13 @@ impl Namespace {
 
     /// Create a [`Component`] in the namespace who's endpoints can be discovered with etcd
     pub fn component(&self, name: impl Into<String>) -> Result<Component> {
-        Ok(ComponentBuilder::from_runtime(self.runtime.clone())
+        let component = ComponentBuilder::from_runtime(self.runtime.clone())
             .name(name)
             .namespace(self.clone())
             .is_static(self.is_static)
-            .build()?)
+            .build()?;
+        component.register_metrics()?; // register a callback to scrape stats and update metrics
+        Ok(component)
     }
 
     /// Create a [`Namespace`] in the parent namespace

lib/runtime/src/component.rs

@@ -106,13 +106,21 @@ impl DistributedRuntime {
             component_registry: component::Registry::new(),
             is_static,
             instance_sources: Arc::new(Mutex::new(HashMap::new())),
-            prometheus_registries_by_prefix: Arc::new(std::sync::Mutex::new(HashMap::<
+            metrics_registry_by_prefix: Arc::new(std::sync::Mutex::new(HashMap::<
                 String,
-                prometheus::Registry,
+                crate::MetricsRegistryEntry,
             >::new())),
             system_health,
         };
 
+        // Register NATS client metrics after creation
+        if let Err(e) = distributed_runtime
+            .nats_client()
+            .register_metrics(&distributed_runtime)
+        {
+            tracing::warn!("Failed to register NATS client metrics: {}", e);
+        }
+
         // Start metrics server if enabled
         if let Some(cancel_token) = cancel_token {
             let host = config.system_host.clone();
@@ -226,11 +234,6 @@ impl DistributedRuntime {
         self.nats_client.clone()
     }
 
-    /// Get metrics server information if available
-    pub fn metrics_server_info(&self) -> Option<Arc<crate::metrics_server::MetricsServerInfo>> {
-        self.metrics_server.get().cloned()
-    }
-
     // todo(ryan): deprecate this as we move to Discovery traits and Component Identifiers
     pub fn etcd_client(&self) -> Option<etcd::Client> {
         self.etcd_client.clone()
@@ -243,6 +246,97 @@ impl DistributedRuntime {
     pub fn instance_sources(&self) -> Arc<Mutex<HashMap<Endpoint, Weak<InstanceSource>>>> {
         self.instance_sources.clone()
     }
+
+    /// Get metrics server information if available
+    pub fn metrics_server_info(&self) -> Option<Arc<crate::metrics_server::MetricsServerInfo>> {
+        self.metrics_server.get().cloned()
+    }
+
+    /// Add a Prometheus metric to a specific prefix's registry
+    pub fn add_prometheus_metric(
+        &self,
+        prefix: &str,
+        prometheus_metric: Box<dyn prometheus::core::Collector>,
+    ) -> anyhow::Result<()> {
+        let mut registries = self.metrics_registry_by_prefix.lock().unwrap();
+        let entry = registries.entry(prefix.to_string()).or_default();
+        entry.prometheus_registry.register(prometheus_metric)?;
+        Ok(())
+    }
+
+    /// Add a callback function to a metrics registry for the given prefix
+    pub fn add_metrics_callback<F>(&self, prefix: &str, callback: F)
+    where
+        F: Fn(&dyn crate::metrics::MetricsRegistry) -> anyhow::Result<String>
+            + Send
+            + Sync
+            + 'static,
+    {
+        let mut registries = self.metrics_registry_by_prefix.lock().unwrap();
+        registries
+            .entry(prefix.to_string())
+            .or_default()
+            .add_callback(self as &dyn crate::metrics::MetricsRegistry, callback);
+    }
+
+    /// Execute all callbacks for a given prefix and return their results
+    pub fn execute_metrics_callbacks(&self, prefix: &str) -> Vec<anyhow::Result<String>> {
+        let registries = self.metrics_registry_by_prefix.lock().unwrap();
+        if let Some(entry) = registries.get(prefix) {
+            entry.execute_callbacks(self as &dyn crate::metrics::MetricsRegistry)
+        } else {
+            Vec::new()
+        }
+    }
+
+    /// Get all registered prefixes. Private because it is only used for testing.
+    fn get_metrics_prefixes(&self) -> Vec<String> {
+        let registries = self.metrics_registry_by_prefix.lock().unwrap();
+        registries.keys().cloned().collect()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_metrics_registry_methods() {
+        let rt = tokio::runtime::Runtime::new().unwrap();
+        let drt = rt.block_on(async {
+            let runtime = Runtime::single_threaded().unwrap();
+            DistributedRuntime::from_settings_without_discovery(runtime)
+                .await
+                .unwrap()
+        });
+
+        // Test adding callbacks
+        drt.add_metrics_callback("test_prefix", |_| Ok("test_callback".to_string()));
+        drt.add_metrics_callback("test_prefix", |_| Ok("test_callback2".to_string()));
+
+        // Test executing callbacks
+        let results = drt.execute_metrics_callbacks("test_prefix");
+        assert_eq!(results.len(), 2);
+        assert_eq!(results[0].as_ref().unwrap(), "test_callback");
+        assert_eq!(results[1].as_ref().unwrap(), "test_callback2");
+
+        // Test getting prefixes
+        let prefixes = drt.get_metrics_prefixes();
+        assert!(prefixes.contains(&"test_prefix".to_string()));
+
+        // Test non-existent prefix
+        let empty_results = drt.execute_metrics_callbacks("non_existent");
+        assert_eq!(empty_results.len(), 0);
+
+        // Test adding a Prometheus metric
+        let counter = prometheus::Counter::new("test_counter", "A test counter").unwrap();
+        drt.add_prometheus_metric("test_prefix", Box::new(counter.clone()))
+            .unwrap();
+
+        // Verify the metric was added by checking if the prefix exists
+        let prefixes = drt.get_metrics_prefixes();
+        assert!(prefixes.contains(&"test_prefix".to_string()));
+    }
 }
 
 #[derive(Dissolve)]

lib/runtime/src/distributed.rs

@@ -147,6 +147,63 @@ impl SystemHealth {
     }
 }
 
+/// Structure to hold Prometheus registries and associated callbacks for a given prefix
+pub struct MetricsRegistryEntry {
+    /// The Prometheus registry for this prefix
+    pub prometheus_registry: prometheus::Registry,
+    /// List of function callbacks that receive a reference to any MetricsRegistry
+    pub runtime_callbacks: Vec<
+        Box<dyn Fn(&dyn crate::metrics::MetricsRegistry) -> anyhow::Result<String> + Send + Sync>,
+    >,
+}
+
+impl MetricsRegistryEntry {
+    /// Create a new metrics registry entry with an empty registry and no callbacks
+    pub fn new() -> Self {
+        Self {
+            prometheus_registry: prometheus::Registry::new(),
+            runtime_callbacks: Vec::new(),
+        }
+    }
+
+    /// Add a callback function that receives a reference to any MetricsRegistry
+    pub fn add_callback<F>(&mut self, _runtime: &dyn crate::metrics::MetricsRegistry, callback: F)
+    where
+        F: Fn(&dyn crate::metrics::MetricsRegistry) -> anyhow::Result<String>
+            + Send
+            + Sync
+            + 'static,
+    {
+        self.runtime_callbacks.push(Box::new(callback));
+    }
+
+    /// Execute all runtime callbacks and return their results
+    pub fn execute_callbacks(
+        &self,
+        runtime: &dyn crate::metrics::MetricsRegistry,
+    ) -> Vec<anyhow::Result<String>> {
+        self.runtime_callbacks
+            .iter()
+            .map(|callback| callback(runtime))
+            .collect()
+    }
+}
+
+impl Default for MetricsRegistryEntry {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Clone for MetricsRegistryEntry {
+    fn clone(&self) -> Self {
+        Self {
+            prometheus_registry: self.prometheus_registry.clone(),
+            runtime_callbacks: Vec::new(), // Callbacks cannot be cloned, so we start with an empty list
+        }
+    }
+}
+
 /// Distributed [Runtime] which provides access to shared resources across the cluster, this includes
 /// communication protocols and transports.
 #[derive(Clone)]
@@ -176,6 +233,6 @@ pub struct DistributedRuntime {
     // Health Status
     system_health: Arc<std::sync::Mutex<SystemHealth>>,
 
-    // This map associates metric prefixes with their corresponding Prometheus registries.
-    prometheus_registries_by_prefix: Arc<std::sync::Mutex<HashMap<String, prometheus::Registry>>>,
+    // This map associates metric prefixes with their corresponding Prometheus registries and callbacks.
+    metrics_registry_by_prefix: Arc<std::sync::Mutex<HashMap<String, MetricsRegistryEntry>>>,
 }