update

Author: kariy

crates/metrics/src/server.rs

@@ -11,6 +11,16 @@ use tracing::info;
 use crate::exporters::Exporter;
 use crate::{Error, Report};
 
+/// 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 shared hook that can be cloned.
+type SharedHook = Arc<dyn Hook<Output = ()>>;
+/// A list of shared hooks.
+type Hooks = Vec<SharedHook>;
+
 /// A handle to the metrics server.
 #[derive(Debug)]
 pub struct MetricsServerHandle {
@@ -45,40 +55,27 @@ impl MetricsServerHandle {
     }
 }
 
-/// 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 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.
 // See, [hyper::server::server::Builder::executor]
-pub struct Server<MetricsExporter> {
+pub struct MetricsServer<E> {
     /// Hooks or callable functions for collecting metrics in the cases where
     /// the metrics are not being collected in the main program flow.
     ///
     /// These are called when metrics are being served through the server.
     hooks: Hooks,
     /// The exporter that is used to export the collected metrics.
-    exporter: MetricsExporter,
+    exporter: E,
 }
 
-impl<MetricsExporter> Server<MetricsExporter>
-where
-    MetricsExporter: Exporter + 'static,
-{
+impl<E: Exporter + 'static> MetricsServer<E> {
     /// Creates a new metrics server using the given exporter.
-    pub fn new(exporter: MetricsExporter) -> Self {
+    pub fn new(exporter: E) -> Self {
         Self { exporter, hooks: Vec::new() }
     }
 
     /// Add new metrics reporter to the server.
-    pub fn with_reports<I>(mut self, reports: I) -> Self
+    pub fn reports<I>(mut self, reports: I) -> Self
     where
         I: IntoIterator<Item = Box<dyn Report>>,
     {
@@ -105,7 +102,7 @@ where
     /// Starts an endpoint at the given address to serve Prometheus metrics.
     ///
     /// 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> {
+    pub 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();
@@ -137,6 +134,7 @@ where
         let actual_addr = server.local_addr();
 
         // Spawn the server with graceful shutdown
+        // TODO: spawn on a task manager
         let task_handle = tokio::spawn(async move {
             server
                 .with_graceful_shutdown(async {
@@ -152,7 +150,7 @@ where
     }
 }
 
-impl<MetricsExporter> fmt::Debug for Server<MetricsExporter> {
+impl<E> fmt::Debug for MetricsServer<E> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("Server")
             .field("hooks", &format_args!("{} hook(s)", self.hooks.len()))

crates/node/src/full/mod.rs

@@ -11,8 +11,9 @@ use katana_chain_spec::ChainSpec;
 use katana_executor::ExecutionFlags;
 use katana_gas_price_oracle::GasPriceOracle;
 use katana_gateway_client::Client as SequencerGateway;
-use katana_metrics::exporters::prometheus::PrometheusRecorder;
-use katana_metrics::{Report, Server as MetricsServer};
+use katana_metrics::exporters::prometheus::{Prometheus, PrometheusRecorder};
+use katana_metrics::sys::DiskReporter;
+use katana_metrics::{MetricsServer, MetricsServerHandle, Report};
 use katana_pipeline::{Pipeline, PipelineHandle};
 use katana_pool::ordering::TipOrdering;
 use katana_provider::providers::db::DbProvider;
@@ -76,6 +77,7 @@ pub struct Node {
     pub pipeline: Pipeline<DbProvider>,
     pub rpc_server: RpcServer,
     pub gateway_client: SequencerGateway,
+    pub metrics_server: Option<MetricsServer<Prometheus>>,
     pub chain_tip_watcher: ChainTipWatcher<SequencerGateway>,
 }
 
@@ -213,29 +215,45 @@ impl Node {
             rpc_server = rpc_server.max_response_body_size(max_response_body_size);
         }
 
+        // --- 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().reports(reports);
+
+            Some(server)
+        } else {
+            None
+        };
+
         Ok(Node {
             db,
             pool,
             pipeline,
             rpc_server,
             task_manager,
             gateway_client,
+            metrics_server,
             chain_tip_watcher,
             config: Arc::new(config),
         })
     }
 
     pub async fn launch(self) -> Result<LaunchedNode> {
-        if let Some(ref cfg) = self.config.metrics {
-            let reports: Vec<Box<dyn Report>> = vec![Box::new(self.db.clone()) as Box<dyn Report>];
-            let exporter = PrometheusRecorder::current().expect("qed; should exist at this point");
+        // --- start the metrics server (if configured)
 
+        let metrics_handle = if let Some(ref server) = self.metrics_server {
+            // safe to unwrap here because metrics_server can only be Some if the metrics config exists
+            let cfg = self.config.metrics.as_ref().expect("qed; must exist");
             let addr = cfg.socket_addr();
-            let server = MetricsServer::new(exporter).with_process_metrics().with_reports(reports);
-
-            // Start the metrics server and discard the handle since full node doesn't track it
-            let _ = server.start(addr).await?;
-        }
+            Some(server.start(addr)?)
+        } else {
+            None
+        };
 
         let pipeline_handle = self.pipeline.handle();
 
@@ -277,6 +295,7 @@ impl Node {
             config: self.config,
             task_manager: self.task_manager,
             pipeline: pipeline_handle,
+            metrics: metrics_handle,
             rpc,
         })
     }
@@ -289,6 +308,8 @@ pub struct LaunchedNode {
     pub config: Arc<Config>,
     pub rpc: RpcServerHandle,
     pub pipeline: PipelineHandle,
+    /// Handle to the metrics server (if enabled).
+    pub metrics: Option<MetricsServerHandle>,
 }
 
 impl LaunchedNode {

crates/node/src/lib.rs

@@ -27,7 +27,7 @@ use katana_gas_price_oracle::{FixedPriceOracle, GasPriceOracle};
 use katana_gateway_server::{GatewayServer, GatewayServerHandle};
 use katana_metrics::exporters::prometheus::{Prometheus, PrometheusRecorder};
 use katana_metrics::sys::DiskReporter;
-use katana_metrics::{MetricsServerHandle, Report, Server as MetricsServer};
+use katana_metrics::{MetricsServer, MetricsServerHandle, Report};
 use katana_pool::ordering::FiFo;
 use katana_pool::TxPool;
 use katana_primitives::env::VersionedConstantsOverrides;
@@ -329,7 +329,7 @@ impl Node {
             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);
+            let server = MetricsServer::new(exporter).with_process_metrics().reports(reports);
 
             Some(server)
         } else {
@@ -359,9 +359,10 @@ impl Node {
         // --- start the metrics server (if configured)
 
         let metrics_handle = if let Some(ref server) = self.metrics_server {
+            // safe to unwrap here because metrics_server can only be Some if the metrics config exists
             let cfg = self.config.metrics.as_ref().expect("qed; must exist");
             let addr = cfg.socket_addr();
-            Some(server.start(addr).await?)
+            Some(server.start(addr)?)
         } else {
             None
         };