enhancement(correctness): add support for checking APM stats in "trace" analysis mode (#1075)

Author: tobz

Cargo.lock

@@ -204,6 +204,7 @@ tracing-appender = { version = "0.2", default-features = false }
 base64 = { version = "0.22.1", default-features = false }
 treediff = { version = "5", default-features = false }
 argh = { version = "0.1", default-features = false }
+rmp-serde = { version = "1.3", default-features = false }
 
 [patch.crates-io]
 # Forked version of `hyper-http-proxy` that removes an unused dependency on `rustls-native-certs`, which transitively depends

Cargo.toml

@@ -32,7 +32,7 @@ ifeq ($(CI),true)
 endif
 export CARGO_TOOL_VERSION_cargo-binstall ?= 1.16.2
 export CARGO_TOOL_VERSION_dd-rust-license-tool ?= 1.0.3
-export CARGO_TOOL_VERSION_cargo-deny ?= 0.18.3
+export CARGO_TOOL_VERSION_cargo-deny ?= 0.18.9
 export CARGO_TOOL_VERSION_cargo-hack ?= 0.6.30
 export CARGO_TOOL_VERSION_cargo-nextest ?= 0.9.99
 export CARGO_TOOL_VERSION_cargo-autoinherit ?= 0.1.5

Makefile

@@ -9,9 +9,10 @@ repository = { workspace = true }
 workspace = true
 
 [dependencies]
-axum = { workspace = true, features = ["http1", "json", "tokio"] }
+axum = { workspace = true, features = ["http1", "json", "tokio", "tracing"] }
 datadog-protos = { workspace = true }
 protobuf = { workspace = true }
+rmp-serde = { workspace = true }
 saluki-error = { workspace = true }
 stele = { workspace = true }
 tokio = { workspace = true, features = [
@@ -23,6 +24,7 @@ tokio = { workspace = true, features = [
 tower-http = { workspace = true, features = [
   "compression-zstd",
   "decompression-deflate",
+  "decompression-gzip",
   "decompression-zstd",
 ] }
 tracing = { workspace = true }

bin/correctness/datadog-intake/Cargo.toml

@@ -2,16 +2,17 @@ use axum::{body::Bytes, extract::State, http::StatusCode, Json};
 use datadog_protos::metrics::{MetricPayload, SketchPayload};
 use protobuf::Message as _;
 use stele::Metric;
-use tracing::{debug, error};
+use tracing::{error, info};
 
 use super::MetricsState;
 
 pub async fn handle_metrics_dump(State(state): State<MetricsState>) -> Json<Vec<Metric>> {
+    info!("Got request to dump metrics.");
     Json(state.dump_metrics())
 }
 
 pub async fn handle_series_v2(State(state): State<MetricsState>, body: Bytes) -> StatusCode {
-    debug!("Received series payload.");
+    info!("Received series payload.");
 
     let payload = match MetricPayload::parse_from_bytes(&body[..]) {
         Ok(payload) => payload,
@@ -23,7 +24,7 @@ pub async fn handle_series_v2(State(state): State<MetricsState>, body: Bytes) ->
 
     match state.merge_series_payload(payload) {
         Ok(()) => {
-            debug!("Processed series payload.");
+            info!("Processed series payload.");
             StatusCode::ACCEPTED
         }
         Err(e) => {
@@ -34,7 +35,7 @@ pub async fn handle_series_v2(State(state): State<MetricsState>, body: Bytes) ->
 }
 
 pub async fn handle_sketch_beta(State(state): State<MetricsState>, body: Bytes) -> StatusCode {
-    debug!("Received sketch payload.");
+    info!("Received sketch payload.");
 
     let payload = match SketchPayload::parse_from_bytes(&body[..]) {
         Ok(payload) => payload,
@@ -46,7 +47,7 @@ pub async fn handle_sketch_beta(State(state): State<MetricsState>, body: Bytes)
 
     match state.merge_sketch_payload(payload) {
         Ok(()) => {
-            debug!("Processed sketch payload.");
+            info!("Processed sketch payload.");
             StatusCode::ACCEPTED
         }
         Err(e) => {

bin/correctness/datadog-intake/src/app/metrics/handlers.rs

@@ -17,7 +17,7 @@ pub fn initialize_app_router() -> Router {
         .merge(misc::build_misc_router())
         .fallback(debug_fallback_handler)
         // Ensure we can handle compressed requests.
-        .route_layer(RequestDecompressionLayer::new().deflate(true).zstd(true))
+        .route_layer(RequestDecompressionLayer::new().deflate(true).gzip(true).zstd(true))
         .route_layer(CompressionLayer::new().zstd(true))
         // Decompressed metrics payloads can be large (~62MB for sketches).
         .route_layer(DefaultBodyLimit::max(64 * 1024 * 1024))

bin/correctness/datadog-intake/src/app/mod.rs

@@ -1,17 +1,23 @@
 use axum::{body::Bytes, extract::State, http::StatusCode, Json};
-use datadog_protos::traces::AgentPayload;
+use datadog_protos::traces::{AgentPayload, StatsPayload};
 use protobuf::Message as _;
-use stele::Span;
-use tracing::{debug, error};
+use stele::{ClientStatisticsAggregator, Span};
+use tracing::{error, info};
 
 use super::TracesState;
 
-pub async fn handle_traces_dump(State(state): State<TracesState>) -> Json<Vec<Span>> {
-    Json(state.dump())
+pub async fn handle_trace_spans_dump(State(state): State<TracesState>) -> Json<Vec<Span>> {
+    info!("Got request to dump trace spans.");
+    Json(state.dump_spans())
+}
+
+pub async fn handle_trace_stats_dump(State(state): State<TracesState>) -> Json<ClientStatisticsAggregator> {
+    info!("Got request to dump trace stats.");
+    Json(state.dump_stats())
 }
 
 pub async fn handle_v02_traces(State(state): State<TracesState>, body: Bytes) -> StatusCode {
-    debug!("Received v0.2 traces payload.");
+    info!("Received v0.2 traces payload.");
 
     let payload = match AgentPayload::parse_from_bytes(&body[..]) {
         Ok(payload) => payload,
@@ -23,7 +29,7 @@ pub async fn handle_v02_traces(State(state): State<TracesState>, body: Bytes) ->
 
     match state.merge_agent_payload(payload) {
         Ok(()) => {
-            debug!("Processed trace payload.");
+            info!("Processed trace payload.");
             StatusCode::ACCEPTED
         }
         Err(e) => {
@@ -32,3 +38,26 @@ pub async fn handle_v02_traces(State(state): State<TracesState>, body: Bytes) ->
         }
     }
 }
+
+pub async fn handle_v02_stats(State(state): State<TracesState>, body: Bytes) -> StatusCode {
+    info!("Received v0.2 stats payload.");
+
+    let payload = match rmp_serde::from_slice::<StatsPayload>(&body[..]) {
+        Ok(payload) => payload,
+        Err(e) => {
+            error!(error = ?e, "Failed to parse stats payload.");
+            return StatusCode::BAD_REQUEST;
+        }
+    };
+
+    match state.merge_stats_payload(payload) {
+        Ok(()) => {
+            info!("Processed stats payload.");
+            StatusCode::ACCEPTED
+        }
+        Err(e) => {
+            error!(error = %e, "Failed to merge stats payload.");
+            StatusCode::BAD_REQUEST
+        }
+    }
+}

bin/correctness/datadog-intake/src/app/traces/handlers.rs

@@ -11,7 +11,9 @@ use self::state::TracesState;
 
 pub fn build_traces_router() -> Router {
     Router::new()
-        .route("/traces/dump", get(handle_traces_dump))
+        .route("/traces/dump_spans", get(handle_trace_spans_dump))
+        .route("/traces/dump_stats", get(handle_trace_stats_dump))
         .route("/api/v0.2/traces", post(handle_v02_traces))
+        .route("/api/v0.2/stats", post(handle_v02_stats))
         .with_state(TracesState::new())
 }

bin/correctness/datadog-intake/src/app/traces/mod.rs

@@ -1,34 +1,52 @@
 use std::sync::{Arc, Mutex};
 
-use datadog_protos::traces::AgentPayload;
+use datadog_protos::traces::{AgentPayload, StatsPayload};
 use saluki_error::GenericError;
-use stele::Span;
+use stele::{ClientStatisticsAggregator, Span};
+
+#[derive(Default)]
+struct Inner {
+    spans: Vec<Span>,
+    stats: ClientStatisticsAggregator,
+}
 
 #[derive(Clone)]
 pub struct TracesState {
-    spans: Arc<Mutex<Vec<Span>>>,
+    inner: Arc<Mutex<Inner>>,
 }
 
 impl TracesState {
     /// Creates a new `TracesState`.
     pub fn new() -> Self {
         Self {
-            spans: Arc::new(Mutex::new(Vec::new())),
+            inner: Arc::new(Mutex::new(Inner::default())),
         }
     }
 
-    /// Dumps the current traces state.
-    pub fn dump(&self) -> Vec<Span> {
-        let data = self.spans.lock().unwrap();
-        data.clone()
+    /// Dumps the current set of spans.
+    pub fn dump_spans(&self) -> Vec<Span> {
+        let inner = self.inner.lock().unwrap();
+        inner.spans.clone()
+    }
+
+    /// Dumps the current stats.
+    pub fn dump_stats(&self) -> ClientStatisticsAggregator {
+        let inner = self.inner.lock().unwrap();
+        inner.stats.clone()
     }
 
     /// Merges the given agent payload into the current traces state.
     pub fn merge_agent_payload(&self, payload: AgentPayload) -> Result<(), GenericError> {
         let new_spans = Span::get_spans_from_agent_payload(&payload);
-        let mut data = self.spans.lock().unwrap();
-        data.extend(new_spans);
+        let mut inner = self.inner.lock().unwrap();
+        inner.spans.extend(new_spans);
 
         Ok(())
     }
+
+    /// Merges the given stats payload into the current traces state.
+    pub fn merge_stats_payload(&self, payload: StatsPayload) -> Result<(), GenericError> {
+        let mut inner = self.inner.lock().unwrap();
+        inner.stats.merge_payload(&payload)
+    }
 }

bin/correctness/datadog-intake/src/app/traces/state.rs

@@ -27,8 +27,6 @@ async fn main() {
         )
         .with_ansi(true)
         .with_target(true)
-        .with_file(true)
-        .with_line_number(true)
         .init();
 
     match run().await {