add sampling

getsentry · Jun 27, 2023 · e04575f · e04575f
1 parent 7b73d0d
commit e04575f
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 {