refactor(node): include metrics server in Node

Author: kariy

Cargo.lock

@@ -8,6 +8,7 @@ version.workspace = true
 [dependencies]
 hyper = { workspace = true, features = [ "http1", "http2", "server", "tcp" ] }
 thiserror.workspace = true
+tokio.workspace = true
 tracing.workspace = true
 
 # Metrics

crates/metrics/Cargo.toml

@@ -29,6 +29,12 @@ pub enum Error {
     #[error("could not bind to address: {addr}")]
     FailedToBindAddress { addr: SocketAddr },
 
+    #[error("metrics server has already been stopped")]
+    AlreadyStopped,
+
+    #[error("failed to join server task: {0}")]
+    JoinError(String),
+
     #[error(transparent)]
     Server(#[from] hyper::Error),
 }

crates/metrics/src/lib.rs

@@ -5,19 +5,55 @@ use std::sync::Arc;
 
 use hyper::service::{make_service_fn, service_fn};
 use hyper::{Body, Request, Response};
+use tokio::sync::oneshot;
+use tracing::info;
 
 use crate::exporters::Exporter;
 use crate::{Error, Report};
 
+/// A handle to the metrics server.
+#[derive(Debug)]
+pub struct MetricsServerHandle {
+    /// The actual address that the server is bound to.
+    addr: SocketAddr,
+    /// The shutdown sender to stop the server.
+    shutdown_tx: Option<oneshot::Sender<()>>,
+    /// The task handle to wait for server completion.
+    task_handle: tokio::task::JoinHandle<Result<(), Error>>,
+}
+
+impl MetricsServerHandle {
+    /// Tell the server to stop without waiting for the server to stop.
+    pub fn stop(&mut self) -> Result<(), Error> {
+        if let Some(tx) = self.shutdown_tx.take() {
+            // Ignore error if receiver already dropped
+            let _ = tx.send(());
+            Ok(())
+        } else {
+            Err(Error::AlreadyStopped)
+        }
+    }
+
+    /// Wait until the server has stopped.
+    pub async fn stopped(self) -> Result<(), Error> {
+        self.task_handle.await.map_err(|e| Error::JoinError(e.to_string()))?
+    }
+
+    /// Returns the socket address the server is listening on.
+    pub fn addr(&self) -> &SocketAddr {
+        &self.addr
+    }
+}
+
 /// A helper trait for defining the type for hooks that are called when the metrics are being
 /// collected by the server.
 trait Hook: Fn() + Send + Sync {}
 impl<T: Fn() + Send + Sync> Hook for T {}
 
-/// A boxed [`Hook`].
-type BoxedHook = Box<dyn Hook<Output = ()>>;
-/// A list of [BoxedHook].
-type Hooks = Vec<BoxedHook>;
+/// A shared hook that can be cloned.
+type SharedHook = Arc<dyn Hook<Output = ()>>;
+/// A list of shared hooks.
+type Hooks = Vec<SharedHook>;
 
 /// Server for serving metrics.
 // TODO: allow configuring the server executor to allow cancelling on invidiual connection tasks.
@@ -46,8 +82,8 @@ where
     where
         I: IntoIterator<Item = Box<dyn Report>>,
     {
-        // convert the report types into callable hooks
-        let hooks = reports.into_iter().map(|r| Box::new(move || r.report()) as BoxedHook);
+        // convert the report types into callable hooks, wrapping in Arc for sharing
+        let hooks = reports.into_iter().map(|r| Arc::new(move || r.report()) as SharedHook);
         self.hooks.extend(hooks);
         self
     }
@@ -60,42 +96,67 @@ where
         describe_memory_stats();
 
         let hooks: Hooks =
-            vec![Box::new(collect_memory_stats), Box::new(move || process.collect())];
+            vec![Arc::new(collect_memory_stats), Arc::new(move || process.collect())];
 
         self.hooks.extend(hooks);
         self
     }
 
     /// Starts an endpoint at the given address to serve Prometheus metrics.
-    pub async fn start(self, addr: SocketAddr) -> Result<(), Error> {
-        let hooks = Arc::new(move || self.hooks.iter().for_each(|hook| hook()));
+    ///
+    /// Returns a handle that can be used to stop the server and wait for it to finish.
+    pub async fn start(&self, addr: SocketAddr) -> Result<MetricsServerHandle, Error> {
+        // Clone the hooks (clones the Arc references, not the closures themselves)
+        let hooks = self.hooks.clone();
+        let exporter = self.exporter.clone();
 
-        hyper::Server::try_bind(&addr)
+        let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
+
+        let server = hyper::Server::try_bind(&addr)
             .map_err(|_| Error::FailedToBindAddress { addr })?
             .serve(make_service_fn(move |_| {
-                let hook = Arc::clone(&hooks);
-                let exporter = self.exporter.clone();
+                let hooks = hooks.clone();
+                let exporter = exporter.clone();
                 async move {
                     Ok::<_, Infallible>(service_fn(move |_: Request<Body>| {
-                        // call the hooks to collect metrics before exporting them
-                        (hook)();
-                        // export the metrics from the installed exporter and send as response
-                        let metrics = Body::from(exporter.export());
-                        async move { Ok::<_, Infallible>(Response::new(metrics)) }
+                        let hooks = hooks.clone();
+                        let exporter = exporter.clone();
+                        async move {
+                            // need to call the hooks to collect metrics before exporting them
+                            for hook in &hooks {
+                                hook();
+                            }
+                            // export the metrics from the installed exporter and send as response
+                            let metrics = Body::from(exporter.export());
+                            Ok::<_, Infallible>(Response::new(metrics))
+                        }
                     }))
                 }
-            }))
-            .await?;
+            }));
+
+        let actual_addr = server.local_addr();
 
-        Ok(())
+        // Spawn the server with graceful shutdown
+        let task_handle = tokio::spawn(async move {
+            server
+                .with_graceful_shutdown(async {
+                    shutdown_rx.await.ok();
+                })
+                .await
+                .map_err(Error::Server)
+        });
+
+        info!(target: "metrics", addr = %actual_addr, "Metrics server started.");
+
+        Ok(MetricsServerHandle { addr: actual_addr, shutdown_tx: Some(shutdown_tx), task_handle })
     }
 }
 
-impl<MetricsExporter> fmt::Debug for Server<MetricsExporter>
-where
-    MetricsExporter: fmt::Debug,
-{
+impl<MetricsExporter> fmt::Debug for Server<MetricsExporter> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("Server").field("hooks", &"...").field("exporter", &self.exporter).finish()
+        f.debug_struct("Server")
+            .field("hooks", &format_args!("{} hook(s)", self.hooks.len()))
+            .field("exporter", &"<exporter>")
+            .finish()
     }
 }

crates/metrics/src/server.rs

@@ -232,9 +232,9 @@ impl Node {
 
             let addr = cfg.socket_addr();
             let server = MetricsServer::new(exporter).with_process_metrics().with_reports(reports);
-            self.task_manager.task_spawner().build_task().spawn(server.start(addr));
 
-            info!(%addr, "Metrics server started.");
+            // Start the metrics server and discard the handle since full node doesn't track it
+            let _ = server.start(addr).await?;
         }
 
         let pipeline_handle = self.pipeline.handle();

crates/node/src/full/mod.rs

@@ -25,9 +25,9 @@ use katana_executor::implementation::blockifier::BlockifierFactory;
 use katana_executor::{ExecutionFlags, ExecutorFactory};
 use katana_gas_price_oracle::{FixedPriceOracle, GasPriceOracle};
 use katana_gateway_server::{GatewayServer, GatewayServerHandle};
-use katana_metrics::exporters::prometheus::PrometheusRecorder;
+use katana_metrics::exporters::prometheus::{Prometheus, PrometheusRecorder};
 use katana_metrics::sys::DiskReporter;
-use katana_metrics::{Report, Server as MetricsServer};
+use katana_metrics::{MetricsServerHandle, Report, Server as MetricsServer};
 use katana_pool::ordering::FiFo;
 use katana_pool::TxPool;
 use katana_primitives::env::VersionedConstantsOverrides;
@@ -65,6 +65,7 @@ pub struct Node {
     backend: Arc<Backend<BlockifierFactory>>,
     block_producer: BlockProducer<BlockifierFactory>,
     gateway_server: Option<GatewayServer<TxPool>>,
+    metrics_server: Option<MetricsServer<Prometheus>>,
 }
 
 impl Node {
@@ -320,13 +321,29 @@ impl Node {
             None
         };
 
+        // --- build metrics server (optional)
+
+        let metrics_server = if config.metrics.is_some() {
+            let db_metrics = Box::new(db.clone()) as Box<dyn Report>;
+            let disk_metrics = Box::new(DiskReporter::new(db.path())?) as Box<dyn Report>;
+            let reports: Vec<Box<dyn Report>> = vec![db_metrics, disk_metrics];
+
+            let exporter = PrometheusRecorder::current().expect("qed; should exist at this point");
+            let server = MetricsServer::new(exporter).with_process_metrics().with_reports(reports);
+
+            Some(server)
+        } else {
+            None
+        };
+
         Ok(Node {
             db,
             pool,
             backend,
             rpc_server,
             gateway_server,
             block_producer,
+            metrics_server,
             config: Arc::new(config),
             task_manager,
         })
@@ -339,19 +356,15 @@ impl Node {
         let chain = self.backend.chain_spec.id();
         info!(%chain, "Starting node.");
 
-        // TODO: maybe move this to the build stage
-        if let Some(ref cfg) = self.config.metrics {
-            let db_metrics = Box::new(self.db.clone()) as Box<dyn Report>;
-            let disk_metrics = Box::new(DiskReporter::new(self.db.path())?) as Box<dyn Report>;
-            let reports: Vec<Box<dyn Report>> = vec![db_metrics, disk_metrics];
-
-            let exporter = PrometheusRecorder::current().expect("qed; should exist at this point");
-            let server = MetricsServer::new(exporter).with_process_metrics().with_reports(reports);
+        // --- start the metrics server (if configured)
 
+        let metrics_handle = if let Some(ref server) = self.metrics_server {
+            let cfg = self.config.metrics.as_ref().expect("qed; must exist");
             let addr = cfg.socket_addr();
-            self.task_manager.task_spawner().build_task().spawn(server.start(addr));
-            info!(%addr, "Metrics server started.");
-        }
+            Some(server.start(addr).await?)
+        } else {
+            None
+        };
 
         let pool = self.pool.clone();
         let backend = self.backend.clone();
@@ -401,7 +414,12 @@ impl Node {
 
         info!(target: "node", "Gas price oracle worker started.");
 
-        Ok(LaunchedNode { node: self, rpc: rpc_handle, gateway: gateway_handle })
+        Ok(LaunchedNode {
+            node: self,
+            rpc: rpc_handle,
+            gateway: gateway_handle,
+            metrics: metrics_handle,
+        })
     }
 
     /// Returns a reference to the node's database environment (if any).
@@ -437,6 +455,8 @@ pub struct LaunchedNode {
     rpc: RpcServerHandle,
     /// Handle to the gateway server (if enabled).
     gateway: Option<GatewayServerHandle>,
+    /// Handle to the metrics server (if enabled).
+    metrics: Option<MetricsServerHandle>,
 }
 
 impl LaunchedNode {
@@ -455,6 +475,11 @@ impl LaunchedNode {
         self.gateway.as_ref()
     }
 
+    /// Returns a reference to the metrics server handle (if enabled).
+    pub fn metrics(&self) -> Option<&MetricsServerHandle> {
+        self.metrics.as_ref()
+    }
+
     /// Stops the node.
     ///
     /// This will instruct the node to stop and wait until it has actually stop.
@@ -467,6 +492,11 @@ impl LaunchedNode {
             handle.stop()?;
         }
 
+        // Stop metrics server if it's running
+        if let Some(mut handle) = self.metrics {
+            handle.stop()?;
+        }
+
         self.node.task_manager.shutdown().await;
         Ok(())
     }