From e3b1bf61646d59e7f6a95e1ab84afd3b1493f8aa Mon Sep 17 00:00:00 2001 From: Sergey Timoshin Date: Sat, 7 Sep 2024 18:29:02 +0700 Subject: [PATCH] Enable optional metrics reporting via CLI flag Added a new CLI flag, `enable_metrics`, to control metrics reporting. Updated configuration initialization and relevant components to respect this flag, ensuring metrics are processed and pushed conditionally based on user input. --- forester/src/cli.rs | 3 +++ forester/src/config.rs | 2 ++ forester/src/epoch_manager.rs | 14 ++++++++------ forester/src/main.rs | 17 +++++++++++------ forester/src/settings.rs | 3 ++- forester/tests/test_utils.rs | 1 + 6 files changed, 27 insertions(+), 13 deletions(-) diff --git a/forester/src/cli.rs b/forester/src/cli.rs index 767e62f549..78a1947932 100644 --- a/forester/src/cli.rs +++ b/forester/src/cli.rs @@ -3,6 +3,9 @@ use clap::{Parser, Subcommand}; #[derive(Parser)] #[clap(author, version, about, long_about=None)] pub struct Cli { + #[arg(long, default_value_t = false, global = true)] + pub enable_metrics: bool, + #[command(subcommand)] pub command: Option, } diff --git a/forester/src/config.rs b/forester/src/config.rs index 530c5b4f5d..5a481b181c 100644 --- a/forester/src/config.rs +++ b/forester/src/config.rs @@ -50,6 +50,7 @@ pub struct ForesterConfig { pub slot_update_interval_seconds: u64, pub address_tree_data: Vec, pub state_tree_data: Vec, + pub enable_metrics: bool, } impl Clone for ForesterConfig { @@ -68,6 +69,7 @@ impl Clone for ForesterConfig { state_tree_data: self.state_tree_data.clone(), address_tree_data: self.address_tree_data.clone(), slot_update_interval_seconds: self.slot_update_interval_seconds, + enable_metrics: self.enable_metrics, } } } diff --git a/forester/src/epoch_manager.rs b/forester/src/epoch_manager.rs index d592342b55..4df91c3d78 100644 --- a/forester/src/epoch_manager.rs +++ b/forester/src/epoch_manager.rs @@ -664,8 +664,9 @@ impl> EpochManager { let num_tx_sent = num_tx_sent.await?; // Prometheus metrics let chunk_duration = start_time.elapsed(); - queue_metric_update(epoch_info.epoch, num_tx_sent, chunk_duration).await; - + if self.config.enable_metrics { + queue_metric_update(epoch_info.epoch, num_tx_sent, chunk_duration).await; + } // TODO: consider do we really need WorkReport self.increment_processed_items_count(epoch_info.epoch, num_tx_sent) .await; @@ -673,10 +674,11 @@ impl> EpochManager { // The forester is not eligible for any more slots in the current epoch break; } - - process_queued_metrics().await; - if let Err(e) = push_metrics(&self.config.external_services.pushgateway_url).await { - error!("Failed to push metrics: {:?}", e); + if self.config.enable_metrics { + process_queued_metrics().await; + if let Err(e) = push_metrics(&self.config.external_services.pushgateway_url).await { + error!("Failed to push metrics: {:?}", e); + } } estimated_slot = self.slot_tracker.estimated_current_slot(); diff --git a/forester/src/main.rs b/forester/src/main.rs index 79362628a0..a83986dff8 100644 --- a/forester/src/main.rs +++ b/forester/src/main.rs @@ -16,17 +16,20 @@ use tracing::{debug, warn}; #[tokio::main] async fn main() -> Result<(), ForesterError> { setup_telemetry(); - register_metrics(); - let config = match init_config() { + let cli = Cli::parse(); + + if cli.enable_metrics { + register_metrics(); + } + + let config = match init_config(cli.enable_metrics) { Ok(config) => Arc::new(config), Err(e) => { panic!("Failed to initialize config: {}", e); } }; - let cli = Cli::parse(); - match &cli.command { Some(Commands::Start) => { let (shutdown_sender, shutdown_receiver) = oneshot::channel(); @@ -65,8 +68,10 @@ async fn main() -> Result<(), ForesterError> { run_queue_info(config.clone(), trees.clone(), TreeType::State).await; run_queue_info(config.clone(), trees.clone(), TreeType::Address).await; - if let Err(e) = push_metrics(&config.external_services.pushgateway_url).await { - warn!("Failed to push metrics: {:?}", e); + if cli.enable_metrics { + if let Err(e) = push_metrics(&config.external_services.pushgateway_url).await { + warn!("Failed to push metrics: {:?}", e); + } } } None => {} diff --git a/forester/src/settings.rs b/forester/src/settings.rs index db46bc2420..4f6f712e8b 100644 --- a/forester/src/settings.rs +++ b/forester/src/settings.rs @@ -90,7 +90,7 @@ fn build_config() -> Result { .map_err(|e| ForesterError::ConfigError(e.to_string())) } -pub fn init_config() -> Result { +pub fn init_config(enable_metrics: bool) -> Result { let settings = build_config()?; let registry_pubkey = light_registry::program::LightRegistry::id().to_string(); @@ -135,6 +135,7 @@ pub fn init_config() -> Result { as u64, address_tree_data: vec![], state_tree_data: vec![], + enable_metrics, }; debug!("Config: {:?}", config); diff --git a/forester/tests/test_utils.rs b/forester/tests/test_utils.rs index 38a6c4e7aa..8642f4611f 100644 --- a/forester/tests/test_utils.rs +++ b/forester/tests/test_utils.rs @@ -97,6 +97,7 @@ pub fn forester_config() -> ForesterConfig { slot_update_interval_seconds: 10, address_tree_data: vec![], state_tree_data: vec![], + enable_metrics: false, } }