From 87c6755daef8b3c08c8edbf1454efde1374b7d02 Mon Sep 17 00:00:00 2001 From: "cliford.shelton" Date: Fri, 9 Oct 2026 13:14:57 -0400 Subject: [PATCH 1/4] [SLES-3039] fix(traces): split trace payloads larger than the batch cap A trace payload larger than the 12 MB batch cap was put back at the front of the trace queue by `get_batch` and never sent, so every later payload in that sandbox queued behind it until the Lambda ran out of memory. Large traces (tens of thousands of spans) exceed the cap once the extension adds its tags to every span. - Split oversized payloads before queuing them: chunks move into separate payloads, and a chunk that is still too large has its spans split into several chunks of the same trace, as tracer partial flushing does. Trace-level `_dd.p.*` tags are copied onto the first span of each piece. A single span larger than the cap is dropped with a warning. - `get_batch` now sends a payload that is over the cap on its own as a batch by itself instead of putting it back. Co-Authored-By: Claude Opus 5.5 --- bottlecap/src/otlp/agent.rs | 14 +- bottlecap/src/traces/mod.rs | 1 + bottlecap/src/traces/payload_split.rs | 311 +++++++++++++++++++++++ bottlecap/src/traces/trace_aggregator.rs | 31 ++- bottlecap/src/traces/trace_processor.rs | 179 ++++++++++--- 5 files changed, 491 insertions(+), 45 deletions(-) create mode 100644 bottlecap/src/traces/payload_split.rs diff --git a/bottlecap/src/otlp/agent.rs b/bottlecap/src/otlp/agent.rs index 20dccc3bc..47083ecc4 100644 --- a/bottlecap/src/otlp/agent.rs +++ b/bottlecap/src/otlp/agent.rs @@ -61,7 +61,7 @@ impl TracePipeline { let lambda_extension_compute_stats = self.config.ext.lambda_extension_compute_stats; // Capture before `tracer_header_tags` is moved into process_traces below. let client_computed_stats = tracer_header_tags.generic.client_computed_stats; - let (send_data_builder, processed_traces) = self.trace_processor.process_traces( + let (send_data_builders, processed_traces) = self.trace_processor.process_traces( self.config.clone(), self.tags_provider.clone(), tracer_header_tags, @@ -70,11 +70,13 @@ impl TracePipeline { None, ); - if let Some(send_data_builder) = send_data_builder { - if let Err(err) = self.trace_tx.send(send_data_builder).await { - return Err(format!( - "Error sending traces to the trace aggregator: {err}" - )); + if !send_data_builders.is_empty() { + for send_data_builder in send_data_builders { + if let Err(err) = self.trace_tx.send(send_data_builder).await { + return Err(format!( + "Error sending traces to the trace aggregator: {err}" + )); + } } debug!("OTLP | Successfully buffered traces to be aggregated."); } diff --git a/bottlecap/src/traces/mod.rs b/bottlecap/src/traces/mod.rs index 0d9981805..29a251801 100644 --- a/bottlecap/src/traces/mod.rs +++ b/bottlecap/src/traces/mod.rs @@ -3,6 +3,7 @@ pub mod data_streams; pub mod http_client; +pub mod payload_split; pub mod propagation; pub mod proxy_aggregator; pub mod proxy_flusher; diff --git a/bottlecap/src/traces/payload_split.rs b/bottlecap/src/traces/payload_split.rs new file mode 100644 index 000000000..fde6ce2ff --- /dev/null +++ b/bottlecap/src/traces/payload_split.rs @@ -0,0 +1,311 @@ +// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Splits trace payloads that are too large to fit in a single outbound batch. +//! +//! A payload larger than the aggregator's batch cap can never be batched, so it is broken +//! up here before it is queued. Chunks are moved into separate payloads, and a chunk that +//! is still too large has its spans moved into several chunks of the same trace, the same +//! shape tracers produce with partial flushing. + +use libdd_trace_protobuf::pb; +use prost::Message; +use prost::encoding::message::encoded_len as field_len; +use tracing::warn; + +/// Protobuf field number of `TracerPayload.chunks`. +const CHUNKS_TAG: u32 = 6; +/// Protobuf field number of `TraceChunk.spans`. +const SPANS_TAG: u32 = 3; +/// Upper bound on the bytes a length-delimited field adds around its contents +/// (1-byte key for these tags plus a varint length of at most 5 bytes). +const FIELD_OVERHEAD: usize = 6; + +/// Splits `payloads` so that each returned payload encodes to at most `max_bytes`. +/// +/// Payloads that already fit are returned unchanged. A single span larger than the +/// limit cannot be split further and is dropped with a warning. +#[must_use] +pub fn split_tracer_payloads( + payloads: Vec, + max_bytes: usize, +) -> Vec { + let mut out = Vec::with_capacity(payloads.len()); + for payload in payloads { + if payload.encoded_len() <= max_bytes { + out.push(payload); + } else { + split_payload(payload, max_bytes, &mut out); + } + } + out +} + +fn split_payload( + mut payload: pb::TracerPayload, + max_bytes: usize, + out: &mut Vec, +) { + let chunks = std::mem::take(&mut payload.chunks); + let budget = max_bytes.saturating_sub(payload.encoded_len()); + + let mut current = Vec::new(); + let mut current_len = 0; + for chunk in chunks { + let pieces = if field_len(CHUNKS_TAG, &chunk) <= budget { + vec![chunk] + } else { + split_chunk(chunk, budget) + }; + for piece in pieces { + let piece_len = field_len(CHUNKS_TAG, &piece); + if !current.is_empty() && current_len + piece_len > budget { + out.push(with_chunks(&payload, std::mem::take(&mut current))); + current_len = 0; + } + current_len += piece_len; + current.push(piece); + } + } + if !current.is_empty() { + out.push(with_chunks(&payload, current)); + } +} + +/// Splits a chunk into chunks of the same trace whose encoded field size fits in `budget`. +/// +/// Trace-level `_dd.p.*` tags (e.g. the upper 64 bits of a 128-bit trace id) are only set on +/// one span of the chunk, so they are copied onto the first span of every piece. +fn split_chunk(mut chunk: pb::TraceChunk, budget: usize) -> Vec { + let spans = std::mem::take(&mut chunk.spans); + let trace_tags = trace_level_tags(&spans); + // Room for the trace tags added to each piece's first span, with slack for map entry + // framing and the span's length varint growing. + let tags_reserve: usize = trace_tags.iter().map(|(k, v)| k.len() + v.len() + 8).sum(); + let span_budget = budget.saturating_sub(chunk.encoded_len() + FIELD_OVERHEAD + tags_reserve); + + let mut pieces = Vec::new(); + let mut current = Vec::new(); + let mut current_len = 0; + let mut dropped = 0; + for span in spans { + let span_len = field_len(SPANS_TAG, &span); + if span_len > span_budget { + dropped += 1; + continue; + } + if !current.is_empty() && current_len + span_len > span_budget { + pieces.push(with_spans(&chunk, std::mem::take(&mut current))); + current_len = 0; + } + current_len += span_len; + current.push(span); + } + if !current.is_empty() { + pieces.push(with_spans(&chunk, current)); + } + if dropped > 0 { + warn!("TRACES | Dropped {dropped} span(s) too large to fit in a trace payload"); + } + for piece in &mut pieces { + if let Some(first) = piece.spans.first_mut() { + for (key, value) in &trace_tags { + first + .meta + .entry(key.clone()) + .or_insert_with(|| value.clone()); + } + } + } + pieces +} + +/// Returns the `_dd.p.*` tags of the chunk's root span, or of the first span that has any. +fn trace_level_tags(spans: &[pb::Span]) -> Vec<(String, String)> { + let has_trace_tags = |s: &&pb::Span| s.meta.keys().any(|k| k.starts_with("_dd.p.")); + spans + .iter() + .filter(has_trace_tags) + .find(|s| s.parent_id == 0) + .or_else(|| spans.iter().find(has_trace_tags)) + .map(|s| { + s.meta + .iter() + .filter(|(k, _)| k.starts_with("_dd.p.")) + .map(|(k, v)| (k.clone(), v.clone())) + .collect() + }) + .unwrap_or_default() +} + +fn with_chunks(template: &pb::TracerPayload, chunks: Vec) -> pb::TracerPayload { + pb::TracerPayload { + chunks, + ..template.clone() + } +} + +fn with_spans(template: &pb::TraceChunk, spans: Vec) -> pb::TraceChunk { + pb::TraceChunk { + spans, + ..template.clone() + } +} + +#[cfg(test)] +#[allow(clippy::unwrap_used)] +mod tests { + use super::*; + use std::collections::HashMap; + + fn span(trace_id: u64, span_id: u64, meta_bytes: usize) -> pb::Span { + pb::Span { + trace_id, + span_id, + name: "policy.evaluate".to_string(), + meta: HashMap::from([("payload".to_string(), "x".repeat(meta_bytes))]), + ..Default::default() + } + } + + fn chunk(trace_id: u64, spans: usize, meta_bytes: usize) -> pb::TraceChunk { + pb::TraceChunk { + priority: 1, + origin: "lambda".to_string(), + spans: (0..spans as u64) + .map(|i| span(trace_id, i + 1, meta_bytes)) + .collect(), + tags: HashMap::from([("_dd.p.dm".to_string(), "-0".to_string())]), + dropped_trace: false, + } + } + + fn payload(chunks: Vec) -> pb::TracerPayload { + pb::TracerPayload { + language_name: "go".to_string(), + tracer_version: "v1.74.6".to_string(), + env: "dev".to_string(), + tags: HashMap::from([("functionname".to_string(), "fn".to_string())]), + chunks, + ..Default::default() + } + } + + fn span_ids(payloads: &[pb::TracerPayload]) -> Vec<(u64, u64)> { + let mut ids: Vec<_> = payloads + .iter() + .flat_map(|p| p.chunks.iter()) + .flat_map(|c| c.spans.iter()) + .map(|s| (s.trace_id, s.span_id)) + .collect(); + ids.sort_unstable(); + ids + } + + #[test] + fn payloads_under_the_limit_are_unchanged() { + let input = vec![ + payload(vec![chunk(1, 10, 100)]), + payload(vec![chunk(2, 5, 100)]), + ]; + let out = split_tracer_payloads(input.clone(), 1_000_000); + assert_eq!(out, input); + } + + #[test] + fn oversized_payload_is_split_by_chunk() { + let input = vec![payload((1..=10).map(|t| chunk(t, 10, 1_000)).collect())]; + let limit = input[0].encoded_len() / 3; + + let out = split_tracer_payloads(input.clone(), limit); + + assert!( + out.len() >= 3, + "expected at least 3 payloads, got {}", + out.len() + ); + for p in &out { + assert!(p.encoded_len() <= limit, "{} > {limit}", p.encoded_len()); + assert_eq!(p.language_name, "go"); + assert_eq!(p.tags, input[0].tags); + } + assert_eq!(span_ids(&out), span_ids(&input)); + // Whole chunks are moved, not split, when each one fits on its own. + assert_eq!(out.iter().map(|p| p.chunks.len()).sum::(), 10); + } + + #[test] + fn oversized_single_chunk_is_split_by_span() { + // One trace with many spans, like a single large invocation. + let input = vec![payload(vec![chunk(7, 2_000, 500)])]; + let limit = input[0].encoded_len() / 5; + + let out = split_tracer_payloads(input.clone(), limit); + + assert!( + out.len() >= 5, + "expected at least 5 payloads, got {}", + out.len() + ); + for p in &out { + assert!(p.encoded_len() <= limit, "{} > {limit}", p.encoded_len()); + for c in &p.chunks { + assert_eq!(c.priority, 1); + assert_eq!(c.origin, "lambda"); + assert_eq!(c.tags, input[0].chunks[0].tags); + assert!(c.spans.iter().all(|s| s.trace_id == 7)); + } + } + assert_eq!(span_ids(&out), span_ids(&input)); + } + + #[test] + fn trace_level_tags_are_copied_to_every_split_chunk() { + let mut trace = chunk(9, 2_000, 500); + let root = &mut trace.spans[0]; + root.meta + .insert("_dd.p.tid".to_string(), "66f1e2a300000000".to_string()); + root.meta.insert("_dd.p.dm".to_string(), "-1".to_string()); + let input = vec![payload(vec![trace])]; + let limit = input[0].encoded_len() / 4; + + let out = split_tracer_payloads(input, limit); + + assert!(out.len() >= 4); + for p in &out { + assert!(p.encoded_len() <= limit, "{} > {limit}", p.encoded_len()); + for c in &p.chunks { + let first = &c.spans[0]; + assert_eq!(first.meta["_dd.p.tid"], "66f1e2a300000000"); + assert_eq!(first.meta["_dd.p.dm"], "-1"); + } + } + } + + #[test] + fn span_larger_than_the_limit_is_dropped() { + let mut big_chunk = chunk(3, 4, 100); + big_chunk.spans.push(span(3, 99, 50_000)); + let input = vec![payload(vec![big_chunk])]; + + let out = split_tracer_payloads(input, 10_000); + + let ids = span_ids(&out); + assert_eq!(ids, vec![(3, 1), (3, 2), (3, 3), (3, 4)]); + assert!(out.iter().all(|p| p.encoded_len() <= 10_000)); + } + + #[test] + fn mixed_payloads_only_split_the_oversized_one() { + let small = payload(vec![chunk(1, 2, 100)]); + let large = payload(vec![chunk(2, 1_000, 500)]); + let limit = large.encoded_len() / 2; + + let out = split_tracer_payloads(vec![small.clone(), large.clone()], limit); + + assert_eq!(out[0], small); + assert!(out.len() >= 3); + assert!(out.iter().all(|p| p.encoded_len() <= limit)); + assert_eq!(span_ids(&out), span_ids(&[small, large])); + } +} diff --git a/bottlecap/src/traces/trace_aggregator.rs b/bottlecap/src/traces/trace_aggregator.rs index 96031d0ed..eb081bb34 100644 --- a/bottlecap/src/traces/trace_aggregator.rs +++ b/bottlecap/src/traces/trace_aggregator.rs @@ -107,12 +107,18 @@ impl TraceAggregator { // Fill the batch while batch_size < self.max_content_size_bytes { if let Some(payload_info) = self.queue.pop_front() { - // TODO(duncanista): revisit if this is bigger than limit let payload_size = payload_info.size; - // Put stats back in the queue if batch_size + payload_size > self.max_content_size_bytes { - self.queue.push_front(payload_info); + // A payload larger than the cap on its own is sent as a batch by itself. + // Putting it back would leave it at the front of the queue forever and + // block every payload behind it. + if self.buffer.is_empty() { + batch_size += payload_size; + self.buffer.push(payload_info); + } else { + self.queue.push_front(payload_info); + } break; } batch_size += payload_size; @@ -221,4 +227,23 @@ mod tests { assert_eq!(second_batch.len(), 1); assert_eq!(aggregator.queue.len(), 0); } + + #[test] + fn test_get_batch_oversized_payload_does_not_block_queue() { + let mut aggregator = TraceAggregator::new(10); + + aggregator.add(make_builder_info(4)); + aggregator.add(make_builder_info(25)); + aggregator.add(make_builder_info(3)); + + // The oversized payload doesn't fit after the first one, so it waits. + assert_eq!(aggregator.get_batch().len(), 1); + // On its own it is sent as a batch by itself instead of being put back. + let oversized = aggregator.get_batch(); + assert_eq!(oversized.len(), 1); + assert_eq!(oversized[0].size, 25); + // Payloads behind it are still sent. + assert_eq!(aggregator.get_batch().len(), 1); + assert_eq!(aggregator.queue.len(), 0); + } } diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs index 9fc07480f..2f1f8ce6f 100644 --- a/bottlecap/src/traces/trace_processor.rs +++ b/bottlecap/src/traces/trace_processor.rs @@ -36,8 +36,11 @@ use tokio::sync::mpsc::Sender; use tokio::sync::mpsc::error::SendError; use tracing::{debug, error}; +use crate::traces::payload_split::split_tracer_payloads; use crate::traces::stats_generator::StatsGenerator; -use crate::traces::trace_aggregator::{OwnedTracerHeaderTags, SendDataBuilderInfo}; +use crate::traces::trace_aggregator::{ + MAX_CONTENT_SIZE_BYTES, OwnedTracerHeaderTags, SendDataBuilderInfo, +}; use libdd_trace_normalization::normalizer::SamplerPriority; /// Which party is responsible for computing trace stats for a trace, derived from @@ -512,7 +515,7 @@ pub trait TraceProcessor { traces: Vec>, body_size: usize, span_pointers: Option>, - ) -> (Option, TracerPayloadCollection); + ) -> (Vec, TracerPayloadCollection); } #[async_trait] @@ -525,7 +528,7 @@ impl TraceProcessor for ServerlessTraceProcessor { traces: Vec>, body_size: usize, span_pointers: Option>, - ) -> (Option, TracerPayloadCollection) { + ) -> (Vec, TracerPayloadCollection) { let mut payload = trace_utils::collect_pb_trace_chunks( traces, &header_tags, @@ -579,7 +582,7 @@ impl TraceProcessor for ServerlessTraceProcessor { { self.drop_sampled_out_chunks(tracer_payloads); if tracer_payloads.is_empty() { - return (None, payloads_for_stats); + return (Vec::new(), payloads_for_stats); } } @@ -593,28 +596,45 @@ impl TraceProcessor for ServerlessTraceProcessor { debug!("TRACES | trace payload size after enrichment: {body_size} bytes"); - let owned_header_tags = OwnedTracerHeaderTags::from(header_tags.clone()); + // A payload larger than the batch cap can never be batched, so split it into + // payloads that each fit. + let payloads = match payload { + TracerPayloadCollection::V07(tracer_payloads) if body_size > MAX_CONTENT_SIZE_BYTES => { + let split = split_tracer_payloads(tracer_payloads, MAX_CONTENT_SIZE_BYTES); + debug!( + "TRACES | split {body_size} byte trace payload into {} payloads", + split.len() + ); + split + .into_iter() + .map(|p| { + let size = prost::Message::encoded_len(&p); + (TracerPayloadCollection::V07(vec![p]), size) + }) + .collect() + } + other => vec![(other, body_size)], + }; - // Move original payload into builder (no clone needed) - let builder = SendDataBuilder::new(body_size, payload, header_tags, &endpoint) - .with_compression(CompressionStrategy::Zstd { - level: config.apm_config_compression_level, + let builders = payloads + .into_iter() + .map(|(payload, size)| { + let owned_header_tags = OwnedTracerHeaderTags::from(header_tags.clone()); + let builder = SendDataBuilder::new(size, payload, header_tags.clone(), &endpoint) + .with_compression(CompressionStrategy::Zstd { + level: config.apm_config_compression_level, + }) + .with_retry_strategy(RetryStrategy::new( + 1, + 100, + RetryBackoffType::Exponential, + None, + )); + SendDataBuilderInfo::new(builder, size, owned_header_tags) }) - .with_retry_strategy(RetryStrategy::new( - 1, - 100, - RetryBackoffType::Exponential, - None, - )); - - ( - Some(SendDataBuilderInfo::new( - builder, - body_size, - owned_header_tags, - )), - payloads_for_stats, - ) + .collect(); + + (builders, payloads_for_stats) } } @@ -697,7 +717,7 @@ impl SendingTraceProcessor { // Capture before `header_tags` is moved into process_traces below. let client_computed_stats = header_tags.generic.client_computed_stats; - let (payload, processed_traces) = self.processor.process_traces( + let (payloads, processed_traces) = self.processor.process_traces( config.clone(), tags_provider, header_tags, @@ -721,7 +741,7 @@ impl SendingTraceProcessor { error!("TRACE_PROCESSOR | Error sending traces to the stats concentrator: {err}"); } - if let Some(payload) = payload { + for payload in payloads { self.trace_tx.send(payload).await?; } Ok(()) @@ -910,7 +930,10 @@ mod tests { 100, None, ); - let tracer_payload = tracer_payload.expect("expected Some payload"); + let tracer_payload = tracer_payload + .into_iter() + .next() + .expect("expected a payload"); let expected_tracer_payload = pb::TracerPayload { container_id: "33".to_string(), @@ -1385,7 +1408,7 @@ mod tests { let (payload_info, stats_collection) = processor.process_traces(config, tags_provider, header_tags, traces, 0, None); - let payload_info = payload_info.expect("expected Some payload"); + let payload_info = payload_info.into_iter().next().expect("expected a payload"); // Stats collection must include all three traces let TracerPayloadCollection::V07(ref stats_payloads) = stats_collection else { @@ -1451,7 +1474,7 @@ mod tests { let (payload_info, _stats) = processor.process_traces(config, tags_provider, header_tags, traces, 0, None); - let payload_info = payload_info.expect("expected Some payload"); + let payload_info = payload_info.into_iter().next().expect("expected a payload"); let backend_send_data = payload_info.builder.build(); let TracerPayloadCollection::V07(backend_payloads) = backend_send_data.get_payloads() else { @@ -1512,7 +1535,10 @@ mod tests { let (payload_info, _stats) = processor.process_traces(config, tags_provider, header_tags, traces, 0, None); - let payload_info = payload_info.expect("errored-child P0 trace rescued"); + let payload_info = payload_info + .into_iter() + .next() + .expect("errored-child P0 trace rescued"); let backend_send_data = payload_info.builder.build(); let TracerPayloadCollection::V07(backend_payloads) = backend_send_data.get_payloads() else { @@ -1556,7 +1582,7 @@ mod tests { let (payload_info, _stats) = processor.process_traces(config, tags_provider, header_tags, traces, 0, None); assert!( - payload_info.is_none(), + payload_info.is_empty(), "errored P0 trace must stay dropped when the error sampler is disabled" ); } @@ -1592,7 +1618,10 @@ mod tests { let (payload_info, stats_collection) = processor.process_traces(config, tags_provider, header_tags, traces, 0, None); - let payload_info = payload_info.expect("kept trace must produce a backend payload"); + let payload_info = payload_info + .into_iter() + .next() + .expect("kept trace must produce a backend payload"); // Stats collection must include both traces, including the dropped errored P0. let TracerPayloadCollection::V07(ref stats_payloads) = stats_collection else { @@ -1658,8 +1687,8 @@ mod tests { processor.process_traces(config, tags_provider, header_tags, traces, 0, None); assert!( - payload.is_none(), - "backend payload must be None when all traces are sampled out" + payload.is_empty(), + "no backend payload must be produced when all traces are sampled out" ); // Stats collection must still include both traces @@ -1713,7 +1742,7 @@ mod tests { let (payload_info, stats_collection) = processor.process_traces(config, tags_provider, header_tags, traces, 999_999, None); - let info = payload_info.expect("expected Some payload"); + let info = payload_info.into_iter().next().expect("expected a payload"); // The reported size must equal the sum of encoded_len() of the kept TracerPayloads. // stats_collection has all 4 traces. Reconstruct the filtered payload (only trace_id=1 @@ -1793,7 +1822,7 @@ mod tests { None, ); - let info = payload_info.expect("expected Some payload"); + let info = payload_info.into_iter().next().expect("expected a payload"); // Nothing is filtered here, so the returned payloads are exactly what gets sent. let TracerPayloadCollection::V07(ref payloads) = processed_payloads else { @@ -1811,6 +1840,84 @@ mod tests { ); } + /// A single trace whose enriched payload exceeds the aggregator batch cap must be + /// split into multiple payloads that each fit, without losing or duplicating spans. + #[test] + fn test_process_traces_splits_oversized_trace() { + const TRACE_ID: u64 = 42; + const SPAN_COUNT: u64 = 40_000; + + let config = Arc::new(Config { + apm_dd_url: "https://trace.agent.datadoghq.com".to_string(), + ..Config::default() + }); + let (tags_provider, processor) = create_test_processor(&config, enabled_error_sampler()); + let header_tags = create_test_header_tags(); + + // Span 1 is the root; every other span is its child. + let make_span = |span_id: u64| -> pb::Span { + pb::Span { + trace_id: TRACE_ID, + span_id, + parent_id: u64::from(span_id != 1), + service: "svc".to_string(), + name: "op".to_string(), + resource: "res".to_string(), + meta: HashMap::from([("payload".to_string(), "x".repeat(300))]), + ..Default::default() + } + }; + let traces = vec![(1..=SPAN_COUNT).map(make_span).collect::>()]; + + let (payload_infos, stats_collection) = + processor.process_traces(config, tags_provider, header_tags, traces, 1, None); + + assert!( + payload_infos.len() > 1, + "oversized trace must be split into multiple payloads" + ); + assert!( + payload_infos + .iter() + .all(|info| info.size <= MAX_CONTENT_SIZE_BYTES), + "every split payload must fit within the batch cap" + ); + + let mut sent_span_ids = Vec::new(); + for info in payload_infos { + let send_data = info.builder.build(); + let TracerPayloadCollection::V07(payloads) = send_data.get_payloads() else { + panic!("expected V07"); + }; + for chunk in payloads.iter().flat_map(|tp| &tp.chunks) { + for span in &chunk.spans { + assert_eq!(span.trace_id, TRACE_ID); + sent_span_ids.push(span.span_id); + } + } + } + sent_span_ids.sort_unstable(); + let expected_span_ids: Vec = (1..=SPAN_COUNT).collect(); + assert_eq!( + sent_span_ids, expected_span_ids, + "split payloads must contain every input span exactly once" + ); + + let TracerPayloadCollection::V07(ref stats_payloads) = stats_collection else { + panic!("expected V07"); + }; + let stats_span_count: usize = stats_payloads + .iter() + .flat_map(|tp| &tp.chunks) + .map(|c| c.spans.len()) + .sum(); + assert_eq!( + stats_span_count, + expected_span_ids.len(), + "stats must include all input spans" + ); + } + fn create_inferred_span() -> pb::Span { let mut meta = HashMap::new(); meta.insert("_inferred_span.tag_source".to_string(), "self".to_string()); From a64cb431a0f8ca1a61f57f2397894ecc647f780c Mon Sep 17 00:00:00 2001 From: "cliford.shelton" Date: Fri, 9 Oct 2026 20:28:27 -0400 Subject: [PATCH 2/4] fix(traces): size trace tags exactly when splitting a chunk The room for copied `_dd.p.*` tags was reserved from every span's budget, although only the first span of each piece gains them and spans that already carry them gain nothing. A span within a few bytes of the limit, such as a large root span, could be dropped even though it fit. Charge each span only for the tags it lacks, and only when it starts a new piece. Co-Authored-By: Claude Opus 5.5 --- bottlecap/src/traces/payload_split.rs | 69 ++++++++++++++++++++++++--- 1 file changed, 62 insertions(+), 7 deletions(-) diff --git a/bottlecap/src/traces/payload_split.rs b/bottlecap/src/traces/payload_split.rs index fde6ce2ff..58a9dbf76 100644 --- a/bottlecap/src/traces/payload_split.rs +++ b/bottlecap/src/traces/payload_split.rs @@ -11,6 +11,8 @@ use libdd_trace_protobuf::pb; use prost::Message; use prost::encoding::message::encoded_len as field_len; +use prost::encoding::{encoded_len_varint, key_len}; +use std::collections::HashMap; use tracing::warn; /// Protobuf field number of `TracerPayload.chunks`. @@ -79,10 +81,7 @@ fn split_payload( fn split_chunk(mut chunk: pb::TraceChunk, budget: usize) -> Vec { let spans = std::mem::take(&mut chunk.spans); let trace_tags = trace_level_tags(&spans); - // Room for the trace tags added to each piece's first span, with slack for map entry - // framing and the span's length varint growing. - let tags_reserve: usize = trace_tags.iter().map(|(k, v)| k.len() + v.len() + 8).sum(); - let span_budget = budget.saturating_sub(chunk.encoded_len() + FIELD_OVERHEAD + tags_reserve); + let span_budget = budget.saturating_sub(chunk.encoded_len() + FIELD_OVERHEAD); let mut pieces = Vec::new(); let mut current = Vec::new(); @@ -90,7 +89,9 @@ fn split_chunk(mut chunk: pb::TraceChunk, budget: usize) -> Vec let mut dropped = 0; for span in spans { let span_len = field_len(SPANS_TAG, &span); - if span_len > span_budget { + // Only a piece's first span gains the trace tags, so this is charged only then. + let extra = trace_tags_growth(&span, span_len, &trace_tags); + if span_len + extra > span_budget { dropped += 1; continue; } @@ -98,7 +99,11 @@ fn split_chunk(mut chunk: pb::TraceChunk, budget: usize) -> Vec pieces.push(with_spans(&chunk, std::mem::take(&mut current))); current_len = 0; } - current_len += span_len; + current_len += if current.is_empty() { + span_len + extra + } else { + span_len + }; current.push(span); } if !current.is_empty() { @@ -120,6 +125,27 @@ fn split_chunk(mut chunk: pb::TraceChunk, budget: usize) -> Vec pieces } +/// Returns how much the encoded `spans` field of `span` (currently `span_len`) grows when +/// the trace tags it lacks are added to its meta. +fn trace_tags_growth(span: &pb::Span, span_len: usize, trace_tags: &[(String, String)]) -> usize { + let missing: HashMap = trace_tags + .iter() + .filter(|(k, _)| !span.meta.contains_key(k)) + .cloned() + .collect(); + if missing.is_empty() { + return 0; + } + // A span holding only the missing tags encodes to exactly the added map entries. + let added = pb::Span { + meta: missing, + ..Default::default() + } + .encoded_len(); + let new_len = span.encoded_len() + added; + key_len(SPANS_TAG) + encoded_len_varint(new_len as u64) + new_len - span_len +} + /// Returns the `_dd.p.*` tags of the chunk's root span, or of the first span that has any. fn trace_level_tags(spans: &[pb::Span]) -> Vec<(String, String)> { let has_trace_tags = |s: &&pb::Span| s.meta.keys().any(|k| k.starts_with("_dd.p.")); @@ -156,7 +182,6 @@ fn with_spans(template: &pb::TraceChunk, spans: Vec) -> pb::TraceChunk #[allow(clippy::unwrap_used)] mod tests { use super::*; - use std::collections::HashMap; fn span(trace_id: u64, span_id: u64, meta_bytes: usize) -> pb::Span { pb::Span { @@ -308,4 +333,34 @@ mod tests { assert!(out.iter().all(|p| p.encoded_len() <= limit)); assert_eq!(span_ids(&out), span_ids(&[small, large])); } + + #[test] + fn near_limit_span_with_trace_tags_is_kept() { + let mut root = span(5, 1, 0); + root.meta + .insert("_dd.p.tid".to_string(), "66f1e2a300000000".to_string()); + root.meta.insert("_dd.p.dm".to_string(), "-1".to_string()); + root.meta.insert("payload".to_string(), "x".repeat(5_000)); + let mut trace = chunk(5, 0, 0); + let empty_chunk_len = trace.encoded_len(); + trace.spans.push(root); + trace.spans.extend((2..=5).map(|i| span(5, i, 100))); + let input = payload(vec![trace]); + let empty_payload_len = { + let mut p = input.clone(); + p.chunks.clear(); + p.encoded_len() + }; + // The root span leaves only a few bytes of the space available to a single span. It + // already carries the trace tags, so it gains nothing when they are propagated. + let root_len = field_len(SPANS_TAG, &input.chunks[0].spans[0]); + let limit = empty_payload_len + empty_chunk_len + FIELD_OVERHEAD + root_len + 10; + assert!(input.encoded_len() > limit); + + let out = split_tracer_payloads(vec![input.clone()], limit); + + assert!(span_ids(&out).contains(&(5, 1)), "root span was dropped"); + assert_eq!(span_ids(&out), span_ids(&[input])); + assert!(out.iter().all(|p| p.encoded_len() <= limit)); + } } From 32a6028808c87e133fa6ea47b3cd1caccdadd6b0 Mon Sep 17 00:00:00 2001 From: "cliford.shelton" Date: Tue, 13 Oct 2026 12:00:10 -0400 Subject: [PATCH 3/4] fix(traces): keep process tags on every split trace payload Tracers set `_dd.tags.process` only on the first span of a payload's first chunk, so after splitting, only the payload holding that span carried the process tags. Promote them to the payload-level tags before splitting, as the Agent does, so every output inherits them and their size is counted. An existing payload-level value takes priority. Co-Authored-By: Claude Opus 5.5 --- bottlecap/src/traces/payload_split.rs | 86 +++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/bottlecap/src/traces/payload_split.rs b/bottlecap/src/traces/payload_split.rs index 58a9dbf76..925a9a319 100644 --- a/bottlecap/src/traces/payload_split.rs +++ b/bottlecap/src/traces/payload_split.rs @@ -22,6 +22,9 @@ const SPANS_TAG: u32 = 3; /// Upper bound on the bytes a length-delimited field adds around its contents /// (1-byte key for these tags plus a varint length of at most 5 bytes). const FIELD_OVERHEAD: usize = 6; +/// Process tags key. Tracers set it only on the first span of a payload's first chunk, and +/// the Agent reads the payload-level tag first, so it is promoted there when splitting. +const PROCESS_TAGS_KEY: &str = "_dd.tags.process"; /// Splits `payloads` so that each returned payload encodes to at most `max_bytes`. /// @@ -48,6 +51,7 @@ fn split_payload( max_bytes: usize, out: &mut Vec, ) { + promote_process_tags(&mut payload); let chunks = std::mem::take(&mut payload.chunks); let budget = max_bytes.saturating_sub(payload.encoded_len()); @@ -74,6 +78,24 @@ fn split_payload( } } +/// Copies the first span's process tags to the payload tags, so every split payload keeps +/// them. An existing payload-level value is left unchanged. +fn promote_process_tags(payload: &mut pb::TracerPayload) { + if payload.tags.contains_key(PROCESS_TAGS_KEY) { + return; + } + let process_tags = payload + .chunks + .first() + .and_then(|c| c.spans.first()) + .and_then(|s| s.meta.get(PROCESS_TAGS_KEY)); + if let Some(value) = process_tags { + payload + .tags + .insert(PROCESS_TAGS_KEY.to_string(), value.clone()); + } +} + /// Splits a chunk into chunks of the same trace whose encoded field size fits in `budget`. /// /// Trace-level `_dd.p.*` tags (e.g. the upper 64 bits of a 128-bit trace id) are only set on @@ -363,4 +385,68 @@ mod tests { assert_eq!(span_ids(&out), span_ids(&[input])); assert!(out.iter().all(|p| p.encoded_len() <= limit)); } + + const PROCESS_TAGS: &str = "entrypoint.name:bootstrap,entrypoint.type:executable"; + + fn with_process_tags_on_first_span(mut input: pb::TracerPayload) -> pb::TracerPayload { + input.chunks[0].spans[0] + .meta + .insert(PROCESS_TAGS_KEY.to_string(), PROCESS_TAGS.to_string()); + input + } + + fn assert_process_tags_on_every_output(input: &pb::TracerPayload, min_outputs: usize) { + let limit = input.encoded_len() / min_outputs; + + let out = split_tracer_payloads(vec![input.clone()], limit); + + assert!(out.len() >= min_outputs, "got {} payloads", out.len()); + for p in &out { + assert!(p.encoded_len() <= limit, "{} > {limit}", p.encoded_len()); + assert_eq!(p.tags[PROCESS_TAGS_KEY], PROCESS_TAGS); + } + assert_eq!(span_ids(&out), span_ids(std::slice::from_ref(input))); + } + + #[test] + fn process_tags_are_promoted_when_splitting_by_chunk() { + let input = with_process_tags_on_first_span(payload( + (1..=10).map(|t| chunk(t, 10, 1_000)).collect(), + )); + assert_process_tags_on_every_output(&input, 3); + } + + #[test] + fn process_tags_are_promoted_when_splitting_by_span() { + let input = with_process_tags_on_first_span(payload(vec![chunk(7, 2_000, 500)])); + assert_process_tags_on_every_output(&input, 5); + } + + #[test] + fn payload_level_process_tags_take_priority() { + let mut input = with_process_tags_on_first_span(payload(vec![chunk(7, 2_000, 500)])); + input + .tags + .insert(PROCESS_TAGS_KEY.to_string(), "payload-level".to_string()); + let limit = input.encoded_len() / 4; + + let out = split_tracer_payloads(vec![input], limit); + + assert!(out.len() >= 4); + for p in &out { + assert!(p.encoded_len() <= limit, "{} > {limit}", p.encoded_len()); + assert_eq!(p.tags[PROCESS_TAGS_KEY], "payload-level"); + } + } + + #[test] + fn payload_without_process_tags_gains_none() { + let input = payload((1..=10).map(|t| chunk(t, 10, 1_000)).collect()); + let limit = input.encoded_len() / 3; + + let out = split_tracer_payloads(vec![input], limit); + + assert!(out.len() >= 3); + assert!(out.iter().all(|p| !p.tags.contains_key(PROCESS_TAGS_KEY))); + } } From cb1806a1cbbd09bfbdcb25dbddb99796e0761a44 Mon Sep 17 00:00:00 2001 From: "cliford.shelton" Date: Tue, 13 Oct 2026 14:39:38 -0400 Subject: [PATCH 4/4] revert(traces): drop trace payload splitting, keep the get_batch fix Match the Go trace-agent, which never splits a trace payload: an oversized payload is sent on its own. Large traces should be split by the tracer with partial flushing (`DD_TRACE_PARTIAL_FLUSH_ENABLED`), which keeps them complete and well under the batch cap. This removes the payload splitting, the trace-level and process tag propagation, and the `process_traces` return type change. What remains is the `get_batch` change: a payload larger than the cap on its own is sent as a batch by itself instead of being put back at the front of the queue, where it blocked every later payload until the Lambda ran out of memory. Co-Authored-By: Claude Opus 5.5 --- bottlecap/src/otlp/agent.rs | 14 +- bottlecap/src/traces/mod.rs | 1 - bottlecap/src/traces/payload_split.rs | 452 ------------------------ bottlecap/src/traces/trace_processor.rs | 179 ++-------- 4 files changed, 42 insertions(+), 604 deletions(-) delete mode 100644 bottlecap/src/traces/payload_split.rs diff --git a/bottlecap/src/otlp/agent.rs b/bottlecap/src/otlp/agent.rs index 47083ecc4..20dccc3bc 100644 --- a/bottlecap/src/otlp/agent.rs +++ b/bottlecap/src/otlp/agent.rs @@ -61,7 +61,7 @@ impl TracePipeline { let lambda_extension_compute_stats = self.config.ext.lambda_extension_compute_stats; // Capture before `tracer_header_tags` is moved into process_traces below. let client_computed_stats = tracer_header_tags.generic.client_computed_stats; - let (send_data_builders, processed_traces) = self.trace_processor.process_traces( + let (send_data_builder, processed_traces) = self.trace_processor.process_traces( self.config.clone(), self.tags_provider.clone(), tracer_header_tags, @@ -70,13 +70,11 @@ impl TracePipeline { None, ); - if !send_data_builders.is_empty() { - for send_data_builder in send_data_builders { - if let Err(err) = self.trace_tx.send(send_data_builder).await { - return Err(format!( - "Error sending traces to the trace aggregator: {err}" - )); - } + if let Some(send_data_builder) = send_data_builder { + if let Err(err) = self.trace_tx.send(send_data_builder).await { + return Err(format!( + "Error sending traces to the trace aggregator: {err}" + )); } debug!("OTLP | Successfully buffered traces to be aggregated."); } diff --git a/bottlecap/src/traces/mod.rs b/bottlecap/src/traces/mod.rs index 29a251801..0d9981805 100644 --- a/bottlecap/src/traces/mod.rs +++ b/bottlecap/src/traces/mod.rs @@ -3,7 +3,6 @@ pub mod data_streams; pub mod http_client; -pub mod payload_split; pub mod propagation; pub mod proxy_aggregator; pub mod proxy_flusher; diff --git a/bottlecap/src/traces/payload_split.rs b/bottlecap/src/traces/payload_split.rs deleted file mode 100644 index 925a9a319..000000000 --- a/bottlecap/src/traces/payload_split.rs +++ /dev/null @@ -1,452 +0,0 @@ -// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/ -// SPDX-License-Identifier: Apache-2.0 - -//! Splits trace payloads that are too large to fit in a single outbound batch. -//! -//! A payload larger than the aggregator's batch cap can never be batched, so it is broken -//! up here before it is queued. Chunks are moved into separate payloads, and a chunk that -//! is still too large has its spans moved into several chunks of the same trace, the same -//! shape tracers produce with partial flushing. - -use libdd_trace_protobuf::pb; -use prost::Message; -use prost::encoding::message::encoded_len as field_len; -use prost::encoding::{encoded_len_varint, key_len}; -use std::collections::HashMap; -use tracing::warn; - -/// Protobuf field number of `TracerPayload.chunks`. -const CHUNKS_TAG: u32 = 6; -/// Protobuf field number of `TraceChunk.spans`. -const SPANS_TAG: u32 = 3; -/// Upper bound on the bytes a length-delimited field adds around its contents -/// (1-byte key for these tags plus a varint length of at most 5 bytes). -const FIELD_OVERHEAD: usize = 6; -/// Process tags key. Tracers set it only on the first span of a payload's first chunk, and -/// the Agent reads the payload-level tag first, so it is promoted there when splitting. -const PROCESS_TAGS_KEY: &str = "_dd.tags.process"; - -/// Splits `payloads` so that each returned payload encodes to at most `max_bytes`. -/// -/// Payloads that already fit are returned unchanged. A single span larger than the -/// limit cannot be split further and is dropped with a warning. -#[must_use] -pub fn split_tracer_payloads( - payloads: Vec, - max_bytes: usize, -) -> Vec { - let mut out = Vec::with_capacity(payloads.len()); - for payload in payloads { - if payload.encoded_len() <= max_bytes { - out.push(payload); - } else { - split_payload(payload, max_bytes, &mut out); - } - } - out -} - -fn split_payload( - mut payload: pb::TracerPayload, - max_bytes: usize, - out: &mut Vec, -) { - promote_process_tags(&mut payload); - let chunks = std::mem::take(&mut payload.chunks); - let budget = max_bytes.saturating_sub(payload.encoded_len()); - - let mut current = Vec::new(); - let mut current_len = 0; - for chunk in chunks { - let pieces = if field_len(CHUNKS_TAG, &chunk) <= budget { - vec![chunk] - } else { - split_chunk(chunk, budget) - }; - for piece in pieces { - let piece_len = field_len(CHUNKS_TAG, &piece); - if !current.is_empty() && current_len + piece_len > budget { - out.push(with_chunks(&payload, std::mem::take(&mut current))); - current_len = 0; - } - current_len += piece_len; - current.push(piece); - } - } - if !current.is_empty() { - out.push(with_chunks(&payload, current)); - } -} - -/// Copies the first span's process tags to the payload tags, so every split payload keeps -/// them. An existing payload-level value is left unchanged. -fn promote_process_tags(payload: &mut pb::TracerPayload) { - if payload.tags.contains_key(PROCESS_TAGS_KEY) { - return; - } - let process_tags = payload - .chunks - .first() - .and_then(|c| c.spans.first()) - .and_then(|s| s.meta.get(PROCESS_TAGS_KEY)); - if let Some(value) = process_tags { - payload - .tags - .insert(PROCESS_TAGS_KEY.to_string(), value.clone()); - } -} - -/// Splits a chunk into chunks of the same trace whose encoded field size fits in `budget`. -/// -/// Trace-level `_dd.p.*` tags (e.g. the upper 64 bits of a 128-bit trace id) are only set on -/// one span of the chunk, so they are copied onto the first span of every piece. -fn split_chunk(mut chunk: pb::TraceChunk, budget: usize) -> Vec { - let spans = std::mem::take(&mut chunk.spans); - let trace_tags = trace_level_tags(&spans); - let span_budget = budget.saturating_sub(chunk.encoded_len() + FIELD_OVERHEAD); - - let mut pieces = Vec::new(); - let mut current = Vec::new(); - let mut current_len = 0; - let mut dropped = 0; - for span in spans { - let span_len = field_len(SPANS_TAG, &span); - // Only a piece's first span gains the trace tags, so this is charged only then. - let extra = trace_tags_growth(&span, span_len, &trace_tags); - if span_len + extra > span_budget { - dropped += 1; - continue; - } - if !current.is_empty() && current_len + span_len > span_budget { - pieces.push(with_spans(&chunk, std::mem::take(&mut current))); - current_len = 0; - } - current_len += if current.is_empty() { - span_len + extra - } else { - span_len - }; - current.push(span); - } - if !current.is_empty() { - pieces.push(with_spans(&chunk, current)); - } - if dropped > 0 { - warn!("TRACES | Dropped {dropped} span(s) too large to fit in a trace payload"); - } - for piece in &mut pieces { - if let Some(first) = piece.spans.first_mut() { - for (key, value) in &trace_tags { - first - .meta - .entry(key.clone()) - .or_insert_with(|| value.clone()); - } - } - } - pieces -} - -/// Returns how much the encoded `spans` field of `span` (currently `span_len`) grows when -/// the trace tags it lacks are added to its meta. -fn trace_tags_growth(span: &pb::Span, span_len: usize, trace_tags: &[(String, String)]) -> usize { - let missing: HashMap = trace_tags - .iter() - .filter(|(k, _)| !span.meta.contains_key(k)) - .cloned() - .collect(); - if missing.is_empty() { - return 0; - } - // A span holding only the missing tags encodes to exactly the added map entries. - let added = pb::Span { - meta: missing, - ..Default::default() - } - .encoded_len(); - let new_len = span.encoded_len() + added; - key_len(SPANS_TAG) + encoded_len_varint(new_len as u64) + new_len - span_len -} - -/// Returns the `_dd.p.*` tags of the chunk's root span, or of the first span that has any. -fn trace_level_tags(spans: &[pb::Span]) -> Vec<(String, String)> { - let has_trace_tags = |s: &&pb::Span| s.meta.keys().any(|k| k.starts_with("_dd.p.")); - spans - .iter() - .filter(has_trace_tags) - .find(|s| s.parent_id == 0) - .or_else(|| spans.iter().find(has_trace_tags)) - .map(|s| { - s.meta - .iter() - .filter(|(k, _)| k.starts_with("_dd.p.")) - .map(|(k, v)| (k.clone(), v.clone())) - .collect() - }) - .unwrap_or_default() -} - -fn with_chunks(template: &pb::TracerPayload, chunks: Vec) -> pb::TracerPayload { - pb::TracerPayload { - chunks, - ..template.clone() - } -} - -fn with_spans(template: &pb::TraceChunk, spans: Vec) -> pb::TraceChunk { - pb::TraceChunk { - spans, - ..template.clone() - } -} - -#[cfg(test)] -#[allow(clippy::unwrap_used)] -mod tests { - use super::*; - - fn span(trace_id: u64, span_id: u64, meta_bytes: usize) -> pb::Span { - pb::Span { - trace_id, - span_id, - name: "policy.evaluate".to_string(), - meta: HashMap::from([("payload".to_string(), "x".repeat(meta_bytes))]), - ..Default::default() - } - } - - fn chunk(trace_id: u64, spans: usize, meta_bytes: usize) -> pb::TraceChunk { - pb::TraceChunk { - priority: 1, - origin: "lambda".to_string(), - spans: (0..spans as u64) - .map(|i| span(trace_id, i + 1, meta_bytes)) - .collect(), - tags: HashMap::from([("_dd.p.dm".to_string(), "-0".to_string())]), - dropped_trace: false, - } - } - - fn payload(chunks: Vec) -> pb::TracerPayload { - pb::TracerPayload { - language_name: "go".to_string(), - tracer_version: "v1.74.6".to_string(), - env: "dev".to_string(), - tags: HashMap::from([("functionname".to_string(), "fn".to_string())]), - chunks, - ..Default::default() - } - } - - fn span_ids(payloads: &[pb::TracerPayload]) -> Vec<(u64, u64)> { - let mut ids: Vec<_> = payloads - .iter() - .flat_map(|p| p.chunks.iter()) - .flat_map(|c| c.spans.iter()) - .map(|s| (s.trace_id, s.span_id)) - .collect(); - ids.sort_unstable(); - ids - } - - #[test] - fn payloads_under_the_limit_are_unchanged() { - let input = vec![ - payload(vec![chunk(1, 10, 100)]), - payload(vec![chunk(2, 5, 100)]), - ]; - let out = split_tracer_payloads(input.clone(), 1_000_000); - assert_eq!(out, input); - } - - #[test] - fn oversized_payload_is_split_by_chunk() { - let input = vec![payload((1..=10).map(|t| chunk(t, 10, 1_000)).collect())]; - let limit = input[0].encoded_len() / 3; - - let out = split_tracer_payloads(input.clone(), limit); - - assert!( - out.len() >= 3, - "expected at least 3 payloads, got {}", - out.len() - ); - for p in &out { - assert!(p.encoded_len() <= limit, "{} > {limit}", p.encoded_len()); - assert_eq!(p.language_name, "go"); - assert_eq!(p.tags, input[0].tags); - } - assert_eq!(span_ids(&out), span_ids(&input)); - // Whole chunks are moved, not split, when each one fits on its own. - assert_eq!(out.iter().map(|p| p.chunks.len()).sum::(), 10); - } - - #[test] - fn oversized_single_chunk_is_split_by_span() { - // One trace with many spans, like a single large invocation. - let input = vec![payload(vec![chunk(7, 2_000, 500)])]; - let limit = input[0].encoded_len() / 5; - - let out = split_tracer_payloads(input.clone(), limit); - - assert!( - out.len() >= 5, - "expected at least 5 payloads, got {}", - out.len() - ); - for p in &out { - assert!(p.encoded_len() <= limit, "{} > {limit}", p.encoded_len()); - for c in &p.chunks { - assert_eq!(c.priority, 1); - assert_eq!(c.origin, "lambda"); - assert_eq!(c.tags, input[0].chunks[0].tags); - assert!(c.spans.iter().all(|s| s.trace_id == 7)); - } - } - assert_eq!(span_ids(&out), span_ids(&input)); - } - - #[test] - fn trace_level_tags_are_copied_to_every_split_chunk() { - let mut trace = chunk(9, 2_000, 500); - let root = &mut trace.spans[0]; - root.meta - .insert("_dd.p.tid".to_string(), "66f1e2a300000000".to_string()); - root.meta.insert("_dd.p.dm".to_string(), "-1".to_string()); - let input = vec![payload(vec![trace])]; - let limit = input[0].encoded_len() / 4; - - let out = split_tracer_payloads(input, limit); - - assert!(out.len() >= 4); - for p in &out { - assert!(p.encoded_len() <= limit, "{} > {limit}", p.encoded_len()); - for c in &p.chunks { - let first = &c.spans[0]; - assert_eq!(first.meta["_dd.p.tid"], "66f1e2a300000000"); - assert_eq!(first.meta["_dd.p.dm"], "-1"); - } - } - } - - #[test] - fn span_larger_than_the_limit_is_dropped() { - let mut big_chunk = chunk(3, 4, 100); - big_chunk.spans.push(span(3, 99, 50_000)); - let input = vec![payload(vec![big_chunk])]; - - let out = split_tracer_payloads(input, 10_000); - - let ids = span_ids(&out); - assert_eq!(ids, vec![(3, 1), (3, 2), (3, 3), (3, 4)]); - assert!(out.iter().all(|p| p.encoded_len() <= 10_000)); - } - - #[test] - fn mixed_payloads_only_split_the_oversized_one() { - let small = payload(vec![chunk(1, 2, 100)]); - let large = payload(vec![chunk(2, 1_000, 500)]); - let limit = large.encoded_len() / 2; - - let out = split_tracer_payloads(vec![small.clone(), large.clone()], limit); - - assert_eq!(out[0], small); - assert!(out.len() >= 3); - assert!(out.iter().all(|p| p.encoded_len() <= limit)); - assert_eq!(span_ids(&out), span_ids(&[small, large])); - } - - #[test] - fn near_limit_span_with_trace_tags_is_kept() { - let mut root = span(5, 1, 0); - root.meta - .insert("_dd.p.tid".to_string(), "66f1e2a300000000".to_string()); - root.meta.insert("_dd.p.dm".to_string(), "-1".to_string()); - root.meta.insert("payload".to_string(), "x".repeat(5_000)); - let mut trace = chunk(5, 0, 0); - let empty_chunk_len = trace.encoded_len(); - trace.spans.push(root); - trace.spans.extend((2..=5).map(|i| span(5, i, 100))); - let input = payload(vec![trace]); - let empty_payload_len = { - let mut p = input.clone(); - p.chunks.clear(); - p.encoded_len() - }; - // The root span leaves only a few bytes of the space available to a single span. It - // already carries the trace tags, so it gains nothing when they are propagated. - let root_len = field_len(SPANS_TAG, &input.chunks[0].spans[0]); - let limit = empty_payload_len + empty_chunk_len + FIELD_OVERHEAD + root_len + 10; - assert!(input.encoded_len() > limit); - - let out = split_tracer_payloads(vec![input.clone()], limit); - - assert!(span_ids(&out).contains(&(5, 1)), "root span was dropped"); - assert_eq!(span_ids(&out), span_ids(&[input])); - assert!(out.iter().all(|p| p.encoded_len() <= limit)); - } - - const PROCESS_TAGS: &str = "entrypoint.name:bootstrap,entrypoint.type:executable"; - - fn with_process_tags_on_first_span(mut input: pb::TracerPayload) -> pb::TracerPayload { - input.chunks[0].spans[0] - .meta - .insert(PROCESS_TAGS_KEY.to_string(), PROCESS_TAGS.to_string()); - input - } - - fn assert_process_tags_on_every_output(input: &pb::TracerPayload, min_outputs: usize) { - let limit = input.encoded_len() / min_outputs; - - let out = split_tracer_payloads(vec![input.clone()], limit); - - assert!(out.len() >= min_outputs, "got {} payloads", out.len()); - for p in &out { - assert!(p.encoded_len() <= limit, "{} > {limit}", p.encoded_len()); - assert_eq!(p.tags[PROCESS_TAGS_KEY], PROCESS_TAGS); - } - assert_eq!(span_ids(&out), span_ids(std::slice::from_ref(input))); - } - - #[test] - fn process_tags_are_promoted_when_splitting_by_chunk() { - let input = with_process_tags_on_first_span(payload( - (1..=10).map(|t| chunk(t, 10, 1_000)).collect(), - )); - assert_process_tags_on_every_output(&input, 3); - } - - #[test] - fn process_tags_are_promoted_when_splitting_by_span() { - let input = with_process_tags_on_first_span(payload(vec![chunk(7, 2_000, 500)])); - assert_process_tags_on_every_output(&input, 5); - } - - #[test] - fn payload_level_process_tags_take_priority() { - let mut input = with_process_tags_on_first_span(payload(vec![chunk(7, 2_000, 500)])); - input - .tags - .insert(PROCESS_TAGS_KEY.to_string(), "payload-level".to_string()); - let limit = input.encoded_len() / 4; - - let out = split_tracer_payloads(vec![input], limit); - - assert!(out.len() >= 4); - for p in &out { - assert!(p.encoded_len() <= limit, "{} > {limit}", p.encoded_len()); - assert_eq!(p.tags[PROCESS_TAGS_KEY], "payload-level"); - } - } - - #[test] - fn payload_without_process_tags_gains_none() { - let input = payload((1..=10).map(|t| chunk(t, 10, 1_000)).collect()); - let limit = input.encoded_len() / 3; - - let out = split_tracer_payloads(vec![input], limit); - - assert!(out.len() >= 3); - assert!(out.iter().all(|p| !p.tags.contains_key(PROCESS_TAGS_KEY))); - } -} diff --git a/bottlecap/src/traces/trace_processor.rs b/bottlecap/src/traces/trace_processor.rs index 2f1f8ce6f..9fc07480f 100644 --- a/bottlecap/src/traces/trace_processor.rs +++ b/bottlecap/src/traces/trace_processor.rs @@ -36,11 +36,8 @@ use tokio::sync::mpsc::Sender; use tokio::sync::mpsc::error::SendError; use tracing::{debug, error}; -use crate::traces::payload_split::split_tracer_payloads; use crate::traces::stats_generator::StatsGenerator; -use crate::traces::trace_aggregator::{ - MAX_CONTENT_SIZE_BYTES, OwnedTracerHeaderTags, SendDataBuilderInfo, -}; +use crate::traces::trace_aggregator::{OwnedTracerHeaderTags, SendDataBuilderInfo}; use libdd_trace_normalization::normalizer::SamplerPriority; /// Which party is responsible for computing trace stats for a trace, derived from @@ -515,7 +512,7 @@ pub trait TraceProcessor { traces: Vec>, body_size: usize, span_pointers: Option>, - ) -> (Vec, TracerPayloadCollection); + ) -> (Option, TracerPayloadCollection); } #[async_trait] @@ -528,7 +525,7 @@ impl TraceProcessor for ServerlessTraceProcessor { traces: Vec>, body_size: usize, span_pointers: Option>, - ) -> (Vec, TracerPayloadCollection) { + ) -> (Option, TracerPayloadCollection) { let mut payload = trace_utils::collect_pb_trace_chunks( traces, &header_tags, @@ -582,7 +579,7 @@ impl TraceProcessor for ServerlessTraceProcessor { { self.drop_sampled_out_chunks(tracer_payloads); if tracer_payloads.is_empty() { - return (Vec::new(), payloads_for_stats); + return (None, payloads_for_stats); } } @@ -596,45 +593,28 @@ impl TraceProcessor for ServerlessTraceProcessor { debug!("TRACES | trace payload size after enrichment: {body_size} bytes"); - // A payload larger than the batch cap can never be batched, so split it into - // payloads that each fit. - let payloads = match payload { - TracerPayloadCollection::V07(tracer_payloads) if body_size > MAX_CONTENT_SIZE_BYTES => { - let split = split_tracer_payloads(tracer_payloads, MAX_CONTENT_SIZE_BYTES); - debug!( - "TRACES | split {body_size} byte trace payload into {} payloads", - split.len() - ); - split - .into_iter() - .map(|p| { - let size = prost::Message::encoded_len(&p); - (TracerPayloadCollection::V07(vec![p]), size) - }) - .collect() - } - other => vec![(other, body_size)], - }; + let owned_header_tags = OwnedTracerHeaderTags::from(header_tags.clone()); - let builders = payloads - .into_iter() - .map(|(payload, size)| { - let owned_header_tags = OwnedTracerHeaderTags::from(header_tags.clone()); - let builder = SendDataBuilder::new(size, payload, header_tags.clone(), &endpoint) - .with_compression(CompressionStrategy::Zstd { - level: config.apm_config_compression_level, - }) - .with_retry_strategy(RetryStrategy::new( - 1, - 100, - RetryBackoffType::Exponential, - None, - )); - SendDataBuilderInfo::new(builder, size, owned_header_tags) + // Move original payload into builder (no clone needed) + let builder = SendDataBuilder::new(body_size, payload, header_tags, &endpoint) + .with_compression(CompressionStrategy::Zstd { + level: config.apm_config_compression_level, }) - .collect(); - - (builders, payloads_for_stats) + .with_retry_strategy(RetryStrategy::new( + 1, + 100, + RetryBackoffType::Exponential, + None, + )); + + ( + Some(SendDataBuilderInfo::new( + builder, + body_size, + owned_header_tags, + )), + payloads_for_stats, + ) } } @@ -717,7 +697,7 @@ impl SendingTraceProcessor { // Capture before `header_tags` is moved into process_traces below. let client_computed_stats = header_tags.generic.client_computed_stats; - let (payloads, processed_traces) = self.processor.process_traces( + let (payload, processed_traces) = self.processor.process_traces( config.clone(), tags_provider, header_tags, @@ -741,7 +721,7 @@ impl SendingTraceProcessor { error!("TRACE_PROCESSOR | Error sending traces to the stats concentrator: {err}"); } - for payload in payloads { + if let Some(payload) = payload { self.trace_tx.send(payload).await?; } Ok(()) @@ -930,10 +910,7 @@ mod tests { 100, None, ); - let tracer_payload = tracer_payload - .into_iter() - .next() - .expect("expected a payload"); + let tracer_payload = tracer_payload.expect("expected Some payload"); let expected_tracer_payload = pb::TracerPayload { container_id: "33".to_string(), @@ -1408,7 +1385,7 @@ mod tests { let (payload_info, stats_collection) = processor.process_traces(config, tags_provider, header_tags, traces, 0, None); - let payload_info = payload_info.into_iter().next().expect("expected a payload"); + let payload_info = payload_info.expect("expected Some payload"); // Stats collection must include all three traces let TracerPayloadCollection::V07(ref stats_payloads) = stats_collection else { @@ -1474,7 +1451,7 @@ mod tests { let (payload_info, _stats) = processor.process_traces(config, tags_provider, header_tags, traces, 0, None); - let payload_info = payload_info.into_iter().next().expect("expected a payload"); + let payload_info = payload_info.expect("expected Some payload"); let backend_send_data = payload_info.builder.build(); let TracerPayloadCollection::V07(backend_payloads) = backend_send_data.get_payloads() else { @@ -1535,10 +1512,7 @@ mod tests { let (payload_info, _stats) = processor.process_traces(config, tags_provider, header_tags, traces, 0, None); - let payload_info = payload_info - .into_iter() - .next() - .expect("errored-child P0 trace rescued"); + let payload_info = payload_info.expect("errored-child P0 trace rescued"); let backend_send_data = payload_info.builder.build(); let TracerPayloadCollection::V07(backend_payloads) = backend_send_data.get_payloads() else { @@ -1582,7 +1556,7 @@ mod tests { let (payload_info, _stats) = processor.process_traces(config, tags_provider, header_tags, traces, 0, None); assert!( - payload_info.is_empty(), + payload_info.is_none(), "errored P0 trace must stay dropped when the error sampler is disabled" ); } @@ -1618,10 +1592,7 @@ mod tests { let (payload_info, stats_collection) = processor.process_traces(config, tags_provider, header_tags, traces, 0, None); - let payload_info = payload_info - .into_iter() - .next() - .expect("kept trace must produce a backend payload"); + let payload_info = payload_info.expect("kept trace must produce a backend payload"); // Stats collection must include both traces, including the dropped errored P0. let TracerPayloadCollection::V07(ref stats_payloads) = stats_collection else { @@ -1687,8 +1658,8 @@ mod tests { processor.process_traces(config, tags_provider, header_tags, traces, 0, None); assert!( - payload.is_empty(), - "no backend payload must be produced when all traces are sampled out" + payload.is_none(), + "backend payload must be None when all traces are sampled out" ); // Stats collection must still include both traces @@ -1742,7 +1713,7 @@ mod tests { let (payload_info, stats_collection) = processor.process_traces(config, tags_provider, header_tags, traces, 999_999, None); - let info = payload_info.into_iter().next().expect("expected a payload"); + let info = payload_info.expect("expected Some payload"); // The reported size must equal the sum of encoded_len() of the kept TracerPayloads. // stats_collection has all 4 traces. Reconstruct the filtered payload (only trace_id=1 @@ -1822,7 +1793,7 @@ mod tests { None, ); - let info = payload_info.into_iter().next().expect("expected a payload"); + let info = payload_info.expect("expected Some payload"); // Nothing is filtered here, so the returned payloads are exactly what gets sent. let TracerPayloadCollection::V07(ref payloads) = processed_payloads else { @@ -1840,84 +1811,6 @@ mod tests { ); } - /// A single trace whose enriched payload exceeds the aggregator batch cap must be - /// split into multiple payloads that each fit, without losing or duplicating spans. - #[test] - fn test_process_traces_splits_oversized_trace() { - const TRACE_ID: u64 = 42; - const SPAN_COUNT: u64 = 40_000; - - let config = Arc::new(Config { - apm_dd_url: "https://trace.agent.datadoghq.com".to_string(), - ..Config::default() - }); - let (tags_provider, processor) = create_test_processor(&config, enabled_error_sampler()); - let header_tags = create_test_header_tags(); - - // Span 1 is the root; every other span is its child. - let make_span = |span_id: u64| -> pb::Span { - pb::Span { - trace_id: TRACE_ID, - span_id, - parent_id: u64::from(span_id != 1), - service: "svc".to_string(), - name: "op".to_string(), - resource: "res".to_string(), - meta: HashMap::from([("payload".to_string(), "x".repeat(300))]), - ..Default::default() - } - }; - let traces = vec![(1..=SPAN_COUNT).map(make_span).collect::>()]; - - let (payload_infos, stats_collection) = - processor.process_traces(config, tags_provider, header_tags, traces, 1, None); - - assert!( - payload_infos.len() > 1, - "oversized trace must be split into multiple payloads" - ); - assert!( - payload_infos - .iter() - .all(|info| info.size <= MAX_CONTENT_SIZE_BYTES), - "every split payload must fit within the batch cap" - ); - - let mut sent_span_ids = Vec::new(); - for info in payload_infos { - let send_data = info.builder.build(); - let TracerPayloadCollection::V07(payloads) = send_data.get_payloads() else { - panic!("expected V07"); - }; - for chunk in payloads.iter().flat_map(|tp| &tp.chunks) { - for span in &chunk.spans { - assert_eq!(span.trace_id, TRACE_ID); - sent_span_ids.push(span.span_id); - } - } - } - sent_span_ids.sort_unstable(); - let expected_span_ids: Vec = (1..=SPAN_COUNT).collect(); - assert_eq!( - sent_span_ids, expected_span_ids, - "split payloads must contain every input span exactly once" - ); - - let TracerPayloadCollection::V07(ref stats_payloads) = stats_collection else { - panic!("expected V07"); - }; - let stats_span_count: usize = stats_payloads - .iter() - .flat_map(|tp| &tp.chunks) - .map(|c| c.spans.len()) - .sum(); - assert_eq!( - stats_span_count, - expected_span_ids.len(), - "stats must include all input spans" - ); - } - fn create_inferred_span() -> pb::Span { let mut meta = HashMap::new(); meta.insert("_inferred_span.tag_source".to_string(), "self".to_string());