chore(runtime): Make nats_client private, refactor NATS stats scraping (#4591)

Signed-off-by: Graham King <grahamk@nvidia.com>

Author: grahamking

lib/runtime/src/component.rs

@@ -269,131 +269,6 @@ impl Component {
         instances.sort();
         Ok(instances)
     }
-
-    /// Scrape ServiceSet, which contains NATS stats as well as user defined stats
-    /// embedded in data field of ServiceInfo.
-    async fn scrape_stats(&self, timeout: Duration) -> anyhow::Result<ServiceSet> {
-        // Debug: scraping stats for component
-        let service_name = self.service_name();
-        let Some(service_client) = self
-            .drt()
-            .nats_client()
-            .map(|nc| ServiceClient::new(nc.clone()))
-        else {
-            anyhow::bail!("ServiceSet is gathered via NATS, do not call this in non-NATS setups.");
-        };
-        service_client
-            .collect_services(&service_name, timeout)
-            .await
-    }
-
-    /// Add Prometheus metrics for this component's NATS service stats.
-    ///
-    /// Starts a background task that periodically requests service statistics from NATS
-    /// and updates the corresponding Prometheus metrics. The first scrape happens immediately,
-    /// then subsequent scrapes occur at a fixed interval of 9.8 seconds (MAX_WAIT_MS),
-    /// which should be near or smaller than typical Prometheus scraping intervals to ensure
-    /// metrics are fresh when Prometheus collects them.
-    fn start_scraping_nats_service_component_metrics(&self) -> anyhow::Result<()> {
-        const MAX_WAIT_MS: std::time::Duration = std::time::Duration::from_millis(9800); // Should be <= Prometheus scrape interval
-
-        // If there is another component with the same service name, this will fail.
-        let component_metrics = ComponentNatsServerPrometheusMetrics::new(self)?;
-
-        let component_clone = self.clone();
-
-        // Start a background task that scrapes stats every 5 seconds
-        let m = component_metrics.clone();
-        let c = component_clone.clone();
-
-        // Use the DRT's runtime handle to spawn the background task.
-        // We cannot use regular `tokio::spawn` here because:
-        // 1. This method may be called from contexts without an active Tokio runtime
-        //    (e.g., tests that create a DRT in a blocking context)
-        // 2. Tests often create a temporary runtime just to build the DRT, then drop it
-        // 3. `tokio::spawn` requires being called from within a runtime context
-        // By using the DRT's own runtime handle, we ensure the task runs in the
-        // correct runtime that will persist for the lifetime of the component.
-        c.drt().runtime().secondary().spawn(async move {
-            let timeout = std::time::Duration::from_millis(500);
-            let mut interval = tokio::time::interval(MAX_WAIT_MS);
-            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
-
-            loop {
-                match c.scrape_stats(timeout).await {
-                    Ok(service_set) => {
-                        m.update_from_service_set(&service_set);
-                    }
-                    Err(err) => {
-                        tracing::error!(
-                            "Background scrape failed for {}: {}",
-                            c.service_name(),
-                            err
-                        );
-                        m.reset_to_zeros();
-                    }
-                }
-
-                interval.tick().await;
-            }
-        });
-
-        Ok(())
-    }
-
-    // Gather NATS metrics
-    async fn add_stats_service(&mut self) -> anyhow::Result<()> {
-        let service_name = self.service_name();
-
-        // Pre-check to save cost of creating the service, but don't hold the lock
-        if self
-            .drt
-            .component_registry()
-            .inner
-            .lock()
-            .await
-            .services
-            .contains_key(&service_name)
-        {
-            // The NATS service is per component, but it is called from `serve_endpoint`, and there
-            // are often multiple endpoints for a component (e.g. `clear_kv_blocks` and `generate`).
-            tracing::trace!("Service {service_name} already exists");
-            return Ok(());
-        }
-
-        let Some(nats_client) = self.drt.nats_client() else {
-            anyhow::bail!("Cannot create NATS service without NATS.");
-        };
-        let description = None;
-        let (nats_service, stats_reg) =
-            service::build_nats_service(nats_client, self, description).await?;
-
-        let mut guard = self.drt.component_registry().inner.lock().await;
-        if !guard.services.contains_key(&service_name) {
-            // Normal case
-            guard.services.insert(service_name.clone(), nats_service);
-            guard.stats_handlers.insert(service_name.clone(), stats_reg);
-
-            tracing::info!("Added NATS / stats service {service_name}");
-
-            drop(guard);
-        } else {
-            drop(guard);
-            let _ = nats_service.stop().await;
-            // The NATS service is per component, but it is called from `serve_endpoint`, and there
-            // are often multiple endpoints for a component (e.g. `clear_kv_blocks` and `generate`).
-            return Ok(());
-        }
-
-        if let Err(err) = self.start_scraping_nats_service_component_metrics() {
-            tracing::debug!(
-                "Metrics registration failed for '{}': {}",
-                self.service_name(),
-                err
-            );
-        }
-        Ok(())
-    }
 }
 
 impl ComponentBuilder {
@@ -404,14 +279,9 @@ impl ComponentBuilder {
     pub fn build(self) -> Result<Component, anyhow::Error> {
         let component = self.build_internal()?;
         // If this component is using NATS, gather it's metrics
-        if component.drt().request_plane().is_nats() {
-            let mut c = component.clone();
-            // Start in the background to isolate the async, and because we don't need it yet
-            component.drt().runtime().secondary().spawn(async move {
-                if let Err(err) = c.add_stats_service().await {
-                    tracing::error!(error = %err, component = c.service_name(), "Failed starting stats service");
-                }
-            });
+        let drt = component.drt();
+        if drt.request_plane().is_nats() {
+            drt.start_stats_service(component.clone());
         }
         Ok(component)
     }

lib/runtime/src/component/component.rs

@@ -32,10 +32,9 @@ impl EventPublisher for Component {
         bytes: Vec<u8>,
     ) -> Result<()> {
         let subject = format!("{}.{}", self.subject(), event_name.as_ref());
-        let Some(nats_client) = self.drt().nats_client() else {
-            anyhow::bail!("KV router's EventPublisher requires NATS");
-        };
-        nats_client.client().publish(subject, bytes.into()).await?;
+        self.drt()
+            .kv_router_nats_publish(subject, bytes.into())
+            .await?;
         Ok(())
     }
 }
@@ -47,10 +46,7 @@ impl EventSubscriber for Component {
         event_name: impl AsRef<str> + Send + Sync,
     ) -> Result<async_nats::Subscriber> {
         let subject = format!("{}.{}", self.subject(), event_name.as_ref());
-        let Some(nats_client) = self.drt().nats_client() else {
-            anyhow::bail!("KV router's EventSubscriber requires NATS");
-        };
-        Ok(nats_client.client().subscribe(subject).await?)
+        Ok(self.drt().kv_router_nats_subscribe(subject).await?)
     }
 
     async fn subscribe_with_type<T: for<'de> Deserialize<'de> + Send + 'static>(

lib/runtime/src/component/namespace.rs

@@ -35,10 +35,9 @@ impl EventPublisher for Namespace {
         bytes: Vec<u8>,
     ) -> Result<()> {
         let subject = format!("{}.{}", self.subject(), event_name.as_ref());
-        let Some(nats_client) = self.drt().nats_client() else {
-            anyhow::bail!("KV router's Namespace EventPublisher requires NATS");
-        };
-        nats_client.client().publish(subject, bytes.into()).await?;
+        self.drt()
+            .kv_router_nats_publish(subject, bytes.into())
+            .await?;
         Ok(())
     }
 }
@@ -50,10 +49,7 @@ impl EventSubscriber for Namespace {
         event_name: impl AsRef<str> + Send + Sync,
     ) -> Result<async_nats::Subscriber> {
         let subject = format!("{}.{}", self.subject(), event_name.as_ref());
-        let Some(nats_client) = self.drt().nats_client() else {
-            anyhow::bail!("KV router's Namespace EventSubscriber requires NATS");
-        };
-        Ok(nats_client.client().subscribe(subject).await?)
+        Ok(self.drt().kv_router_nats_subscribe(subject).await?)
     }
 
     async fn subscribe_with_type<T: for<'de> Deserialize<'de> + Send + 'static>(