From 4aed84b110663c353242f1b582f4b41a4c28f29b Mon Sep 17 00:00:00 2001 From: Andrew Glaude Date: Tue, 28 Jul 2026 08:57:29 -0400 Subject: [PATCH] adding some tests to validate fwd compatability --- .../src/msgpack_decoder/v1/mod.rs | 213 ++++++++ libdd-trace-utils/src/span/trace_utils_v1.rs | 461 ++++++++++++++++++ 2 files changed, 674 insertions(+) create mode 100644 libdd-trace-utils/src/span/trace_utils_v1.rs diff --git a/libdd-trace-utils/src/msgpack_decoder/v1/mod.rs b/libdd-trace-utils/src/msgpack_decoder/v1/mod.rs index b7ee7c45d7..4297888d12 100644 --- a/libdd-trace-utils/src/msgpack_decoder/v1/mod.rs +++ b/libdd-trace-utils/src/msgpack_decoder/v1/mod.rs @@ -693,4 +693,217 @@ mod tests { }, ); } + + // --------------------------------------------------------------------------------------------- + // Forward-compatibility: unknown map keys must be skipped for every V1 map type. These tests + // hand-build wire bytes (the encoder never emits unknown keys) with `rmp::encode`, injecting a + // future/unknown key that older decoders don't recognize, and assert the surrounding known + // fields still decode correctly. + // --------------------------------------------------------------------------------------------- + + use rmp::encode::{self, ByteBuf}; + + /// Writes a `u8` msgpack map key. + fn wkey(buf: &mut ByteBuf, k: u8) { + encode::write_uint(buf, k as u64).unwrap(); + } + + /// Wraps pre-encoded span map bytes into a full single-chunk V1 payload so the span decoder is + /// exercised through the real entry point. + fn payload_wrapping_span(span_bytes: &[u8]) -> Vec { + let mut buf = ByteBuf::new(); + encode::write_map_len(&mut buf, 1).unwrap(); // payload: chunks only + wkey(&mut buf, trace_key::CHUNKS); + encode::write_array_len(&mut buf, 1).unwrap(); + encode::write_map_len(&mut buf, 2).unwrap(); // chunk: trace_id + spans + wkey(&mut buf, chunk_key::TRACE_ID); + encode::write_bin(&mut buf, &[0u8; 16]).unwrap(); + wkey(&mut buf, chunk_key::SPANS); + encode::write_array_len(&mut buf, 1).unwrap(); + let mut out = buf.into_vec(); + out.extend_from_slice(span_bytes); + out + } + + #[test] + fn unknown_payload_key_is_skipped() { + let mut buf = ByteBuf::new(); + encode::write_map_len(&mut buf, 3).unwrap(); + // known field before the unknown key + wkey(&mut buf, trace_key::ENV_REF); + encode::write_str(&mut buf, "prod").unwrap(); + // future/unknown payload key -> arbitrary scalar + wkey(&mut buf, 99); + encode::write_uint(&mut buf, 123_456).unwrap(); + // required chunks field after the unknown key + wkey(&mut buf, trace_key::CHUNKS); + encode::write_array_len(&mut buf, 0).unwrap(); + let buf = buf.into_vec(); + + let (decoded, consumed) = + from_bytes(Bytes::from(buf.clone())).expect("unknown payload key must be skipped"); + assert_eq!( + consumed, + buf.len(), + "decoder must consume the skipped value" + ); + assert_eq!(decoded.env.as_str(), "prod"); + assert!(decoded.chunks.is_empty()); + } + + #[test] + fn unknown_payload_key_with_nested_value_is_skipped() { + // The skipped value is a nested array containing a map, exercising rmpv's recursive skip. + let mut buf = ByteBuf::new(); + encode::write_map_len(&mut buf, 2).unwrap(); + wkey(&mut buf, 120); + encode::write_array_len(&mut buf, 3).unwrap(); + encode::write_uint(&mut buf, 1).unwrap(); + encode::write_str(&mut buf, "x").unwrap(); + encode::write_map_len(&mut buf, 1).unwrap(); + encode::write_uint(&mut buf, 5).unwrap(); + encode::write_bool(&mut buf, true).unwrap(); + wkey(&mut buf, trace_key::CHUNKS); + encode::write_array_len(&mut buf, 0).unwrap(); + let buf = buf.into_vec(); + + let (decoded, consumed) = from_bytes(Bytes::from(buf.clone())) + .expect("nested unknown value must be fully skipped"); + assert_eq!(consumed, buf.len()); + assert!(decoded.chunks.is_empty()); + } + + #[test] + fn unknown_chunk_key_is_skipped() { + let mut buf = ByteBuf::new(); + encode::write_map_len(&mut buf, 1).unwrap(); + wkey(&mut buf, trace_key::CHUNKS); + encode::write_array_len(&mut buf, 1).unwrap(); + encode::write_map_len(&mut buf, 3).unwrap(); + // unknown chunk key before the required fields + wkey(&mut buf, 77); + encode::write_bool(&mut buf, true).unwrap(); + wkey(&mut buf, chunk_key::TRACE_ID); + encode::write_bin(&mut buf, &[9u8; 16]).unwrap(); + wkey(&mut buf, chunk_key::SPANS); + encode::write_array_len(&mut buf, 0).unwrap(); + let buf = buf.into_vec(); + + let (decoded, _) = from_bytes(Bytes::from(buf)).expect("unknown chunk key must be skipped"); + assert_eq!(decoded.chunks.len(), 1); + assert_eq!(decoded.chunks[0].trace_id, [9u8; 16]); + assert!(decoded.chunks[0].spans.is_empty()); + } + + #[test] + fn unknown_span_key_is_skipped() { + let mut span = ByteBuf::new(); + encode::write_map_len(&mut span, 4).unwrap(); + wkey(&mut span, span_key::SPAN_ID); + encode::write_uint(&mut span, 42).unwrap(); + wkey(&mut span, span_key::START); + encode::write_uint(&mut span, 100).unwrap(); + // unknown span key between required and known optional fields + wkey(&mut span, 88); + encode::write_f64(&mut span, 2.5).unwrap(); + wkey(&mut span, span_key::SERVICE); + encode::write_str(&mut span, "svc").unwrap(); + + let buf = payload_wrapping_span(&span.into_vec()); + let (decoded, _) = from_bytes(Bytes::from(buf)).expect("unknown span key must be skipped"); + let span = &decoded.chunks[0].spans[0]; + assert_eq!(span.span_id, 42); + assert_eq!(span.start, 100); + assert_eq!(span.service.as_str(), "svc"); + } + + #[test] + fn unknown_span_link_key_is_skipped() { + let mut span = ByteBuf::new(); + encode::write_map_len(&mut span, 3).unwrap(); + wkey(&mut span, span_key::SPAN_ID); + encode::write_uint(&mut span, 1).unwrap(); + wkey(&mut span, span_key::START); + encode::write_uint(&mut span, 1).unwrap(); + wkey(&mut span, span_key::SPAN_LINKS); + encode::write_array_len(&mut span, 1).unwrap(); + // span link map with an unknown key preceding a known one + encode::write_map_len(&mut span, 2).unwrap(); + wkey(&mut span, 55); + encode::write_bool(&mut span, false).unwrap(); + wkey(&mut span, span_link_key::SPAN_ID); + encode::write_uint(&mut span, 777).unwrap(); + + let buf = payload_wrapping_span(&span.into_vec()); + let (decoded, _) = + from_bytes(Bytes::from(buf)).expect("unknown span_link key must be skipped"); + let links = &decoded.chunks[0].spans[0].span_links; + assert_eq!(links.len(), 1); + assert_eq!(links[0].span_id, 777); + } + + #[test] + fn unknown_span_event_key_is_skipped() { + let mut span = ByteBuf::new(); + encode::write_map_len(&mut span, 3).unwrap(); + wkey(&mut span, span_key::SPAN_ID); + encode::write_uint(&mut span, 1).unwrap(); + wkey(&mut span, span_key::START); + encode::write_uint(&mut span, 1).unwrap(); + wkey(&mut span, span_key::SPAN_EVENTS); + encode::write_array_len(&mut span, 1).unwrap(); + // span event map with an unknown key preceding a known one + encode::write_map_len(&mut span, 2).unwrap(); + wkey(&mut span, 66); + encode::write_uint(&mut span, 999).unwrap(); + wkey(&mut span, span_event_key::TIME); + encode::write_uint(&mut span, 123).unwrap(); + + let buf = payload_wrapping_span(&span.into_vec()); + let (decoded, _) = + from_bytes(Bytes::from(buf)).expect("unknown span_event key must be skipped"); + let events = &decoded.chunks[0].spans[0].span_events; + assert_eq!(events.len(), 1); + assert_eq!(events[0].time_unix_nano, 123); + } + + /// Forward-compatibility boundary. Skipping an unknown key advances the buffer correctly, but + /// it does NOT feed inline strings into the shared streaming intern table. So if a future field + /// carries a *first-occurrence* interned string that a later field back-references by ID, the + /// decoder's table is shifted and the reference no longer resolves. This pins the current + /// behavior — a loud error rather than silent corruption. If forward-compat is ever hardened to + /// record skipped strings, this test should be updated deliberately. + #[test] + fn unknown_key_with_interned_string_desyncs_table() { + // Simulate a future encoder whose intern table is: ""=0, "ghost"=1, "prod"=2. + let mut buf = ByteBuf::new(); + encode::write_map_len(&mut buf, 3).unwrap(); + // Future field 99: first occurrence of "ghost" (encoder id 1). The decoder skips it and + // never records it, so its table stays one entry short. + wkey(&mut buf, 99); + encode::write_str(&mut buf, "ghost").unwrap(); + // ENV first occurrence "prod": encoder id 2, but the decoder records it as id 1. + wkey(&mut buf, trace_key::ENV_REF); + encode::write_str(&mut buf, "prod").unwrap(); + // A span whose service references the encoder's id 2 ("prod"). + wkey(&mut buf, trace_key::CHUNKS); + encode::write_array_len(&mut buf, 1).unwrap(); + encode::write_map_len(&mut buf, 2).unwrap(); + wkey(&mut buf, chunk_key::TRACE_ID); + encode::write_bin(&mut buf, &[0u8; 16]).unwrap(); + wkey(&mut buf, chunk_key::SPANS); + encode::write_array_len(&mut buf, 1).unwrap(); + encode::write_map_len(&mut buf, 3).unwrap(); + wkey(&mut buf, span_key::SPAN_ID); + encode::write_uint(&mut buf, 1).unwrap(); + wkey(&mut buf, span_key::START); + encode::write_uint(&mut buf, 1).unwrap(); + wkey(&mut buf, span_key::SERVICE); + encode::write_uint(&mut buf, 2).unwrap(); // reference to encoder id 2 = "prod" + let buf = buf.into_vec(); + + let err = from_bytes(Bytes::from(buf)) + .expect_err("skipping an interned string desyncs the table, so the ref must fail"); + assert!(matches!(err, DecodeError::InvalidFormat(_))); + } } diff --git a/libdd-trace-utils/src/span/trace_utils_v1.rs b/libdd-trace-utils/src/span/trace_utils_v1.rs new file mode 100644 index 0000000000..d47f2d792e --- /dev/null +++ b/libdd-trace-utils/src/span/trace_utils_v1.rs @@ -0,0 +1,461 @@ +// Copyright 2026-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Trace-utils functionalities implementation for V1 spans. + +use crate::span::trace_utils::DroppedP0Stats; +use crate::span::v1::{AttributeValue, Span, TraceChunk}; +use crate::span::{SpanText, TraceData}; +use std::collections::{HashMap, HashSet}; +use tracing::debug; + +/// Span metric the mini agent must set for the backend to recognize top level span +const TOP_LEVEL_KEY: &str = "_top_level"; +/// Span metric the tracer sets to denote a top level span +const TRACER_TOP_LEVEL_KEY: &str = "_dd.top_level"; +const MEASURED_KEY: &str = "_dd.measured"; +const PARTIAL_VERSION_KEY: &str = "_dd.partial_version"; +const SAMPLING_SINGLE_SPAN_MECHANISM: &str = "_dd.span_sampling.mechanism"; +const SAMPLING_ANALYTICS_RATE_KEY: &str = "_dd1.sr.eausr"; + +/// Reads a numeric attribute (`Float` or `Int`), mirroring how v0.4's `metrics` map is read. +fn attribute_as_f64(value: &AttributeValue) -> Option { + match value { + AttributeValue::Float(v) => Some(*v), + AttributeValue::Int(v) => Some(*v as f64), + _ => None, + } +} + +fn set_top_level_span(span: &mut Span) { + span.attributes.insert( + T::Text::from_static_str(TOP_LEVEL_KEY), + AttributeValue::Float(1.0), + ); +} + +/// Updates all the spans top-level attribute. +/// A span is considered top-level if: +/// - it's a root span +/// - OR its parent is unknown (other part of the code, distributed trace) +/// - OR its parent belongs to another service (in that case it's a "local root" being the highest +/// ancestor of other spans belonging to this service and attached to it). +pub fn compute_top_level_span(trace: &mut [Span]) { + let span_id_idx: HashMap = trace + .iter() + .enumerate() + .map(|(i, span)| (span.span_id, i)) + .collect(); + for span_idx in 0..trace.len() { + let parent_id = trace[span_idx].parent_id; + if parent_id == 0 { + set_top_level_span(&mut trace[span_idx]); + continue; + } + match span_id_idx.get(&parent_id).map(|i| &trace[*i].service) { + Some(parent_span_service) => { + if !(parent_span_service == &trace[span_idx].service) { + // parent is not in the same service + set_top_level_span(&mut trace[span_idx]) + } + } + None => { + // span has no parent in chunk + set_top_level_span(&mut trace[span_idx]) + } + } + } +} + +/// Returns the index of the root span in `trace`. +pub fn get_root_span_index(trace: &[Span]) -> anyhow::Result { + if trace.is_empty() { + anyhow::bail!("Cannot find root span index in an empty trace."); + } + + // Do a first pass to find if we have an obvious root span (starting from the end) since some + // clients put the root span last. + for (i, span) in trace.iter().enumerate().rev() { + if span.parent_id == 0 { + return Ok(i); + } + } + + let span_ids: HashSet<_> = trace.iter().map(|span| span.span_id).collect(); + + let mut root_span_id = None; + for (i, span) in trace.iter().enumerate() { + // If a span's parent is not in the trace, it is a root + if !span_ids.contains(&span.parent_id) { + if root_span_id.is_some() { + debug!("trace has multiple root spans"); + } + root_span_id = Some(i); + } + } + Ok(match root_span_id { + Some(i) => i, + None => { + debug!("Could not find the root span for trace"); + trace.len() - 1 + } + }) +} + +/// Return true if the span has a top level key set +pub fn has_top_level(span: &Span) -> bool { + span.attributes + .get(TRACER_TOP_LEVEL_KEY) + .and_then(attribute_as_f64) + .is_some_and(|v| v == 1.0) + || span + .attributes + .get(TOP_LEVEL_KEY) + .and_then(attribute_as_f64) + .is_some_and(|v| v == 1.0) +} + +/// Returns true if a span should be measured (i.e., it should get trace metrics calculated). +pub fn is_measured(span: &Span) -> bool { + span.attributes + .get(MEASURED_KEY) + .and_then(attribute_as_f64) + .is_some_and(|v| v == 1.0) +} + +/// Returns true if the span is a partial snapshot. +/// This kind of spans are partial images of long-running spans. +/// When incomplete, a partial snapshot has a metric _dd.partial_version which is a positive +/// integer. The metric usually increases each time a new version of the same span is sent by +/// the tracer +pub fn is_partial_snapshot(span: &Span) -> bool { + span.attributes + .get(PARTIAL_VERSION_KEY) + .and_then(attribute_as_f64) + .is_some_and(|v| v >= 0.0) +} + +/// Remove spans and chunks, only keeping the ones that may be sampled by the agent. +/// +/// Unlike v0.4 (where sampling priority is a per-span metric), v1's `priority` is already a +/// direct field on [`TraceChunk`], so it is read once per chunk instead of being searched for +/// across the chunk's spans. +/// +/// # Returns +/// +/// A tuple containing the dropped p0 stats, the first value correspond the amount of traces +/// dropped and the latter to the spans dropped. +pub fn drop_chunks(traces: &mut Vec>) -> DroppedP0Stats { + let mut dropped_p0_traces = 0; + let mut dropped_p0_spans = 0; + + traces.retain_mut(|chunk| { + // ErrorSampler + if chunk.spans.iter().any(|s| s.error) { + // We send chunks containing an error + return true; + } + + // PrioritySampler and NoPrioritySampler + if chunk.priority.is_none_or(|p| p > 0) { + // We send chunks with positive priority or no priority + return true; + } + + // SingleSpanSampler and AnalyzedSpansSampler + // List of spans to keep even if the chunk is dropped + let mut sampled_indexes = Vec::new(); + for (index, span) in chunk.spans.iter().enumerate() { + if span + .attributes + .get(SAMPLING_SINGLE_SPAN_MECHANISM) + .and_then(attribute_as_f64) + .is_some_and(|m| m == 8.0) + || span.attributes.contains_key(SAMPLING_ANALYTICS_RATE_KEY) + { + // We send spans sampled by single-span sampling or analyzed spans + sampled_indexes.push(index); + } + } + dropped_p0_spans += chunk.spans.len() - sampled_indexes.len(); + if sampled_indexes.is_empty() { + // If no spans were sampled we can drop the whole chunk + dropped_p0_traces += 1; + return false; + } + let sampled_spans = sampled_indexes + .iter() + .map(|i| std::mem::take(&mut chunk.spans[*i])) + .collect(); + chunk.spans = sampled_spans; + true + }); + + DroppedP0Stats { + dropped_p0_traces, + dropped_p0_spans, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::span::v1::{SpanBytes, TraceChunkBytes}; + + fn create_test_span(is_top_level: bool) -> SpanBytes { + let mut span = SpanBytes { + service: "test-service".into(), + name: "test_name".into(), + resource: "test-resource".into(), + ..Default::default() + }; + if is_top_level { + span.attributes + .insert("_top_level".into(), AttributeValue::Float(1.0)); + } + span + } + + fn create_test_span_with_ids(span_id: u64, parent_id: u64) -> SpanBytes { + SpanBytes { + service: "test-service".into(), + name: "test_name".into(), + resource: "test-resource".into(), + span_id, + parent_id, + ..Default::default() + } + } + + #[test] + fn test_has_top_level() { + let top_level_span = create_test_span(true); + let not_top_level_span = create_test_span(false); + assert!(has_top_level(&top_level_span)); + assert!(!has_top_level(¬_top_level_span)); + } + + #[test] + fn test_is_measured() { + let mut measured_span = create_test_span(true); + measured_span + .attributes + .insert(MEASURED_KEY.into(), AttributeValue::Float(1.0)); + let not_measured_span = create_test_span(true); + assert!(is_measured(&measured_span)); + assert!(!is_measured(¬_measured_span)); + } + + #[test] + fn test_is_partial_snapshot() { + let mut partial_span = create_test_span(false); + partial_span + .attributes + .insert(PARTIAL_VERSION_KEY.into(), AttributeValue::Int(2)); + let not_partial_span = create_test_span(false); + assert!(is_partial_snapshot(&partial_span)); + assert!(!is_partial_snapshot(¬_partial_span)); + } + + #[test] + fn test_compute_top_level() { + let mut span_with_different_service = create_test_span_with_ids(5, 2); + span_with_different_service.service = "another_service".into(); + let mut trace = vec![ + // Root span, should be marked as top-level + create_test_span_with_ids(1, 0), + // Should not be marked as top-level + create_test_span_with_ids(2, 1), + // No parent in local trace, should be marked as top-level + create_test_span_with_ids(4, 3), + // Parent belongs to another service, should be marked as top-level + span_with_different_service, + ]; + + compute_top_level_span(trace.as_mut_slice()); + + let spans_marked_as_top_level: Vec = trace + .iter() + .filter_map(|span| has_top_level(span).then_some(span.span_id)) + .collect(); + assert_eq!(spans_marked_as_top_level, [1, 4, 5]); + } + + #[test] + fn test_get_root_span_index_from_complete_trace() { + let trace = vec![ + create_test_span_with_ids(1, 0), + create_test_span_with_ids(2, 1), + create_test_span_with_ids(3, 1), + ]; + assert_eq!(get_root_span_index(&trace).unwrap(), 0); + } + + #[test] + fn test_get_root_span_index_root_last() { + let trace = vec![ + create_test_span_with_ids(2, 1), + create_test_span_with_ids(3, 1), + create_test_span_with_ids(1, 0), + ]; + assert_eq!(get_root_span_index(&trace).unwrap(), 2); + } + + #[test] + fn test_get_root_span_index_from_partial_trace() { + // No span has parent_id == 0, but span 1's parent (99) isn't in the trace. + let trace = vec![ + create_test_span_with_ids(1, 99), + create_test_span_with_ids(2, 1), + ]; + assert_eq!(get_root_span_index(&trace).unwrap(), 0); + } + + #[test] + fn test_get_root_span_index_empty_trace_errors() { + let trace: Vec = vec![]; + assert!(get_root_span_index(&trace).is_err()); + } + + fn chunk_with_spans(priority: Option, spans: Vec) -> TraceChunkBytes { + TraceChunkBytes { + priority, + spans, + ..Default::default() + } + } + + #[test] + fn test_drop_chunks() { + let chunk_with_priority = chunk_with_spans( + Some(1), + vec![ + SpanBytes { + span_id: 1, + ..Default::default() + }, + SpanBytes { + span_id: 2, + parent_id: 1, + ..Default::default() + }, + ], + ); + let chunk_with_null_priority = chunk_with_spans( + Some(0), + vec![ + SpanBytes { + span_id: 1, + ..Default::default() + }, + SpanBytes { + span_id: 2, + parent_id: 1, + ..Default::default() + }, + ], + ); + let chunk_without_priority = chunk_with_spans( + None, + vec![ + SpanBytes { + span_id: 1, + ..Default::default() + }, + SpanBytes { + span_id: 2, + parent_id: 1, + ..Default::default() + }, + ], + ); + let chunk_with_negative_priority = chunk_with_spans( + Some(-1), + vec![ + SpanBytes { + span_id: 1, + ..Default::default() + }, + SpanBytes { + span_id: 2, + parent_id: 1, + ..Default::default() + }, + ], + ); + let chunk_with_error = chunk_with_spans( + Some(0), + vec![ + SpanBytes { + span_id: 1, + error: true, + ..Default::default() + }, + SpanBytes { + span_id: 2, + parent_id: 1, + ..Default::default() + }, + ], + ); + let chunk_with_a_single_span = chunk_with_spans( + Some(0), + vec![ + SpanBytes { + span_id: 1, + ..Default::default() + }, + SpanBytes { + span_id: 2, + parent_id: 1, + attributes: vec![( + SAMPLING_SINGLE_SPAN_MECHANISM.into(), + AttributeValue::Float(8.0), + )] + .into(), + ..Default::default() + }, + ], + ); + let chunk_with_analyzed_span = chunk_with_spans( + Some(0), + vec![ + SpanBytes { + span_id: 1, + ..Default::default() + }, + SpanBytes { + span_id: 2, + parent_id: 1, + attributes: vec![( + SAMPLING_ANALYTICS_RATE_KEY.into(), + AttributeValue::Float(1.0), + )] + .into(), + ..Default::default() + }, + ], + ); + + let chunks_and_expected_sampled_spans = vec![ + (chunk_with_priority, 2), + (chunk_with_null_priority, 0), + (chunk_without_priority, 2), + (chunk_with_negative_priority, 0), + (chunk_with_error, 2), + (chunk_with_a_single_span, 1), + (chunk_with_analyzed_span, 1), + ]; + + for (chunk, expected_count) in chunks_and_expected_sampled_spans.into_iter() { + let mut traces = vec![chunk]; + drop_chunks(&mut traces); + + if expected_count == 0 { + assert!(traces.is_empty()); + } else { + assert_eq!(traces[0].spans.len(), expected_count); + } + } + } +}