diff --git a/Cargo.lock b/Cargo.lock index a8d1b467..f42fb6ed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -585,6 +585,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" @@ -664,8 +685,8 @@ dependencies = [ "async-trait", "bytes", "datadog-fips", + "datadog-mock-intake", "duplicate", - "flate2", "http-body-util", "hyper", "hyper-http-proxy", @@ -3015,9 +3036,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.45" +version = "0.23.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" +checksum = "6725596c3f2c3a0aef021139e145d4eafe314a6623e4680ca83852b2c67ab2ba" dependencies = [ "aws-lc-rs", "once_cell", diff --git a/crates/datadog-mock-intake/Cargo.toml b/crates/datadog-mock-intake/Cargo.toml new file mode 100644 index 00000000..894fe183 --- /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 = ["io-util", "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..0b7ccb24 --- /dev/null +++ b/crates/datadog-mock-intake/src/lib.rs @@ -0,0 +1,2590 @@ +// 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 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. +//! +//! 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; + +/// 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", default)] + pub env: String, + #[serde(rename = "Service", default)] + pub service: String, + #[serde(rename = "TracerVersion", default)] + pub tracer_version: String, + #[serde(rename = "Version", default)] + pub version: String, + #[serde(rename = "Tags", default)] + pub tags: Vec, + #[serde(rename = "Stats", default)] + pub stats: Vec, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct PipelineStatsBucket { + #[serde(rename = "Stats", default)] + pub stats: Vec, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct PipelineStatsPoint { + #[serde(rename = "Hash", default)] + pub hash: u64, + #[serde(rename = "ParentHash", default)] + pub parent_hash: u64, + #[serde(rename = "EdgeTags", default)] + 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, 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, + pub path: String, + pub headers: Vec<(String, String)>, + 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 { + 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. + 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 +/// 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>, + 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 { + /// 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::new()), + options, + request_counter: AtomicU64::new(0), + stats_attempts: AtomicU64::new(0), + }); + + 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 { + tokio::select! { + result = listener.accept() => { + 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}"); + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + continue; + } + }; + + let io = TokioIo::new(stream); + let state = std::sync::Arc::clone(&state); + 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 } + }); + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection(io, service) + .await; + }); + if let Ok(mut conns) = task_connections.lock() { + 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 => { + break; + } + } + } + }); + + Ok(Self { + base_url, + state, + shutdown_tx: Some(shutdown_tx), + task: Some(task), + connections, + }) + } + + /// 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, 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 + .captured + .lock() + .expect("mock_intake: captured 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: captured 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: captured 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: captured 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(); + } + // Abort any accepted-but-still-running connection tasks so a + // 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() { + conns.shutting_down = true; + for handle in conns.handles.drain(..) { + handle.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(), + String::from_utf8_lossy(v.as_bytes()).into_owned(), + ) + }) + .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, 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 { + 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); + 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), + PipelineStats(PipelineStatsPayload), +} + +/// 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, + request_id: u64, + headers: &hyper::HeaderMap, + body: &Bytes, +) -> (StatusCode, Option) { + 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 + // failure-injection budget: the budget counts attempts that reach + // payload decoding. + Err(e @ DecompressError::UnsupportedEncoding(_)) => { + eprintln!("{e}"); + HandledRequest { + request_id, + 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()); + + 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 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(serde_json::Value::Null), + ); + } + + let accepted = (handled.status == StatusCode::ACCEPTED) + .then_some(handled.decoded) + .flatten() + .map(Accepted::Stats); + (handled.status, accepted) +} + +fn handle_traces( + state: &std::sync::Arc, + request_id: u64, + headers: &hyper::HeaderMap, + body: &Bytes, +) -> (StatusCode, Option) { + 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, + } + } + }, + Err(e) => { + eprintln!("{e}"); + HandledRequest { + request_id, + status: decompress_failure_status(&e, false), + 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 state.options.dump_dir.is_some() + && let Some(payload) = &handled.decoded + { + dump_request( + state, + handled.request_id, + "/api/v0.2/traces", + headers, + handled.status, + agent_payload_to_json(payload), + ); + } + + 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, + request_id: u64, + headers: &hyper::HeaderMap, + body: &Bytes, +) -> (StatusCode, Option) { + 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, + } + } + }, + Err(e) => { + eprintln!("{e}"); + HandledRequest { + request_id, + status: decompress_failure_status(&e, false), + decoded: None, + } + } + }; + + 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, + payload_count, + None, + ); + } + + 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(serde_json::Value::Null), + ); + } + + 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 +/// 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 +} + +/// 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") + .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 +// --------------------------------------------------------------------------- + +/// 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 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(); + + 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| { + DecompressError::Decode(format!("mock_intake: gzip decode failed: {e}")) + }) + } + "zstd" => { + 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())), + } +} + +/// 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 { + + 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() + } + } + + 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 + /// 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 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: captured 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: captured mutex poisoned") + .requests + .len(); + assert_eq!( + captured_after, 1, + "request sent after Drop must not be captured" + ); + } + + #[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 + ); + } + + #[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; + 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 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; + 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}")); + } + + 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()); + } +} 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/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(()); - } - } -} 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..4c4ea902 100644 --- a/crates/datadog-trace-agent/tests/integration_test.rs +++ b/crates/datadog-trace-agent/tests/integration_test.rs @@ -5,11 +5,12 @@ 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, PipelineStatsBucket, PipelineStatsPayload, PipelineStatsPoint, +}; use datadog_trace_agent::{ config::{Config, Tags, test_helpers::create_tcp_test_config}, mini_agent::MiniAgent, @@ -39,15 +40,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 +115,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 +156,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 +197,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_requests = mock_intake.requests_for_path("/api/v0.2/stats"); assert!( - stats_reqs.is_empty(), - "Expected no stats request to mock server, received {} request(s)", - stats_reqs.len() + stats_requests.is_empty(), + "Expected no stats request to mock intake, received {} request(s)", + stats_requests.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 +235,24 @@ pub async fn verify_dsm_request( "Expected DSM payload body to be forwarded unchanged" ); + 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" + ); + let api_key = dsm_req .headers .iter() @@ -469,11 +503,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 +545,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 +595,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![PipelineStatsBucket { + stats: vec![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 +630,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 +651,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 +703,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 +743,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 +756,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 +798,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 +808,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 +872,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 +894,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 +938,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 +960,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 +1009,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 +1048,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 +1097,15 @@ 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 + .first() + .expect("at least one decoded stats payload"); let all_peer_tags: Vec<&str> = payload .stats .iter() @@ -1090,12 +1124,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 +1166,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 +1247,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 +1281,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; }