From 2ff48c47f77dd7b647c220f80e5abf3b9604e6a6 Mon Sep 17 00:00:00 2001
From: Dalitso Banda <dalitso.banda@sentry.io>
Date: Mon, 22 May 2023 23:22:30 -0700
Subject: [PATCH] add sampling

---
 rust_snuba/rust_arroyo/src/utils/metrics.rs   | 19 +++++++++++++------
 .../src/utils/metrics/backends/datadog.rs     | 16 +++++++++++++---
 .../src/utils/metrics/backends/testing.rs     | 17 ++++++++++++++---
 3 files changed, 40 insertions(+), 12 deletions(-)

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<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
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<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
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<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 {