chore(dogstatsd): enforce a minimum sample rate on metrics (#1021)

Author: tobz

lib/saluki-core/src/data_model/event/metric/mod.rs

@@ -190,6 +190,11 @@ impl SampleRate {
         Self(1.0)
     }
 
+    /// Returns the sample rate.
+    pub const fn rate(&self) -> f64 {
+        self.0
+    }
+
     /// Returns the weight of the sample rate.
     pub fn weight(&self) -> u64 {
         (1.0 / self.0) as u64

lib/saluki-io/src/deser/codec/dogstatsd/metric.rs

@@ -9,6 +9,7 @@ use nom::{
 };
 use saluki_context::{origin::OriginTagCardinality, tags::RawTags};
 use saluki_core::data_model::event::metric::*;
+use tracing::warn;
 
 use super::{helpers::*, DogstatsdCodecConfiguration, NomParserError};
 
@@ -72,7 +73,9 @@ pub fn parse_dogstatsd_metric<'a>(
                 // point downstream to calculate the true metric value.
                 b'@' => {
                     let (_, sample_rate) =
-                        all_consuming(preceded(tag("@"), map_res(double, SampleRate::try_from))).parse(chunk)?;
+                        all_consuming(preceded(tag("@"), map_res(double, sample_rate(metric_name, config))))
+                            .parse(chunk)?;
+
                     maybe_sample_rate = Some(sample_rate);
                 }
                 // Tags: additional tags to be added to the metric.
@@ -156,6 +159,24 @@ fn permissive_metric_name(input: &[u8]) -> IResult<&[u8], &str> {
     .parse(input)
 }
 
+#[inline]
+fn sample_rate<'a>(
+    metric_name: &'a str, config: &'a DogstatsdCodecConfiguration,
+) -> impl Fn(f64) -> Result<SampleRate, &'static str> + 'a {
+    let minimum_sample_rate = config.minimum_sample_rate;
+
+    move |mut raw_sample_rate| {
+        if raw_sample_rate < minimum_sample_rate {
+            raw_sample_rate = minimum_sample_rate;
+            warn!(
+                "Sample rate for metric '{}' is below minimum of {}. Clamping to minimum.",
+                metric_name, minimum_sample_rate
+            );
+        }
+        SampleRate::try_from(raw_sample_rate)
+    }
+}
+
 #[inline]
 fn raw_metric_values(input: &[u8]) -> IResult<&[u8], (MetricType, &[u8])> {
     let (remaining, raw_values) = terminated(take_while1(|b| b != b'|'), tag("|")).parse(input)?;
@@ -568,6 +589,47 @@ mod tests {
         }
     }
 
+    #[test]
+    fn minimum_sample_rate() {
+        // Sample rate of 0.01 should lead to a count of 100 when handling a single value.
+        let minimum_sample_rate = SampleRate::try_from(0.01).unwrap();
+        let config = DogstatsdCodecConfiguration::default().with_minimum_sample_rate(minimum_sample_rate.rate());
+
+        let cases = [
+            // Worst case scenario: sample rate of zero, or "infinitely sampled".
+            "test:1|d|@0".to_string(),
+            // Bunch of values with different sample rates all below the minimum sample rate.
+            "test:1|d|@0.001".to_string(),
+            "test:1|d|@0.0005".to_string(),
+            "test:1|d|@0.00001".to_string(),
+            // Control: use the minimum sample rate.
+            format!("test:1|d|@{}", minimum_sample_rate.rate()),
+            // Bunch of values with _greater_ sampling rates than the minimum.
+            "test:1|d|@0.1".to_string(),
+            "test:1|d|@0.5".to_string(),
+            "test:1|d".to_string(),
+        ];
+
+        for input in cases {
+            let metric = parse_dsd_metric_with_conf(input.as_bytes(), &config)
+                .expect("Should not fail to parse metric.")
+                .expect("Metric should be present.");
+
+            let sketch = match metric.values() {
+                MetricValues::Distribution(points) => {
+                    points
+                        .into_iter()
+                        .next()
+                        .expect("Should have at least one sketch point.")
+                        .1
+                }
+                _ => panic!("Unexpected metric type."),
+            };
+
+            assert!(sketch.count() as u64 <= minimum_sample_rate.weight());
+        }
+    }
+
     proptest! {
         #![proptest_config(ProptestConfig::with_cases(1000))]
         #[test]

lib/saluki-io/src/deser/codec/dogstatsd/mod.rs

@@ -14,6 +14,16 @@ pub use self::service_check::ServiceCheckPacket;
 
 type NomParserError<'a> = nom::Err<nom::error::Error<&'a [u8]>>;
 
+// This is the lowest sample rate that we consider to be "safe" with typical DogStatsD default settings.
+//
+// Our logic here is:
+// - DogStatsD payloads are limited to 8KiB by default
+// - a valid distribution metric could have a multi-value payload with ~4093 values (value of `1`, when factoring for protocol overhead)
+// - to avoid overflow in resulting sketch, total count of all values must be less than or equal to 2^32
+// - 2^32 / 4093 = 1,049,344... or 1,000,000 to be safe
+// - 1 / 1,000,000 = 0.000001
+const MINIMUM_SAFE_DEFAULT_SAMPLE_RATE: f64 = 0.000001;
+
 /// Parser error.
 #[derive(Debug)]
 pub struct ParseError {
@@ -64,6 +74,7 @@ pub struct DogstatsdCodecConfiguration {
     maximum_tag_length: usize,
     maximum_tag_count: usize,
     timestamps: bool,
+    minimum_sample_rate: f64,
 }
 
 impl DogstatsdCodecConfiguration {
@@ -115,6 +126,17 @@ impl DogstatsdCodecConfiguration {
         self.timestamps = timestamps;
         self
     }
+
+    /// Sets the minimum sample rate.
+    ///
+    /// This is the minimum sample rate that is allowed for a metric payload. If the sample rate is less than this limit,
+    /// the sample rate is clamped to this value and a log message is emitted.
+    ///
+    /// Defaults to `0.000001`.
+    pub fn with_minimum_sample_rate(mut self, minimum_sample_rate: f64) -> Self {
+        self.minimum_sample_rate = minimum_sample_rate;
+        self
+    }
 }
 
 impl Default for DogstatsdCodecConfiguration {
@@ -124,6 +146,7 @@ impl Default for DogstatsdCodecConfiguration {
             maximum_tag_count: usize::MAX,
             timestamps: true,
             permissive: false,
+            minimum_sample_rate: MINIMUM_SAFE_DEFAULT_SAMPLE_RATE,
         }
     }
 }