diff --git a/grpc-server/src/statsd.rs b/grpc-server/src/statsd.rs index 99c95e456fe0adb96cfddf237bb65050eb2141fd..04722922fc30a9db35df604c66a8817f7eaf7120 100644 --- a/grpc-server/src/statsd.rs +++ b/grpc-server/src/statsd.rs @@ -2,12 +2,14 @@ // See the AUTHORS file at the top-level directory of this distribution // License: GNU General Public License version 3, or any later version // See top-level LICENSE file for more information -use anyhow::{anyhow, Context, Result}; + +use std::collections::HashMap; + +use anyhow::{anyhow, bail, Context, Result}; use std::net::{ToSocketAddrs, UdpSocket}; use std::str::FromStr; use cadence::{StatsdClient, UdpMetricSink}; - macro_rules! resolve_host { ($host:expr) => { $host @@ -20,16 +22,33 @@ macro_rules! resolve_host { #[allow(clippy::comparison_to_empty)] /// Parses the `$STATSD_TAGS` environment variable as a comma-separated list of `key=value`. -fn parse_statsd_tags(tags: &str) -> Result<Vec<(&str, &str)>> { +fn parse_statsd_tags<'a>( + tags: &'a str, + env: &'a HashMap<String, String>, +) -> Result<Vec<(&'a str, &'a str)>> { if tags == "" { return Ok(Vec::new()); } + // TODO: Add support for variable expansion, like swh.core.statsd does. tags.split(",") .map(|tag| { - tag.split_once(':').ok_or(anyhow!( - "STATSD_TAGS needs to be in 'key:value' format, not {tag}" - )) + let Some((k, mut v)) = tag.split_once(':') else { + bail!("STATSD_TAGS needs to be in 'key:value' format, not {tag}"); + }; + if v.starts_with('$') { + // Replace the tag value with an environment variable. + // Mimics behavior of swh/core/statsd.py, which also only supports substitution, + // and not expansion in general. + v = &v[1..]; + if v.starts_with('{') && v.ends_with('}') { + v = &v[1..v.len() - 1]; + } + v = env.get(v).ok_or(anyhow!( + "STATSD_TAGS references env variable {v}, which is undefined" + ))?; + } + Ok((k, v)) }) .collect() } @@ -50,46 +69,90 @@ pub fn statsd_client(host: Option<String>) -> Result<StatsdClient> { None => resolve_host!(default_host), }; let sink = UdpMetricSink::from(host, socket).unwrap(); - let client = parse_statsd_tags(&std::env::var("STATSD_TAGS").unwrap_or("".to_string()))? - .into_iter() - .fold( - StatsdClient::builder("swh_graph_grpc_server", sink), - |client_builder, (k, v)| client_builder.with_tag(k, v), - ) - .with_error_handler(|e| log::error!("Could not update Statsd metric: {e}")) - .build(); + let client = parse_statsd_tags( + &std::env::var("STATSD_TAGS").unwrap_or("".to_string()), + &std::env::vars().collect(), + )? + .into_iter() + .fold( + StatsdClient::builder("swh_graph_grpc_server", sink), + |client_builder, (k, v)| client_builder.with_tag(k, v), + ) + .with_error_handler(|e| log::error!("Could not update Statsd metric: {e}")) + .build(); Ok(client) } #[test] fn test_parse_statsd_tags() -> Result<()> { + let env = HashMap::new(); + assert_eq!( - parse_statsd_tags("").context("Could not parse empty string")?, + parse_statsd_tags("", &env).context("Could not parse empty string")?, Vec::new() ); assert_eq!( - parse_statsd_tags("foo:bar").context("Could not parse foo:bar")?, + parse_statsd_tags("foo:bar", &env).context("Could not parse foo:bar")?, vec![("foo", "bar")] ); assert_eq!( - parse_statsd_tags("foo:").context("Could not parse foo:bar")?, + parse_statsd_tags("foo:", &env).context("Could not parse foo:bar")?, vec![("foo", "")] ); assert_eq!( - parse_statsd_tags("foo:bar,baz:qux").context("Could not parse foo:bar,baz:qux")?, + parse_statsd_tags("foo:bar,baz:qux", &env).context("Could not parse foo:bar,baz:qux")?, vec![("foo", "bar"), ("baz", "qux")] ); assert_eq!( - parse_statsd_tags("foo:bar:bar2,baz:qux").context("Could not parse foo:bar,baz:qux")?, + parse_statsd_tags("foo:bar:bar2,baz:qux", &env) + .context("Could not parse foo:bar,baz:qux")?, vec![("foo", "bar:bar2"), ("baz", "qux")] ); - assert!(parse_statsd_tags("foo").is_err()); + Ok(()) +} + +#[test] +fn test_parse_statsd_tags_fail() -> Result<()> { + let env = HashMap::new(); + + assert!(parse_statsd_tags("foo", &env).is_err()); + + Ok(()) +} + +#[test] +fn test_parse_statsd_tags_var_substitution() -> Result<()> { + let mut env = HashMap::new(); + env.insert("ENV_VAR".to_string(), "ENV_VALUE".to_string()); + + assert_eq!( + parse_statsd_tags("foo:ENV_VAR", &env).context("Could not parse foo:bar")?, + vec![("foo", "ENV_VAR")] // not substituted + ); + + assert_eq!( + parse_statsd_tags("foo:$ENV_VAR", &env).context("Could not parse foo:bar")?, + vec![("foo", "ENV_VALUE")] + ); + + assert_eq!( + parse_statsd_tags("foo:${ENV_VAR}", &env).context("Could not parse foo:bar")?, + vec![("foo", "ENV_VALUE")] + ); + + assert_eq!( + parse_statsd_tags("foo:${ENV_VAR},bar:baz", &env).context("Could not parse foo:bar")?, + vec![("foo", "ENV_VALUE"), ("bar", "baz")] + ); + + assert!(parse_statsd_tags("foo:$NOT_ENV_VAR", &env).is_err()); + assert!(parse_statsd_tags("foo:${ENV_VAR", &env).is_err()); Ok(()) }