Introduce DogStatsD payload generation (#432)

* Introduce DogStatsD payload generation

This commit introduces DogStatsD payload generation into lading. It's not quite
done yet but the basic arbitrary definitions are in place.

Signed-off-by: Brian L. Troutwine <brian.troutwine@datadoghq.com>

* Actually serialize

Signed-off-by: Brian L. Troutwine <brian.troutwine@datadoghq.com>

* `dog_stats_d` -> `dogstatsd`

Signed-off-by: Brian L. Troutwine <brian.troutwine@datadoghq.com>

* intro changelog

Signed-off-by: Brian L. Troutwine <brian.troutwine@datadoghq.com>

* feedback on Arbitrary use

Signed-off-by: Brian L. Troutwine <brian.troutwine@datadoghq.com>

Signed-off-by: Brian L. Troutwine <brian.troutwine@datadoghq.com>

Author: blt

CHANGELOG.md

@@ -8,6 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## Fixed
 - Fix OpenTelemetry message length calculation for some messages.
 
+## [0.12.0-rc3]
+## Changed
+- **Breaking change:** Added support for DogStatsD payload.
+
 ## [0.12.0]
 ## Changed
 - **Breaking change:** Support for Kafka generator is removed.

src/block.rs

@@ -131,6 +131,12 @@ where
             block_chunks,
             labels,
         ),
+        payload::Config::DogStatsD => construct_block_cache_inner(
+            &mut rng,
+            &payload::DogStatsD::default(),
+            block_chunks,
+            labels,
+        ),
         payload::Config::Fluent => {
             construct_block_cache_inner(&mut rng, &payload::Fluent::default(), block_chunks, labels)
         }

src/payload.rs

@@ -9,6 +9,7 @@ use serde::Deserialize;
 pub(crate) use apache_common::ApacheCommon;
 pub(crate) use ascii::Ascii;
 pub(crate) use datadog_logs::DatadogLog;
+pub(crate) use dogstatsd::DogStatsD;
 pub(crate) use fluent::Fluent;
 pub(crate) use foundationdb::FoundationDb;
 pub(crate) use json::Json;
@@ -23,6 +24,7 @@ mod apache_common;
 mod ascii;
 mod common;
 mod datadog_logs;
+mod dogstatsd;
 mod fluent;
 mod foundationdb;
 mod json;
@@ -115,6 +117,9 @@ pub enum Config {
     OpentelemetryLogs,
     /// Generates OpenTelemetry metrics
     OpentelemetryMetrics,
+    /// Generates DogStatsD
+    #[serde(rename = "dogstatsd")]
+    DogStatsD,
 }
 
 #[derive(Debug)]
@@ -136,6 +141,7 @@ pub(crate) enum Payload {
     OtelTraces(OpentelemetryTraces),
     OtelLogs(OpentelemetryLogs),
     OtelMetrics(OpentelemetryMetrics),
+    DogStatsdD(DogStatsD),
 }
 
 impl Payload {}
@@ -159,6 +165,7 @@ impl Serialize for Payload {
             Payload::OtelTraces(ser) => ser.to_bytes(rng, max_bytes, writer),
             Payload::OtelLogs(ser) => ser.to_bytes(rng, max_bytes, writer),
             Payload::OtelMetrics(ser) => ser.to_bytes(rng, max_bytes, writer),
+            Payload::DogStatsdD(ser) => ser.to_bytes(rng, max_bytes, writer),
         }
     }
 }

src/payload/datadog_logs.rs

@@ -1,32 +1,19 @@
 use std::io::Write;
 
-use arbitrary::{size_hint, Unstructured};
+use arbitrary::{size_hint, Arbitrary, Unstructured};
 use rand::Rng;
 
 use crate::payload::{common::AsciiStr, Error, Serialize};
 
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, serde::Serialize, serde::Deserialize, Arbitrary)]
 #[serde(rename_all = "lowercase")]
 enum Status {
     Notice,
     Info,
     Warning,
 }
 
-impl<'a> arbitrary::Arbitrary<'a> for Status {
-    fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
-        let choice = u.arbitrary::<u8>()?;
-        let res = match choice % 3 {
-            0 => Status::Notice,
-            1 => Status::Info,
-            2 => Status::Warning,
-            _ => unreachable!(),
-        };
-        Ok(res)
-    }
-}
-
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, serde::Serialize, serde::Deserialize, Arbitrary)]
 #[serde(rename_all = "lowercase")]
 enum Hostname {
     Alpha,
@@ -35,42 +22,15 @@ enum Hostname {
     Localhost,
 }
 
-impl<'a> arbitrary::Arbitrary<'a> for Hostname {
-    fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
-        let choice = u.arbitrary::<u8>()?;
-        let res = match choice % 4 {
-            0 => Hostname::Alpha,
-            1 => Hostname::Beta,
-            2 => Hostname::Gamma,
-            3 => Hostname::Localhost,
-            _ => unreachable!(),
-        };
-        Ok(res)
-    }
-}
-
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, serde::Serialize, serde::Deserialize, Arbitrary)]
 #[serde(rename_all = "lowercase")]
 enum Service {
     Vector,
     Lading,
     Cernan,
 }
 
-impl<'a> arbitrary::Arbitrary<'a> for Service {
-    fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
-        let choice = u.arbitrary::<u8>()?;
-        let res = match choice % 3 {
-            0 => Service::Vector,
-            1 => Service::Lading,
-            2 => Service::Cernan,
-            _ => unreachable!(),
-        };
-        Ok(res)
-    }
-}
-
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, serde::Serialize, serde::Deserialize, Arbitrary)]
 #[serde(rename_all = "lowercase")]
 enum Source {
     Bergman,
@@ -81,22 +41,6 @@ enum Source {
     Tarkovsky,
 }
 
-impl<'a> arbitrary::Arbitrary<'a> for Source {
-    fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
-        let choice = u.arbitrary::<u8>()?;
-        let res = match choice % 6 {
-            0 => Source::Bergman,
-            1 => Source::Keaton,
-            2 => Source::Kurosawa,
-            3 => Source::Lynch,
-            4 => Source::Waters,
-            5 => Source::Tarkovsky,
-            _ => unreachable!(),
-        };
-        Ok(res)
-    }
-}
-
 const TAG_OPTIONS: [&str; 4] = ["", "env:prod", "env:dev", "env:prod,version:1.1"];
 
 #[derive(Debug, serde::Serialize, serde::Deserialize)]

src/payload/dogstatsd.rs

@@ -0,0 +1,87 @@
+use std::{fmt, io::Write};
+
+use arbitrary::{Arbitrary, Unstructured};
+use rand::Rng;
+
+use crate::payload::{Error, Serialize};
+
+mod common;
+mod event;
+mod metric;
+mod service_check;
+
+// https://docs.datadoghq.com/developers/dogstatsd/datagram_shell/
+#[derive(Arbitrary)]
+enum Member {
+    Metric(metric::Metric),
+    Event(event::Event),
+    ServiceCheck(service_check::ServiceCheck),
+}
+
+impl fmt::Display for Member {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Metric(ref m) => write!(f, "{m}"),
+            Self::Event(ref e) => write!(f, "{e}"),
+            Self::ServiceCheck(ref sc) => write!(f, "{sc}"),
+        }
+    }
+}
+
+#[derive(Debug, Default, Clone, Copy)]
+#[allow(clippy::module_name_repetitions)]
+#[cfg_attr(test, derive(proptest_derive::Arbitrary))]
+pub(crate) struct DogStatsD {}
+
+impl Serialize for DogStatsD {
+    fn to_bytes<W, R>(&self, mut rng: R, max_bytes: usize, writer: &mut W) -> Result<(), Error>
+    where
+        R: Rng + Sized,
+        W: Write,
+    {
+        let mut entropy: Vec<u8> = vec![0; max_bytes];
+        rng.fill_bytes(&mut entropy);
+        let mut unstructured = Unstructured::new(&entropy);
+
+        let mut bytes_remaining = max_bytes;
+        while let Ok(member) = unstructured.arbitrary::<Member>() {
+            let encoding = format!("{member}");
+            let line_length = encoding.len() + 1; // add one for the newline
+            match bytes_remaining.checked_sub(line_length) {
+                Some(remainder) => {
+                    writeln!(writer, "{encoding}")?;
+                    bytes_remaining = remainder;
+                }
+                None => break,
+            }
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use proptest::prelude::*;
+    use rand::{rngs::SmallRng, SeedableRng};
+
+    use crate::payload::{DogStatsD, Serialize};
+
+    // We want to be sure that the serialized size of the payload does not
+    // exceed `max_bytes`.
+    proptest! {
+        #[test]
+        fn payload_not_exceed_max_bytes(seed: u64, max_bytes: u16) {
+            let max_bytes = max_bytes as usize;
+            let rng = SmallRng::seed_from_u64(seed);
+            let dogstatsd = DogStatsD::default();
+
+            let mut bytes = Vec::with_capacity(max_bytes);
+            dogstatsd.to_bytes(rng, max_bytes, &mut bytes).unwrap();
+            debug_assert!(
+                bytes.len() <= max_bytes,
+                "{:?}",
+                std::str::from_utf8(&bytes).unwrap()
+            );
+        }
+    }
+}