diff --git a/rust_snuba/rust_arroyo/src/utils/metrics.rs b/rust_snuba/rust_arroyo/src/utils/metrics.rs index c85628fff4..9aee1197c3 100644 --- a/rust_snuba/rust_arroyo/src/utils/metrics.rs +++ b/rust_snuba/rust_arroyo/src/utils/metrics.rs @@ -9,6 +9,11 @@ use std::net::{ToSocketAddrs, UdpSocket}; use std::sync::Arc; pub trait MetricsClientTrait: Send + Sync { + + fn should_sample(&self, sample_rate: Option) -> bool { + rand::thread_rng().gen_range(0.0..=1.0) < sample_rate.unwrap_or(1.0) + } + fn counter( &self, key: &str, @@ -34,7 +39,6 @@ pub trait MetricsClientTrait: Send + Sync { ); } -// #[derive( Clone)] pub struct MetricsClient { statsd_client: StatsdClient, prefix: String, @@ -120,9 +124,6 @@ impl MetricsClientTrait for MetricsClient { } impl MetricsClient { - fn should_sample(&self, sample_rate: Option) -> bool { - rand::thread_rng().gen_range(0.0..=1.0) < sample_rate.unwrap_or(1.0) - } fn send_with_tags<'t, T: cadence::Metric + From>( &self, @@ -210,11 +211,17 @@ mod tests { fn test_metrics() { init("my_host", "0.0.0.0:8125"); + assert!(!METRICS_CLIENT + .read() + .clone() + .unwrap() + .should_sample(Some(0.0)),); + assert!(METRICS_CLIENT .read() .clone() - .is_some() - ); + .unwrap() + .should_sample(Some(1.0)),); increment( "a", diff --git a/rust_snuba/src/utils/metrics/backends/datadog.rs b/rust_snuba/src/utils/metrics/backends/datadog.rs index 3a5c2fbe6c..da12a8aa27 100644 --- a/rust_snuba/src/utils/metrics/backends/datadog.rs +++ b/rust_snuba/src/utils/metrics/backends/datadog.rs @@ -43,8 +43,12 @@ impl MetricsClientTrait for DatadogMetricsBackend { key: &str, value: Option, tags: Option>, - _sample_rate: Option, + sample_rate: Option, ) { + if !self.should_sample(sample_rate) { + return; + } + let tags_str: Vec = tags.unwrap().iter().map(|(k, v)| format!("{k}:{v}")).collect(); match value { @@ -70,8 +74,11 @@ impl MetricsClientTrait for DatadogMetricsBackend { key: &str, value: u64, tags: Option>, - _sample_rate: Option, + sample_rate: Option, ) { + if !self.should_sample(sample_rate) { + return; + } let tags_str: Vec = tags.unwrap().iter().map(|(k, v)| format!("{k}:{v}")).collect(); self.client_sd.gauge(key, value.to_string(), tags_str).unwrap(); } @@ -81,8 +88,11 @@ impl MetricsClientTrait for DatadogMetricsBackend { key: &str, value: u64, tags: Option>, - _sample_rate: Option, + sample_rate: Option, ) { + if !self.should_sample(sample_rate) { + return; + } let tags_str: Vec = tags.unwrap().iter().map(|(k, v)| format!("{k}:{v}")).collect(); self.client_sd.timing(key, value.try_into().unwrap(), tags_str).unwrap(); } diff --git a/rust_snuba/src/utils/metrics/backends/testing.rs b/rust_snuba/src/utils/metrics/backends/testing.rs index a28b609784..3a61d38559 100644 --- a/rust_snuba/src/utils/metrics/backends/testing.rs +++ b/rust_snuba/src/utils/metrics/backends/testing.rs @@ -35,7 +35,11 @@ impl MetricsBackend for TestingMetricsBackend { } impl MetricsClientTrait for TestingMetricsBackend{ - fn counter(&self, name: &str, value: Option, tags: Option>, _sample_rate: Option) { + fn counter(&self, name: &str, value: Option, tags: Option>, sample_rate: Option) { + if !self.should_sample(sample_rate) { + return; + } + let mut tags_vec = Vec::new(); if let Some(tags) = tags { for (k, v) in tags { @@ -50,7 +54,10 @@ impl MetricsClientTrait for TestingMetricsBackend{ }); } - fn gauge(&self, name: &str, value: u64, tags: Option>, _sample_rate: Option) { + fn gauge(&self, name: &str, value: u64, tags: Option>, sample_rate: Option) { + if !self.should_sample(sample_rate) { + return; + } let mut tags_vec = Vec::new(); if let Some(tags) = tags { for (k, v) in tags { @@ -66,7 +73,11 @@ impl MetricsClientTrait for TestingMetricsBackend{ }); } - fn time(&self, name: &str, value: u64, tags: Option>, _sample_rate: Option) { + fn time(&self, name: &str, value: u64, tags: Option>, sample_rate: Option) { + if !self.should_sample(sample_rate) { + return; + } + let mut tags_vec = Vec::new(); if let Some(tags) = tags { for (k, v) in tags {