kona_cli/
metrics_args.rs

1//! Utility module to house implementation and declaration of MetricsArgs since it's being used in
2//! multiple places, it's just being referenced from this module.
3
4use crate::init_prometheus_server;
5use clap::{Parser, arg};
6use std::net::IpAddr;
7
8/// Configuration for Prometheus metrics.
9#[derive(Debug, Clone, Parser)]
10#[command(next_help_heading = "Metrics")]
11pub struct MetricsArgs {
12    /// Controls whether Prometheus metrics are enabled. Disabled by default.
13    #[arg(
14        long = "metrics.enabled",
15        global = true,
16        default_value_t = false,
17        env = "KONA_METRICS_ENABLED"
18    )]
19    pub enabled: bool,
20
21    /// The port to serve Prometheus metrics on.
22    #[arg(long = "metrics.port", global = true, default_value = "9090", env = "KONA_METRICS_PORT")]
23    pub port: u16,
24
25    /// The IP address to use for Prometheus metrics.
26    #[arg(
27        long = "metrics.addr",
28        global = true,
29        default_value = "0.0.0.0",
30        env = "KONA_METRICS_ADDR"
31    )]
32    pub addr: IpAddr,
33}
34
35impl Default for MetricsArgs {
36    fn default() -> Self {
37        Self::parse_from::<[_; 0], &str>([])
38    }
39}
40
41impl MetricsArgs {
42    /// Initialize the tracing stack and Prometheus metrics recorder.
43    ///
44    /// This function should be called at the beginning of the program.
45    pub fn init_metrics(&self) -> anyhow::Result<()> {
46        if self.enabled {
47            init_prometheus_server(self.addr, self.port)?;
48        }
49
50        Ok(())
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57    use clap::Parser;
58    use std::net::{IpAddr, Ipv4Addr};
59
60    /// Helper struct to parse MetricsArgs within a test CLI structure.
61    #[derive(Parser, Debug)]
62    struct TestCli {
63        #[command(flatten)]
64        metrics: MetricsArgs,
65    }
66
67    #[test]
68    fn test_default_metrics_args() {
69        let cli = TestCli::parse_from(["test_app"]);
70        assert!(!cli.metrics.enabled, "Default for metrics.enabled should be false.");
71        assert_eq!(cli.metrics.port, 9090, "Default for metrics.port should be 9090.");
72        assert_eq!(
73            cli.metrics.addr,
74            IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
75            "Default for metrics.addr should be 0.0.0.0."
76        );
77    }
78
79    #[test]
80    fn test_metrics_args_from_cli() {
81        let cli = TestCli::parse_from([
82            "test_app",
83            "--metrics.enabled",
84            "--metrics.port",
85            "9999",
86            "--metrics.addr",
87            "127.0.0.1",
88        ]);
89        assert!(cli.metrics.enabled, "metrics.enabled should be true.");
90        assert_eq!(cli.metrics.port, 9999, "metrics.port should be parsed from CLI.");
91        assert_eq!(
92            cli.metrics.addr,
93            IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
94            "metrics.addr should be parsed from CLI."
95        );
96    }
97}