diff --git a/crates/shared/cli-utils/src/lib.rs b/crates/shared/cli-utils/src/lib.rs index fe92be91..9ffeeff0 100644 --- a/crates/shared/cli-utils/src/lib.rs +++ b/crates/shared/cli-utils/src/lib.rs @@ -12,6 +12,9 @@ pub use backtrace::Backtracing; mod args; pub use args::GlobalArgs; +mod metrics; +pub use metrics::MetricsArgs; + mod logging; pub use logging::{LogFormat, LogRotation, LoggingArgs}; diff --git a/crates/shared/cli-utils/src/metrics.rs b/crates/shared/cli-utils/src/metrics.rs new file mode 100644 index 00000000..81390d18 --- /dev/null +++ b/crates/shared/cli-utils/src/metrics.rs @@ -0,0 +1,101 @@ +//! Utility module to house implementation and declaration of MetricsArgs since it's being used in +//! multiple places, it's just being referenced from this module. + +use std::net::IpAddr; + +use clap::Parser; + +use crate::{CliResult, PrometheusServer}; + +/// Configuration for Prometheus metrics. +#[derive(Debug, Clone, Parser)] +#[command(next_help_heading = "Metrics")] +pub struct MetricsArgs { + /// Controls whether Prometheus metrics are enabled. Disabled by default. + #[arg( + long = "metrics.enabled", + global = true, + default_value_t = false, + env = "BASE_METRICS_ENABLED" + )] + pub enabled: bool, + + /// The port to serve Prometheus metrics on. + #[arg(long = "metrics.port", global = true, default_value = "9090", env = "BASE_METRICS_PORT")] + pub port: u16, + + /// The IP address to use for Prometheus metrics. + #[arg( + long = "metrics.addr", + global = true, + default_value = "0.0.0.0", + env = "BASE_METRICS_ADDR" + )] + pub addr: IpAddr, +} + +impl Default for MetricsArgs { + fn default() -> Self { + Self::parse_from::<[_; 0], &str>([]) + } +} + +impl MetricsArgs { + /// Initialize the tracing stack and Prometheus metrics recorder. + /// + /// This function should be called at the beginning of the program. + pub fn init_metrics(&self) -> CliResult<()> { + if self.enabled { + PrometheusServer::init(self.addr, self.port)?; + } + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::net::{IpAddr, Ipv4Addr}; + + use clap::Parser; + + use super::*; + + /// Helper struct to parse MetricsArgs within a test CLI structure. + #[derive(Parser, Debug)] + struct TestCli { + #[command(flatten)] + metrics: MetricsArgs, + } + + #[test] + fn test_default_metrics_args() { + let cli = TestCli::parse_from(["test_app"]); + assert!(!cli.metrics.enabled, "Default for metrics.enabled should be false."); + assert_eq!(cli.metrics.port, 9090, "Default for metrics.port should be 9090."); + assert_eq!( + cli.metrics.addr, + IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), + "Default for metrics.addr should be 0.0.0.0." + ); + } + + #[test] + fn test_metrics_args_from_cli() { + let cli = TestCli::parse_from([ + "test_app", + "--metrics.enabled", + "--metrics.port", + "9999", + "--metrics.addr", + "127.0.0.1", + ]); + assert!(cli.metrics.enabled, "metrics.enabled should be true."); + assert_eq!(cli.metrics.port, 9999, "metrics.port should be parsed from CLI."); + assert_eq!( + cli.metrics.addr, + IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), + "metrics.addr should be parsed from CLI." + ); + } +}