From f5c0208bea13a70f215c7f24b9ac415f1ebc98e1 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 25 Sep 2026 12:05:41 -0400 Subject: [PATCH 01/32] Add datadog-mock-intake crate for APM payload-level integration tests In-process mock Datadog intake that accepts the APM endpoints serverless components flush to (/api/v0.2/stats, /api/v0.2/traces, /api/v0.1/pipeline_stats), decodes msgpack/protobuf payloads on arrival, and exposes typed query methods plus raw request capture for test assertions. Ported from the bottlecap fake intake and served with hyper; preserves the 2 MiB wire-body limit, gzip/zstd decoding, stats failure injection, request summaries, and JSON dumps. --- Cargo.lock | 23 +- crates/datadog-mock-intake/Cargo.toml | 30 + crates/datadog-mock-intake/src/lib.rs | 2067 +++++++++++++++++ .../tests/common/mock_server.rs | 157 -- 4 files changed, 2119 insertions(+), 158 deletions(-) create mode 100644 crates/datadog-mock-intake/Cargo.toml create mode 100644 crates/datadog-mock-intake/src/lib.rs delete mode 100644 crates/datadog-trace-agent/tests/common/mock_server.rs diff --git a/Cargo.lock b/Cargo.lock index f13b00ed..beb28362 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -606,6 +606,27 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "datadog-mock-intake" +version = "0.1.0" +dependencies = [ + "bytes", + "flate2", + "http-body-util", + "hyper", + "hyper-util", + "libdd-trace-protobuf", + "prost 0.14.3", + "reqwest", + "rmp-serde", + "serde", + "serde_json", + "tempfile", + "thiserror 1.0.69", + "tokio", + "zstd", +] + [[package]] name = "datadog-opentelemetry" version = "0.5.2" @@ -692,8 +713,8 @@ dependencies = [ "async-trait", "bytes", "datadog-fips", + "datadog-mock-intake", "duplicate", - "flate2", "http-body-util", "hyper", "hyper-http-proxy", diff --git a/crates/datadog-mock-intake/Cargo.toml b/crates/datadog-mock-intake/Cargo.toml new file mode 100644 index 00000000..db1cf336 --- /dev/null +++ b/crates/datadog-mock-intake/Cargo.toml @@ -0,0 +1,30 @@ +# Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "datadog-mock-intake" +version = "0.1.0" +license.workspace = true +edition.workspace = true + +[lib] +bench = false + +[dependencies] +bytes = "1.10.1" +flate2 = { version = "1", default-features = false, features = ["rust_backend"] } +http-body-util = "0.1" +hyper = { version = "1.6", features = ["http1", "server"] } +hyper-util = { version = "0.1", features = ["tokio"] } +libdd-trace-protobuf = "4.0.1" +prost = "0.14.1" +rmp-serde = { version = "1.3.1", default-features = false } +serde = { version = "1.0.145", features = ["derive"] } +serde_json = "1.0" +thiserror = { version = "1.0.58", default-features = false } +tokio = { version = "1", features = ["macros", "net", "rt-multi-thread", "sync", "time"] } +zstd = { version = "0.13.3", default-features = false } + +[dev-dependencies] +reqwest = { version = "0.12.23", default-features = false } +tempfile = "3.3.0" diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs new file mode 100644 index 00000000..8f477e6e --- /dev/null +++ b/crates/datadog-mock-intake/src/lib.rs @@ -0,0 +1,2067 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! In-process mock Datadog intake for APM payload-level integration tests. +//! +//! Spawns a hyper server on a local port that accepts the same APM endpoints +//! serverless components flush to, decodes msgpack / protobuf payloads on +//! arrival, and stores the decoded structs. Callers then use typed query +//! methods to assert on payload contents. +//! +//! Endpoints supported: +//! +//! - `POST /api/v0.2/stats`: msgpack, gzip-compressed, `pb::StatsPayload` +//! - `POST /api/v0.2/traces`: protobuf (optionally zstd-compressed), `pb::AgentPayload` +//! - `POST /api/v0.1/pipeline_stats`: msgpack (struct-as-map), gzip-compressed, DSM pipeline stats +//! +//! Every completed POST attempt to one of these endpoints, including decode +//! failures and failure-injection rejections, is also recorded as a raw +//! [`CapturedRequest`], queryable via [`MockIntake::requests_for_path`]. +//! Typed query methods return only accepted, successfully decoded payloads. +//! +//! `MockIntake::start()` binds `127.0.0.1:0` with defaults and is used by the +//! trace-agent integration tests in this workspace and by the APM / DSM +//! integration tests in `bottlecap`. The standalone `fake-intake` binary in +//! the `datadog-lambda-extension` repository wraps this crate with request +//! summaries, optional stats failure injection +//! (`FAKE_INTAKE_FAIL_STATS_FIRST_N`), and optional JSON dumps +//! (`FAKE_INTAKE_DUMP_DIR`) for local debugging against a live tracer. +//! +//! This crate is self-contained (no host config, logging framework, or trace +//! processing dependencies) so it stays usable from multiple repositories. +//! +//! DSM JSON dumps contain only the fields this fixture decodes +//! ([`PipelineStatsPayload`] below); serde ignores the rest of the wire +//! payload, including the `serde_bytes` latency sketches. + +use std::collections::BTreeMap; +use std::convert::Infallible; +use std::fmt::Write as _; +use std::io::Read; +use std::path::PathBuf; +use std::process; +use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use bytes::Bytes; +use http_body_util::{BodyExt, Full, Limited}; +use hyper::body::Incoming; +use hyper::header; +use hyper::{Method, Request, Response, StatusCode, service::service_fn}; +use hyper_util::rt::TokioIo; +use libdd_trace_protobuf::pb; +use prost::Message; +use serde::{Deserialize, Serialize}; +use tokio::net::TcpListener; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; + +/// Maximum accepted request body size, matching the 2 MiB default extractor +/// limit the intake was originally served with. Enforced on the wire body, +/// before decompression. +const MAX_BODY_SIZE: usize = 2_097_152; + +/// A DSM pipeline-stats payload as it lands on `/api/v0.1/pipeline_stats`. +/// Only the fields tests assert on are decoded; serde ignores the rest +/// (including the `serde_bytes` latency sketches). JSON dumps therefore +/// contain only these fields, not the full wire payload. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct PipelineStatsPayload { + #[serde(rename = "Env")] + pub env: String, + #[serde(rename = "Service")] + pub service: String, + #[serde(rename = "TracerVersion")] + pub tracer_version: String, + #[serde(rename = "Version")] + pub version: String, + #[serde(rename = "Tags")] + pub tags: Vec, + #[serde(rename = "Stats")] + pub stats: Vec, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct PipelineStatsBucket { + #[serde(rename = "Stats")] + pub stats: Vec, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct PipelineStatsPoint { + #[serde(rename = "Hash")] + pub hash: u64, + #[serde(rename = "ParentHash")] + pub parent_hash: u64, + #[serde(rename = "EdgeTags")] + pub edge_tags: Vec, +} + +/// Options controlling a mock-intake server. The default +/// (`MockIntakeOptions::default()`) disables all diagnostics. +#[derive(Clone, Debug, Default)] +pub struct MockIntakeOptions { + /// Port to bind on `127.0.0.1`. `0` lets the OS assign a free port; the + /// actual address is available via `MockIntake::base_url()`. + pub port: u16, + /// Emit one summary line per handled request to stderr. + pub request_summaries: bool, + /// Return HTTP 500 for the first N stats request attempts. Attempts are + /// still decoded so they can be summarized and dumped, but rejected + /// payloads are not captured by `stats_payloads()`. + pub fail_stats_first_n: usize, + /// Write one JSON envelope per successfully decoded request attempt + /// (including rejected stats attempts) into this directory. + pub dump_dir: Option, +} + +/// Errors from `MockIntake::start_with_options`. +#[derive(Debug, thiserror::Error)] +pub enum MockIntakeError { + #[error("mock_intake: failed to bind listener on 127.0.0.1:{port}: {source}")] + Bind { + port: u16, + #[source] + source: std::io::Error, + }, + #[error( + "mock_intake: failed to initialize dump directory {}: {source}", + .dir.display() + )] + DumpDir { + dir: PathBuf, + #[source] + source: std::io::Error, + }, +} + +/// A raw request captured by the intake: method, path, headers in arrival +/// order, and the original wire body (pre-decompression). Recorded for every +/// completed POST attempt to a supported endpoint, including attempts that +/// were rejected by failure injection or failed to decode. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CapturedRequest { + pub method: String, + pub path: String, + pub headers: Vec<(String, String)>, + pub body: Vec, +} + +/// Captured, decoded APM payloads for a single test run. +#[derive(Debug, Default)] +struct Captured { + requests: Vec, + stats: Vec, + traces: Vec, + pipeline_stats: Vec, +} + +/// Shared server state. The request handler writes to the mutex; callers read +/// via `MockIntake::stats_payloads()` / `trace_payloads()` and friends. +#[derive(Debug)] +struct SharedState { + captured: Mutex, + options: MockIntakeOptions, + /// Monotonic request identity for summaries and dumps. + request_counter: AtomicU64, + /// Stats request attempts, used for `fail_stats_first_n` injection. + stats_attempts: AtomicU64, +} + +/// The APM endpoints the intake decodes. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Endpoint { + Stats, + Traces, + PipelineStats, +} + +impl Endpoint { + fn from_path(path: &str) -> Option { + match path { + "/api/v0.2/stats" => Some(Self::Stats), + "/api/v0.2/traces" => Some(Self::Traces), + "/api/v0.1/pipeline_stats" => Some(Self::PipelineStats), + _ => None, + } + } +} + +/// A running mock-intake server. Drop shuts it down. +#[derive(Debug)] +pub struct MockIntake { + base_url: String, + state: std::sync::Arc, + shutdown_tx: Option>, + task: Option>, +} + +impl MockIntake { + /// Bind to `127.0.0.1` on an OS-assigned port and start serving. + pub async fn start() -> Self { + Self::start_with_options(MockIntakeOptions::default()) + .await + .expect("mock_intake: failed to start server on 127.0.0.1:0") + } + + /// Start serving with explicit options. Fails on bind errors or when the + /// dump directory cannot be created; nothing is served otherwise. + pub async fn start_with_options(options: MockIntakeOptions) -> Result { + if let Some(dir) = &options.dump_dir { + std::fs::create_dir_all(dir).map_err(|source| MockIntakeError::DumpDir { + dir: dir.clone(), + source, + })?; + } + + let listener = TcpListener::bind(("127.0.0.1", options.port)) + .await + .map_err(|source| MockIntakeError::Bind { + port: options.port, + source, + })?; + let addr = listener + .local_addr() + .map_err(|source| MockIntakeError::Bind { + port: options.port, + source, + })?; + let base_url = format!("http://{addr}"); + + let state = std::sync::Arc::new(SharedState { + captured: Mutex::new(Captured::default()), + options, + request_counter: AtomicU64::new(0), + stats_attempts: AtomicU64::new(0), + }); + + let (shutdown_tx, mut shutdown_rx) = oneshot::channel::<()>(); + let task_state = std::sync::Arc::clone(&state); + let task = tokio::spawn(async move { + let state = task_state; + loop { + tokio::select! { + result = listener.accept() => { + let (stream, _) = match result { + Ok(conn) => conn, + Err(e) => { + eprintln!("mock_intake: accept error: {e}"); + break; + } + }; + + let io = TokioIo::new(stream); + let state = std::sync::Arc::clone(&state); + tokio::spawn(async move { + let service = service_fn(move |req: Request| { + let state = std::sync::Arc::clone(&state); + async move { handle_request(state, req).await } + }); + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection(io, service) + .await; + }); + } + _ = &mut shutdown_rx => { + break; + } + } + } + }); + + Ok(Self { + base_url, + state, + shutdown_tx: Some(shutdown_tx), + task: Some(task), + }) + } + + /// Full URL for the stats endpoint. + #[must_use] + pub fn stats_url(&self) -> String { + format!("{}/api/v0.2/stats", self.base_url) + } + + /// Full URL for the traces endpoint. + #[must_use] + pub fn traces_url(&self) -> String { + format!("{}/api/v0.2/traces", self.base_url) + } + + /// Full URL for the DSM pipeline-stats endpoint. + #[must_use] + pub fn pipeline_stats_url(&self) -> String { + format!("{}/api/v0.1/pipeline_stats", self.base_url) + } + + /// Base URL (scheme + host + port, no path). Use as the `apm_dd_url` for + /// components that build their own endpoint path (e.g. `DsmProcessor`). + #[must_use] + pub fn base_url(&self) -> String { + self.base_url.clone() + } + + /// Raw requests captured for `path`, in arrival order. Includes attempts + /// that were rejected by failure injection or failed to decode; excludes + /// requests rejected at the transport level (unknown path, unsupported + /// method, or oversized body). + #[must_use] + pub fn requests_for_path(&self, path: &str) -> Vec { + self.state + .captured + .lock() + .expect("mock_intake: requests mutex poisoned") + .requests + .iter() + .filter(|req| req.path == path) + .cloned() + .collect() + } + + /// All DSM pipeline-stats payloads captured so far, in arrival order. + #[must_use] + pub fn pipeline_stats_payloads(&self) -> Vec { + self.state + .captured + .lock() + .expect("mock_intake: pipeline_stats mutex poisoned") + .pipeline_stats + .clone() + } + + /// All `StatsPayload`s captured so far, in arrival order. Rejected stats + /// attempts (failure injection) are excluded. + #[must_use] + pub fn stats_payloads(&self) -> Vec { + self.state + .captured + .lock() + .expect("mock_intake: stats mutex poisoned") + .stats + .clone() + } + + /// All `AgentPayload`s captured so far, in arrival order. + #[must_use] + pub fn trace_payloads(&self) -> Vec { + self.state + .captured + .lock() + .expect("mock_intake: traces mutex poisoned") + .traces + .clone() + } +} + +impl Drop for MockIntake { + fn drop(&mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + if let Some(task) = self.task.take() { + task.abort(); + } + } +} + +/// Build an empty-bodied response with the given status. +fn response(status: StatusCode) -> Response> { + Response::builder() + .status(status) + .body(Full::new(Bytes::new())) + .expect("mock_intake: static response construction cannot fail") +} + +/// Handle one incoming request: route it, enforce the body size limit, then +/// capture and dispatch it to the endpoint-specific handler. +async fn handle_request( + state: std::sync::Arc, + req: Request, +) -> Result>, Infallible> { + let path = req.uri().path().to_string(); + let Some(endpoint) = Endpoint::from_path(&path) else { + return Ok(response(StatusCode::NOT_FOUND)); + }; + if req.method() != Method::POST { + let mut resp = response(StatusCode::METHOD_NOT_ALLOWED); + resp.headers_mut() + .insert(header::ALLOW, header::HeaderValue::from_static("POST")); + return Ok(resp); + } + + let headers = req.headers().clone(); + let body = match Limited::new(req.into_body(), MAX_BODY_SIZE).collect().await { + Ok(collected) => collected.to_bytes(), + Err(err) => { + if err + .downcast_ref::() + .is_some() + { + return Ok(response(StatusCode::PAYLOAD_TOO_LARGE)); + } + eprintln!("mock_intake: failed to read request body: {err}"); + return Ok(response(StatusCode::BAD_REQUEST)); + } + }; + + let captured = CapturedRequest { + method: Method::POST.as_str().to_string(), + path: path.clone(), + headers: headers + .iter() + .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("").to_string())) + .collect(), + body: body.to_vec(), + }; + state + .captured + .lock() + .expect("mock_intake: requests mutex poisoned") + .requests + .push(captured); + + let status = match endpoint { + Endpoint::Stats => handle_stats(&state, &headers, &body), + Endpoint::Traces => handle_traces(&state, &headers, &body), + Endpoint::PipelineStats => handle_pipeline_stats(&state, &headers, &body), + }; + Ok(response(status)) +} + +/// Result of handling one intake request, before summary and dump emission. +struct HandledRequest { + request_id: u64, + status: StatusCode, + decoded: Option, +} + +fn handle_stats( + state: &std::sync::Arc, + headers: &hyper::HeaderMap, + body: &Bytes, +) -> StatusCode { + let request_id = next_request_id(state); + let attempt = state.stats_attempts.fetch_add(1, Ordering::SeqCst) + 1; + let inject_failure = attempt <= state.options.fail_stats_first_n as u64; + + let handled: HandledRequest = match decompress(headers, body) { + Ok(d) => match rmp_serde::from_slice::(&d) { + Ok(payload) => { + let status = if inject_failure { + StatusCode::INTERNAL_SERVER_ERROR + } else { + state + .captured + .lock() + .expect("mock_intake: stats mutex poisoned") + .stats + .push(payload.clone()); + StatusCode::ACCEPTED + }; + HandledRequest { + request_id, + status, + decoded: Some(payload), + } + } + Err(err) => { + eprintln!("mock_intake: failed to decode StatsPayload msgpack: {err}"); + HandledRequest { + request_id, + status: failure_status(inject_failure), + decoded: None, + } + } + }, + Err(e) => { + eprintln!("{e}"); + HandledRequest { + request_id, + status: failure_status(inject_failure), + decoded: None, + } + } + }; + + let payload_count = handled.decoded.as_ref().map_or(0, |p| p.stats.len()); + + if state.options.request_summaries { + let groups = handled.decoded.as_ref().map(stats_hits_by_key); + log_summary( + handled.request_id, + "/api/v0.2/stats", + headers, + handled.status, + payload_count, + groups.as_ref(), + ); + } + + if let Some(payload) = &handled.decoded { + dump_request( + state, + handled.request_id, + "/api/v0.2/stats", + headers, + handled.status, + serde_json::to_value(payload).unwrap_or_else(|_| serde_json::Value::Null), + ); + } + + handled.status +} + +fn handle_traces( + state: &std::sync::Arc, + headers: &hyper::HeaderMap, + body: &Bytes, +) -> StatusCode { + let request_id = next_request_id(state); + let handled: HandledRequest = match decompress(headers, body) { + Ok(d) => match pb::AgentPayload::decode(d.as_slice()) { + Ok(payload) => { + state + .captured + .lock() + .expect("mock_intake: traces mutex poisoned") + .traces + .push(payload.clone()); + HandledRequest { + request_id, + status: StatusCode::ACCEPTED, + decoded: Some(payload), + } + } + Err(err) => { + eprintln!("mock_intake: failed to decode AgentPayload protobuf: {err}"); + HandledRequest { + request_id, + status: StatusCode::BAD_REQUEST, + decoded: None, + } + } + }, + Err(e) => { + eprintln!("{e}"); + HandledRequest { + request_id, + status: StatusCode::BAD_REQUEST, + decoded: None, + } + } + }; + + // Tracer payload count across the ordinary and indexed collections. + let payload_count = handled + .decoded + .as_ref() + .map_or(0, |p| p.tracer_payloads.len() + p.idx_tracer_payloads.len()); + + if state.options.request_summaries { + log_summary( + handled.request_id, + "/api/v0.2/traces", + headers, + handled.status, + payload_count, + None, + ); + } + + if let Some(payload) = &handled.decoded { + dump_request( + state, + handled.request_id, + "/api/v0.2/traces", + headers, + handled.status, + agent_payload_to_json(payload), + ); + } + + handled.status +} + +fn handle_pipeline_stats( + state: &std::sync::Arc, + headers: &hyper::HeaderMap, + body: &Bytes, +) -> StatusCode { + let request_id = next_request_id(state); + let handled: HandledRequest = match decompress(headers, body) { + Ok(d) => match rmp_serde::from_slice::(&d) { + Ok(payload) => { + state + .captured + .lock() + .expect("mock_intake: pipeline_stats mutex poisoned") + .pipeline_stats + .push(payload.clone()); + HandledRequest { + request_id, + status: StatusCode::ACCEPTED, + decoded: Some(payload), + } + } + Err(err) => { + eprintln!("mock_intake: failed to decode pipeline stats msgpack: {err}"); + HandledRequest { + request_id, + status: StatusCode::BAD_REQUEST, + decoded: None, + } + } + }, + Err(e) => { + eprintln!("{e}"); + HandledRequest { + request_id, + status: StatusCode::BAD_REQUEST, + decoded: None, + } + } + }; + + if state.options.request_summaries { + log_summary( + handled.request_id, + "/api/v0.1/pipeline_stats", + headers, + handled.status, + handled.decoded.as_ref().map_or(0, |_| 1), + None, + ); + } + + if let Some(payload) = &handled.decoded { + dump_request( + state, + handled.request_id, + "/api/v0.1/pipeline_stats", + headers, + handled.status, + serde_json::to_value(payload).unwrap_or_else(|_| serde_json::Value::Null), + ); + } + + handled.status +} + +/// Status used for a stats request that failed to decode or was rejected by +/// failure injection. Decode failures outside the injection window keep the +/// historical `400 Bad Request` behavior. +fn failure_status(inject_failure: bool) -> StatusCode { + if inject_failure { + StatusCode::INTERNAL_SERVER_ERROR + } else { + StatusCode::BAD_REQUEST + } +} + +fn next_request_id(state: &SharedState) -> u64 { + state.request_counter.fetch_add(1, Ordering::SeqCst) + 1 +} + +fn content_encoding(headers: &hyper::HeaderMap) -> String { + headers + .get("content-encoding") + .and_then(|v| v.to_str().ok()) + .map_or_else(|| "identity".to_string(), ToString::to_string) +} + +/// Emit one summary line per handled request. Payload contents stay out of +/// logs; for stats, only aggregated hits per full grouping key are shown. +fn log_summary( + request_id: u64, + endpoint: &str, + headers: &hyper::HeaderMap, + status: StatusCode, + payload_count: usize, + groups: Option<&BTreeMap>, +) { + let mut line = format!( + "mock-intake: request {request_id} | POST {endpoint} | encoding={} | status={} | payloads={payload_count}", + content_encoding(headers), + status.as_u16(), + ); + if let Some(groups) = groups { + let rendered: Vec = groups + .iter() + .map(|(key, hits)| format!("hits={hits} {}", key.render())) + .collect(); + let _ = write!(line, " | groups={} | {}", groups.len(), rendered.join("; ")); + } + eprintln!("{line}"); +} + +/// Serialize a JSON envelope for one decoded request attempt and write it to +/// the dump directory. Diagnostic only: failures are reported but never +/// change the intake response. No capture lock is held during serialization +/// or filesystem I/O. +fn dump_request( + state: &SharedState, + request_id: u64, + endpoint: &str, + headers: &hyper::HeaderMap, + status: StatusCode, + payload: serde_json::Value, +) { + let Some(dump_dir) = &state.options.dump_dir else { + return; + }; + + let envelope = serde_json::json!({ + "request_id": request_id, + "endpoint": endpoint, + "encoding": content_encoding(headers), + "status": status.as_u16(), + "payload": payload, + }); + let Ok(mut json) = serde_json::to_vec_pretty(&envelope) else { + eprintln!("mock_intake: failed to serialize dump envelope for request {request_id}"); + return; + }; + json.push(b'\n'); + + // Collision-resistant filename: slug + request id + pid + nanos. A final + // uniqueness check with create_new guarantees no earlier dump is ever + // overwritten, even across process restarts. + let slug = endpoint.replace('/', "_"); + let base = format!( + "{slug}-{request_id:04}-{}-{}", + process::id(), + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |d| d.as_nanos()) + ); + let mut suffix = 0; + loop { + let path = if suffix == 0 { + dump_dir.join(format!("{base}.json")) + } else { + dump_dir.join(format!("{base}-{suffix}.json")) + }; + match std::fs::OpenOptions::new() + .create_new(true) + .write(true) + .open(&path) + { + Ok(mut file) => { + if let Err(e) = std::io::Write::write_all(&mut file, &json) { + eprintln!("mock_intake: failed to write dump {}: {e}", path.display()); + } + return; + } + Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => { + suffix += 1; + if suffix > 1000 { + eprintln!( + "mock_intake: giving up on dump filename for request {request_id} after 1000 collisions" + ); + return; + } + } + Err(e) => { + eprintln!("mock_intake: failed to create dump {}: {e}", path.display()); + return; + } + } + } +} + +/// Manual JSON serialization for `pb::AgentPayload`, which does not implement +/// `Serialize` in `libdd-trace-protobuf` 4.0.1. Its tracer-payload collections +/// do, so the fields are embedded directly. Fields mirror the protobuf tags. +fn agent_payload_to_json(payload: &pb::AgentPayload) -> serde_json::Value { + serde_json::json!({ + "host_name": payload.host_name, + "env": payload.env, + "tracer_payloads": payload.tracer_payloads, + "tags": payload.tags, + "agent_version": payload.agent_version, + "target_tps": payload.target_tps, + "error_tps": payload.error_tps, + "rare_sampler_enabled": payload.rare_sampler_enabled, + "idx_tracer_payloads": payload.idx_tracer_payloads, + }) +} + +// --------------------------------------------------------------------------- +// Stats grouping +// --------------------------------------------------------------------------- + +/// Full aggregation-dimension key for one `ClientGroupedStats` entry. +/// +/// Includes every dimension stats aggregation distinguishes: client identity +/// and origin context plus the grouped-stat dimensions. Excludes +/// measurements and delivery metadata (hits, errors, duration, sketches, +/// top-level hits, runtime ID, sequence, tracer version/language) and time +/// bucket fields, so identical dimensions combine across time windows. +/// +/// Tag lists are normalized (sorted) here for deterministic summaries; +/// captured payloads are never modified. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +struct StatsGroupKey { + // Client dimensions (from `ClientStatsPayload`). + client_service: String, + client_env: String, + client_version: String, + client_hostname: String, + client_container_id: String, + client_tags: Vec, + client_git_commit_sha: String, + client_image_tag: String, + client_process_tags: String, + client_process_tags_hash: u64, + client_agent_aggregation: String, + + // Grouped dimensions (from `ClientGroupedStats`). + stats_service: String, + stats_name: String, + stats_resource: String, + stats_type: String, + stats_db_type: String, + stats_http_status_code: u32, + stats_grpc_status_code: String, + stats_http_method: String, + stats_http_endpoint: String, + stats_span_kind: String, + stats_synthetics: bool, + stats_is_trace_root: i32, + stats_service_source: String, + stats_peer_tags: Vec, + stats_span_derived_primary_tags: Vec, + stats_additional_metric_tags: Vec, +} + +impl StatsGroupKey { + fn new(client: &pb::ClientStatsPayload, grouped: &pb::ClientGroupedStats) -> Self { + let sorted = |mut tags: Vec| { + tags.sort(); + tags + }; + Self { + client_service: client.service.clone(), + client_env: client.env.clone(), + client_version: client.version.clone(), + client_hostname: client.hostname.clone(), + client_container_id: client.container_id.clone(), + client_tags: sorted(client.tags.clone()), + client_git_commit_sha: client.git_commit_sha.clone(), + client_image_tag: client.image_tag.clone(), + client_process_tags: client.process_tags.clone(), + client_process_tags_hash: client.process_tags_hash, + client_agent_aggregation: client.agent_aggregation.clone(), + + stats_service: grouped.service.clone(), + stats_name: grouped.name.clone(), + stats_resource: grouped.resource.clone(), + stats_type: grouped.r#type.clone(), + stats_db_type: grouped.db_type.clone(), + stats_http_status_code: grouped.http_status_code, + stats_grpc_status_code: grouped.grpc_status_code.clone(), + stats_http_method: grouped.http_method.clone(), + stats_http_endpoint: grouped.http_endpoint.clone(), + stats_span_kind: grouped.span_kind.clone(), + stats_synthetics: grouped.synthetics, + stats_is_trace_root: grouped.is_trace_root, + stats_service_source: grouped.service_source.clone(), + stats_peer_tags: sorted(grouped.peer_tags.clone()), + stats_span_derived_primary_tags: sorted(grouped.span_derived_primary_tags.clone()), + stats_additional_metric_tags: sorted(grouped.additional_metric_tags.clone()), + } + } + + /// Stable, escaped one-line representation in a fixed field order. + fn render(&self) -> String { + let mut fields: Vec = Vec::new(); + fields.push(format!("client.service={}", escape(&self.client_service))); + fields.push(format!("client.env={}", escape(&self.client_env))); + fields.push(format!("client.version={}", escape(&self.client_version))); + fields.push(format!("client.hostname={}", escape(&self.client_hostname))); + fields.push(format!( + "client.container_id={}", + escape(&self.client_container_id) + )); + fields.push(format!("client.tags={}", render_tags(&self.client_tags))); + fields.push(format!( + "client.git_commit_sha={}", + escape(&self.client_git_commit_sha) + )); + fields.push(format!( + "client.image_tag={}", + escape(&self.client_image_tag) + )); + fields.push(format!( + "client.process_tags={}", + escape(&self.client_process_tags) + )); + fields.push(format!( + "client.process_tags_hash={}", + self.client_process_tags_hash + )); + fields.push(format!( + "client.agent_aggregation={}", + escape(&self.client_agent_aggregation) + )); + + fields.push(format!("stats.service={}", escape(&self.stats_service))); + fields.push(format!("stats.name={}", escape(&self.stats_name))); + fields.push(format!("stats.resource={}", escape(&self.stats_resource))); + fields.push(format!("stats.type={}", escape(&self.stats_type))); + fields.push(format!("stats.db_type={}", escape(&self.stats_db_type))); + fields.push(format!( + "stats.http_status_code={}", + self.stats_http_status_code + )); + fields.push(format!( + "stats.grpc_status_code={}", + escape(&self.stats_grpc_status_code) + )); + fields.push(format!( + "stats.http_method={}", + escape(&self.stats_http_method) + )); + fields.push(format!( + "stats.http_endpoint={}", + escape(&self.stats_http_endpoint) + )); + fields.push(format!("stats.span_kind={}", escape(&self.stats_span_kind))); + fields.push(format!("stats.synthetics={}", self.stats_synthetics)); + fields.push(format!( + "stats.is_trace_root={}", + escape(&trilean_name(self.stats_is_trace_root)) + )); + fields.push(format!( + "stats.service_source={}", + escape(&self.stats_service_source) + )); + fields.push(format!( + "stats.peer_tags={}", + render_tags(&self.stats_peer_tags) + )); + fields.push(format!( + "stats.span_derived_primary_tags={}", + render_tags(&self.stats_span_derived_primary_tags) + )); + fields.push(format!( + "stats.additional_metric_tags={}", + render_tags(&self.stats_additional_metric_tags) + )); + fields.join(" ") + } +} + +/// JSON-escape a value for stable output in summaries. +fn escape(value: &str) -> String { + serde_json::to_string(value).unwrap_or_else(|_| "\"\"".to_string()) +} + +/// Render a sorted tag list as `["a:1","b:2"]` for summaries. +fn render_tags(tags: &[String]) -> String { + let rendered: Vec = tags.iter().map(|t| escape(t)).collect(); + format!("[{}]", rendered.join(",")) +} + +fn trilean_name(value: i32) -> String { + match pb::Trilean::try_from(value) { + Ok(t) => t.as_str_name().to_string(), + Err(_) => value.to_string(), + } +} + +/// Sum `hits` across client payloads and time buckets, grouped by the full +/// aggregation key. Returned as a `BTreeMap` for deterministic ordering. +fn stats_hits_by_key(payload: &pb::StatsPayload) -> BTreeMap { + let mut groups: BTreeMap = BTreeMap::new(); + for client in &payload.stats { + for bucket in &client.stats { + for grouped in &bucket.stats { + *groups + .entry(StatsGroupKey::new(client, grouped)) + .or_insert(0) += grouped.hits; + } + } + } + groups +} + +// --------------------------------------------------------------------------- +// Body decoding +// --------------------------------------------------------------------------- + +/// Decompress a request body based on its `Content-Encoding` header. +/// Supports `gzip` and `zstd`. An unknown or absent encoding is treated as +/// identity: the body is returned unchanged. +fn decompress(headers: &hyper::HeaderMap, body: &Bytes) -> Result, String> { + let encoding = headers + .get("content-encoding") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_ascii_lowercase(); + + match encoding.as_str() { + "gzip" => { + let mut decoder = flate2::read::GzDecoder::new(body.as_ref()); + let mut out = Vec::new(); + decoder + .read_to_end(&mut out) + .map_err(|e| format!("mock_intake: gzip decode failed: {e}"))?; + Ok(out) + } + "zstd" => zstd::stream::decode_all(body.as_ref()) + .map_err(|e| format!("mock_intake: zstd decode failed: {e}")), + _ => { + if !encoding.is_empty() { + eprintln!( + "mock_intake: unrecognized Content-Encoding '{encoding}', treating as identity" + ); + } + Ok(body.to_vec()) + } + } +} + +#[cfg(test)] +mod tests { + + use super::*; + + /// POST a raw body to a path on the intake with an optional + /// Content-Encoding header. Returns the response status code. + async fn post(base_url: &str, path: &str, encoding: Option<&str>, body: Vec) -> StatusCode { + let client = reqwest::Client::new(); + let url = format!("{base_url}{path}"); + let mut request = client.post(&url); + if let Some(enc) = encoding { + request = request.header("content-encoding", enc); + } + request + .header("content-type", "application/msgpack") + .body(body) + .send() + .await + .expect("test: request to mock intake failed") + .status() + } + + fn gzip(data: Vec) -> Vec { + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + std::io::Write::write_all(&mut encoder, &data).expect("test: gzip compression failed"); + encoder.finish().expect("test: gzip compression failed") + } + + /// Build a `pb::StatsPayload` with one client payload, one bucket, and the + /// given grouped stats entries. + fn stats_payload( + client: pb::ClientStatsPayload, + start: u64, + grouped: Vec, + ) -> pb::StatsPayload { + pb::StatsPayload { + stats: vec![pb::ClientStatsPayload { + stats: vec![pb::ClientStatsBucket { + start, + duration: 10_000_000_000, + stats: grouped, + agent_time_shift: 0, + }], + ..client + }], + ..pb::StatsPayload::default() + } + } + + fn grouped_stats( + service: &str, + resource: &str, + hits: u64, + overrides: impl FnOnce(&mut pb::ClientGroupedStats), + ) -> pb::ClientGroupedStats { + let mut g = pb::ClientGroupedStats { + service: service.to_string(), + name: "smoke.request".to_string(), + resource: resource.to_string(), + r#type: "web".to_string(), + span_kind: "server".to_string(), + http_status_code: 200, + is_trace_root: pb::Trilean::True as i32, + hits, + ..pb::ClientGroupedStats::default() + }; + overrides(&mut g); + g + } + + fn client_payload(env: &str) -> pb::ClientStatsPayload { + pb::ClientStatsPayload { + service: "mock-intake-smoke".to_string(), + env: env.to_string(), + version: "smoke".to_string(), + ..pb::ClientStatsPayload::default() + } + } + + /// Serialize a stats payload whose wire size is exactly `target` bytes by + /// padding the client hostname. Above the msgpack str16 boundary the + /// encoding grows linearly with the string, so a fixed-point adjustment + /// converges; fails the test if it cannot hit the target exactly. + fn stats_payload_wire_of_size(target: usize) -> Vec { + let base = stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + ); + let base_len = rmp_serde::to_vec_named(&base) + .expect("test: msgpack encode failed") + .len(); + assert!( + target > base_len, + "test: target {} must exceed base payload size {base_len}", + target + ); + + let mut pad = target - base_len; + for _ in 0..8 { + let mut client = client_payload("local"); + client.hostname = "x".repeat(pad); + let body = rmp_serde::to_vec_named(&stats_payload( + client, + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + match body.len().cmp(&target) { + std::cmp::Ordering::Equal => return body, + std::cmp::Ordering::Greater => { + pad -= body.len() - target; + } + std::cmp::Ordering::Less => { + pad += target - body.len(); + } + } + } + panic!("test: could not build a payload of exactly {target} bytes"); + } + + async fn start_default() -> MockIntake { + MockIntake::start().await + } + + #[tokio::test] + async fn default_startup_binds_port_zero_and_exposes_url_helpers() { + let intake = start_default().await; + assert!(intake.base_url().starts_with("http://127.0.0.1:")); + assert!(intake.stats_url().ends_with("/api/v0.2/stats")); + assert!(intake.traces_url().ends_with("/api/v0.2/traces")); + assert!( + intake + .pipeline_stats_url() + .ends_with("/api/v0.1/pipeline_stats") + ); + } + + #[tokio::test] + async fn instances_are_independent() { + let a = start_default().await; + let b = start_default().await; + assert_ne!(a.base_url(), b.base_url()); + + let payload = stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + ); + let body = rmp_serde::to_vec_named(&payload).expect("test: msgpack encode failed"); + assert_eq!( + post(&a.base_url(), "/api/v0.2/stats", None, body).await, + StatusCode::ACCEPTED + ); + assert_eq!(a.stats_payloads().len(), 1); + assert!(b.stats_payloads().is_empty()); + + // Dropping `a` must not shut down `b`. + drop(a); + let client = reqwest::Client::new(); + // Any HTTP response proves `b` is still listening; an unmatched route + // answers 404. + let response = client + .get(format!("{}/", b.base_url())) + .send() + .await + .expect("test: intake b should still be listening"); + assert_eq!(response.status(), reqwest::StatusCode::NOT_FOUND); + } + + #[tokio::test] + async fn drop_shuts_down_server() { + let intake = start_default().await; + let url = intake.base_url(); + drop(intake); + // The listener socket is released on shutdown; a fresh bind on the + // same port must succeed shortly afterwards. + let port = url.rsplit(':').next().and_then(|p| p.parse::().ok()); + let port = port.expect("test: no port in base url"); + let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(5); + let mut bound = false; + while tokio::time::Instant::now() < deadline { + if TcpListener::bind(("127.0.0.1", port)).await.is_ok() { + bound = true; + break; + } + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + } + assert!(bound, "listener port was not released after Drop"); + } + + #[tokio::test] + async fn stats_endpoint_decodes_msgpack_and_gzip() { + let intake = start_default().await; + let payload = stats_payload( + client_payload("local"), + 42, + vec![grouped_stats("svc", "GET /a", 3, |_| {})], + ); + let raw = rmp_serde::to_vec_named(&payload).expect("test: msgpack encode failed"); + + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + Some("gzip"), + gzip(raw.clone()) + ) + .await, + StatusCode::ACCEPTED + ); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, raw).await, + StatusCode::ACCEPTED + ); + + let captured = intake.stats_payloads(); + assert_eq!(captured.len(), 2); + assert_eq!(captured[0].stats[0].stats[0].stats[0].hits, 3); + assert_eq!(captured[0].stats[0].stats[0].start, 42); + } + + #[tokio::test] + async fn traces_endpoint_decodes_protobuf() { + let intake = start_default().await; + let payload = pb::AgentPayload { + env: "local".to_string(), + tracer_payloads: vec![pb::TracerPayload::default()], + ..pb::AgentPayload::default() + }; + let body = payload.encode_to_vec(); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/traces", None, body).await, + StatusCode::ACCEPTED + ); + let captured = intake.trace_payloads(); + assert_eq!(captured.len(), 1); + assert_eq!(captured[0].env, "local"); + assert_eq!(captured[0].tracer_payloads.len(), 1); + } + + #[tokio::test] + async fn traces_endpoint_decodes_zstd_compressed_protobuf() { + let intake = start_default().await; + let payload = pb::AgentPayload { + env: "local".to_string(), + tracer_payloads: vec![pb::TracerPayload::default()], + ..pb::AgentPayload::default() + }; + let body = zstd::stream::encode_all(payload.encode_to_vec().as_slice(), 0) + .expect("test: zstd compression failed"); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/traces", Some("zstd"), body).await, + StatusCode::ACCEPTED + ); + let captured = intake.trace_payloads(); + assert_eq!(captured.len(), 1); + assert_eq!(captured[0].env, "local"); + } + + #[tokio::test] + async fn dsm_endpoint_decodes_msgpack_gzip() { + let intake = start_default().await; + let payload = PipelineStatsPayload { + env: "local".to_string(), + service: "svc".to_string(), + tracer_version: "1.0".to_string(), + version: "2.0".to_string(), + tags: vec!["a:b".to_string()], + stats: vec![PipelineStatsBucket { + stats: vec![PipelineStatsPoint { + hash: 7, + parent_hash: 0, + edge_tags: vec!["direction:out".to_string()], + }], + }], + }; + let body = gzip(rmp_serde::to_vec_named(&payload).expect("test: msgpack encode failed")); + assert_eq!( + post( + &intake.base_url(), + "/api/v0.1/pipeline_stats", + Some("gzip"), + body + ) + .await, + StatusCode::ACCEPTED + ); + let captured = intake.pipeline_stats_payloads(); + assert_eq!(captured.len(), 1); + assert_eq!(captured[0].stats[0].stats[0].hash, 7); + } + + #[tokio::test] + async fn malformed_bodies_return_400_and_are_not_typed_captured() { + let intake = start_default().await; + for path in [ + "/api/v0.2/stats", + "/api/v0.2/traces", + "/api/v0.1/pipeline_stats", + ] { + assert_eq!( + post(&intake.base_url(), path, None, vec![0xFF; 16]).await, + StatusCode::BAD_REQUEST, + "malformed body on {path} should return 400" + ); + } + assert!(intake.stats_payloads().is_empty()); + assert!(intake.trace_payloads().is_empty()); + assert!(intake.pipeline_stats_payloads().is_empty()); + // Raw capture records completed attempts even when decoding failed. + for path in [ + "/api/v0.2/stats", + "/api/v0.2/traces", + "/api/v0.1/pipeline_stats", + ] { + assert_eq!( + intake.requests_for_path(path).len(), + 1, + "malformed POST on {path} should still be raw-captured" + ); + } + } + + #[tokio::test] + async fn fail_stats_first_n_rejects_exactly_n_then_accepts() { + let intake = MockIntake::start_with_options(MockIntakeOptions { + fail_stats_first_n: 2, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + let payload = stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + ); + let body = rmp_serde::to_vec_named(&payload).expect("test: msgpack encode failed"); + + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body.clone()).await, + StatusCode::INTERNAL_SERVER_ERROR + ); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body.clone()).await, + StatusCode::INTERNAL_SERVER_ERROR + ); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body).await, + StatusCode::ACCEPTED + ); + + // Rejected attempts are not typed-captured; only the successful retry is. + let captured = intake.stats_payloads(); + assert_eq!(captured.len(), 1); + assert_eq!(captured[0].stats[0].stats[0].stats[0].hits, 1); + // Raw capture records every completed attempt, including rejections. + assert_eq!(intake.requests_for_path("/api/v0.2/stats").len(), 3); + } + + #[tokio::test] + async fn concurrent_stats_requests_respect_rejection_limit() { + let intake = MockIntake::start_with_options(MockIntakeOptions { + fail_stats_first_n: 3, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + let payload = stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + ); + let body = rmp_serde::to_vec_named(&payload).expect("test: msgpack encode failed"); + + let client = reqwest::Client::new(); + let url = format!("{}/api/v0.2/stats", intake.base_url()); + let mut handles = Vec::new(); + for _ in 0..10 { + let client = client.clone(); + let url = url.clone(); + let body = body.clone(); + handles.push(tokio::spawn(async move { + client + .post(&url) + .body(body) + .send() + .await + .expect("test: request failed") + .status() + .as_u16() + })); + } + let mut rejected = 0; + let mut accepted = 0; + for handle in handles { + match handle.await.expect("test: task panicked") { + 500 => rejected += 1, + 202 => accepted += 1, + other => panic!("unexpected status {other}"), + } + } + assert_eq!(rejected, 3, "exactly the first 3 attempts must be rejected"); + assert_eq!(accepted, 7); + assert_eq!(intake.stats_payloads().len(), 7); + } + + #[tokio::test] + async fn traces_and_dsm_do_not_consume_stats_rejections() { + let intake = MockIntake::start_with_options(MockIntakeOptions { + fail_stats_first_n: 1, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + let trace = pb::AgentPayload::default().encode_to_vec(); + let dsm = gzip( + rmp_serde::to_vec_named(&PipelineStatsPayload { + env: "local".to_string(), + service: "svc".to_string(), + tracer_version: "1.0".to_string(), + version: "2.0".to_string(), + tags: Vec::new(), + stats: Vec::new(), + }) + .expect("test: msgpack encode failed"), + ); + + assert_eq!( + post(&intake.base_url(), "/api/v0.2/traces", None, trace).await, + StatusCode::ACCEPTED + ); + assert_eq!( + post( + &intake.base_url(), + "/api/v0.1/pipeline_stats", + Some("gzip"), + dsm + ) + .await, + StatusCode::ACCEPTED + ); + // The stats rejection budget is untouched by the two requests above. + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + None, + rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})] + )) + .expect("test: msgpack encode failed") + ) + .await, + StatusCode::INTERNAL_SERVER_ERROR + ); + } + + #[tokio::test] + async fn unknown_path_returns_404_and_is_not_captured() { + let intake = start_default().await; + let body = rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + assert_eq!( + post(&intake.base_url(), "/api/v2/profile", None, body).await, + StatusCode::NOT_FOUND + ); + assert!(intake.stats_payloads().is_empty()); + assert!(intake.trace_payloads().is_empty()); + assert!(intake.pipeline_stats_payloads().is_empty()); + assert!(intake.requests_for_path("/api/v2/profile").is_empty()); + } + + #[tokio::test] + async fn non_post_methods_on_known_paths_return_405_and_are_not_captured() { + let intake = start_default().await; + let client = reqwest::Client::new(); + for path in [ + "/api/v0.2/stats", + "/api/v0.2/traces", + "/api/v0.1/pipeline_stats", + ] { + let response = client + .get(format!("{}{path}", intake.base_url())) + .send() + .await + .expect("test: request failed"); + assert_eq!(response.status(), reqwest::StatusCode::METHOD_NOT_ALLOWED); + assert!( + response.headers().get("allow").is_some_and(|v| v == "POST"), + "405 response must advertise Allow: POST" + ); + assert!(intake.requests_for_path(path).is_empty()); + } + assert!(intake.stats_payloads().is_empty()); + assert!(intake.trace_payloads().is_empty()); + assert!(intake.pipeline_stats_payloads().is_empty()); + } + + #[tokio::test] + async fn oversized_bodies_return_413_and_are_not_captured_or_counted() { + let intake = MockIntake::start_with_options(MockIntakeOptions { + fail_stats_first_n: 1, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + // One byte over the wire-body limit. + let oversized = vec![0u8; MAX_BODY_SIZE + 1]; + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, oversized).await, + StatusCode::PAYLOAD_TOO_LARGE + ); + assert!(intake.requests_for_path("/api/v0.2/stats").is_empty()); + assert!(intake.stats_payloads().is_empty()); + + // The rejected oversized attempt must not consume the failure-injection + // budget: the next valid stats request is rejected as attempt 1 of 1. + let body = rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body).await, + StatusCode::INTERNAL_SERVER_ERROR + ); + assert!(intake.stats_payloads().is_empty()); + } + + #[tokio::test] + async fn body_at_exact_size_limit_is_accepted() { + let intake = start_default().await; + let body = stats_payload_wire_of_size(MAX_BODY_SIZE); + assert_eq!(body.len(), MAX_BODY_SIZE); + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body).await, + StatusCode::ACCEPTED + ); + assert_eq!(intake.stats_payloads().len(), 1); + } + + #[tokio::test] + async fn grouping_combines_across_time_buckets_and_clients() { + let payload = pb::StatsPayload { + stats: vec![ + pb::ClientStatsPayload { + stats: vec![ + pb::ClientStatsBucket { + start: 100, + duration: 10, + stats: vec![ + grouped_stats("svc", "GET /a", 2, |_| {}), + grouped_stats("svc", "GET /b", 1, |_| {}), + ], + agent_time_shift: 0, + }, + pb::ClientStatsBucket { + start: 200, + duration: 10, + stats: vec![grouped_stats("svc", "GET /a", 3, |_| {})], + agent_time_shift: 0, + }, + ], + ..client_payload("prod") + }, + // Same dimensions in a second client payload combine too. + pb::ClientStatsPayload { + stats: vec![pb::ClientStatsBucket { + start: 300, + duration: 10, + stats: vec![grouped_stats("svc", "GET /a", 4, |_| {})], + agent_time_shift: 0, + }], + ..client_payload("prod") + }, + ], + ..pb::StatsPayload::default() + }; + + let groups = stats_hits_by_key(&payload); + assert_eq!(groups.len(), 2, "matching keys must combine across buckets"); + let hits: Vec = groups.values().copied().collect(); + assert_eq!(hits, vec![2 + 3 + 4, 1]); + } + + type PayloadMutation = Box; + + #[test] + #[allow(clippy::too_many_lines)] + fn grouping_separates_changed_dimensions() { + let base = || pb::StatsPayload { + stats: vec![pb::ClientStatsPayload { + stats: vec![pb::ClientStatsBucket { + start: 1, + duration: 1, + stats: vec![grouped_stats("svc", "GET /a", 1, |_| {})], + agent_time_shift: 0, + }], + ..client_payload("prod") + }], + ..pb::StatsPayload::default() + }; + + let mut dimensions: Vec<(&str, PayloadMutation)> = Vec::new(); + dimensions.push(( + "resource", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].resource = "GET /b".to_string(); + }), + )); + dimensions.push(( + "peer_tags", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].peer_tags = vec!["grpc.target:other".to_string()]; + }), + )); + dimensions.push(( + "additional_metric_tags", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].additional_metric_tags = vec!["x:y".to_string()]; + }), + )); + dimensions.push(( + "is_trace_root", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].is_trace_root = pb::Trilean::False as i32; + }), + )); + dimensions.push(( + "http_status_code", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].http_status_code = 500; + }), + )); + dimensions.push(( + "grpc_status_code", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].grpc_status_code = "14".to_string(); + }), + )); + dimensions.push(( + "client_env", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].env = "staging".to_string(); + }), + )); + dimensions.push(( + "client_version", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].version = "2.0".to_string(); + }), + )); + dimensions.push(( + "span_kind", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].span_kind = "client".to_string(); + }), + )); + dimensions.push(( + "http_method", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].http_method = "POST".to_string(); + }), + )); + dimensions.push(( + "http_endpoint", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].http_endpoint = "/api".to_string(); + }), + )); + dimensions.push(( + "service_source", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].service_source = "dd.trace".to_string(); + }), + )); + dimensions.push(( + "span_derived_primary_tags", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].span_derived_primary_tags = vec!["t:1".to_string()]; + }), + )); + dimensions.push(( + "db_type", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].db_type = "postgres".to_string(); + }), + )); + dimensions.push(( + "synthetics", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].synthetics = true; + }), + )); + dimensions.push(( + "name", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].name = "other.request".to_string(); + }), + )); + dimensions.push(( + "type", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].r#type = "db".to_string(); + }), + )); + dimensions.push(( + "client_hostname", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].hostname = "host-1".to_string(); + }), + )); + dimensions.push(( + "client_container_id", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].container_id = "abc".to_string(); + }), + )); + dimensions.push(( + "client_tags", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].tags = vec!["k:v".to_string()]; + }), + )); + dimensions.push(( + "client_git_commit_sha", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].git_commit_sha = "deadbeef".to_string(); + }), + )); + dimensions.push(( + "client_image_tag", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].image_tag = "v1".to_string(); + }), + )); + dimensions.push(( + "client_process_tags", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].process_tags = "proc:a".to_string(); + }), + )); + dimensions.push(( + "client_process_tags_hash", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].process_tags_hash = 99; + }), + )); + dimensions.push(( + "client_agent_aggregation", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].agent_aggregation = "counts".to_string(); + }), + )); + dimensions.push(( + "stats_service", + Box::new(|p: &mut pb::StatsPayload| { + p.stats[0].stats[0].stats[0].service = "other-svc".to_string(); + }), + )); + + let baseline = stats_hits_by_key(&base()); + assert_eq!(baseline.len(), 1); + for (name, mutate) in dimensions { + let mut payload = base(); + mutate(&mut payload); + let groups = stats_hits_by_key(&payload); + assert_eq!( + groups.len(), + 1, + "{name} change must not merge with baseline" + ); + assert_ne!( + groups.keys().next(), + baseline.keys().next(), + "{name} must change the grouping key" + ); + } + } + + #[test] + fn grouping_ignores_measurements_and_delivery_metadata() { + let base = grouped_stats("svc", "GET /a", 1, |_| {}); + let variant = grouped_stats("svc", "GET /a", 7, |g| { + g.errors = 2; + g.duration = 500; + g.top_level_hits = 3; + g.ok_summary = vec![1, 2, 3]; + g.error_summary = vec![4, 5]; + }); + let client_a = client_payload("prod"); + let mut client_b = client_payload("prod"); + client_b.runtime_id = "other".to_string(); + client_b.sequence = 12; + client_b.lang = "rust".to_string(); + client_b.tracer_version = "9.9".to_string(); + + let key_base = StatsGroupKey::new(&client_a, &base); + let key_variant = StatsGroupKey::new(&client_b, &variant); + assert_eq!( + key_base, key_variant, + "measurement or delivery changes must not split groups" + ); + } + + #[test] + fn summary_output_is_deterministic_and_normalizes_tag_order() { + let mut client = client_payload("prod"); + client.tags = vec!["z:1".to_string(), "a:2".to_string()]; + let mut grouped = grouped_stats("svc", "GET /a", 5, |_| {}); + grouped.peer_tags = vec!["b:x".to_string(), "a:y".to_string()]; + + let key = StatsGroupKey::new(&client, &grouped); + let first = key.render(); + let second = StatsGroupKey::new(&client, &grouped).render(); + assert_eq!(first, second, "rendering must be deterministic"); + assert!( + first.contains("client.tags=[\"a:2\",\"z:1\"]"), + "tag lists must be sorted: {first}" + ); + assert!( + first.contains("stats.peer_tags=[\"a:y\",\"b:x\"]"), + "tag lists must be sorted: {first}" + ); + assert!(!first.contains("hits="), "summary key must exclude hits"); + assert!(first.starts_with("client.service=")); + } + + #[tokio::test] + async fn json_dumps_parse_for_all_endpoints_and_rejections() { + let tmp = tempfile::tempdir().expect("test: tempdir failed"); + let intake = MockIntake::start_with_options(MockIntakeOptions { + dump_dir: Some(tmp.path().to_path_buf()), + fail_stats_first_n: 1, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + let stats = stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 2, |_| {})], + ); + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + None, + rmp_serde::to_vec_named(&stats).expect("test: msgpack encode failed") + ) + .await, + StatusCode::INTERNAL_SERVER_ERROR + ); + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + None, + rmp_serde::to_vec_named(&stats).expect("test: msgpack encode failed") + ) + .await, + StatusCode::ACCEPTED + ); + + let trace = pb::AgentPayload { + env: "local".to_string(), + tracer_payloads: vec![pb::TracerPayload::default()], + ..pb::AgentPayload::default() + }; + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/traces", + None, + trace.encode_to_vec() + ) + .await, + StatusCode::ACCEPTED + ); + + let dsm = PipelineStatsPayload { + env: "local".to_string(), + service: "svc".to_string(), + tracer_version: "1.0".to_string(), + version: "2.0".to_string(), + tags: Vec::new(), + stats: Vec::new(), + }; + assert_eq!( + post( + &intake.base_url(), + "/api/v0.1/pipeline_stats", + None, + rmp_serde::to_vec_named(&dsm).expect("test: msgpack encode failed") + ) + .await, + StatusCode::ACCEPTED + ); + + let mut files: Vec<_> = std::fs::read_dir(tmp.path()) + .expect("test: read dump dir failed") + .filter_map(std::result::Result::ok) + .map(|e| e.path()) + .collect(); + files.sort(); + assert_eq!(files.len(), 4, "one dump per decoded attempt: {files:?}"); + + let mut seen = std::collections::HashSet::new(); + for path in &files { + let raw = std::fs::read(path).expect("test: read dump failed"); + let envelope: serde_json::Value = + serde_json::from_slice(&raw).expect("dump envelope must parse as JSON"); + let dumped_endpoint = envelope["endpoint"].as_str().expect("endpoint key missing"); + let dumped_status = envelope["status"].as_u64().expect("status key missing"); + assert!( + envelope["payload"].is_object(), + "decoded payload must be present" + ); + assert!( + envelope["request_id"].is_u64(), + "request id must be present" + ); + assert!( + envelope["encoding"].as_str() == Some("identity"), + "absent content-encoding must report identity" + ); + seen.insert(format!("{dumped_endpoint}/{dumped_status}")); + } + // Rejected stats attempt (500) and its successful retry (202) are + // both visible with their own identities and statuses. + assert!(seen.contains("/api/v0.2/stats/500")); + assert!(seen.contains("/api/v0.2/stats/202")); + assert!(seen.contains("/api/v0.2/traces/202")); + assert!(seen.contains("/api/v0.1/pipeline_stats/202")); + } + + #[tokio::test] + async fn dumps_never_overwrite_earlier_attempts() { + let tmp = tempfile::tempdir().expect("test: tempdir failed"); + let intake = MockIntake::start_with_options(MockIntakeOptions { + dump_dir: Some(tmp.path().to_path_buf()), + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + let body = rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + for _ in 0..5 { + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body.clone()).await, + StatusCode::ACCEPTED + ); + } + + let files: Vec<_> = std::fs::read_dir(tmp.path()) + .expect("test: read dump dir failed") + .filter_map(std::result::Result::ok) + .map(|e| e.path()) + .collect(); + assert_eq!(files.len(), 5, "every attempt gets its own file"); + let mut contents: Vec = files + .iter() + .map(|p| std::fs::read_to_string(p).expect("test: read dump failed")) + .collect(); + contents.sort(); + let unique: std::collections::HashSet = contents.iter().cloned().collect(); + assert_eq!( + unique.len(), + 5, + "each dump must have distinct request identity" + ); + } + + #[tokio::test] + async fn invalid_options_fail_clearly() { + // A dump directory whose parent is a file cannot be created. + let file = tempfile::NamedTempFile::new().expect("test: tempfile failed"); + let err = MockIntake::start_with_options(MockIntakeOptions { + dump_dir: Some(file.path().join("nested")), + ..MockIntakeOptions::default() + }) + .await + .expect_err("dump dir under a file must fail"); + assert!(matches!(err, MockIntakeError::DumpDir { .. })); + + // Binding to a port already in use must fail with a bind error. + let listener = TcpListener::bind(("127.0.0.1", 0)) + .await + .expect("test: bind failed"); + let port = listener + .local_addr() + .expect("test: local addr failed") + .port(); + let err = MockIntake::start_with_options(MockIntakeOptions { + port, + ..MockIntakeOptions::default() + }) + .await + .expect_err("bind conflict must fail"); + assert!(matches!(err, MockIntakeError::Bind { .. })); + } + + #[tokio::test] + async fn start_with_options_uses_configured_port() { + let listener = TcpListener::bind(("127.0.0.1", 0)) + .await + .expect("test: bind failed"); + let port = listener + .local_addr() + .expect("test: local addr failed") + .port(); + drop(listener); + + let intake = MockIntake::start_with_options(MockIntakeOptions { + port, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + assert_eq!(intake.base_url(), format!("http://127.0.0.1:{port}")); + } +} diff --git a/crates/datadog-trace-agent/tests/common/mock_server.rs b/crates/datadog-trace-agent/tests/common/mock_server.rs deleted file mode 100644 index cd0cd6b9..00000000 --- a/crates/datadog-trace-agent/tests/common/mock_server.rs +++ /dev/null @@ -1,157 +0,0 @@ -// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/ -// SPDX-License-Identifier: Apache-2.0 - -//! Simple mock HTTP server for testing flushers - -use http_body_util::BodyExt; -use hyper::{Request, Response, StatusCode, body::Incoming}; -use hyper_util::rt::TokioIo; -use libdd_common::http_common; -use std::net::SocketAddr; -use std::sync::{Arc, Mutex}; -use tokio::net::TcpListener; - -#[derive(Clone, Debug)] -pub struct ReceivedRequest { - pub method: String, - pub path: String, - pub headers: Vec<(String, String)>, - pub body: Vec, -} - -pub struct MockServer { - pub addr: SocketAddr, - pub received_requests: Arc>>, - shutdown_tx: Option>, -} - -impl MockServer { - /// Start a mock HTTP server on a random port - pub async fn start() -> Self { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("Failed to bind mock server"); - let addr = listener.local_addr().expect("Failed to get local addr"); - - let received_requests = Arc::new(Mutex::new(Vec::new())); - let requests_clone = received_requests.clone(); - - let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel(); - - tokio::spawn(async move { - loop { - tokio::select! { - result = listener.accept() => { - let (stream, _) = match result { - Ok(conn) => conn, - Err(e) => { - eprintln!("Mock server accept error: {}", e); - break; - } - }; - - let io = TokioIo::new(stream); - let requests = requests_clone.clone(); - - tokio::spawn(async move { - let service = hyper::service::service_fn(move |req: Request| { - let requests = requests.clone(); - async move { - // Capture the request - let method = req.method().to_string(); - let path = req.uri().path().to_string(); - let is_stats_intake = path.ends_with("/stats"); - let headers: Vec<(String, String)> = req - .headers() - .iter() - .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string())) - .collect(); - - // Read the body - let body_bytes = req - .into_body() - .collect() - .await - .map(|collected| collected.to_bytes().to_vec()) - .unwrap_or_default(); - - // Store the request - requests.lock().unwrap().push(ReceivedRequest { - method, - path, - headers, - body: body_bytes, - }); - - // Trace intake accepts 2xx - // Stats intake accepts 202 - // see `libdd_trace_utils::stats_utils::send_stats_payload_with_client` - let (status, body) = if is_stats_intake { - (StatusCode::ACCEPTED, http_common::Body::empty()) - } else { - (StatusCode::OK, http_common::Body::from(r#"{"ok":true}"#)) - }; - Ok::<_, hyper::http::Error>( - Response::builder() - .status(status) - .body(body) - .unwrap(), - ) - } - }); - - let _ = hyper::server::conn::http1::Builder::new() - .serve_connection(io, service) - .await; - }); - } - _ = &mut shutdown_rx => { - break; - } - } - } - }); - - MockServer { - addr, - received_requests, - shutdown_tx: Some(shutdown_tx), - } - } - - /// Get the base URL of the mock server - pub fn url(&self) -> String { - format!("http://{}", self.addr) - } - - /// Get all received requests - #[allow(dead_code)] - pub fn get_requests(&self) -> Vec { - self.received_requests.lock().unwrap().clone() - } - - /// Get requests matching a path - pub fn get_requests_for_path(&self, path: &str) -> Vec { - self.received_requests - .lock() - .unwrap() - .iter() - .filter(|req| req.path == path) - .cloned() - .collect() - } - - /// Clear all received requests - #[allow(dead_code)] - pub fn clear_requests(&self) { - self.received_requests.lock().unwrap().clear(); - } -} - -impl Drop for MockServer { - fn drop(&mut self) { - if let Some(shutdown_tx) = self.shutdown_tx.take() { - let _ = shutdown_tx.send(()); - } - } -} From a21468b2d68fd0b779ad9328865f36e40562ef5f Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 25 Sep 2026 12:07:55 -0400 Subject: [PATCH 02/32] Serve trace-agent integration tests with datadog-mock-intake Replace the test-local MockServer with the shared datadog-mock-intake crate. Payload assertions now use the intake's decoded capture instead of re-decoding captured bytes, positive assertions wait on decoded payload arrival, and the DSM forwarding test sends a valid msgpack pipeline-stats payload so it verifies both unchanged byte forwarding and intake acceptance. --- crates/datadog-trace-agent/Cargo.toml | 2 +- .../tests/common/helpers.rs | 13 - .../datadog-trace-agent/tests/common/mod.rs | 1 - .../tests/integration_test.rs | 239 ++++++++++-------- 4 files changed, 133 insertions(+), 122 deletions(-) diff --git a/crates/datadog-trace-agent/Cargo.toml b/crates/datadog-trace-agent/Cargo.toml index 291ff87f..4e868937 100644 --- a/crates/datadog-trace-agent/Cargo.toml +++ b/crates/datadog-trace-agent/Cargo.toml @@ -42,7 +42,7 @@ bytes = "1.10.1" prost = "0.14.1" [dev-dependencies] -flate2 = "1" +datadog-mock-intake = { path = "../datadog-mock-intake" } rmp-serde = "1.1.1" serial_test = "2.0.0" duplicate = "2.0.1" diff --git a/crates/datadog-trace-agent/tests/common/helpers.rs b/crates/datadog-trace-agent/tests/common/helpers.rs index 32cd4736..8650cb48 100644 --- a/crates/datadog-trace-agent/tests/common/helpers.rs +++ b/crates/datadog-trace-agent/tests/common/helpers.rs @@ -3,14 +3,12 @@ //! Helper functions for integration tests -use flate2::read::GzDecoder; use hyper::{Request, Response}; use hyper_util::rt::TokioIo; use libdd_common::http_common; use libdd_trace_protobuf::pb; use libdd_trace_utils::test_utils::create_test_json_span; use serde_json::json; -use std::io::Read; use std::time::{Duration, UNIX_EPOCH}; use tokio::time::timeout; @@ -70,17 +68,6 @@ pub fn create_test_client_stats_payload(service: &str) -> Vec { rmp_serde::to_vec_named(&payload).expect("Failed to serialize client stats payload") } -/// Decompress a gzip+msgpack stats payload and deserialize it into a `StatsPayload`. -/// The stats flusher encodes payloads as `gzip(rmp_serde::to_vec_named(StatsPayload))`. -pub fn decode_stats_payload(body: &[u8]) -> pb::StatsPayload { - let mut decoder = GzDecoder::new(body); - let mut decompressed = Vec::new(); - decoder - .read_to_end(&mut decompressed) - .expect("Failed to decompress stats payload"); - rmp_serde::from_slice(&decompressed).expect("Failed to deserialize stats payload") -} - /// Send an HTTP request over TCP and return the response pub async fn send_tcp_request( port: u16, diff --git a/crates/datadog-trace-agent/tests/common/mod.rs b/crates/datadog-trace-agent/tests/common/mod.rs index 447946a7..3d83bbdb 100644 --- a/crates/datadog-trace-agent/tests/common/mod.rs +++ b/crates/datadog-trace-agent/tests/common/mod.rs @@ -4,5 +4,4 @@ //! Common test utilities, mocks, and helpers for integration tests pub mod helpers; -pub mod mock_server; pub mod mocks; diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index cb0117ba..0f5722ef 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -5,11 +5,10 @@ mod common; use common::helpers::{ create_client_span_with_peer_tag_payload, create_test_client_stats_payload, - create_test_trace_payload, create_trace_with_span_kind_children_payload, decode_stats_payload, - send_tcp_request, + create_test_trace_payload, create_trace_with_span_kind_children_payload, send_tcp_request, }; -use common::mock_server::MockServer; use common::mocks::{MockEnvVerifier, MockStatsFlusher, MockStatsProcessor, MockTraceFlusher}; +use datadog_mock_intake::{MockIntake, PipelineStatsPayload}; use datadog_trace_agent::{ config::{Config, Tags, test_helpers::create_tcp_test_config}, mini_agent::MiniAgent, @@ -39,15 +38,15 @@ const FLUSH_WAIT_DURATION: Duration = Duration::from_millis(1500); // the timeout exists to fail loudly rather than hang indefinitely. const VERIFY_REQUEST_TIMEOUT: Duration = Duration::from_secs(60); -async fn wait_for_request_at_path(mock_server: &common::mock_server::MockServer, path: &str) { +async fn wait_for_capture(ready: impl Fn() -> bool, what: &str) { let deadline = tokio::time::Instant::now() + VERIFY_REQUEST_TIMEOUT; while tokio::time::Instant::now() < deadline { - if !mock_server.get_requests_for_path(path).is_empty() { + if ready() { return; } tokio::time::sleep(Duration::from_millis(50)).await; } - panic!("Timed out after {VERIFY_REQUEST_TIMEOUT:?} waiting for request at {path}"); + panic!("Timed out after {VERIFY_REQUEST_TIMEOUT:?} waiting for {what}"); } /// Helper to configure a config with mock server endpoints @@ -114,14 +113,18 @@ pub fn create_mini_agent_with_real_flushers( (mini_agent, stats_concentrator_service_handle) } -/// Helper to verify trace request sent to mock server -pub async fn verify_trace_request(mock_server: &common::mock_server::MockServer) { - wait_for_request_at_path(mock_server, "/api/v0.2/traces").await; - let trace_reqs = mock_server.get_requests_for_path("/api/v0.2/traces"); +/// Helper to verify trace request sent to the mock intake +pub async fn verify_trace_request(mock_intake: &MockIntake) { + wait_for_capture( + || !mock_intake.trace_payloads().is_empty(), + "a trace request at /api/v0.2/traces", + ) + .await; + let trace_reqs = mock_intake.requests_for_path("/api/v0.2/traces"); assert!( !trace_reqs.is_empty(), - "Expected at least one trace request to mock server" + "Expected at least one trace request to mock intake" ); let trace_req = &trace_reqs[0]; @@ -151,14 +154,18 @@ pub async fn verify_trace_request(mock_server: &common::mock_server::MockServer) ); } -/// Helper to verify stats request sent to mock server -pub async fn verify_stats_request(mock_server: &common::mock_server::MockServer) { - wait_for_request_at_path(mock_server, "/api/v0.2/stats").await; - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); +/// Helper to verify stats request sent to the mock intake +pub async fn verify_stats_request(mock_intake: &MockIntake) { + wait_for_capture( + || !mock_intake.stats_payloads().is_empty(), + "a stats request at /api/v0.2/stats", + ) + .await; + let stats_reqs = mock_intake.requests_for_path("/api/v0.2/stats"); assert!( !stats_reqs.is_empty(), - "Expected at least one stats request to mock server" + "Expected at least one stats request to mock intake" ); let stats_req = &stats_reqs[0]; @@ -188,28 +195,35 @@ pub async fn verify_stats_request(mock_server: &common::mock_server::MockServer) ); } -/// Helper to verify stats request was not sent to mock server -pub fn verify_no_stats_request(mock_server: &common::mock_server::MockServer) { - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); +/// Helper to verify stats request was not sent to mock intake +pub fn verify_no_stats_request(mock_intake: &MockIntake) { + let stats_payloads = mock_intake.stats_payloads(); assert!( - stats_reqs.is_empty(), - "Expected no stats request to mock server, received {} request(s)", - stats_reqs.len() + stats_payloads.is_empty(), + "Expected no stats request to mock intake, received {} request(s)", + stats_payloads.len() ); } -/// Helper to verify a DSM request sent to the mock server +/// Helper to verify a DSM request sent to the mock intake. The forwarded body +/// must arrive byte-for-byte unchanged AND be decodable by the intake: the +/// decoded capture proves the proxy forwards a payload the real intake accepts. pub async fn verify_dsm_request( - mock_server: &common::mock_server::MockServer, + mock_intake: &MockIntake, expected_body: &[u8], + expected_edge_tags: &[&str], expected_additional_tags: &[&str], ) { - wait_for_request_at_path(mock_server, "/api/v0.1/pipeline_stats").await; - let dsm_reqs = mock_server.get_requests_for_path("/api/v0.1/pipeline_stats"); + wait_for_capture( + || !mock_intake.pipeline_stats_payloads().is_empty(), + "a DSM request at /api/v0.1/pipeline_stats", + ) + .await; + let dsm_reqs = mock_intake.requests_for_path("/api/v0.1/pipeline_stats"); assert!( !dsm_reqs.is_empty(), - "Expected at least one DSM request to mock server" + "Expected at least one DSM request to mock intake" ); let dsm_req = &dsm_reqs[0]; @@ -219,6 +233,17 @@ pub async fn verify_dsm_request( "Expected DSM payload body to be forwarded unchanged" ); + let decoded = &mock_intake.pipeline_stats_payloads()[0]; + let edge_tags: Vec<&str> = decoded.stats[0].stats[0] + .edge_tags + .iter() + .map(String::as_str) + .collect(); + assert_eq!( + edge_tags, expected_edge_tags, + "Expected decoded DSM payload to carry the fixture edge tags" + ); + let api_key = dsm_req .headers .iter() @@ -469,11 +494,10 @@ async fn test_mini_agent_named_pipe_handles_requests() { #[tokio::test] #[serial] async fn test_mini_agent_tcp_with_real_flushers() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8127); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = true; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -512,23 +536,22 @@ async fn test_mini_agent_tcp_with_real_flushers() { .expect("Failed to send /v0.4/traces request"); assert_eq!(trace_response.status(), StatusCode::OK); - verify_trace_request(&mock_server).await; + verify_trace_request(&mock_intake).await; // Trigger shutdown to force flush in progress concentrator buckets let _ = shutdown_tx.send(true); let _ = agent_handle.await; - verify_stats_request(&mock_server).await; // Stats generator should generate stats from trace payload + verify_stats_request(&mock_intake).await; // Stats generator should generate stats from trace payload } #[cfg(test)] #[tokio::test] #[serial] async fn test_mini_agent_tcp_proxies_dsm_requests() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8133); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.tags = Tags::from_env_string("env:test,service:payments"); let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -563,7 +586,23 @@ async fn test_mini_agent_tcp_proxies_dsm_requests() { "Mini agent server failed to start within timeout" ); - let dsm_payload = br#"{"Stats":[{"EdgeTags":["direction:out","type:kafka"]}]}"#.to_vec(); + // A valid DSM pipeline-stats payload: the proxy must forward it unchanged + // and the intake must accept and decode it. + let dsm_payload = rmp_serde::to_vec_named(&PipelineStatsPayload { + env: "local".to_string(), + service: "svc".to_string(), + tracer_version: "1.0".to_string(), + version: "2.0".to_string(), + tags: Vec::new(), + stats: vec![datadog_mock_intake::PipelineStatsBucket { + stats: vec![datadog_mock_intake::PipelineStatsPoint { + hash: 7, + parent_hash: 0, + edge_tags: vec!["direction:out".to_string(), "type:kafka".to_string()], + }], + }], + }) + .expect("Failed to serialize DSM pipeline stats payload"); let response = send_tcp_request( test_port, "/v0.1/pipeline_stats", @@ -582,8 +621,9 @@ async fn test_mini_agent_tcp_proxies_dsm_requests() { assert_eq!(response.status(), StatusCode::OK); verify_dsm_request( - &mock_server, + &mock_intake, &dsm_payload, + &["direction:out", "type:kafka"], &[ "host:worker-1", "default_env:prod", @@ -602,11 +642,10 @@ async fn test_mini_agent_tcp_proxies_dsm_requests() { #[tokio::test] #[serial] async fn test_concentrator_task_death_shuts_down_mini_agent() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8129); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -655,11 +694,10 @@ async fn test_concentrator_task_death_shuts_down_mini_agent() { #[tokio::test] #[serial] async fn test_tracer_and_agent_stats_disabled_produces_no_stats() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8128); // use different port to avoid race condition with other tests - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = false; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -696,10 +734,10 @@ async fn test_tracer_and_agent_stats_disabled_produces_no_stats() { .expect("Failed to send /v0.4/traces request"); assert_eq!(trace_response.status(), StatusCode::OK); - verify_trace_request(&mock_server).await; + verify_trace_request(&mock_intake).await; // Bounded wait to confirm absence of stats request — neither side computed stats. tokio::time::sleep(FLUSH_WAIT_DURATION).await; - verify_no_stats_request(&mock_server); + verify_no_stats_request(&mock_intake); // Clean up agent_handle.abort(); @@ -709,11 +747,10 @@ async fn test_tracer_and_agent_stats_disabled_produces_no_stats() { #[tokio::test] #[serial] async fn test_tracer_disabled_agent_stats_enabled_uses_agent_stats() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8136); // use different port to avoid race condition with other tests - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = true; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -752,8 +789,8 @@ async fn test_tracer_disabled_agent_stats_enabled_uses_agent_stats() { .expect("Failed to send /v0.4/traces request"); assert_eq!(trace_response.status(), StatusCode::OK); - verify_trace_request(&mock_server).await; - verify_stats_request(&mock_server).await; // Agent should compute stats + verify_trace_request(&mock_intake).await; + verify_stats_request(&mock_intake).await; // Agent should compute stats // Clean up agent_handle.abort(); @@ -762,11 +799,10 @@ async fn test_tracer_disabled_agent_stats_enabled_uses_agent_stats() { #[tokio::test] #[serial] async fn test_tracer_and_agent_stats_enabled_uses_agent_stats_no_duplicates() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8135); // use different port to avoid race condition with other tests - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = true; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -827,20 +863,16 @@ async fn test_tracer_and_agent_stats_enabled_uses_agent_stats_no_duplicates() { .expect("Failed to send /v0.4/traces request"); assert_eq!(trace_response.status(), StatusCode::OK); - verify_trace_request(&mock_server).await; - verify_stats_request(&mock_server).await; + verify_trace_request(&mock_intake).await; + verify_stats_request(&mock_intake).await; // The tracer computed stats should never reach the backend. Only the agent computed stats should reach the backend. - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); - let has_marker = stats_reqs - .iter() - .map(|req| decode_stats_payload(&req.body)) - .any(|payload| { - payload - .stats - .iter() - .any(|csp| csp.service == "tracer-marker-stats") - }); + let has_marker = mock_intake.stats_payloads().iter().any(|payload| { + payload + .stats + .iter() + .any(|csp| csp.service == "tracer-marker-stats") + }); assert!( !has_marker, "Expected tracer computed stats to be dropped, not forwarded to the backend" @@ -853,11 +885,10 @@ async fn test_tracer_and_agent_stats_enabled_uses_agent_stats_no_duplicates() { #[tokio::test] #[serial] async fn test_tracer_stats_enabled_agent_stats_disabled_forwards_tracer_stats() { - let mock_server: MockServer = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8134); // use different port to avoid race condition with other tests - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = false; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -898,17 +929,13 @@ async fn test_tracer_stats_enabled_agent_stats_disabled_forwards_tracer_stats() .expect("Failed to send /v0.6/stats request"); assert_eq!(stats_response.status(), StatusCode::ACCEPTED); - verify_stats_request(&mock_server).await; - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); - let has_marker = stats_reqs - .iter() - .map(|req| decode_stats_payload(&req.body)) - .any(|payload| { - payload - .stats - .iter() - .any(|csp| csp.service == "tracer-marker-stats") - }); + verify_stats_request(&mock_intake).await; + let has_marker = mock_intake.stats_payloads().iter().any(|payload| { + payload + .stats + .iter() + .any(|csp| csp.service == "tracer-marker-stats") + }); assert!( has_marker, "Expected tracer computed stats to be forwarded to the backend" @@ -924,11 +951,10 @@ async fn test_tracer_stats_enabled_agent_stats_disabled_forwards_tracer_stats() #[tokio::test] #[serial] async fn test_internal_span_kind_does_not_produce_stats() { - let mock_server = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8132); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = true; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -974,22 +1000,20 @@ async fn test_internal_span_kind_does_not_produce_stats() { let _ = shutdown_tx.send(true); let _ = agent_handle.await; - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); + let stats_payloads = mock_intake.stats_payloads(); assert!( - !stats_reqs.is_empty(), + !stats_payloads.is_empty(), "Expected a stats request from the root span" ); - let all_groups: Vec<_> = stats_reqs + let all_groups: Vec<_> = stats_payloads .iter() - .map(|req| decode_stats_payload(&req.body)) .flat_map(|payload| { payload .stats - .into_iter() - .flat_map(|csp| csp.stats.into_iter()) - .flat_map(|bucket| bucket.stats.into_iter()) - .collect::>() + .iter() + .flat_map(|csp| csp.stats.iter()) + .flat_map(|bucket| bucket.stats.iter()) }) .collect(); @@ -1015,11 +1039,10 @@ async fn test_internal_span_kind_does_not_produce_stats() { #[tokio::test] #[serial] async fn test_peer_tags_in_flushed_stats() { - let mock_server = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let mut config = create_tcp_test_config(8131); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.agent_stats_computation_enabled = true; let config = Arc::new(config); let test_port = config.dd_apm_receiver_port; @@ -1065,13 +1088,13 @@ async fn test_peer_tags_in_flushed_stats() { let _ = shutdown_tx.send(true); let _ = agent_handle.await; - let stats_reqs = mock_server.get_requests_for_path("/api/v0.2/stats"); + let stats_payloads = mock_intake.stats_payloads(); assert!( - !stats_reqs.is_empty(), + !stats_payloads.is_empty(), "Expected at least one stats request" ); - let payload = decode_stats_payload(&stats_reqs[0].body); + let payload = &stats_payloads[0]; let all_peer_tags: Vec<&str> = payload .stats .iter() @@ -1090,12 +1113,11 @@ async fn test_peer_tags_in_flushed_stats() { #[tokio::test] #[serial] async fn test_mini_agent_named_pipe_with_real_flushers() { - let mock_server = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let pipe_name = r"\\.\pipe\dd_trace_real_flusher_test"; let mut config = create_tcp_test_config(0); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.dd_apm_windows_pipe_name = Some(pipe_name.to_string()); config.dd_apm_receiver_port = 0; config.agent_stats_computation_enabled = true; @@ -1133,26 +1155,25 @@ async fn test_mini_agent_named_pipe_with_real_flushers() { .expect("Failed to send /v0.4/traces request over named pipe"); assert_eq!(trace_response.status(), StatusCode::OK); - verify_trace_request(&mock_server).await; + verify_trace_request(&mock_intake).await; // Trigger shutdown to force flush in progress concentrator buckets let _ = shutdown_tx.send(true); let _ = agent_handle.await; - verify_stats_request(&mock_server).await; + verify_stats_request(&mock_intake).await; } #[cfg(all(test, windows, feature = "windows-pipes"))] #[tokio::test] #[serial] async fn test_mini_agent_dual_transport_with_real_flushers() { - let mock_server = MockServer::start().await; - tokio::time::sleep(Duration::from_millis(50)).await; + let mock_intake = MockIntake::start().await; let pipe_name = r"\\.\pipe\dd_trace_dual_transport_test"; let tcp_port: u16 = 8130; let mut config = create_tcp_test_config(tcp_port); - configure_mock_endpoints(&mut config, &mock_server.url()); + configure_mock_endpoints(&mut config, &mock_intake.base_url()); config.dd_apm_windows_pipe_name = Some(pipe_name.to_string()); // Both transports are deliberately set on the same agent: a non-zero TCP // port AND a pipe name. They must come up concurrently. @@ -1215,12 +1236,16 @@ async fn test_mini_agent_dual_transport_with_real_flushers() { .expect("Failed to send /v0.4/traces request over named pipe"); assert_eq!(pipe_response.status(), StatusCode::OK); - wait_for_request_at_path(&mock_server, "/api/v0.2/traces").await; + wait_for_capture( + || !mock_intake.trace_payloads().is_empty(), + "a trace request at /api/v0.2/traces", + ) + .await; // Both payloads must reach the same backend through the shared flusher // pipeline. The flusher may batch them into one POST or two; either is // fine, what matters is that both service-name needles show up. - let trace_reqs = mock_server.get_requests_for_path("/api/v0.2/traces"); + let trace_reqs = mock_intake.requests_for_path("/api/v0.2/traces"); assert!( !trace_reqs.is_empty(), "no trace POST reached backend; expected traces from both transports" @@ -1245,5 +1270,5 @@ async fn test_mini_agent_dual_transport_with_real_flushers() { // a transport, agent_handle would hang or stats wouldn't arrive. let _ = shutdown_tx.send(true); let _ = agent_handle.await; - verify_stats_request(&mock_server).await; + verify_stats_request(&mock_intake).await; } From 11fe447f1357db5160a1d2a4dc62d0aeb3798cb5 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 25 Sep 2026 13:59:34 -0400 Subject: [PATCH 03/32] Rename the standalone binary and env vars to mock-intake in crate docs The datadog-lambda-extension standalone binary and its environment variables are now named mock-intake / MOCK_INTAKE_*; align the crate documentation. --- crates/datadog-mock-intake/src/lib.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 8f477e6e..a360cb5f 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -21,11 +21,11 @@ //! //! `MockIntake::start()` binds `127.0.0.1:0` with defaults and is used by the //! trace-agent integration tests in this workspace and by the APM / DSM -//! integration tests in `bottlecap`. The standalone `fake-intake` binary in +//! integration tests in `bottlecap`. The standalone `mock-intake` binary in //! the `datadog-lambda-extension` repository wraps this crate with request //! summaries, optional stats failure injection -//! (`FAKE_INTAKE_FAIL_STATS_FIRST_N`), and optional JSON dumps -//! (`FAKE_INTAKE_DUMP_DIR`) for local debugging against a live tracer. +//! (`MOCK_INTAKE_FAIL_STATS_FIRST_N`), and optional JSON dumps +//! (`MOCK_INTAKE_DUMP_DIR`) for local debugging against a live tracer. //! //! This crate is self-contained (no host config, logging framework, or trace //! processing dependencies) so it stays usable from multiple repositories. From a893f8bbd01d8d2bf5fcfc6c6778fb9197d6f070 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Fri, 25 Sep 2026 18:54:18 -0400 Subject: [PATCH 04/32] Cap decompressed body size and skip unused dump serialization in mock-intake MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Guards decompress() against decompression-bomb-style memory blow-up by capping gzip/zstd output at 64 MiB. Also skips building the diagnostic JSON dump when no dump_dir is configured (the default), avoiding a wasted allocation on every request, and marks DSM pipeline-stats payload fields #[serde(default)] so a payload missing an optional field still decodes. 🤖 --- crates/datadog-mock-intake/src/lib.rs | 78 +++++++++++++++++++-------- 1 file changed, 55 insertions(+), 23 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index a360cb5f..51ace914 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -62,39 +62,45 @@ use tokio::task::JoinHandle; /// before decompression. const MAX_BODY_SIZE: usize = 2_097_152; +/// Maximum accepted size for a decompressed body. `MAX_BODY_SIZE` caps the +/// wire body, but gzip/zstd can still expand a small compressed body into a +/// much larger one; this bounds that expansion so a pathological or +/// malformed payload cannot exhaust memory. +const MAX_DECOMPRESSED_SIZE: usize = 64 * 1024 * 1024; + /// A DSM pipeline-stats payload as it lands on `/api/v0.1/pipeline_stats`. /// Only the fields tests assert on are decoded; serde ignores the rest /// (including the `serde_bytes` latency sketches). JSON dumps therefore /// contain only these fields, not the full wire payload. #[derive(Clone, Debug, Deserialize, Serialize)] pub struct PipelineStatsPayload { - #[serde(rename = "Env")] + #[serde(rename = "Env", default)] pub env: String, - #[serde(rename = "Service")] + #[serde(rename = "Service", default)] pub service: String, - #[serde(rename = "TracerVersion")] + #[serde(rename = "TracerVersion", default)] pub tracer_version: String, - #[serde(rename = "Version")] + #[serde(rename = "Version", default)] pub version: String, - #[serde(rename = "Tags")] + #[serde(rename = "Tags", default)] pub tags: Vec, - #[serde(rename = "Stats")] + #[serde(rename = "Stats", default)] pub stats: Vec, } #[derive(Clone, Debug, Deserialize, Serialize)] pub struct PipelineStatsBucket { - #[serde(rename = "Stats")] + #[serde(rename = "Stats", default)] pub stats: Vec, } #[derive(Clone, Debug, Deserialize, Serialize)] pub struct PipelineStatsPoint { - #[serde(rename = "Hash")] + #[serde(rename = "Hash", default)] pub hash: u64, - #[serde(rename = "ParentHash")] + #[serde(rename = "ParentHash", default)] pub parent_hash: u64, - #[serde(rename = "EdgeTags")] + #[serde(rename = "EdgeTags", default)] pub edge_tags: Vec, } @@ -499,14 +505,16 @@ fn handle_stats( ); } - if let Some(payload) = &handled.decoded { + if state.options.dump_dir.is_some() + && let Some(payload) = &handled.decoded + { dump_request( state, handled.request_id, "/api/v0.2/stats", headers, handled.status, - serde_json::to_value(payload).unwrap_or_else(|_| serde_json::Value::Null), + serde_json::to_value(payload).unwrap_or(serde_json::Value::Null), ); } @@ -570,7 +578,9 @@ fn handle_traces( ); } - if let Some(payload) = &handled.decoded { + if state.options.dump_dir.is_some() + && let Some(payload) = &handled.decoded + { dump_request( state, handled.request_id, @@ -635,14 +645,16 @@ fn handle_pipeline_stats( ); } - if let Some(payload) = &handled.decoded { + if state.options.dump_dir.is_some() + && let Some(payload) = &handled.decoded + { dump_request( state, handled.request_id, "/api/v0.1/pipeline_stats", headers, handled.status, - serde_json::to_value(payload).unwrap_or_else(|_| serde_json::Value::Null), + serde_json::to_value(payload).unwrap_or(serde_json::Value::Null), ); } @@ -1004,15 +1016,16 @@ fn decompress(headers: &hyper::HeaderMap, body: &Bytes) -> Result, Strin match encoding.as_str() { "gzip" => { - let mut decoder = flate2::read::GzDecoder::new(body.as_ref()); - let mut out = Vec::new(); - decoder - .read_to_end(&mut out) - .map_err(|e| format!("mock_intake: gzip decode failed: {e}"))?; - Ok(out) + let decoder = flate2::read::GzDecoder::new(body.as_ref()); + read_capped(decoder, MAX_DECOMPRESSED_SIZE) + .map_err(|e| format!("mock_intake: gzip decode failed: {e}")) + } + "zstd" => { + let decoder = zstd::stream::read::Decoder::new(body.as_ref()) + .map_err(|e| format!("mock_intake: zstd decoder init failed: {e}"))?; + read_capped(decoder, MAX_DECOMPRESSED_SIZE) + .map_err(|e| format!("mock_intake: zstd decode failed: {e}")) } - "zstd" => zstd::stream::decode_all(body.as_ref()) - .map_err(|e| format!("mock_intake: zstd decode failed: {e}")), _ => { if !encoding.is_empty() { eprintln!( @@ -1024,6 +1037,25 @@ fn decompress(headers: &hyper::HeaderMap, body: &Bytes) -> Result, Strin } } +/// Read `reader` to the end, rejecting output past `cap` bytes instead of +/// buffering it. Guards decompression against small inputs that expand into +/// an unbounded amount of memory (a decompression bomb). +fn read_capped(reader: impl Read, cap: usize) -> Result, String> { + let mut out = Vec::new(); + // Ask for one byte more than the cap so an exactly-at-cap payload isn't + // mistaken for a truncated, over-cap one. + reader + .take(cap as u64 + 1) + .read_to_end(&mut out) + .map_err(|e| format!("read failed: {e}"))?; + if out.len() > cap { + return Err(format!( + "decompressed body exceeds {cap} byte cap, rejecting as a likely decompression bomb" + )); + } + Ok(out) +} + #[cfg(test)] mod tests { From 35e10433e3fde65a61485e4c74b33061e16f2190 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 00:34:18 -0400 Subject: [PATCH 05/32] Abort accepted connections when mock intake is dropped --- crates/datadog-mock-intake/Cargo.toml | 2 +- crates/datadog-mock-intake/src/lib.rs | 101 +++++++++++++++++++++++++- 2 files changed, 101 insertions(+), 2 deletions(-) diff --git a/crates/datadog-mock-intake/Cargo.toml b/crates/datadog-mock-intake/Cargo.toml index db1cf336..894fe183 100644 --- a/crates/datadog-mock-intake/Cargo.toml +++ b/crates/datadog-mock-intake/Cargo.toml @@ -22,7 +22,7 @@ rmp-serde = { version = "1.3.1", default-features = false } serde = { version = "1.0.145", features = ["derive"] } serde_json = "1.0" thiserror = { version = "1.0.58", default-features = false } -tokio = { version = "1", features = ["macros", "net", "rt-multi-thread", "sync", "time"] } +tokio = { version = "1", features = ["io-util", "macros", "net", "rt-multi-thread", "sync", "time"] } zstd = { version = "0.13.3", default-features = false } [dev-dependencies] diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 51ace914..4ba0f4e4 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -201,6 +201,7 @@ pub struct MockIntake { state: std::sync::Arc, shutdown_tx: Option>, task: Option>, + connections: std::sync::Arc>>>, } impl MockIntake { @@ -243,7 +244,10 @@ impl MockIntake { }); let (shutdown_tx, mut shutdown_rx) = oneshot::channel::<()>(); + let connections: std::sync::Arc>>> = + std::sync::Arc::default(); let task_state = std::sync::Arc::clone(&state); + let task_connections = std::sync::Arc::clone(&connections); let task = tokio::spawn(async move { let state = task_state; loop { @@ -259,7 +263,7 @@ impl MockIntake { let io = TokioIo::new(stream); let state = std::sync::Arc::clone(&state); - tokio::spawn(async move { + let handle = tokio::spawn(async move { let service = service_fn(move |req: Request| { let state = std::sync::Arc::clone(&state); async move { handle_request(state, req).await } @@ -268,6 +272,10 @@ impl MockIntake { .serve_connection(io, service) .await; }); + if let Ok(mut conns) = task_connections.lock() { + conns.retain(|h| !h.is_finished()); + conns.push(handle); + } } _ = &mut shutdown_rx => { break; @@ -281,6 +289,7 @@ impl MockIntake { state, shutdown_tx: Some(shutdown_tx), task: Some(task), + connections, }) } @@ -369,6 +378,13 @@ impl Drop for MockIntake { if let Some(task) = self.task.take() { task.abort(); } + // Abort any accepted-but-still-running connection tasks so a + // keep-alive client cannot keep sending requests after Drop. + if let Ok(mut conns) = self.connections.lock() { + for handle in conns.drain(..) { + handle.abort(); + } + } } } @@ -1248,6 +1264,89 @@ mod tests { assert!(bound, "listener port was not released after Drop"); } + #[tokio::test] + async fn drop_aborts_keep_alive_connections() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let intake = start_default().await; + let state = std::sync::Arc::clone(&intake.state); + let addr = intake + .base_url() + .trim_start_matches("http://") + .to_string(); + + // Raw TCP so the connection stays open (keep-alive) across Drop. + let mut stream = tokio::net::TcpStream::connect(&addr) + .await + .expect("test: connect to mock intake failed"); + let body = b"not-msgpack"; + let request = format!( + "POST /api/v0.2/stats HTTP/1.1\r\nHost: {addr}\r\nContent-Type: application/msgpack\r\nContent-Length: {}\r\n\r\n", + body.len() + ); + stream + .write_all(request.as_bytes()) + .await + .expect("test: write request head failed"); + stream + .write_all(body.as_slice()) + .await + .expect("test: write request body failed"); + + // The response has an empty body, so it ends at the header + // terminator; read until we have seen it. Do not read to EOF: the + // server keeps the connection open for keep-alive reuse. + let mut buf = Vec::new(); + let mut chunk = [0u8; 1024]; + loop { + let n = stream + .read(&mut chunk) + .await + .expect("test: read response failed"); + assert!(n > 0, "test: connection closed before response complete"); + buf.extend_from_slice(&chunk[..n]); + if buf.windows(4).any(|w| w == b"\r\n\r\n") { + break; + } + } + assert!( + buf.starts_with(b"HTTP/1.1"), + "test: expected an HTTP response, got {buf:?}" + ); + let captured_before = state + .captured + .lock() + .expect("test: requests mutex poisoned") + .requests + .len(); + assert_eq!(captured_before, 1, "test: first request should be captured"); + + drop(intake); + + // Reuse the still-open keep-alive connection. The connection task + // was aborted on Drop, so this request must not be captured. + stream + .write_all(request.as_bytes()) + .await + .expect("test: write second request failed"); + stream + .write_all(body.as_slice()) + .await + .expect("test: write second body failed"); + tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + + let captured_after = state + .captured + .lock() + .expect("test: requests mutex poisoned") + .requests + .len(); + assert_eq!( + captured_after, 1, + "request sent after Drop must not be captured" + ); + } + #[tokio::test] async fn stats_endpoint_decodes_msgpack_and_gzip() { let intake = start_default().await; From 4c2ee5c112f80f927893412c8562f1edae01fb35 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 00:39:51 -0400 Subject: [PATCH 06/32] Assert no raw stats requests in no-stats integration test --- crates/datadog-trace-agent/tests/integration_test.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 0f5722ef..08cbdcec 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -197,11 +197,11 @@ pub async fn verify_stats_request(mock_intake: &MockIntake) { /// Helper to verify stats request was not sent to mock intake pub fn verify_no_stats_request(mock_intake: &MockIntake) { - let stats_payloads = mock_intake.stats_payloads(); + let stats_requests = mock_intake.requests_for_path("/api/v0.2/stats"); assert!( - stats_payloads.is_empty(), + stats_requests.is_empty(), "Expected no stats request to mock intake, received {} request(s)", - stats_payloads.len() + stats_requests.len() ); } From c89bc9228f6ca0e4a87b5c9c59a12efc8e32a3a9 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 11:03:27 -0400 Subject: [PATCH 07/32] Keep mock intake serving after accept errors --- crates/datadog-mock-intake/src/lib.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 4ba0f4e4..6aba7cbc 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -256,8 +256,13 @@ impl MockIntake { let (stream, _) = match result { Ok(conn) => conn, Err(e) => { + // Accept errors (e.g. EMFILE, ECONNABORTED) are + // usually transient; keep serving instead of + // silently stopping the intake. Back off briefly + // so a persistent error does not spin the loop. eprintln!("mock_intake: accept error: {e}"); - break; + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + continue; } }; From b311294150653291f7b073d6f2453611a8539f47 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 11:05:58 -0400 Subject: [PATCH 08/32] Record raw and typed intake captures atomically --- crates/datadog-mock-intake/src/lib.rs | 75 +++++++++++++++------------ 1 file changed, 43 insertions(+), 32 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 6aba7cbc..13c39a68 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -442,21 +442,38 @@ async fn handle_request( .collect(), body: body.to_vec(), }; - state - .captured - .lock() - .expect("mock_intake: requests mutex poisoned") - .requests - .push(captured); - - let status = match endpoint { + let (status, accepted) = match endpoint { Endpoint::Stats => handle_stats(&state, &headers, &body), Endpoint::Traces => handle_traces(&state, &headers, &body), Endpoint::PipelineStats => handle_pipeline_stats(&state, &headers, &body), }; + + // Record the raw request and its typed payload in one critical section so + // readers never observe a raw capture whose typed capture is still missing. + { + let mut guard = state + .captured + .lock() + .expect("mock_intake: requests mutex poisoned"); + guard.requests.push(captured); + match accepted { + Some(Accepted::Stats(payload)) => guard.stats.push(payload), + Some(Accepted::Traces(payload)) => guard.traces.push(payload), + Some(Accepted::PipelineStats(payload)) => guard.pipeline_stats.push(payload), + None => {} + } + } Ok(response(status)) } +/// A decoded payload the intake accepted and should expose through its typed +/// query methods. +enum Accepted { + Stats(pb::StatsPayload), + Traces(pb::AgentPayload), + PipelineStats(PipelineStatsPayload), +} + /// Result of handling one intake request, before summary and dump emission. struct HandledRequest { request_id: u64, @@ -468,7 +485,7 @@ fn handle_stats( state: &std::sync::Arc, headers: &hyper::HeaderMap, body: &Bytes, -) -> StatusCode { +) -> (StatusCode, Option) { let request_id = next_request_id(state); let attempt = state.stats_attempts.fetch_add(1, Ordering::SeqCst) + 1; let inject_failure = attempt <= state.options.fail_stats_first_n as u64; @@ -479,12 +496,6 @@ fn handle_stats( let status = if inject_failure { StatusCode::INTERNAL_SERVER_ERROR } else { - state - .captured - .lock() - .expect("mock_intake: stats mutex poisoned") - .stats - .push(payload.clone()); StatusCode::ACCEPTED }; HandledRequest { @@ -539,24 +550,22 @@ fn handle_stats( ); } - handled.status + let accepted = (handled.status == StatusCode::ACCEPTED) + .then_some(handled.decoded) + .flatten() + .map(Accepted::Stats); + (handled.status, accepted) } fn handle_traces( state: &std::sync::Arc, headers: &hyper::HeaderMap, body: &Bytes, -) -> StatusCode { +) -> (StatusCode, Option) { let request_id = next_request_id(state); let handled: HandledRequest = match decompress(headers, body) { Ok(d) => match pb::AgentPayload::decode(d.as_slice()) { Ok(payload) => { - state - .captured - .lock() - .expect("mock_intake: traces mutex poisoned") - .traces - .push(payload.clone()); HandledRequest { request_id, status: StatusCode::ACCEPTED, @@ -612,24 +621,22 @@ fn handle_traces( ); } - handled.status + let accepted = (handled.status == StatusCode::ACCEPTED) + .then_some(handled.decoded) + .flatten() + .map(Accepted::Traces); + (handled.status, accepted) } fn handle_pipeline_stats( state: &std::sync::Arc, headers: &hyper::HeaderMap, body: &Bytes, -) -> StatusCode { +) -> (StatusCode, Option) { let request_id = next_request_id(state); let handled: HandledRequest = match decompress(headers, body) { Ok(d) => match rmp_serde::from_slice::(&d) { Ok(payload) => { - state - .captured - .lock() - .expect("mock_intake: pipeline_stats mutex poisoned") - .pipeline_stats - .push(payload.clone()); HandledRequest { request_id, status: StatusCode::ACCEPTED, @@ -679,7 +686,11 @@ fn handle_pipeline_stats( ); } - handled.status + let accepted = (handled.status == StatusCode::ACCEPTED) + .then_some(handled.decoded) + .flatten() + .map(Accepted::PipelineStats); + (handled.status, accepted) } /// Status used for a stats request that failed to decode or was rejected by From 0a66d33f4fbeadb6a778182c3c22040d262d31db Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 11:07:16 -0400 Subject: [PATCH 09/32] Decode intake requests off the async runtime --- crates/datadog-mock-intake/src/lib.rs | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 13c39a68..1e549a1e 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -442,10 +442,22 @@ async fn handle_request( .collect(), body: body.to_vec(), }; - let (status, accepted) = match endpoint { - Endpoint::Stats => handle_stats(&state, &headers, &body), - Endpoint::Traces => handle_traces(&state, &headers, &body), - Endpoint::PipelineStats => handle_pipeline_stats(&state, &headers, &body), + // Decompression (up to `MAX_DECOMPRESSED_SIZE`) and optional JSON dumps + // are blocking work; run them off the async worker so a large payload + // cannot stall a single-threaded test runtime. + let handler_state = std::sync::Arc::clone(&state); + let handled = tokio::task::spawn_blocking(move || match endpoint { + Endpoint::Stats => handle_stats(&handler_state, &headers, &body), + Endpoint::Traces => handle_traces(&handler_state, &headers, &body), + Endpoint::PipelineStats => handle_pipeline_stats(&handler_state, &headers, &body), + }) + .await; + let (status, accepted) = match handled { + Ok(result) => result, + Err(err) => { + eprintln!("mock_intake: request handler failed: {err}"); + (StatusCode::INTERNAL_SERVER_ERROR, None) + } }; // Record the raw request and its typed payload in one critical section so From 2d9882ab94004b4c245073028d6caedd1d397f33 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 11:12:06 -0400 Subject: [PATCH 10/32] Keep non-UTF-8 header values in captured requests --- crates/datadog-mock-intake/src/lib.rs | 58 ++++++++++++++++++++++++++- 1 file changed, 56 insertions(+), 2 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 1e549a1e..aa7ef9a5 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -145,7 +145,9 @@ pub enum MockIntakeError { /// A raw request captured by the intake: method, path, headers in arrival /// order, and the original wire body (pre-decompression). Recorded for every /// completed POST attempt to a supported endpoint, including attempts that -/// were rejected by failure injection or failed to decode. +/// were rejected by failure injection or failed to decode. Header values +/// that are not valid UTF-8 are converted lossily (invalid bytes become +/// U+FFFD) rather than dropped. #[derive(Clone, Debug, PartialEq, Eq)] pub struct CapturedRequest { pub method: String, @@ -438,7 +440,12 @@ async fn handle_request( path: path.clone(), headers: headers .iter() - .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("").to_string())) + .map(|(k, v)| { + ( + k.as_str().to_string(), + String::from_utf8_lossy(v.as_bytes()).into_owned(), + ) + }) .collect(), body: body.to_vec(), }; @@ -1647,6 +1654,53 @@ mod tests { ); } + #[tokio::test] + async fn non_utf8_header_values_are_captured_lossily() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let intake = start_default().await; + let addr = intake + .base_url() + .trim_start_matches("http://") + .to_string(); + + // Raw TCP so the header value can carry a non-UTF-8 (obs-text) byte. + let mut stream = tokio::net::TcpStream::connect(&addr) + .await + .expect("test: connect to mock intake failed"); + let body = b"not-msgpack"; + let mut request = format!( + "POST /api/v0.2/stats HTTP/1.1\r\nHost: {addr}\r\nConnection: close\r\nContent-Length: {}\r\nx-test: a", + body.len() + ) + .into_bytes(); + request.push(0xFF); + request.extend_from_slice(b"b\r\n\r\n"); + request.extend_from_slice(body); + stream + .write_all(&request) + .await + .expect("test: write request failed"); + let mut response = Vec::new(); + stream + .read_to_end(&mut response) + .await + .expect("test: read response failed"); + assert!( + response.starts_with(b"HTTP/1.1 400"), + "test: expected a 400 for the undecodable body, got {response:?}" + ); + + let requests = intake.requests_for_path("/api/v0.2/stats"); + assert_eq!(requests.len(), 1); + let value = requests[0] + .headers + .iter() + .find(|(k, _)| k == "x-test") + .map(|(_, v)| v.as_str()); + assert_eq!(value, Some("a\u{FFFD}b")); + } + #[tokio::test] async fn unknown_path_returns_404_and_is_not_captured() { let intake = start_default().await; From 04ac689caea3e5c1ca33d2313248d2425a116182 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 11:12:42 -0400 Subject: [PATCH 11/32] Name the single capture mutex in panic messages --- crates/datadog-mock-intake/src/lib.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index aa7ef9a5..fc583700 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -334,7 +334,7 @@ impl MockIntake { self.state .captured .lock() - .expect("mock_intake: requests mutex poisoned") + .expect("mock_intake: captured mutex poisoned") .requests .iter() .filter(|req| req.path == path) @@ -348,7 +348,7 @@ impl MockIntake { self.state .captured .lock() - .expect("mock_intake: pipeline_stats mutex poisoned") + .expect("mock_intake: captured mutex poisoned") .pipeline_stats .clone() } @@ -360,7 +360,7 @@ impl MockIntake { self.state .captured .lock() - .expect("mock_intake: stats mutex poisoned") + .expect("mock_intake: captured mutex poisoned") .stats .clone() } @@ -371,7 +371,7 @@ impl MockIntake { self.state .captured .lock() - .expect("mock_intake: traces mutex poisoned") + .expect("mock_intake: captured mutex poisoned") .traces .clone() } @@ -473,7 +473,7 @@ async fn handle_request( let mut guard = state .captured .lock() - .expect("mock_intake: requests mutex poisoned"); + .expect("mock_intake: captured mutex poisoned"); guard.requests.push(captured); match accepted { Some(Accepted::Stats(payload)) => guard.stats.push(payload), @@ -1351,7 +1351,7 @@ mod tests { let captured_before = state .captured .lock() - .expect("test: requests mutex poisoned") + .expect("test: captured mutex poisoned") .requests .len(); assert_eq!(captured_before, 1, "test: first request should be captured"); @@ -1373,7 +1373,7 @@ mod tests { let captured_after = state .captured .lock() - .expect("test: requests mutex poisoned") + .expect("test: captured mutex poisoned") .requests .len(); assert_eq!( From 40f7c1afaf57dd6d5462bb7981033c6209cb0421 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 11:14:50 -0400 Subject: [PATCH 12/32] Describe only this crate's API in mock-intake docs --- crates/datadog-mock-intake/src/lib.rs | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index fc583700..27766f39 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -19,13 +19,10 @@ //! [`CapturedRequest`], queryable via [`MockIntake::requests_for_path`]. //! Typed query methods return only accepted, successfully decoded payloads. //! -//! `MockIntake::start()` binds `127.0.0.1:0` with defaults and is used by the -//! trace-agent integration tests in this workspace and by the APM / DSM -//! integration tests in `bottlecap`. The standalone `mock-intake` binary in -//! the `datadog-lambda-extension` repository wraps this crate with request -//! summaries, optional stats failure injection -//! (`MOCK_INTAKE_FAIL_STATS_FIRST_N`), and optional JSON dumps -//! (`MOCK_INTAKE_DUMP_DIR`) for local debugging against a live tracer. +//! [`MockIntake::start`] binds `127.0.0.1:0` with all diagnostics disabled. +//! [`MockIntake::start_with_options`] accepts [`MockIntakeOptions`] to pick a +//! port and enable per-request summaries, stats failure injection, and JSON +//! dumps of decoded payloads. //! //! This crate is self-contained (no host config, logging framework, or trace //! processing dependencies) so it stays usable from multiple repositories. From 9c24dd0fe1ea6dec6573e61a0e599e094070a862 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 11:23:55 -0400 Subject: [PATCH 13/32] Apply rustfmt to mock-intake --- crates/datadog-mock-intake/src/lib.rs | 37 ++++++++++----------------- 1 file changed, 13 insertions(+), 24 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 27766f39..6c8f3c4b 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -243,8 +243,7 @@ impl MockIntake { }); let (shutdown_tx, mut shutdown_rx) = oneshot::channel::<()>(); - let connections: std::sync::Arc>>> = - std::sync::Arc::default(); + let connections: std::sync::Arc>>> = std::sync::Arc::default(); let task_state = std::sync::Arc::clone(&state); let task_connections = std::sync::Arc::clone(&connections); let task = tokio::spawn(async move { @@ -581,13 +580,11 @@ fn handle_traces( let request_id = next_request_id(state); let handled: HandledRequest = match decompress(headers, body) { Ok(d) => match pb::AgentPayload::decode(d.as_slice()) { - Ok(payload) => { - HandledRequest { - request_id, - status: StatusCode::ACCEPTED, - decoded: Some(payload), - } - } + Ok(payload) => HandledRequest { + request_id, + status: StatusCode::ACCEPTED, + decoded: Some(payload), + }, Err(err) => { eprintln!("mock_intake: failed to decode AgentPayload protobuf: {err}"); HandledRequest { @@ -652,13 +649,11 @@ fn handle_pipeline_stats( let request_id = next_request_id(state); let handled: HandledRequest = match decompress(headers, body) { Ok(d) => match rmp_serde::from_slice::(&d) { - Ok(payload) => { - HandledRequest { - request_id, - status: StatusCode::ACCEPTED, - decoded: Some(payload), - } - } + Ok(payload) => HandledRequest { + request_id, + status: StatusCode::ACCEPTED, + decoded: Some(payload), + }, Err(err) => { eprintln!("mock_intake: failed to decode pipeline stats msgpack: {err}"); HandledRequest { @@ -1302,10 +1297,7 @@ mod tests { let intake = start_default().await; let state = std::sync::Arc::clone(&intake.state); - let addr = intake - .base_url() - .trim_start_matches("http://") - .to_string(); + let addr = intake.base_url().trim_start_matches("http://").to_string(); // Raw TCP so the connection stays open (keep-alive) across Drop. let mut stream = tokio::net::TcpStream::connect(&addr) @@ -1656,10 +1648,7 @@ mod tests { use tokio::io::{AsyncReadExt, AsyncWriteExt}; let intake = start_default().await; - let addr = intake - .base_url() - .trim_start_matches("http://") - .to_string(); + let addr = intake.base_url().trim_start_matches("http://").to_string(); // Raw TCP so the header value can carry a non-UTF-8 (obs-text) byte. let mut stream = tokio::net::TcpStream::connect(&addr) From 664fb6cb653ec32fab67ed4094185c46b86a43af Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 11:33:20 -0400 Subject: [PATCH 14/32] Abort connections racing with MockIntake shutdown --- crates/datadog-mock-intake/src/lib.rs | 31 +++++++++++++++++++++------ 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 6c8f3c4b..1eb77745 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -200,7 +200,17 @@ pub struct MockIntake { state: std::sync::Arc, shutdown_tx: Option>, task: Option>, - connections: std::sync::Arc>>>, + connections: std::sync::Arc>, +} + +/// Tracks live connection tasks plus whether shutdown has begun. Keeping the +/// flag in the same mutex as the handles makes registration and shutdown +/// mutually exclusive: a connection spawned during `Drop` is aborted at +/// registration time instead of being parked in a registry nobody drains. +#[derive(Debug, Default)] +struct ConnectionRegistry { + shutting_down: bool, + handles: Vec>, } impl MockIntake { @@ -243,7 +253,7 @@ impl MockIntake { }); let (shutdown_tx, mut shutdown_rx) = oneshot::channel::<()>(); - let connections: std::sync::Arc>>> = std::sync::Arc::default(); + let connections: std::sync::Arc> = std::sync::Arc::default(); let task_state = std::sync::Arc::clone(&state); let task_connections = std::sync::Arc::clone(&connections); let task = tokio::spawn(async move { @@ -276,8 +286,14 @@ impl MockIntake { .await; }); if let Ok(mut conns) = task_connections.lock() { - conns.retain(|h| !h.is_finished()); - conns.push(handle); + if conns.shutting_down { + // Drop already ran; abort this connection now + // so it cannot keep serving after shutdown. + handle.abort(); + } else { + conns.handles.retain(|h| !h.is_finished()); + conns.handles.push(handle); + } } } _ = &mut shutdown_rx => { @@ -382,9 +398,12 @@ impl Drop for MockIntake { task.abort(); } // Abort any accepted-but-still-running connection tasks so a - // keep-alive client cannot keep sending requests after Drop. + // keep-alive client cannot keep sending requests after Drop. The + // flag is set in the same lock as the drain, so a connection whose + // registration races with Drop is aborted at registration time. if let Ok(mut conns) = self.connections.lock() { - for handle in conns.drain(..) { + conns.shutting_down = true; + for handle in conns.handles.drain(..) { handle.abort(); } } From f4441d7d7b8e7e4b4d003bc8ee676472e10f5358 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 13:14:14 -0400 Subject: [PATCH 15/32] fix(mock-intake): reject unsupported content-encoding with 415 Unknown Content-Encoding values were silently decoded as identity, masking mismatches with the real intake, which rejects unsupported encodings. Rejected attempts are still raw-captured but do not consume the stats failure-injection budget; explicit identity encoding is still accepted. --- crates/datadog-mock-intake/src/lib.rs | 267 +++++++++++++++++++++----- 1 file changed, 215 insertions(+), 52 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 1eb77745..56e4f4dd 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -142,9 +142,9 @@ pub enum MockIntakeError { /// A raw request captured by the intake: method, path, headers in arrival /// order, and the original wire body (pre-decompression). Recorded for every /// completed POST attempt to a supported endpoint, including attempts that -/// were rejected by failure injection or failed to decode. Header values -/// that are not valid UTF-8 are converted lossily (invalid bytes become -/// U+FFFD) rather than dropped. +/// were rejected by failure injection, failed to decode, or used an +/// unsupported `Content-Encoding`. Header values that are not valid UTF-8 are +/// converted lossily (invalid bytes become U+FFFD) rather than dropped. #[derive(Clone, Debug, PartialEq, Eq)] pub struct CapturedRequest { pub method: String, @@ -338,9 +338,9 @@ impl MockIntake { } /// Raw requests captured for `path`, in arrival order. Includes attempts - /// that were rejected by failure injection or failed to decode; excludes - /// requests rejected at the transport level (unknown path, unsupported - /// method, or oversized body). + /// that were rejected by failure injection, failed to decode, or used an + /// unsupported `Content-Encoding`; excludes requests rejected at the + /// transport level (unknown path, unsupported method, or oversized body). #[must_use] pub fn requests_for_path(&self, path: &str) -> Vec { self.state @@ -521,40 +521,55 @@ fn handle_stats( body: &Bytes, ) -> (StatusCode, Option) { let request_id = next_request_id(state); - let attempt = state.stats_attempts.fetch_add(1, Ordering::SeqCst) + 1; - let inject_failure = attempt <= state.options.fail_stats_first_n as u64; - let handled: HandledRequest = match decompress(headers, body) { - Ok(d) => match rmp_serde::from_slice::(&d) { - Ok(payload) => { - let status = if inject_failure { - StatusCode::INTERNAL_SERVER_ERROR - } else { - StatusCode::ACCEPTED - }; - HandledRequest { - request_id, - status, - decoded: Some(payload), - } - } - Err(err) => { - eprintln!("mock_intake: failed to decode StatsPayload msgpack: {err}"); - HandledRequest { - request_id, - status: failure_status(inject_failure), - decoded: None, - } - } - }, - Err(e) => { + // An unsupported Content-Encoding is a client error (415). Like the + // transport-level rejections (404/405/413), it does not consume the + // failure-injection budget: the budget counts attempts that reach + // payload decoding. + Err(e @ DecompressError::UnsupportedEncoding(_)) => { eprintln!("{e}"); HandledRequest { request_id, - status: failure_status(inject_failure), + status: StatusCode::UNSUPPORTED_MEDIA_TYPE, decoded: None, } } + result => { + let attempt = state.stats_attempts.fetch_add(1, Ordering::SeqCst) + 1; + let inject_failure = attempt <= state.options.fail_stats_first_n as u64; + match result { + Ok(d) => match rmp_serde::from_slice::(&d) { + Ok(payload) => { + let status = if inject_failure { + StatusCode::INTERNAL_SERVER_ERROR + } else { + StatusCode::ACCEPTED + }; + HandledRequest { + request_id, + status, + decoded: Some(payload), + } + } + Err(err) => { + eprintln!("mock_intake: failed to decode StatsPayload msgpack: {err}"); + HandledRequest { + request_id, + status: failure_status(inject_failure), + decoded: None, + } + } + }, + Err(e) => { + eprintln!("{e}"); + HandledRequest { + request_id, + status: failure_status(inject_failure), + decoded: None, + } + } + } + } }; let payload_count = handled.decoded.as_ref().map_or(0, |p| p.stats.len()); @@ -617,7 +632,7 @@ fn handle_traces( eprintln!("{e}"); HandledRequest { request_id, - status: StatusCode::BAD_REQUEST, + status: decompress_failure_status(&e, false), decoded: None, } } @@ -686,7 +701,7 @@ fn handle_pipeline_stats( eprintln!("{e}"); HandledRequest { request_id, - status: StatusCode::BAD_REQUEST, + status: decompress_failure_status(&e, false), decoded: None, } } @@ -738,6 +753,19 @@ fn next_request_id(state: &SharedState) -> u64 { state.request_counter.fetch_add(1, Ordering::SeqCst) + 1 } +/// Status for a request whose body could not be decompressed. An unsupported +/// `Content-Encoding` is a client error (`415 Unsupported Media Type`, the +/// same status the real intake stack uses for unsupported media types) +/// regardless of stats failure injection; an actual decompression failure +/// keeps the endpoint's historical status (`400`, or `500` inside the stats +/// injection window). +fn decompress_failure_status(err: &DecompressError, inject_failure: bool) -> StatusCode { + match err { + DecompressError::UnsupportedEncoding(_) => StatusCode::UNSUPPORTED_MEDIA_TYPE, + DecompressError::Decode(_) => failure_status(inject_failure), + } +} + fn content_encoding(headers: &hyper::HeaderMap) -> String { headers .get("content-encoding") @@ -1066,10 +1094,34 @@ fn stats_hits_by_key(payload: &pb::StatsPayload) -> BTreeMap // Body decoding // --------------------------------------------------------------------------- +/// An error from decoding a request body's `Content-Encoding`. +#[derive(Debug)] +enum DecompressError { + /// The `Content-Encoding` header names an encoding the intake does not + /// support. Rejected with `415`, matching the intake's treatment of + /// unsupported media types, so a test cannot silently pass against the + /// mock while the real intake would reject the payload. + UnsupportedEncoding(String), + /// The body could not be decompressed with its declared encoding. + Decode(String), +} + +impl std::fmt::Display for DecompressError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::UnsupportedEncoding(encoding) => { + write!(f, "mock_intake: unsupported Content-Encoding '{encoding}'") + } + Self::Decode(message) => write!(f, "{message}"), + } + } +} + /// Decompress a request body based on its `Content-Encoding` header. -/// Supports `gzip` and `zstd`. An unknown or absent encoding is treated as -/// identity: the body is returned unchanged. -fn decompress(headers: &hyper::HeaderMap, body: &Bytes) -> Result, String> { +/// Supports `gzip` and `zstd`; an absent or `identity` encoding returns the +/// body unchanged. Any other encoding is rejected as unsupported rather than +/// silently decoded as identity. +fn decompress(headers: &hyper::HeaderMap, body: &Bytes) -> Result, DecompressError> { let encoding = headers .get("content-encoding") .and_then(|v| v.to_str().ok()) @@ -1077,25 +1129,22 @@ fn decompress(headers: &hyper::HeaderMap, body: &Bytes) -> Result, Strin .to_ascii_lowercase(); match encoding.as_str() { + "" | "identity" => Ok(body.to_vec()), "gzip" => { let decoder = flate2::read::GzDecoder::new(body.as_ref()); - read_capped(decoder, MAX_DECOMPRESSED_SIZE) - .map_err(|e| format!("mock_intake: gzip decode failed: {e}")) + read_capped(decoder, MAX_DECOMPRESSED_SIZE).map_err(|e| { + DecompressError::Decode(format!("mock_intake: gzip decode failed: {e}")) + }) } "zstd" => { - let decoder = zstd::stream::read::Decoder::new(body.as_ref()) - .map_err(|e| format!("mock_intake: zstd decoder init failed: {e}"))?; - read_capped(decoder, MAX_DECOMPRESSED_SIZE) - .map_err(|e| format!("mock_intake: zstd decode failed: {e}")) - } - _ => { - if !encoding.is_empty() { - eprintln!( - "mock_intake: unrecognized Content-Encoding '{encoding}', treating as identity" - ); - } - Ok(body.to_vec()) + let decoder = zstd::stream::read::Decoder::new(body.as_ref()).map_err(|e| { + DecompressError::Decode(format!("mock_intake: zstd decoder init failed: {e}")) + })?; + read_capped(decoder, MAX_DECOMPRESSED_SIZE).map_err(|e| { + DecompressError::Decode(format!("mock_intake: zstd decode failed: {e}")) + }) } + other => Err(DecompressError::UnsupportedEncoding(other.to_string())), } } @@ -1198,6 +1247,14 @@ mod tests { } } + fn trace_payload() -> pb::AgentPayload { + pb::AgentPayload { + env: "local".to_string(), + tracer_payloads: vec![pb::TracerPayload::default()], + ..pb::AgentPayload::default() + } + } + /// Serialize a stats payload whose wire size is exactly `target` bytes by /// padding the client hostname. Above the msgpack str16 boundary the /// encoding grows linearly with the string, so a fixed-point adjustment @@ -1390,6 +1447,112 @@ mod tests { ); } + #[tokio::test] + async fn unsupported_content_encoding_returns_415_and_is_not_typed_captured() { + let intake = start_default().await; + let body = rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + + // An unknown Content-Encoding is rejected like the real intake rejects + // unsupported media types, not silently decoded as identity. + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + Some("deflate"), + body.clone() + ) + .await, + StatusCode::UNSUPPORTED_MEDIA_TYPE + ); + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/traces", + Some("br"), + trace_payload().encode_to_vec() + ) + .await, + StatusCode::UNSUPPORTED_MEDIA_TYPE + ); + + // Not typed-captured: the payload was never accepted. + assert!(intake.stats_payloads().is_empty()); + assert!(intake.trace_payloads().is_empty()); + // But raw-captured as a completed POST attempt. + assert_eq!(intake.requests_for_path("/api/v0.2/stats").len(), 1); + assert_eq!(intake.requests_for_path("/api/v0.2/traces").len(), 1); + + // The explicit identity encoding is still accepted. + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + Some("identity"), + body + ) + .await, + StatusCode::ACCEPTED + ); + assert_eq!(intake.stats_payloads().len(), 1); + } + + #[tokio::test] + async fn corrupted_compressed_bodies_return_400() { + let intake = start_default().await; + // A gzip-declared body that is not valid gzip is a decode failure + // (400), distinct from an unsupported encoding (415). + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + Some("gzip"), + b"not-gzip".to_vec() + ) + .await, + StatusCode::BAD_REQUEST + ); + } + + #[tokio::test] + async fn unsupported_content_encoding_does_not_consume_stats_rejection_budget() { + let intake = MockIntake::start_with_options(MockIntakeOptions { + fail_stats_first_n: 1, + ..MockIntakeOptions::default() + }) + .await + .expect("test: start failed"); + + let body = rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + + assert_eq!( + post( + &intake.base_url(), + "/api/v0.2/stats", + Some("deflate"), + body.clone() + ) + .await, + StatusCode::UNSUPPORTED_MEDIA_TYPE + ); + // The 415 attempt is a client error, not an injected rejection: the + // next valid request is rejected as attempt 1 of 1, proving the + // unsupported-encoding attempt did not consume the budget. + assert_eq!( + post(&intake.base_url(), "/api/v0.2/stats", None, body).await, + StatusCode::INTERNAL_SERVER_ERROR + ); + } + #[tokio::test] async fn stats_endpoint_decodes_msgpack_and_gzip() { let intake = start_default().await; From a7191dd83741184125df3f4cd3056dd9c4493560 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 13:18:39 -0400 Subject: [PATCH 16/32] fix(mock-intake): assign request ids at arrival time Request ids were assigned inside the blocking handlers, so under concurrency dump and summary ids could interleave out of arrival order relative to the captured requests. --- crates/datadog-mock-intake/src/lib.rs | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 56e4f4dd..89a8ee5c 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -464,14 +464,20 @@ async fn handle_request( .collect(), body: body.to_vec(), }; + // Request identity is assigned at arrival time (before the blocking + // handler runs) so summaries and dumps stay in arrival order even when + // requests complete out of order. + let request_id = next_request_id(&state); // Decompression (up to `MAX_DECOMPRESSED_SIZE`) and optional JSON dumps // are blocking work; run them off the async worker so a large payload // cannot stall a single-threaded test runtime. let handler_state = std::sync::Arc::clone(&state); let handled = tokio::task::spawn_blocking(move || match endpoint { - Endpoint::Stats => handle_stats(&handler_state, &headers, &body), - Endpoint::Traces => handle_traces(&handler_state, &headers, &body), - Endpoint::PipelineStats => handle_pipeline_stats(&handler_state, &headers, &body), + Endpoint::Stats => handle_stats(&handler_state, request_id, &headers, &body), + Endpoint::Traces => handle_traces(&handler_state, request_id, &headers, &body), + Endpoint::PipelineStats => { + handle_pipeline_stats(&handler_state, request_id, &headers, &body) + } }) .await; let (status, accepted) = match handled { @@ -517,10 +523,10 @@ struct HandledRequest { fn handle_stats( state: &std::sync::Arc, + request_id: u64, headers: &hyper::HeaderMap, body: &Bytes, ) -> (StatusCode, Option) { - let request_id = next_request_id(state); let handled: HandledRequest = match decompress(headers, body) { // An unsupported Content-Encoding is a client error (415). Like the // transport-level rejections (404/405/413), it does not consume the @@ -608,10 +614,10 @@ fn handle_stats( fn handle_traces( state: &std::sync::Arc, + request_id: u64, headers: &hyper::HeaderMap, body: &Bytes, ) -> (StatusCode, Option) { - let request_id = next_request_id(state); let handled: HandledRequest = match decompress(headers, body) { Ok(d) => match pb::AgentPayload::decode(d.as_slice()) { Ok(payload) => HandledRequest { @@ -677,10 +683,10 @@ fn handle_traces( fn handle_pipeline_stats( state: &std::sync::Arc, + request_id: u64, headers: &hyper::HeaderMap, body: &Bytes, ) -> (StatusCode, Option) { - let request_id = next_request_id(state); let handled: HandledRequest = match decompress(headers, body) { Ok(d) => match rmp_serde::from_slice::(&d) { Ok(payload) => HandledRequest { From eff6ff490f6e053e0dcb3cf0c36f0b57529d8d52 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 13:26:00 -0400 Subject: [PATCH 17/32] fix(mock-intake): count pipeline stats points in summaries Summaries reported a fixed payload count of 1 per pipeline-stats request; report the number of decoded points like the other endpoints. --- crates/datadog-mock-intake/src/lib.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 89a8ee5c..6fc32e9f 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -714,12 +714,16 @@ fn handle_pipeline_stats( }; if state.options.request_summaries { + let payload_count = handled + .decoded + .as_ref() + .map_or(0, |p| p.stats.iter().map(|b| b.stats.len()).sum()); log_summary( handled.request_id, "/api/v0.1/pipeline_stats", headers, handled.status, - handled.decoded.as_ref().map_or(0, |_| 1), + payload_count, None, ); } From c92ba9b4df2ce137f35f7c5a070bda8a9f1a0290 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 13:27:47 -0400 Subject: [PATCH 18/32] test(mock-intake): pin AgentPayload JSON dump fields The dump serializer mirrors pb::AgentPayload fields by hand; an exhaustive struct literal in the test forces it to be updated whenever the protobuf struct gains a field, instead of silently dropping it from dumps. --- crates/datadog-mock-intake/src/lib.rs | 49 ++++++++++++++++++++++++++- 1 file changed, 48 insertions(+), 1 deletion(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 6fc32e9f..2aa02d02 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -1126,7 +1126,6 @@ impl std::fmt::Display for DecompressError { } } } - /// Decompress a request body based on its `Content-Encoding` header. /// Supports `gzip` and `zstd`; an absent or `identity` encoding returns the /// body unchanged. Any other encoding is rejected as unsupported rather than @@ -1563,6 +1562,54 @@ mod tests { ); } + #[test] + fn agent_payload_json_covers_all_proto_fields() { + // Exhaustive struct literal: adding a field to `pb::AgentPayload` + // breaks this test's compilation, which forces `agent_payload_to_json` + // to be updated in the same change instead of silently dropping the + // new field from JSON dumps. + let payload = pb::AgentPayload { + host_name: "host-1".to_string(), + env: "prod".to_string(), + tracer_payloads: vec![pb::TracerPayload::default()], + tags: [("key".to_string(), "value".to_string())] + .into_iter() + .collect(), + agent_version: "1.2.3".to_string(), + target_tps: 10.0, + error_tps: 20.0, + rare_sampler_enabled: true, + idx_tracer_payloads: vec![pb::idx::TracerPayload::default()], + }; + + let json = agent_payload_to_json(&payload) + .as_object() + .expect("test: JSON dumps must be objects") + .clone(); + + let expected: &[(&str, serde_json::Value)] = &[ + ("host_name", "host-1".into()), + ("env", "prod".into()), + ( + "tracer_payloads", + serde_json::json!([pb::TracerPayload::default()]), + ), + ("tags", serde_json::json!({"key": "value"})), + ("agent_version", "1.2.3".into()), + ("target_tps", 10.0.into()), + ("error_tps", 20.0.into()), + ("rare_sampler_enabled", true.into()), + ( + "idx_tracer_payloads", + serde_json::json!([pb::idx::TracerPayload::default()]), + ), + ]; + assert_eq!(json.len(), expected.len(), "JSON dump field set changed"); + for (key, value) in expected { + assert_eq!(json.get(*key), Some(value), "unexpected value for {key}"); + } + } + #[tokio::test] async fn stats_endpoint_decodes_msgpack_and_gzip() { let intake = start_default().await; From c9a0ebdc5befb1117af10da9db5545d3991034b2 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 13:29:03 -0400 Subject: [PATCH 19/32] test(trace-agent): readable failures for empty decoded captures Indexing decoded payload collections panicked with bare out-of-bounds errors; use first() with expects that name the failed assertion. --- .../tests/integration_test.rs | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 08cbdcec..c68f555e 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -233,12 +233,19 @@ pub async fn verify_dsm_request( "Expected DSM payload body to be forwarded unchanged" ); - let decoded = &mock_intake.pipeline_stats_payloads()[0]; - let edge_tags: Vec<&str> = decoded.stats[0].stats[0] - .edge_tags - .iter() - .map(String::as_str) - .collect(); + let payloads = mock_intake.pipeline_stats_payloads(); + let decoded = payloads + .first() + .expect("intake must have decoded at least one DSM payload"); + let bucket = decoded + .stats + .first() + .expect("decoded DSM payload must have at least one bucket"); + let point = bucket + .stats + .first() + .expect("decoded DSM bucket must have at least one point"); + let edge_tags: Vec<&str> = point.edge_tags.iter().map(String::as_str).collect(); assert_eq!( edge_tags, expected_edge_tags, "Expected decoded DSM payload to carry the fixture edge tags" @@ -1094,7 +1101,9 @@ async fn test_peer_tags_in_flushed_stats() { "Expected at least one stats request" ); - let payload = &stats_payloads[0]; + let payload = stats_payloads + .first() + .expect("at least one decoded stats payload"); let all_peer_tags: Vec<&str> = payload .stats .iter() From af94d9852e0588c5d5314b479e6ce19fba8d7829 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 13:59:43 -0400 Subject: [PATCH 20/32] test(trace-agent): import PipelineStats helper types Drop the inline datadog_mock_intake qualifications in favor of the top-level import used for the other helper types. --- crates/datadog-trace-agent/tests/integration_test.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index c68f555e..4c4ea902 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -8,7 +8,9 @@ use common::helpers::{ create_test_trace_payload, create_trace_with_span_kind_children_payload, send_tcp_request, }; use common::mocks::{MockEnvVerifier, MockStatsFlusher, MockStatsProcessor, MockTraceFlusher}; -use datadog_mock_intake::{MockIntake, PipelineStatsPayload}; +use datadog_mock_intake::{ + MockIntake, PipelineStatsBucket, PipelineStatsPayload, PipelineStatsPoint, +}; use datadog_trace_agent::{ config::{Config, Tags, test_helpers::create_tcp_test_config}, mini_agent::MiniAgent, @@ -601,8 +603,8 @@ async fn test_mini_agent_tcp_proxies_dsm_requests() { tracer_version: "1.0".to_string(), version: "2.0".to_string(), tags: Vec::new(), - stats: vec![datadog_mock_intake::PipelineStatsBucket { - stats: vec![datadog_mock_intake::PipelineStatsPoint { + stats: vec![PipelineStatsBucket { + stats: vec![PipelineStatsPoint { hash: 7, parent_hash: 0, edge_tags: vec!["direction:out".to_string(), "type:kafka".to_string()], From d84a05e7ee1fd44848cef3210dcb373e52e5f577 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Tue, 29 Sep 2026 14:11:26 -0400 Subject: [PATCH 21/32] fix(mock-intake): capture concurrent requests in arrival order Captured requests were appended after the blocking handler completed, so under concurrent clients a later-arriving request that decompressed or decoded faster could be recorded before an earlier one, breaking the arrival-order contract of the query APIs. Completed requests are now buffered by request id and flushed as a contiguous prefix in one critical section. --- crates/datadog-mock-intake/src/lib.rs | 117 ++++++++++++++++++++++---- 1 file changed, 101 insertions(+), 16 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 2aa02d02..0b7ccb24 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -153,6 +153,15 @@ pub struct CapturedRequest { pub body: Vec, } +/// A completed request held back until every earlier request has been +/// flushed to the visible capture vectors, so readers only ever observe a +/// contiguous, arrival-ordered prefix of captures. +#[derive(Debug)] +struct PendingCapture { + request: CapturedRequest, + accepted: Option, +} + /// Captured, decoded APM payloads for a single test run. #[derive(Debug, Default)] struct Captured { @@ -160,6 +169,43 @@ struct Captured { stats: Vec, traces: Vec, pipeline_stats: Vec, + /// Requests that finished handling out of arrival order and are waiting + /// for earlier requests before they can be exposed to readers. + pending: BTreeMap, + /// The request id the next flush expects, so the visible vectors stay in + /// arrival order even when concurrent requests complete out of order. + next_flush: u64, +} + +impl Captured { + fn new() -> Self { + // Request ids start at 1 (`next_request_id` adds 1 after the fetch), + // so the first flush expects id 1. + Self { + next_flush: 1, + ..Default::default() + } + } + + /// Record a completed request. Out-of-order completions are buffered + /// until all earlier requests arrive, then a contiguous prefix is + /// appended to the visible vectors in one pass so readers never observe + /// a raw capture whose typed capture is missing, a gap, or a later + /// request before an earlier one. + fn record(&mut self, request_id: u64, request: CapturedRequest, accepted: Option) { + self.pending + .insert(request_id, PendingCapture { request, accepted }); + while let Some(pending) = self.pending.remove(&self.next_flush) { + self.requests.push(pending.request); + match pending.accepted { + Some(Accepted::Stats(payload)) => self.stats.push(payload), + Some(Accepted::Traces(payload)) => self.traces.push(payload), + Some(Accepted::PipelineStats(payload)) => self.pipeline_stats.push(payload), + None => {} + } + self.next_flush += 1; + } + } } /// Shared server state. The request handler writes to the mutex; callers read @@ -246,7 +292,7 @@ impl MockIntake { let base_url = format!("http://{addr}"); let state = std::sync::Arc::new(SharedState { - captured: Mutex::new(Captured::default()), + captured: Mutex::new(Captured::new()), options, request_counter: AtomicU64::new(0), stats_attempts: AtomicU64::new(0), @@ -488,26 +534,21 @@ async fn handle_request( } }; - // Record the raw request and its typed payload in one critical section so - // readers never observe a raw capture whose typed capture is still missing. - { - let mut guard = state - .captured - .lock() - .expect("mock_intake: captured mutex poisoned"); - guard.requests.push(captured); - match accepted { - Some(Accepted::Stats(payload)) => guard.stats.push(payload), - Some(Accepted::Traces(payload)) => guard.traces.push(payload), - Some(Accepted::PipelineStats(payload)) => guard.pipeline_stats.push(payload), - None => {} - } - } + // Record the raw request and its typed payload in one critical section. + // Out-of-order completions are buffered so readers only ever observe a + // contiguous, arrival-ordered prefix: no raw capture without its typed + // payload, no gaps, and no later request before an earlier one. + state + .captured + .lock() + .expect("mock_intake: captured mutex poisoned") + .record(request_id, captured, accepted); Ok(response(status)) } /// A decoded payload the intake accepted and should expose through its typed /// query methods. +#[derive(Debug)] enum Accepted { Stats(pb::StatsPayload), Traces(pb::AgentPayload), @@ -2502,4 +2543,48 @@ mod tests { .expect("test: start failed"); assert_eq!(intake.base_url(), format!("http://127.0.0.1:{port}")); } + + fn captured_request(body: &[u8]) -> CapturedRequest { + CapturedRequest { + method: "POST".to_string(), + path: "/api/v0.2/stats".to_string(), + headers: Vec::new(), + body: body.to_vec(), + } + } + + #[test] + fn record_flushes_captures_in_arrival_order() { + let mut captured = Captured::new(); + + // Request 2 completes before request 1: nothing becomes visible yet, + // because the visible vectors must stay in arrival order. + captured.record(2, captured_request(b"two"), None); + assert!(captured.requests.is_empty()); + + // A later request completing early is also held back. + captured.record(4, captured_request(b"four"), None); + assert!(captured.requests.is_empty()); + + // Request 1 completes: the contiguous prefix (1, 2) flushes in + // arrival order; request 4 stays buffered. + captured.record(1, captured_request(b"one"), None); + let bodies: Vec<_> = captured.requests.iter().map(|r| r.body.clone()).collect(); + assert_eq!(bodies, vec![b"one".to_vec(), b"two".to_vec()]); + + // Request 3 completes: the remaining prefix (2, 3, 4) flushes in + // arrival order, including the out-of-order request 4. + captured.record(3, captured_request(b"three"), None); + let bodies: Vec<_> = captured.requests.iter().map(|r| r.body.clone()).collect(); + assert_eq!( + bodies, + vec![ + b"one".to_vec(), + b"two".to_vec(), + b"three".to_vec(), + b"four".to_vec() + ] + ); + assert!(captured.pending.is_empty()); + } } From d8fd44143f629901ffec846b1a0b48b95aaa1321 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 13:00:29 -0400 Subject: [PATCH 22/32] fix(mock-intake): don't stall on client disconnect A request canceled by a client disconnect could leave its request id unrecorded, holding back every later capture from the query APIs. Captures are now recorded inside the blocking handler task, which always runs to completion, with handler panics recorded as rejected requests. --- crates/datadog-mock-intake/src/lib.rs | 60 +++++++++++++++++---------- 1 file changed, 37 insertions(+), 23 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 0b7ccb24..6728aca0 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -517,32 +517,46 @@ async fn handle_request( // Decompression (up to `MAX_DECOMPRESSED_SIZE`) and optional JSON dumps // are blocking work; run them off the async worker so a large payload // cannot stall a single-threaded test runtime. + // + // The capture is recorded inside the blocking task rather than after + // awaiting it: Hyper drops this future if the client disconnects, but a + // blocking task always runs to completion. Recording here guarantees every + // allocated request id is recorded, so a canceled request cannot leave a + // gap that holds back all later captures. let handler_state = std::sync::Arc::clone(&state); - let handled = tokio::task::spawn_blocking(move || match endpoint { - Endpoint::Stats => handle_stats(&handler_state, request_id, &headers, &body), - Endpoint::Traces => handle_traces(&handler_state, request_id, &headers, &body), - Endpoint::PipelineStats => { - handle_pipeline_stats(&handler_state, request_id, &headers, &body) - } + let handled = tokio::task::spawn_blocking(move || { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| match endpoint { + Endpoint::Stats => handle_stats(&handler_state, request_id, &headers, &body), + Endpoint::Traces => handle_traces(&handler_state, request_id, &headers, &body), + Endpoint::PipelineStats => { + handle_pipeline_stats(&handler_state, request_id, &headers, &body) + } + })); + let (status, accepted) = match result { + Ok(result) => result, + Err(_) => { + eprintln!("mock_intake: request handler panicked for request {request_id}"); + (StatusCode::INTERNAL_SERVER_ERROR, None) + } + }; + + // Record the raw request and its typed payload in one critical + // section. Out-of-order completions are buffered so readers only ever + // observe a contiguous, arrival-ordered prefix: no raw capture without + // its typed payload, no gaps, and no later request before an earlier + // one. + handler_state + .captured + .lock() + .expect("mock_intake: captured mutex poisoned") + .record(request_id, captured, accepted); + status }) .await; - let (status, accepted) = match handled { - Ok(result) => result, - Err(err) => { - eprintln!("mock_intake: request handler failed: {err}"); - (StatusCode::INTERNAL_SERVER_ERROR, None) - } - }; - - // Record the raw request and its typed payload in one critical section. - // Out-of-order completions are buffered so readers only ever observe a - // contiguous, arrival-ordered prefix: no raw capture without its typed - // payload, no gaps, and no later request before an earlier one. - state - .captured - .lock() - .expect("mock_intake: captured mutex poisoned") - .record(request_id, captured, accepted); + let status = handled.unwrap_or_else(|err| { + eprintln!("mock_intake: request handler failed: {err}"); + StatusCode::INTERNAL_SERVER_ERROR + }); Ok(response(status)) } From b882fe5739bb0539a0c9eb03a7a0fcae769f8a68 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 13:01:36 -0400 Subject: [PATCH 23/32] fix(mock-intake): log connection errors Connection failures such as malformed HTTP or requests cut off mid-message were silently discarded, making a test timeout look like a missing payload. They are now logged to stderr. --- crates/datadog-mock-intake/src/lib.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 6728aca0..1a790c23 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -327,9 +327,15 @@ impl MockIntake { let state = std::sync::Arc::clone(&state); async move { handle_request(state, req).await } }); - let _ = hyper::server::conn::http1::Builder::new() + // Log connection failures (malformed HTTP, resets, + // messages cut off mid-request) so a test timeout is + // not mistaken for a payload that was never sent. + if let Err(e) = hyper::server::conn::http1::Builder::new() .serve_connection(io, service) - .await; + .await + { + eprintln!("mock_intake: connection error: {e}"); + } }); if let Ok(mut conns) = task_connections.lock() { if conns.shutting_down { From e70d1acd01e1b10ba818e5c2e87bfdc22212bce5 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 13:20:36 -0400 Subject: [PATCH 24/32] test(mock-intake): accept closed socket after Drop The keep-alive shutdown test required writes after Drop to succeed, but the server may already have closed the socket. Broken pipe and reset errors are now accepted as evidence of shutdown; the capture-count assertion is unchanged. --- crates/datadog-mock-intake/src/lib.rs | 29 ++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 1a790c23..80756227 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -1494,15 +1494,26 @@ mod tests { drop(intake); // Reuse the still-open keep-alive connection. The connection task - // was aborted on Drop, so this request must not be captured. - stream - .write_all(request.as_bytes()) - .await - .expect("test: write second request failed"); - stream - .write_all(body.as_slice()) - .await - .expect("test: write second body failed"); + // was aborted on Drop, so this request must not be captured. The + // server may already have closed its side of the socket, in which + // case a write fails with a reset or broken pipe: that is also valid + // evidence of shutdown. Any other write error is a test failure. + let write_result = async { + stream.write_all(request.as_bytes()).await?; + stream.write_all(body.as_slice()).await + } + .await; + if let Err(e) = write_result { + assert!( + matches!( + e.kind(), + std::io::ErrorKind::BrokenPipe + | std::io::ErrorKind::ConnectionReset + | std::io::ErrorKind::ConnectionAborted + ), + "test: unexpected error writing after Drop: {e}" + ); + } tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; let captured_after = state From a4d92b9f52852f81a473de51556a8fa97f14bf9b Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 13:29:59 -0400 Subject: [PATCH 25/32] fix(mock-intake): order captures by request start Request ids were assigned after the body was fully read, so a slow upload could be ordered after a later request. Ids are now assigned when the request head arrives. Requests that end without a capture (oversized or unreadable body, client disconnect) release their id so they cannot hold back later captures. --- crates/datadog-mock-intake/src/lib.rs | 201 ++++++++++++++++++++++---- 1 file changed, 176 insertions(+), 25 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 80756227..d5b4a34d 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -170,8 +170,11 @@ struct Captured { traces: Vec, pipeline_stats: Vec, /// Requests that finished handling out of arrival order and are waiting - /// for earlier requests before they can be exposed to readers. - pending: BTreeMap, + /// for earlier requests before they can be exposed to readers. `None` + /// marks a request id that finished without a capture (e.g. an oversized + /// body or a client that disconnected mid-request), so later captures are + /// not held back waiting for it. + pending: BTreeMap>, /// The request id the next flush expects, so the visible vectors stay in /// arrival order even when concurrent requests complete out of order. next_flush: u64, @@ -193,21 +196,78 @@ impl Captured { /// a raw capture whose typed capture is missing, a gap, or a later /// request before an earlier one. fn record(&mut self, request_id: u64, request: CapturedRequest, accepted: Option) { - self.pending - .insert(request_id, PendingCapture { request, accepted }); - while let Some(pending) = self.pending.remove(&self.next_flush) { - self.requests.push(pending.request); - match pending.accepted { - Some(Accepted::Stats(payload)) => self.stats.push(payload), - Some(Accepted::Traces(payload)) => self.traces.push(payload), - Some(Accepted::PipelineStats(payload)) => self.pipeline_stats.push(payload), - None => {} + self.complete(request_id, Some(PendingCapture { request, accepted })); + } + + /// Mark a request id as finished without a capture, so the flush does not + /// wait for it forever. + fn skip(&mut self, request_id: u64) { + self.complete(request_id, None); + } + + fn complete(&mut self, request_id: u64, capture: Option) { + self.pending.insert(request_id, capture); + while let Some(entry) = self.pending.remove(&self.next_flush) { + if let Some(pending) = entry { + self.requests.push(pending.request); + match pending.accepted { + Some(Accepted::Stats(payload)) => self.stats.push(payload), + Some(Accepted::Traces(payload)) => self.traces.push(payload), + Some(Accepted::PipelineStats(payload)) => self.pipeline_stats.push(payload), + None => {} + } } self.next_flush += 1; } } } +/// An allocated request id that must be completed exactly once. Every id +/// handed out by `next_request_id` has to reach `Captured` (recorded or +/// skipped), otherwise the arrival-ordered flush waits for it forever and all +/// later captures stay hidden. Dropping an unrecorded slot skips its id, which +/// covers early returns, handler panics, and Hyper canceling the request +/// future when a client disconnects mid-request. +struct RequestSlot { + state: std::sync::Arc, + request_id: u64, + completed: bool, +} + +impl RequestSlot { + fn allocate(state: &std::sync::Arc) -> Self { + Self { + state: std::sync::Arc::clone(state), + request_id: next_request_id(state), + completed: false, + } + } + + fn record(mut self, request: CapturedRequest, accepted: Option) { + self.completed = true; + self.state + .captured + .lock() + .expect("mock_intake: captured mutex poisoned") + .record(self.request_id, request, accepted); + } +} + +impl Drop for RequestSlot { + fn drop(&mut self) { + if self.completed { + return; + } + // Tolerate a poisoned mutex: panicking in Drop during an unwind would + // abort the process, and skipping the id is still correct. + self.state + .captured + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .skip(self.request_id); + } +} + /// Shared server state. The request handler writes to the mutex; callers read /// via `MockIntake::stats_payloads()` / `trace_payloads()` and friends. #[derive(Debug)] @@ -487,6 +547,15 @@ async fn handle_request( return Ok(resp); } + // Request identity is assigned when the request head arrives, before the + // body is read, so captures, summaries, and dumps follow arrival order + // even when concurrent requests upload at different speeds or complete + // out of order. The slot skips its id if the request ends without a + // capture (oversized or unreadable body, or the client disconnecting and + // Hyper dropping this future), so it cannot hold back later captures. + let slot = RequestSlot::allocate(&state); + let request_id = slot.request_id; + let headers = req.headers().clone(); let body = match Limited::new(req.into_body(), MAX_BODY_SIZE).collect().await { Ok(collected) => collected.to_bytes(), @@ -516,19 +585,14 @@ async fn handle_request( .collect(), body: body.to_vec(), }; - // Request identity is assigned at arrival time (before the blocking - // handler runs) so summaries and dumps stay in arrival order even when - // requests complete out of order. - let request_id = next_request_id(&state); // Decompression (up to `MAX_DECOMPRESSED_SIZE`) and optional JSON dumps // are blocking work; run them off the async worker so a large payload // cannot stall a single-threaded test runtime. // - // The capture is recorded inside the blocking task rather than after - // awaiting it: Hyper drops this future if the client disconnects, but a - // blocking task always runs to completion. Recording here guarantees every - // allocated request id is recorded, so a canceled request cannot leave a - // gap that holds back all later captures. + // The slot moves into the blocking task and the capture is recorded there + // rather than after awaiting it: Hyper drops this future if the client + // disconnects, but a blocking task always runs to completion, so a + // request that was fully read is still captured. let handler_state = std::sync::Arc::clone(&state); let handled = tokio::task::spawn_blocking(move || { let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| match endpoint { @@ -551,11 +615,7 @@ async fn handle_request( // observe a contiguous, arrival-ordered prefix: no raw capture without // its typed payload, no gaps, and no later request before an earlier // one. - handler_state - .captured - .lock() - .expect("mock_intake: captured mutex poisoned") - .record(request_id, captured, accepted); + slot.record(captured, accepted); status }) .await; @@ -2076,6 +2136,71 @@ mod tests { assert!(intake.stats_payloads().is_empty()); } + #[tokio::test] + async fn requests_ending_without_capture_do_not_hold_back_later_captures() { + use tokio::io::AsyncWriteExt; + + let intake = start_default().await; + let base_url = intake.base_url(); + let addr = base_url.trim_start_matches("http://").to_string(); + + // An oversized body is rejected after its request id is allocated, so + // the id must be skipped rather than left as a gap. + assert_eq!( + post( + &base_url, + "/api/v0.2/stats", + None, + vec![0u8; MAX_BODY_SIZE + 1] + ) + .await, + StatusCode::PAYLOAD_TOO_LARGE + ); + + // A client that disconnects mid-body also ends its request without a + // capture. + { + let mut stream = tokio::net::TcpStream::connect(&addr) + .await + .expect("test: connect to mock intake failed"); + let head = format!( + "POST /api/v0.2/stats HTTP/1.1\r\nHost: {addr}\r\nContent-Type: application/msgpack\r\nContent-Length: 1000\r\n\r\n" + ); + stream + .write_all(head.as_bytes()) + .await + .expect("test: write request head failed"); + stream + .write_all(b"partial") + .await + .expect("test: write partial body failed"); + } + + let body = rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + assert_eq!( + post(&base_url, "/api/v0.2/stats", None, body).await, + StatusCode::ACCEPTED + ); + + // The server notices the disconnect asynchronously, so poll until the + // valid request becomes visible. + let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(5); + while intake.requests_for_path("/api/v0.2/stats").is_empty() { + assert!( + tokio::time::Instant::now() < deadline, + "valid request was held back by a request that ended without a capture" + ); + tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; + } + assert_eq!(intake.requests_for_path("/api/v0.2/stats").len(), 1); + assert_eq!(intake.stats_payloads().len(), 1); + } + #[tokio::test] async fn body_at_exact_size_limit_is_accepted() { let intake = start_default().await; @@ -2618,4 +2743,30 @@ mod tests { ); assert!(captured.pending.is_empty()); } + + #[test] + fn skip_releases_later_captures() { + let mut captured = Captured::new(); + + // Request 2 is held back until request 1 completes. + captured.record(2, captured_request(b"two"), None); + assert!(captured.requests.is_empty()); + + // Request 1 ends without a capture: skipping it flushes request 2. + captured.skip(1); + let bodies: Vec<_> = captured.requests.iter().map(|r| r.body.clone()).collect(); + assert_eq!(bodies, vec![b"two".to_vec()]); + + // A skip that arrives out of order is buffered like a capture. + captured.skip(4); + captured.record(5, captured_request(b"five"), None); + assert_eq!(captured.requests.len(), 1); + captured.record(3, captured_request(b"three"), None); + let bodies: Vec<_> = captured.requests.iter().map(|r| r.body.clone()).collect(); + assert_eq!( + bodies, + vec![b"two".to_vec(), b"three".to_vec(), b"five".to_vec()] + ); + assert!(captured.pending.is_empty()); + } } From daf34d4e521cd2e4d61f56ac58d33a1017766ab8 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 14:14:22 -0400 Subject: [PATCH 26/32] refactor(mock-intake): simplify payload handler plumbing Drop the HandledRequest request_id field, which always echoed the handler's request_id parameter, and use the parameter directly at summary and dump call sites. Replace the three then_some/flatten accepted-payload constructs with Option::filter, and remove the always-false inject_failure parameter from decompress_failure_status. No behavior change. --- crates/datadog-mock-intake/src/lib.rs | 61 +++++++++++---------------- 1 file changed, 25 insertions(+), 36 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index d5b4a34d..09156658 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -637,7 +637,6 @@ enum Accepted { /// Result of handling one intake request, before summary and dump emission. struct HandledRequest { - request_id: u64, status: StatusCode, decoded: Option, } @@ -656,7 +655,6 @@ fn handle_stats( Err(e @ DecompressError::UnsupportedEncoding(_)) => { eprintln!("{e}"); HandledRequest { - request_id, status: StatusCode::UNSUPPORTED_MEDIA_TYPE, decoded: None, } @@ -673,7 +671,6 @@ fn handle_stats( StatusCode::ACCEPTED }; HandledRequest { - request_id, status, decoded: Some(payload), } @@ -681,7 +678,6 @@ fn handle_stats( Err(err) => { eprintln!("mock_intake: failed to decode StatsPayload msgpack: {err}"); HandledRequest { - request_id, status: failure_status(inject_failure), decoded: None, } @@ -690,7 +686,6 @@ fn handle_stats( Err(e) => { eprintln!("{e}"); HandledRequest { - request_id, status: failure_status(inject_failure), decoded: None, } @@ -704,7 +699,7 @@ fn handle_stats( if state.options.request_summaries { let groups = handled.decoded.as_ref().map(stats_hits_by_key); log_summary( - handled.request_id, + request_id, "/api/v0.2/stats", headers, handled.status, @@ -718,7 +713,7 @@ fn handle_stats( { dump_request( state, - handled.request_id, + request_id, "/api/v0.2/stats", headers, handled.status, @@ -726,9 +721,9 @@ fn handle_stats( ); } - let accepted = (handled.status == StatusCode::ACCEPTED) - .then_some(handled.decoded) - .flatten() + let accepted = handled + .decoded + .filter(|_| handled.status == StatusCode::ACCEPTED) .map(Accepted::Stats); (handled.status, accepted) } @@ -742,14 +737,12 @@ fn handle_traces( let handled: HandledRequest = match decompress(headers, body) { Ok(d) => match pb::AgentPayload::decode(d.as_slice()) { Ok(payload) => HandledRequest { - request_id, status: StatusCode::ACCEPTED, decoded: Some(payload), }, Err(err) => { eprintln!("mock_intake: failed to decode AgentPayload protobuf: {err}"); HandledRequest { - request_id, status: StatusCode::BAD_REQUEST, decoded: None, } @@ -758,8 +751,7 @@ fn handle_traces( Err(e) => { eprintln!("{e}"); HandledRequest { - request_id, - status: decompress_failure_status(&e, false), + status: decompress_failure_status(&e), decoded: None, } } @@ -773,7 +765,7 @@ fn handle_traces( if state.options.request_summaries { log_summary( - handled.request_id, + request_id, "/api/v0.2/traces", headers, handled.status, @@ -787,7 +779,7 @@ fn handle_traces( { dump_request( state, - handled.request_id, + request_id, "/api/v0.2/traces", headers, handled.status, @@ -795,9 +787,9 @@ fn handle_traces( ); } - let accepted = (handled.status == StatusCode::ACCEPTED) - .then_some(handled.decoded) - .flatten() + let accepted = handled + .decoded + .filter(|_| handled.status == StatusCode::ACCEPTED) .map(Accepted::Traces); (handled.status, accepted) } @@ -811,14 +803,12 @@ fn handle_pipeline_stats( let handled: HandledRequest = match decompress(headers, body) { Ok(d) => match rmp_serde::from_slice::(&d) { Ok(payload) => HandledRequest { - request_id, status: StatusCode::ACCEPTED, decoded: Some(payload), }, Err(err) => { eprintln!("mock_intake: failed to decode pipeline stats msgpack: {err}"); HandledRequest { - request_id, status: StatusCode::BAD_REQUEST, decoded: None, } @@ -827,8 +817,7 @@ fn handle_pipeline_stats( Err(e) => { eprintln!("{e}"); HandledRequest { - request_id, - status: decompress_failure_status(&e, false), + status: decompress_failure_status(&e), decoded: None, } } @@ -840,7 +829,7 @@ fn handle_pipeline_stats( .as_ref() .map_or(0, |p| p.stats.iter().map(|b| b.stats.len()).sum()); log_summary( - handled.request_id, + request_id, "/api/v0.1/pipeline_stats", headers, handled.status, @@ -854,7 +843,7 @@ fn handle_pipeline_stats( { dump_request( state, - handled.request_id, + request_id, "/api/v0.1/pipeline_stats", headers, handled.status, @@ -862,9 +851,9 @@ fn handle_pipeline_stats( ); } - let accepted = (handled.status == StatusCode::ACCEPTED) - .then_some(handled.decoded) - .flatten() + let accepted = handled + .decoded + .filter(|_| handled.status == StatusCode::ACCEPTED) .map(Accepted::PipelineStats); (handled.status, accepted) } @@ -884,16 +873,16 @@ fn next_request_id(state: &SharedState) -> u64 { state.request_counter.fetch_add(1, Ordering::SeqCst) + 1 } -/// Status for a request whose body could not be decompressed. An unsupported -/// `Content-Encoding` is a client error (`415 Unsupported Media Type`, the -/// same status the real intake stack uses for unsupported media types) -/// regardless of stats failure injection; an actual decompression failure -/// keeps the endpoint's historical status (`400`, or `500` inside the stats -/// injection window). -fn decompress_failure_status(err: &DecompressError, inject_failure: bool) -> StatusCode { +/// Status for a traces or pipeline-stats request whose body could not be +/// decompressed. An unsupported `Content-Encoding` is a client error (`415 +/// Unsupported Media Type`, the same status the real intake stack uses for +/// unsupported media types); an actual decompression failure keeps the +/// endpoint's historical `400 Bad Request` status. (Stats handles failure +/// injection separately in `handle_stats`.) +fn decompress_failure_status(err: &DecompressError) -> StatusCode { match err { DecompressError::UnsupportedEncoding(_) => StatusCode::UNSUPPORTED_MEDIA_TYPE, - DecompressError::Decode(_) => failure_status(inject_failure), + DecompressError::Decode(_) => StatusCode::BAD_REQUEST, } } From 108eda490d0cd31282b772fa1fe86df6d5eb9a15 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 15:57:38 -0400 Subject: [PATCH 27/32] refactor: clean up mock intake and test duplication Replace hand-written Display impl with thiserror, reuse the existing content-encoding helper, inline the one-use request-id helper, drop a redundant state clone, and only compute payload counts when summaries are enabled. Extract a shared helper for the duplicated marker-stats assertion in the trace-agent tests. --- crates/datadog-mock-intake/src/lib.rs | 103 ++++++++---------- .../tests/integration_test.rs | 22 ++-- 2 files changed, 55 insertions(+), 70 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 09156658..38577fab 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -37,8 +37,8 @@ use std::fmt::Write as _; use std::io::Read; use std::path::PathBuf; use std::process; -use std::sync::Mutex; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; use std::time::{SystemTime, UNIX_EPOCH}; use bytes::Bytes; @@ -182,8 +182,6 @@ struct Captured { impl Captured { fn new() -> Self { - // Request ids start at 1 (`next_request_id` adds 1 after the fetch), - // so the first flush expects id 1. Self { next_flush: 1, ..Default::default() @@ -223,22 +221,23 @@ impl Captured { } /// An allocated request id that must be completed exactly once. Every id -/// handed out by `next_request_id` has to reach `Captured` (recorded or +/// handed out by the request counter has to reach `Captured` (recorded or /// skipped), otherwise the arrival-ordered flush waits for it forever and all /// later captures stay hidden. Dropping an unrecorded slot skips its id, which /// covers early returns, handler panics, and Hyper canceling the request /// future when a client disconnects mid-request. struct RequestSlot { - state: std::sync::Arc, + state: Arc, request_id: u64, completed: bool, } impl RequestSlot { - fn allocate(state: &std::sync::Arc) -> Self { + fn allocate(state: &Arc) -> Self { Self { - state: std::sync::Arc::clone(state), - request_id: next_request_id(state), + state: Arc::clone(state), + // Ids start at 1 so the first flush expects id 1. + request_id: state.request_counter.fetch_add(1, Ordering::SeqCst) + 1, completed: false, } } @@ -274,7 +273,8 @@ impl Drop for RequestSlot { struct SharedState { captured: Mutex, options: MockIntakeOptions, - /// Monotonic request identity for summaries and dumps. + /// Monotonic request identity for summaries and dumps. Ids start at 1 + /// (`Captured::new` initializes `next_flush` to match). request_counter: AtomicU64, /// Stats request attempts, used for `fail_stats_first_n` injection. stats_attempts: AtomicU64, @@ -303,10 +303,10 @@ impl Endpoint { #[derive(Debug)] pub struct MockIntake { base_url: String, - state: std::sync::Arc, + state: Arc, shutdown_tx: Option>, task: Option>, - connections: std::sync::Arc>, + connections: Arc>, } /// Tracks live connection tasks plus whether shutdown has begun. Keeping the @@ -351,7 +351,7 @@ impl MockIntake { })?; let base_url = format!("http://{addr}"); - let state = std::sync::Arc::new(SharedState { + let state = Arc::new(SharedState { captured: Mutex::new(Captured::new()), options, request_counter: AtomicU64::new(0), @@ -359,9 +359,9 @@ impl MockIntake { }); let (shutdown_tx, mut shutdown_rx) = oneshot::channel::<()>(); - let connections: std::sync::Arc> = std::sync::Arc::default(); - let task_state = std::sync::Arc::clone(&state); - let task_connections = std::sync::Arc::clone(&connections); + let connections: Arc> = Arc::default(); + let task_state = Arc::clone(&state); + let task_connections = Arc::clone(&connections); let task = tokio::spawn(async move { let state = task_state; loop { @@ -381,10 +381,10 @@ impl MockIntake { }; let io = TokioIo::new(stream); - let state = std::sync::Arc::clone(&state); + let state = Arc::clone(&state); let handle = tokio::spawn(async move { let service = service_fn(move |req: Request| { - let state = std::sync::Arc::clone(&state); + let state = Arc::clone(&state); async move { handle_request(state, req).await } }); // Log connection failures (malformed HTTP, resets, @@ -533,7 +533,7 @@ fn response(status: StatusCode) -> Response> { /// Handle one incoming request: route it, enforce the body size limit, then /// capture and dispatch it to the endpoint-specific handler. async fn handle_request( - state: std::sync::Arc, + state: Arc, req: Request, ) -> Result>, Infallible> { let path = req.uri().path().to_string(); @@ -593,14 +593,13 @@ async fn handle_request( // rather than after awaiting it: Hyper drops this future if the client // disconnects, but a blocking task always runs to completion, so a // request that was fully read is still captured. - let handler_state = std::sync::Arc::clone(&state); + // `state` is moved into the blocking task; the slot holds its own Arc + // clone, so the handler still shares the same SharedState. let handled = tokio::task::spawn_blocking(move || { let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| match endpoint { - Endpoint::Stats => handle_stats(&handler_state, request_id, &headers, &body), - Endpoint::Traces => handle_traces(&handler_state, request_id, &headers, &body), - Endpoint::PipelineStats => { - handle_pipeline_stats(&handler_state, request_id, &headers, &body) - } + Endpoint::Stats => handle_stats(&state, request_id, &headers, &body), + Endpoint::Traces => handle_traces(&state, request_id, &headers, &body), + Endpoint::PipelineStats => handle_pipeline_stats(&state, request_id, &headers, &body), })); let (status, accepted) = match result { Ok(result) => result, @@ -642,7 +641,7 @@ struct HandledRequest { } fn handle_stats( - state: &std::sync::Arc, + state: &Arc, request_id: u64, headers: &hyper::HeaderMap, body: &Bytes, @@ -694,9 +693,8 @@ fn handle_stats( } }; - let payload_count = handled.decoded.as_ref().map_or(0, |p| p.stats.len()); - if state.options.request_summaries { + let payload_count = handled.decoded.as_ref().map_or(0, |p| p.stats.len()); let groups = handled.decoded.as_ref().map(stats_hits_by_key); log_summary( request_id, @@ -729,7 +727,7 @@ fn handle_stats( } fn handle_traces( - state: &std::sync::Arc, + state: &Arc, request_id: u64, headers: &hyper::HeaderMap, body: &Bytes, @@ -757,13 +755,12 @@ fn handle_traces( } }; - // Tracer payload count across the ordinary and indexed collections. - let payload_count = handled - .decoded - .as_ref() - .map_or(0, |p| p.tracer_payloads.len() + p.idx_tracer_payloads.len()); - if state.options.request_summaries { + // Tracer payload count across the ordinary and indexed collections. + let payload_count = handled + .decoded + .as_ref() + .map_or(0, |p| p.tracer_payloads.len() + p.idx_tracer_payloads.len()); log_summary( request_id, "/api/v0.2/traces", @@ -795,7 +792,7 @@ fn handle_traces( } fn handle_pipeline_stats( - state: &std::sync::Arc, + state: &Arc, request_id: u64, headers: &hyper::HeaderMap, body: &Bytes, @@ -869,10 +866,6 @@ fn failure_status(inject_failure: bool) -> StatusCode { } } -fn next_request_id(state: &SharedState) -> u64 { - state.request_counter.fetch_add(1, Ordering::SeqCst) + 1 -} - /// Status for a traces or pipeline-stats request whose body could not be /// decompressed. An unsupported `Content-Encoding` is a client error (`415 /// Unsupported Media Type`, the same status the real intake stack uses for @@ -930,6 +923,8 @@ fn dump_request( status: StatusCode, payload: serde_json::Value, ) { + use std::io::Write as _; + let Some(dump_dir) = &state.options.dump_dir else { return; }; @@ -971,7 +966,7 @@ fn dump_request( .open(&path) { Ok(mut file) => { - if let Err(e) = std::io::Write::write_all(&mut file, &json) { + if let Err(e) = file.write_all(&json) { eprintln!("mock_intake: failed to write dump {}: {e}", path.display()); } return; @@ -1215,37 +1210,25 @@ fn stats_hits_by_key(payload: &pb::StatsPayload) -> BTreeMap // --------------------------------------------------------------------------- /// An error from decoding a request body's `Content-Encoding`. -#[derive(Debug)] +#[derive(Debug, thiserror::Error)] enum DecompressError { /// The `Content-Encoding` header names an encoding the intake does not /// support. Rejected with `415`, matching the intake's treatment of /// unsupported media types, so a test cannot silently pass against the /// mock while the real intake would reject the payload. + #[error("mock_intake: unsupported Content-Encoding '{0}'")] UnsupportedEncoding(String), /// The body could not be decompressed with its declared encoding. + #[error("{0}")] Decode(String), } -impl std::fmt::Display for DecompressError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::UnsupportedEncoding(encoding) => { - write!(f, "mock_intake: unsupported Content-Encoding '{encoding}'") - } - Self::Decode(message) => write!(f, "{message}"), - } - } -} /// Decompress a request body based on its `Content-Encoding` header. /// Supports `gzip` and `zstd`; an absent or `identity` encoding returns the /// body unchanged. Any other encoding is rejected as unsupported rather than /// silently decoded as identity. fn decompress(headers: &hyper::HeaderMap, body: &Bytes) -> Result, DecompressError> { - let encoding = headers - .get("content-encoding") - .and_then(|v| v.to_str().ok()) - .unwrap_or("") - .to_ascii_lowercase(); + let encoding = content_encoding(headers).to_ascii_lowercase(); match encoding.as_str() { "" | "identity" => Ok(body.to_vec()), @@ -1310,8 +1293,12 @@ mod tests { } fn gzip(data: Vec) -> Vec { + use std::io::Write as _; + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); - std::io::Write::write_all(&mut encoder, &data).expect("test: gzip compression failed"); + encoder + .write_all(&data) + .expect("test: gzip compression failed"); encoder.finish().expect("test: gzip compression failed") } @@ -1491,7 +1478,7 @@ mod tests { use tokio::io::{AsyncReadExt, AsyncWriteExt}; let intake = start_default().await; - let state = std::sync::Arc::clone(&intake.state); + let state = Arc::clone(&intake.state); let addr = intake.base_url().trim_start_matches("http://").to_string(); // Raw TCP so the connection stays open (keep-alive) across Drop. diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 4c4ea902..691d9f95 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -207,6 +207,14 @@ pub fn verify_no_stats_request(mock_intake: &MockIntake) { ); } +/// Check whether any captured stats payload carries stats from the given service. +fn has_stats_from_service(mock_intake: &MockIntake, service: &str) -> bool { + mock_intake + .stats_payloads() + .iter() + .any(|payload| payload.stats.iter().any(|csp| csp.service == service)) +} + /// Helper to verify a DSM request sent to the mock intake. The forwarded body /// must arrive byte-for-byte unchanged AND be decodable by the intake: the /// decoded capture proves the proxy forwards a payload the real intake accepts. @@ -876,12 +884,7 @@ async fn test_tracer_and_agent_stats_enabled_uses_agent_stats_no_duplicates() { verify_stats_request(&mock_intake).await; // The tracer computed stats should never reach the backend. Only the agent computed stats should reach the backend. - let has_marker = mock_intake.stats_payloads().iter().any(|payload| { - payload - .stats - .iter() - .any(|csp| csp.service == "tracer-marker-stats") - }); + let has_marker = has_stats_from_service(&mock_intake, "tracer-marker-stats"); assert!( !has_marker, "Expected tracer computed stats to be dropped, not forwarded to the backend" @@ -939,12 +942,7 @@ async fn test_tracer_stats_enabled_agent_stats_disabled_forwards_tracer_stats() assert_eq!(stats_response.status(), StatusCode::ACCEPTED); verify_stats_request(&mock_intake).await; - let has_marker = mock_intake.stats_payloads().iter().any(|payload| { - payload - .stats - .iter() - .any(|csp| csp.service == "tracer-marker-stats") - }); + let has_marker = has_stats_from_service(&mock_intake, "tracer-marker-stats"); assert!( has_marker, "Expected tracer computed stats to be forwarded to the backend" From 54d0c8eb3b89bec1ba7e027fd61b42a3b3a87d7a Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 16:52:44 -0400 Subject: [PATCH 28/32] Show mock intake captures once responded Captures were only exposed once every earlier request finished, so a 2xx did not guarantee the capture was visible, and a stalled request hid all later captures indefinitely. Captures are now visible as soon as the intake responds, still ordered by arrival. --- crates/datadog-mock-intake/src/lib.rs | 334 +++++++++----------------- 1 file changed, 117 insertions(+), 217 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index 38577fab..cb43f1e9 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -153,118 +153,48 @@ pub struct CapturedRequest { pub body: Vec, } -/// A completed request held back until every earlier request has been -/// flushed to the visible capture vectors, so readers only ever observe a -/// contiguous, arrival-ordered prefix of captures. -#[derive(Debug)] -struct PendingCapture { - request: CapturedRequest, - accepted: Option, -} - -/// Captured, decoded APM payloads for a single test run. +/// Captured, decoded APM payloads for a single test run. Each entry is +/// tagged with its request id and kept sorted by it, so readers see captures +/// in arrival order (when the request head arrived) even when concurrent +/// requests finish handling out of order. #[derive(Debug, Default)] struct Captured { - requests: Vec, - stats: Vec, - traces: Vec, - pipeline_stats: Vec, - /// Requests that finished handling out of arrival order and are waiting - /// for earlier requests before they can be exposed to readers. `None` - /// marks a request id that finished without a capture (e.g. an oversized - /// body or a client that disconnected mid-request), so later captures are - /// not held back waiting for it. - pending: BTreeMap>, - /// The request id the next flush expects, so the visible vectors stay in - /// arrival order even when concurrent requests complete out of order. - next_flush: u64, + requests: Vec<(u64, CapturedRequest)>, + stats: Vec<(u64, pb::StatsPayload)>, + traces: Vec<(u64, pb::AgentPayload)>, + pipeline_stats: Vec<(u64, PipelineStatsPayload)>, } impl Captured { - fn new() -> Self { - Self { - next_flush: 1, - ..Default::default() - } - } - - /// Record a completed request. Out-of-order completions are buffered - /// until all earlier requests arrive, then a contiguous prefix is - /// appended to the visible vectors in one pass so readers never observe - /// a raw capture whose typed capture is missing, a gap, or a later - /// request before an earlier one. + /// Record a completed request and make it visible immediately. The raw + /// capture and its typed payload are inserted in one call (under the + /// caller's lock), so readers never observe one without the other. A + /// request is never held back waiting for earlier requests: once the + /// intake has responded, its capture is visible, and a stalled or + /// abandoned request cannot hide later ones. fn record(&mut self, request_id: u64, request: CapturedRequest, accepted: Option) { - self.complete(request_id, Some(PendingCapture { request, accepted })); - } - - /// Mark a request id as finished without a capture, so the flush does not - /// wait for it forever. - fn skip(&mut self, request_id: u64) { - self.complete(request_id, None); - } - - fn complete(&mut self, request_id: u64, capture: Option) { - self.pending.insert(request_id, capture); - while let Some(entry) = self.pending.remove(&self.next_flush) { - if let Some(pending) = entry { - self.requests.push(pending.request); - match pending.accepted { - Some(Accepted::Stats(payload)) => self.stats.push(payload), - Some(Accepted::Traces(payload)) => self.traces.push(payload), - Some(Accepted::PipelineStats(payload)) => self.pipeline_stats.push(payload), - None => {} - } + insert_sorted(&mut self.requests, request_id, request); + match accepted { + Some(Accepted::Stats(payload)) => insert_sorted(&mut self.stats, request_id, payload), + Some(Accepted::Traces(payload)) => insert_sorted(&mut self.traces, request_id, payload), + Some(Accepted::PipelineStats(payload)) => { + insert_sorted(&mut self.pipeline_stats, request_id, payload); } - self.next_flush += 1; + None => {} } } } -/// An allocated request id that must be completed exactly once. Every id -/// handed out by the request counter has to reach `Captured` (recorded or -/// skipped), otherwise the arrival-ordered flush waits for it forever and all -/// later captures stay hidden. Dropping an unrecorded slot skips its id, which -/// covers early returns, handler panics, and Hyper canceling the request -/// future when a client disconnects mid-request. -struct RequestSlot { - state: Arc, - request_id: u64, - completed: bool, +/// Insert `item` keeping `entries` sorted by request id. Captures usually +/// complete in order, so this is normally an append. +fn insert_sorted(entries: &mut Vec<(u64, T)>, request_id: u64, item: T) { + let index = entries.partition_point(|(id, _)| *id < request_id); + entries.insert(index, (request_id, item)); } -impl RequestSlot { - fn allocate(state: &Arc) -> Self { - Self { - state: Arc::clone(state), - // Ids start at 1 so the first flush expects id 1. - request_id: state.request_counter.fetch_add(1, Ordering::SeqCst) + 1, - completed: false, - } - } - - fn record(mut self, request: CapturedRequest, accepted: Option) { - self.completed = true; - self.state - .captured - .lock() - .expect("mock_intake: captured mutex poisoned") - .record(self.request_id, request, accepted); - } -} - -impl Drop for RequestSlot { - fn drop(&mut self) { - if self.completed { - return; - } - // Tolerate a poisoned mutex: panicking in Drop during an unwind would - // abort the process, and skipping the id is still correct. - self.state - .captured - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .skip(self.request_id); - } +/// Clone the items of an id-tagged capture vector, in request-id order. +fn items(entries: &[(u64, T)]) -> Vec { + entries.iter().map(|(_, item)| item.clone()).collect() } /// Shared server state. The request handler writes to the mutex; callers read @@ -273,8 +203,8 @@ impl Drop for RequestSlot { struct SharedState { captured: Mutex, options: MockIntakeOptions, - /// Monotonic request identity for summaries and dumps. Ids start at 1 - /// (`Captured::new` initializes `next_flush` to match). + /// Monotonic request identity for summaries, dumps, and capture order. + /// Ids start at 1. request_counter: AtomicU64, /// Stats request attempts, used for `fail_stats_first_n` injection. stats_attempts: AtomicU64, @@ -352,7 +282,7 @@ impl MockIntake { let base_url = format!("http://{addr}"); let state = Arc::new(SharedState { - captured: Mutex::new(Captured::new()), + captured: Mutex::new(Captured::default()), options, request_counter: AtomicU64::new(0), stats_attempts: AtomicU64::new(0), @@ -461,43 +391,49 @@ impl MockIntake { .expect("mock_intake: captured mutex poisoned") .requests .iter() - .filter(|req| req.path == path) - .cloned() + .filter(|(_, req)| req.path == path) + .map(|(_, req)| req.clone()) .collect() } /// All DSM pipeline-stats payloads captured so far, in arrival order. #[must_use] pub fn pipeline_stats_payloads(&self) -> Vec { - self.state - .captured - .lock() - .expect("mock_intake: captured mutex poisoned") - .pipeline_stats - .clone() + items( + &self + .state + .captured + .lock() + .expect("mock_intake: captured mutex poisoned") + .pipeline_stats, + ) } /// All `StatsPayload`s captured so far, in arrival order. Rejected stats /// attempts (failure injection) are excluded. #[must_use] pub fn stats_payloads(&self) -> Vec { - self.state - .captured - .lock() - .expect("mock_intake: captured mutex poisoned") - .stats - .clone() + items( + &self + .state + .captured + .lock() + .expect("mock_intake: captured mutex poisoned") + .stats, + ) } /// All `AgentPayload`s captured so far, in arrival order. #[must_use] pub fn trace_payloads(&self) -> Vec { - self.state - .captured - .lock() - .expect("mock_intake: captured mutex poisoned") - .traces - .clone() + items( + &self + .state + .captured + .lock() + .expect("mock_intake: captured mutex poisoned") + .traces, + ) } } @@ -550,11 +486,9 @@ async fn handle_request( // Request identity is assigned when the request head arrives, before the // body is read, so captures, summaries, and dumps follow arrival order // even when concurrent requests upload at different speeds or complete - // out of order. The slot skips its id if the request ends without a - // capture (oversized or unreadable body, or the client disconnecting and - // Hyper dropping this future), so it cannot hold back later captures. - let slot = RequestSlot::allocate(&state); - let request_id = slot.request_id; + // out of order. An id whose request ends without a capture (oversized or + // unreadable body, or a client disconnecting) is simply never recorded. + let request_id = state.request_counter.fetch_add(1, Ordering::SeqCst) + 1; let headers = req.headers().clone(); let body = match Limited::new(req.into_body(), MAX_BODY_SIZE).collect().await { @@ -589,12 +523,10 @@ async fn handle_request( // are blocking work; run them off the async worker so a large payload // cannot stall a single-threaded test runtime. // - // The slot moves into the blocking task and the capture is recorded there - // rather than after awaiting it: Hyper drops this future if the client - // disconnects, but a blocking task always runs to completion, so a - // request that was fully read is still captured. - // `state` is moved into the blocking task; the slot holds its own Arc - // clone, so the handler still shares the same SharedState. + // The capture is recorded inside the blocking task rather than after + // awaiting it: Hyper drops this future if the client disconnects, but a + // blocking task always runs to completion, so a request that was fully + // read is still captured. let handled = tokio::task::spawn_blocking(move || { let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| match endpoint { Endpoint::Stats => handle_stats(&state, request_id, &headers, &body), @@ -610,11 +542,13 @@ async fn handle_request( }; // Record the raw request and its typed payload in one critical - // section. Out-of-order completions are buffered so readers only ever - // observe a contiguous, arrival-ordered prefix: no raw capture without - // its typed payload, no gaps, and no later request before an earlier - // one. - slot.record(captured, accepted); + // section, before the response is sent, so a client that has seen + // the response always finds its capture. + state + .captured + .lock() + .expect("mock_intake: captured mutex poisoned") + .record(request_id, captured, accepted); status }) .await; @@ -2133,24 +2067,26 @@ mod tests { StatusCode::PAYLOAD_TOO_LARGE ); - // A client that disconnects mid-body also ends its request without a - // capture. - { - let mut stream = tokio::net::TcpStream::connect(&addr) - .await - .expect("test: connect to mock intake failed"); - let head = format!( - "POST /api/v0.2/stats HTTP/1.1\r\nHost: {addr}\r\nContent-Type: application/msgpack\r\nContent-Length: 1000\r\n\r\n" - ); - stream - .write_all(head.as_bytes()) - .await - .expect("test: write request head failed"); - stream - .write_all(b"partial") - .await - .expect("test: write partial body failed"); - } + // A client that stalls mid-body keeps its request in flight: the + // connection stays open, so the request never completes. It must not + // hide the captures of later requests. + let mut stalled = tokio::net::TcpStream::connect(&addr) + .await + .expect("test: connect to mock intake failed"); + let head = format!( + "POST /api/v0.2/stats HTTP/1.1\r\nHost: {addr}\r\nContent-Type: application/msgpack\r\nContent-Length: 1000\r\n\r\n" + ); + stalled + .write_all(head.as_bytes()) + .await + .expect("test: write request head failed"); + stalled + .write_all(b"partial") + .await + .expect("test: write partial body failed"); + // Give the server time to read the head and assign the stalled + // request an id ahead of the valid request below. + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; let body = rmp_serde::to_vec_named(&stats_payload( client_payload("local"), @@ -2163,18 +2099,15 @@ mod tests { StatusCode::ACCEPTED ); - // The server notices the disconnect asynchronously, so poll until the - // valid request becomes visible. - let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(5); - while intake.requests_for_path("/api/v0.2/stats").is_empty() { - assert!( - tokio::time::Instant::now() < deadline, - "valid request was held back by a request that ended without a capture" - ); - tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; - } - assert_eq!(intake.requests_for_path("/api/v0.2/stats").len(), 1); + // Once the client has seen the response, the capture is visible + // without polling, even though the stalled request is still open. + assert_eq!( + intake.requests_for_path("/api/v0.2/stats").len(), + 1, + "valid request was held back by an incomplete earlier request" + ); assert_eq!(intake.stats_payloads().len(), 1); + drop(stalled); } #[tokio::test] @@ -2686,63 +2619,30 @@ mod tests { } #[test] - fn record_flushes_captures_in_arrival_order() { - let mut captured = Captured::new(); + fn record_is_visible_immediately_in_arrival_order() { + let mut captured = Captured::default(); + let bodies = |captured: &Captured| -> Vec> { + captured + .requests + .iter() + .map(|(_, r)| r.body.clone()) + .collect() + }; - // Request 2 completes before request 1: nothing becomes visible yet, - // because the visible vectors must stay in arrival order. + // Request 2 completes before request 1: it is visible right away. captured.record(2, captured_request(b"two"), None); - assert!(captured.requests.is_empty()); + assert_eq!(bodies(&captured), vec![b"two".to_vec()]); - // A later request completing early is also held back. + // Request 4 completes early too; request 3 never completes (e.g. a + // stalled or disconnected client) and must not hide anything. captured.record(4, captured_request(b"four"), None); - assert!(captured.requests.is_empty()); + assert_eq!(bodies(&captured), vec![b"two".to_vec(), b"four".to_vec()]); - // Request 1 completes: the contiguous prefix (1, 2) flushes in - // arrival order; request 4 stays buffered. + // Request 1 completes last but is ordered by arrival (request id). captured.record(1, captured_request(b"one"), None); - let bodies: Vec<_> = captured.requests.iter().map(|r| r.body.clone()).collect(); - assert_eq!(bodies, vec![b"one".to_vec(), b"two".to_vec()]); - - // Request 3 completes: the remaining prefix (2, 3, 4) flushes in - // arrival order, including the out-of-order request 4. - captured.record(3, captured_request(b"three"), None); - let bodies: Vec<_> = captured.requests.iter().map(|r| r.body.clone()).collect(); - assert_eq!( - bodies, - vec![ - b"one".to_vec(), - b"two".to_vec(), - b"three".to_vec(), - b"four".to_vec() - ] - ); - assert!(captured.pending.is_empty()); - } - - #[test] - fn skip_releases_later_captures() { - let mut captured = Captured::new(); - - // Request 2 is held back until request 1 completes. - captured.record(2, captured_request(b"two"), None); - assert!(captured.requests.is_empty()); - - // Request 1 ends without a capture: skipping it flushes request 2. - captured.skip(1); - let bodies: Vec<_> = captured.requests.iter().map(|r| r.body.clone()).collect(); - assert_eq!(bodies, vec![b"two".to_vec()]); - - // A skip that arrives out of order is buffered like a capture. - captured.skip(4); - captured.record(5, captured_request(b"five"), None); - assert_eq!(captured.requests.len(), 1); - captured.record(3, captured_request(b"three"), None); - let bodies: Vec<_> = captured.requests.iter().map(|r| r.body.clone()).collect(); assert_eq!( - bodies, - vec![b"two".to_vec(), b"three".to_vec(), b"five".to_vec()] + bodies(&captured), + vec![b"one".to_vec(), b"two".to_vec(), b"four".to_vec()] ); - assert!(captured.pending.is_empty()); } } From 695d4469d252aa1b250eeb0518fbcbfa88412f99 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 16:55:24 -0400 Subject: [PATCH 29/32] Wait for stats captures in trace-agent tests test_internal_span_kind_does_not_produce_stats and test_peer_tags_in_flushed_stats read the intake capture immediately after shutdown; wait for it like the other tests to avoid timing-dependent failures. --- .../tests/integration_test.rs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 691d9f95..7dd35a3d 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -1007,11 +1007,12 @@ async fn test_internal_span_kind_does_not_produce_stats() { let _ = shutdown_tx.send(true); let _ = agent_handle.await; + wait_for_capture( + || !mock_intake.stats_payloads().is_empty(), + "a stats request from the root span", + ) + .await; let stats_payloads = mock_intake.stats_payloads(); - assert!( - !stats_payloads.is_empty(), - "Expected a stats request from the root span" - ); let all_groups: Vec<_> = stats_payloads .iter() @@ -1095,11 +1096,12 @@ async fn test_peer_tags_in_flushed_stats() { let _ = shutdown_tx.send(true); let _ = agent_handle.await; + wait_for_capture( + || !mock_intake.stats_payloads().is_empty(), + "a stats request at /api/v0.2/stats", + ) + .await; let stats_payloads = mock_intake.stats_payloads(); - assert!( - !stats_payloads.is_empty(), - "Expected at least one stats request" - ); let payload = stats_payloads .first() From 098d30d2ea15fbbd952658fa1d13bc57ab0921c0 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 16:56:50 -0400 Subject: [PATCH 30/32] Fail fast when the mock intake rejects a payload The trace-agent verify helpers waited for decoded payloads, so a payload the intake rejected surfaced as a 60s timeout claiming no request arrived. Wait for the raw request, then assert it was accepted. --- .../tests/integration_test.rs | 33 ++++++++++--------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 7dd35a3d..61e49dc5 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -117,16 +117,17 @@ pub fn create_mini_agent_with_real_flushers( /// Helper to verify trace request sent to the mock intake pub async fn verify_trace_request(mock_intake: &MockIntake) { + // Wait on the raw capture, not the decoded one, so a payload the intake + // rejects fails fast with an accurate message instead of timing out. wait_for_capture( - || !mock_intake.trace_payloads().is_empty(), + || !mock_intake.requests_for_path("/api/v0.2/traces").is_empty(), "a trace request at /api/v0.2/traces", ) .await; let trace_reqs = mock_intake.requests_for_path("/api/v0.2/traces"); - assert!( - !trace_reqs.is_empty(), - "Expected at least one trace request to mock intake" + !mock_intake.trace_payloads().is_empty(), + "Trace request reached the mock intake but was not accepted; see mock_intake stderr" ); let trace_req = &trace_reqs[0]; @@ -158,16 +159,17 @@ pub async fn verify_trace_request(mock_intake: &MockIntake) { /// Helper to verify stats request sent to the mock intake pub async fn verify_stats_request(mock_intake: &MockIntake) { + // Wait on the raw capture, not the decoded one, so a payload the intake + // rejects fails fast with an accurate message instead of timing out. wait_for_capture( - || !mock_intake.stats_payloads().is_empty(), + || !mock_intake.requests_for_path("/api/v0.2/stats").is_empty(), "a stats request at /api/v0.2/stats", ) .await; let stats_reqs = mock_intake.requests_for_path("/api/v0.2/stats"); - assert!( - !stats_reqs.is_empty(), - "Expected at least one stats request to mock intake" + !mock_intake.stats_payloads().is_empty(), + "Stats request reached the mock intake but was not accepted; see mock_intake stderr" ); let stats_req = &stats_reqs[0]; @@ -224,18 +226,19 @@ pub async fn verify_dsm_request( expected_edge_tags: &[&str], expected_additional_tags: &[&str], ) { + // Wait on the raw capture, not the decoded one, so a payload the intake + // rejects fails fast with an accurate message instead of timing out. wait_for_capture( - || !mock_intake.pipeline_stats_payloads().is_empty(), + || { + !mock_intake + .requests_for_path("/api/v0.1/pipeline_stats") + .is_empty() + }, "a DSM request at /api/v0.1/pipeline_stats", ) .await; let dsm_reqs = mock_intake.requests_for_path("/api/v0.1/pipeline_stats"); - assert!( - !dsm_reqs.is_empty(), - "Expected at least one DSM request to mock intake" - ); - let dsm_req = &dsm_reqs[0]; assert_eq!(dsm_req.method, "POST", "Expected POST method"); assert_eq!( @@ -246,7 +249,7 @@ pub async fn verify_dsm_request( let payloads = mock_intake.pipeline_stats_payloads(); let decoded = payloads .first() - .expect("intake must have decoded at least one DSM payload"); + .expect("DSM request reached the mock intake but was not accepted; see mock_intake stderr"); let bucket = decoded .stats .first() From 45afb6cbaafece7c2b3f94883ca41aa1bec1fb12 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 16:59:17 -0400 Subject: [PATCH 31/32] Reject non-UTF-8 Content-Encoding in mock intake A Content-Encoding value that was not valid UTF-8 was treated as absent, so the body was silently decoded as identity. It is now rejected with 415 like any other unsupported encoding. --- crates/datadog-mock-intake/src/lib.rs | 56 +++++++++++++++++++++++++-- 1 file changed, 52 insertions(+), 4 deletions(-) diff --git a/crates/datadog-mock-intake/src/lib.rs b/crates/datadog-mock-intake/src/lib.rs index cb43f1e9..88e2a29a 100644 --- a/crates/datadog-mock-intake/src/lib.rs +++ b/crates/datadog-mock-intake/src/lib.rs @@ -813,11 +813,14 @@ fn decompress_failure_status(err: &DecompressError) -> StatusCode { } } +/// The request's `Content-Encoding`, or `identity` when absent. A value that +/// is not valid UTF-8 is converted lossily rather than treated as absent, so +/// it is rejected as unsupported instead of silently decoded as identity. fn content_encoding(headers: &hyper::HeaderMap) -> String { - headers - .get("content-encoding") - .and_then(|v| v.to_str().ok()) - .map_or_else(|| "identity".to_string(), ToString::to_string) + headers.get("content-encoding").map_or_else( + || "identity".to_string(), + |v| String::from_utf8_lossy(v.as_bytes()).into_owned(), + ) } /// Emit one summary line per handled request. Payload contents stay out of @@ -1552,6 +1555,51 @@ mod tests { assert_eq!(intake.stats_payloads().len(), 1); } + #[tokio::test] + async fn non_utf8_content_encoding_returns_415_and_is_not_typed_captured() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let intake = start_default().await; + let addr = intake.base_url().trim_start_matches("http://").to_string(); + // A valid identity body: if the non-UTF-8 encoding were treated as + // absent, this would be accepted. + let body = rmp_serde::to_vec_named(&stats_payload( + client_payload("local"), + 1, + vec![grouped_stats("svc", "GET /a", 1, |_| {})], + )) + .expect("test: msgpack encode failed"); + + // Raw TCP so the header value can carry a non-UTF-8 (obs-text) byte. + let mut stream = tokio::net::TcpStream::connect(&addr) + .await + .expect("test: connect to mock intake failed"); + let mut request = format!( + "POST /api/v0.2/stats HTTP/1.1\r\nHost: {addr}\r\nConnection: close\r\nContent-Length: {}\r\nContent-Encoding: gz", + body.len() + ) + .into_bytes(); + request.push(0xFF); + request.extend_from_slice(b"ip\r\n\r\n"); + request.extend_from_slice(&body); + stream + .write_all(&request) + .await + .expect("test: write request failed"); + let mut response = Vec::new(); + stream + .read_to_end(&mut response) + .await + .expect("test: read response failed"); + assert!( + response.starts_with(b"HTTP/1.1 415"), + "test: expected a 415 for the non-UTF-8 encoding, got {response:?}" + ); + + assert!(intake.stats_payloads().is_empty()); + assert_eq!(intake.requests_for_path("/api/v0.2/stats").len(), 1); + } + #[tokio::test] async fn corrupted_compressed_bodies_return_400() { let intake = start_default().await; From d1070122bd306bd2061b10156f208f95decc7f59 Mon Sep 17 00:00:00 2001 From: Lucas Pimentel Date: Wed, 30 Sep 2026 17:26:03 -0400 Subject: [PATCH 32/32] Assert decoded services in dual-transport test The dual-transport test byte-searched raw request bodies for the service names and checked only after the first trace POST arrived. It now waits until spans from both services appear in the decoded trace payloads. --- .../tests/integration_test.rs | 40 ++++++++----------- 1 file changed, 16 insertions(+), 24 deletions(-) diff --git a/crates/datadog-trace-agent/tests/integration_test.rs b/crates/datadog-trace-agent/tests/integration_test.rs index 61e49dc5..57de50ef 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -1250,34 +1250,26 @@ async fn test_mini_agent_dual_transport_with_real_flushers() { .expect("Failed to send /v0.4/traces request over named pipe"); assert_eq!(pipe_response.status(), StatusCode::OK); + // Both payloads must reach the same backend through the shared flusher + // pipeline. The flusher may batch them into one POST or two; either is + // fine, what matters is that spans from both services show up in the + // decoded trace payloads. + let has_service = |service: &str| { + mock_intake.trace_payloads().iter().any(|payload| { + payload + .tracer_payloads + .iter() + .flat_map(|tp| tp.chunks.iter()) + .flat_map(|chunk| chunk.spans.iter()) + .any(|span| span.service == service) + }) + }; wait_for_capture( - || !mock_intake.trace_payloads().is_empty(), - "a trace request at /api/v0.2/traces", + || has_service("dual-tcp-svc") && has_service("dual-pipe-svc"), + "traces from both the TCP and pipe transports at /api/v0.2/traces", ) .await; - // Both payloads must reach the same backend through the shared flusher - // pipeline. The flusher may batch them into one POST or two; either is - // fine, what matters is that both service-name needles show up. - let trace_reqs = mock_intake.requests_for_path("/api/v0.2/traces"); - assert!( - !trace_reqs.is_empty(), - "no trace POST reached backend; expected traces from both transports" - ); - let mut all_bytes = Vec::new(); - for req in &trace_reqs { - assert_eq!(req.method, "POST"); - all_bytes.extend_from_slice(&req.body); - } - assert!( - all_bytes.windows(12).any(|w| w == b"dual-tcp-svc"), - "TCP-side trace did not reach backend" - ); - assert!( - all_bytes.windows(13).any(|w| w == b"dual-pipe-svc"), - "pipe-side trace did not reach backend" - ); - // Trigger graceful shutdown. The watch fan-out must reach BOTH accept // loops; each drains its in-flight handlers; the supervisor then // signals the stats flusher to do its final emit. If fan-out skipped