add sampling

getsentry · May 23, 2023 · 2ff48c4 · 2ff48c4
1 parent 506ca85
commit 2ff48c4
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 12 deletions.
diff --git 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<f64>) -> 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<f64>) -> 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<String>>(
         &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
@@ -43,8 +43,12 @@ impl MetricsClientTrait for DatadogMetricsBackend {
         key: &str,
         value: Option<i64>,
         tags: Option<std::collections::HashMap<&str, &str>>,
-        _sample_rate: Option<f64>,
+        sample_rate: Option<f64>,
     ) {
+        if !self.should_sample(sample_rate) {
+            return;
+        }
+
         let tags_str: Vec<String> = 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<std::collections::HashMap<&str, &str>>,
-        _sample_rate: Option<f64>,
+        sample_rate: Option<f64>,
     ) {
+        if !self.should_sample(sample_rate) {
+            return;
+        }
         let tags_str: Vec<String> = 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<std::collections::HashMap<&str, &str>>,
-        _sample_rate: Option<f64>,
+        sample_rate: Option<f64>,
     ) {
+        if !self.should_sample(sample_rate) {
+            return;
+        }
         let tags_str: Vec<String> = 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
@@ -35,7 +35,11 @@ impl MetricsBackend for TestingMetricsBackend {
 }
 
 impl MetricsClientTrait for TestingMetricsBackend{
-    fn counter(&self, name: &str, value: Option<i64>, tags: Option<HashMap<&str, &str>>, _sample_rate: Option<f64>) {
+    fn counter(&self, name: &str, value: Option<i64>, tags: Option<HashMap<&str, &str>>, sample_rate: Option<f64>) {
+        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<HashMap<&str, &str>>, _sample_rate: Option<f64>) {
+    fn gauge(&self, name: &str, value: u64, tags: Option<HashMap<&str, &str>>, sample_rate: Option<f64>) {
+        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<HashMap<&str, &str>>, _sample_rate: Option<f64>) {
+    fn time(&self, name: &str, value: u64, tags: Option<HashMap<&str, &str>>, sample_rate: Option<f64>) {
+        if !self.should_sample(sample_rate) {
+            return;
+        }
+
         let mut tags_vec = Vec::new();
         if let Some(tags) = tags {
             for (k, v) in tags {