chore(datadog_traces sink): Compute APM stats (#12806)

* +code path for APM stats computation

* +aggregator skeleton

* +placeholders for full sketches

* +switch to msgpack for stats

* +minor adjustment

* +adjust encoding to exactly match intake expectation

* +fix sketch offset

* +fmt

* +adjusments & test

* +address reviews & fix ci

* +make check-event happy

* +side cases & inline doc

* +additional doc

* +additional comments

Co-authored-by: Jesse Szwedko <[email protected]>

Author: prognant

Cargo.toml

@@ -673,7 +673,7 @@ sinks-datadog_archives = ["sinks-aws_s3", "sinks-azure_blob", "sinks-gcp"]
 sinks-datadog_events = []
 sinks-datadog_logs = []
 sinks-datadog_metrics = ["protobuf-build", "sinks-azure_blob"]
-sinks-datadog_traces = ["protobuf-build"]
+sinks-datadog_traces = ["protobuf-build", "rmpv", "rmp-serde", "serde_bytes"]
 sinks-elasticsearch = ["aws-core", "aws-sigv4", "transforms-metric_to_log"]
 sinks-file = ["async-compression"]
 sinks-gcp = ["base64", "gcp", "gouth"]

build.rs

@@ -102,6 +102,7 @@ fn main() {
         println!("cargo:rerun-if-changed=proto/dd_trace.proto");
         println!("cargo:rerun-if-changed=proto/dnstap.proto");
         println!("cargo:rerun-if-changed=proto/ddsketch.proto");
+        println!("cargo:rerun-if-changed=proto/ddsketch_full.proto");
         println!("cargo:rerun-if-changed=proto/google/pubsub/v1/pubsub.proto");
         println!("cargo:rerun-if-changed=proto/vector.proto");
 
@@ -115,6 +116,7 @@ fn main() {
                     "lib/vector-core/proto/event.proto",
                     "proto/dnstap.proto",
                     "proto/ddsketch.proto",
+                    "proto/ddsketch_full.proto",
                     "proto/dd_trace.proto",
                     "proto/google/pubsub/v1/pubsub.proto",
                     "proto/vector.proto",

lib/vector-core/src/metrics/ddsketch.rs

@@ -259,6 +259,14 @@ impl AgentDDSketch {
         })
     }
 
+    pub fn gamma(&self) -> f64 {
+        self.config.gamma_v
+    }
+
+    pub fn bin_index_offset(&self) -> i32 {
+        self.config.norm_bias
+    }
+
     #[allow(dead_code)]
     fn bin_count(&self) -> usize {
         self.bins.len()

proto/ddsketch_full.proto

@@ -0,0 +1,63 @@
+// Copied from https://github.com/DataDog/sketches-go/blob/0a92170/ddsketch/pb/ddsketch.proto
+
+syntax = "proto3";
+
+package ddsketch_full;
+
+// A DDSketch is essentially a histogram that partitions the range of positive values into an infinite number of
+// indexed bins whose size grows exponentially. It keeps track of the number of values (or possibly floating-point
+// weights) added to each bin. Negative values are partitioned like positive values, symmetrically to zero.
+// The value zero as well as its close neighborhood that would be mapped to extreme bin indexes is mapped to a specific
+// counter.
+message DDSketch {
+    // The mapping between positive values and the bin indexes they belong to.
+    IndexMapping mapping = 1;
+
+    // The store for keeping track of positive values.
+    Store positiveValues = 2;
+
+    // The store for keeping track of negative values. A negative value v is mapped using its positive opposite -v.
+    Store negativeValues = 3;
+
+    // The count for the value zero and its close neighborhood (whose width depends on the mapping).
+    double zeroCount = 4;
+  }
+
+  // How to map positive values to the bins they belong to.
+  message IndexMapping {
+    // The gamma parameter of the mapping, such that bin index that a value v belongs to is roughly equal to
+    // log(v)/log(gamma).
+    double gamma = 1;
+
+    // An offset that can be used to shift all bin indexes.
+    double indexOffset = 2;
+
+    // To speed up the computation of the index a value belongs to, the computation of the log may be approximated using
+    // the fact that the log to the base 2 of powers of 2 can be computed at a low cost from the binary representation of
+    // the input value. Other values can be approximated by interpolating between successive powers of 2 (linearly,
+    // quadratically or cubically).
+    // NONE means that the log is to be computed exactly (no interpolation).
+    Interpolation interpolation = 3;
+    enum Interpolation {
+      NONE = 0;
+      LINEAR = 1;
+      QUADRATIC = 2;
+      CUBIC = 3;
+    }
+  }
+
+  // A Store maps bin indexes to their respective counts.
+  // Counts can be encoded sparsely using binCounts, but also in a contiguous way using contiguousBinCounts and
+  // contiguousBinIndexOffset. Given that non-empty bins are in practice usually contiguous or close to one another, the
+  // latter contiguous encoding method is usually more efficient than the sparse one.
+  // Both encoding methods can be used conjointly. If a bin appears in both the sparse and the contiguous encodings, its
+  // count value is the sum of the counts in each encodings.
+  message Store {
+    // The bin counts, encoded sparsely.
+    map<sint32, double> binCounts = 1;
+
+    // The bin counts, encoded contiguously. The values of contiguousBinCounts are the counts for the bins of indexes
+    // o, o+1, o+2, etc., where o is contiguousBinIndexOffset.
+    repeated double contiguousBinCounts = 2 [packed = true];
+    sint32 contiguousBinIndexOffset = 3;
+  }

src/sinks/datadog/traces/config.rs

@@ -98,7 +98,6 @@ pub enum DatadogTracesEndpoint {
 /// Store traces & APM stats endpoints actual URIs.
 pub struct DatadogTracesEndpointConfiguration {
     traces_endpoint: Uri,
-    // Unused so far
     stats_endpoint: Uri,
 }
 
@@ -186,7 +185,6 @@ impl SinkConfig for DatadogTracesConfig {
     }
 
     fn input(&self) -> Input {
-        // Metric will be accepted as soon as APM stats will be handled
         Input::trace()
     }
 

src/sinks/datadog/traces/integration_tests.rs

@@ -27,7 +27,7 @@ async fn to_real_traces_endpoint() {
     let (batch, receiver) = BatchNotifier::new_with_receiver();
 
     let trace = vec![Event::Trace(
-        simple_trace_event().with_batch_notifier(&batch),
+        simple_trace_event("a_trace".to_string()).with_batch_notifier(&batch),
     )];
 
     let stream = map_event_batch_stream(stream::iter(trace), Some(batch));

src/sinks/datadog/traces/mod.rs

@@ -15,6 +15,15 @@ mod config;
 mod request_builder;
 mod service;
 mod sink;
+mod stats;
+
+pub(crate) mod ddsketch_full {
+    include!(concat!(env!("OUT_DIR"), "/ddsketch_full.rs"));
+}
+
+pub(crate) mod dd_proto {
+    include!(concat!(env!("OUT_DIR"), "/dd_trace.rs"));
+}
 
 use crate::{config::SinkDescription, sinks::datadog::traces::config::DatadogTracesConfig};
 

src/sinks/datadog/traces/request_builder.rs

@@ -2,23 +2,21 @@ use std::{collections::BTreeMap, io::Write, sync::Arc};
 
 use bytes::Bytes;
 use prost::Message;
+use rmp_serde;
 use snafu::Snafu;
 use vector_core::event::{EventFinalizers, Finalizable};
 
 use super::{
     config::{DatadogTracesEndpoint, DatadogTracesEndpointConfiguration},
+    dd_proto,
     service::TraceApiRequest,
+    sink::PartitionKey,
+    stats,
 };
 use crate::{
     event::{Event, TraceEvent, Value},
-    sinks::{
-        datadog::traces::sink::PartitionKey,
-        util::{Compression, Compressor, IncrementalRequestBuilder},
-    },
+    sinks::util::{Compression, Compressor, IncrementalRequestBuilder},
 };
-mod dd_proto {
-    include!(concat!(env!("OUT_DIR"), "/dd_trace.rs"));
-}
 
 #[derive(Debug, Snafu)]
 pub enum RequestBuilderError {
@@ -79,6 +77,7 @@ pub struct RequestMetadata {
     endpoint: DatadogTracesEndpoint,
     finalizers: EventFinalizers,
     uncompressed_size: usize,
+    content_type: String,
 }
 
 impl IncrementalRequestBuilder<(PartitionKey, Vec<Event>)> for DatadogTracesRequestBuilder {
@@ -91,14 +90,21 @@ impl IncrementalRequestBuilder<(PartitionKey, Vec<Event>)> for DatadogTracesRequ
         &mut self,
         input: (PartitionKey, Vec<Event>),
     ) -> Vec<Result<(Self::Metadata, Self::Payload), Self::Error>> {
-        let (mut key, events) = input;
+        let (key, events) = input;
         let mut results = Vec::new();
         let n = events.len();
-
         let traces_event = events
             .into_iter()
             .filter_map(|e| e.try_into_trace())
-            .collect();
+            .collect::<Vec<TraceEvent>>();
+
+        results.push(build_apm_stats_request(
+            &key,
+            &traces_event,
+            self.compression,
+            &self.api_key,
+        ));
+
         self.trace_encoder
             .encode_trace(&key, traces_event)
             .into_iter()
@@ -108,12 +114,13 @@ impl IncrementalRequestBuilder<(PartitionKey, Vec<Event>)> for DatadogTracesRequ
                     let metadata = RequestMetadata {
                         api_key: key
                             .api_key
-                            .take()
+                            .clone()
                             .unwrap_or_else(|| Arc::clone(&self.api_key)),
                         batch_size: n,
                         endpoint: DatadogTracesEndpoint::Traces,
                         finalizers: processed.take_finalizers(),
                         uncompressed_size,
+                        content_type: "application/x-protobuf".to_string(),
                     };
                     let mut compressor = Compressor::from(self.compression);
                     match compressor.write_all(&payload) {
@@ -136,10 +143,7 @@ impl IncrementalRequestBuilder<(PartitionKey, Vec<Event>)> for DatadogTracesRequ
 
     fn build_request(&mut self, metadata: Self::Metadata, payload: Self::Payload) -> Self::Request {
         let mut headers = BTreeMap::<String, String>::new();
-        headers.insert(
-            "Content-Type".to_string(),
-            "application/x-protobuf".to_string(),
-        );
+        headers.insert("Content-Type".to_string(), metadata.content_type);
         headers.insert("DD-API-KEY".to_string(), metadata.api_key.to_string());
         if let Some(ce) = self.compression.content_encoding() {
             headers.insert("Content-Encoding".to_string(), ce.to_string());
@@ -391,3 +395,39 @@ impl DatadogTracesEncoder {
         }
     }
 }
+
+fn build_apm_stats_request(
+    key: &PartitionKey,
+    events: &[TraceEvent],
+    compression: Compression,
+    default_api_key: &Arc<str>,
+) -> Result<(RequestMetadata, Bytes), RequestBuilderError> {
+    let payload = stats::compute_apm_stats(key, events);
+    let encoded_payload =
+        rmp_serde::to_vec_named(&payload).map_err(|e| RequestBuilderError::FailedToEncode {
+            message: "APM stats encoding failed.",
+            reason: e.to_string(),
+            dropped_events: 0,
+        })?;
+    let uncompressed_size = encoded_payload.len();
+    let metadata = RequestMetadata {
+        api_key: key
+            .api_key
+            .clone()
+            .unwrap_or_else(|| Arc::clone(default_api_key)),
+        batch_size: 0,
+        endpoint: DatadogTracesEndpoint::APMStats,
+        finalizers: EventFinalizers::default(),
+        uncompressed_size,
+        content_type: "application/msgpack".to_string(),
+    };
+    let mut compressor = Compressor::from(compression);
+    match compressor.write_all(&encoded_payload) {
+        Ok(()) => Ok((metadata, compressor.into_inner().freeze())),
+        Err(e) => Err(RequestBuilderError::FailedToEncode {
+            message: "APM stats payload compression failed.",
+            reason: e.to_string(),
+            dropped_events: 0,
+        }),
+    }
+}