libdd_trace_utils/
tracer_payload.rs

1// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::span::{v04, v05, v1, BytesData, SharedDictBytes, TraceData};
5use crate::trace_utils::convert_trace_chunks_v04_to_v05;
6use crate::{msgpack_decoder, trace_utils::cmp_send_data_payloads};
7use libdd_trace_protobuf::pb;
8use std::cmp::Ordering;
9use std::iter::Iterator;
10use tracing::warn;
11
12pub type TracerPayloadV04 = Vec<v04::SpanBytes>;
13pub type TracerPayloadV05 = Vec<v05::Span>;
14
15#[derive(Debug, Clone, Copy)]
16/// Enumerates the different encoding types.
17pub enum TraceEncoding {
18    /// v0.4 encoding (TracerPayloadV04).
19    V04,
20    /// v0.5 encoding (TracerPayloadV05).
21    V05,
22    /// v1 encoding (TracerPayloadV1).
23    V1,
24}
25
26#[derive(Debug)]
27pub enum TraceChunks<T: TraceData> {
28    /// Collection of TraceChunkSpan.
29    V04(Vec<Vec<v04::Span<T>>>),
30    /// Collection of TraceChunkSpan with de-duplicated strings.
31    ///
32    /// The dictionary always owns its strings ([`SharedDictBytes`]) because the v0.5
33    /// conversion interns dynamically-built JSON (span links / events) alongside the
34    /// (possibly borrowed) span text.
35    V05((SharedDictBytes, Vec<Vec<v05::Span>>)),
36    /// Collection of v0.4 spans to be serialized as a V1 msgpack payload.
37    V1(Box<v1::TracerPayload<BytesData>>),
38}
39
40impl TraceChunks<BytesData> {
41    pub fn into_tracer_payload_collection(self) -> TracerPayloadCollection {
42        match self {
43            TraceChunks::V04(traces) => TracerPayloadCollection::V04(traces),
44            TraceChunks::V05(traces) => TracerPayloadCollection::V05(traces),
45            TraceChunks::V1(traces) => TracerPayloadCollection::V1(traces),
46        }
47    }
48}
49
50impl<T: TraceData> TraceChunks<T> {
51    /// Returns the number of traces in the chunk
52    pub fn size(&self) -> usize {
53        match self {
54            TraceChunks::V04(traces) => traces.len(),
55            TraceChunks::V05((_, traces)) => traces.len(),
56            TraceChunks::V1(trace) => trace.chunks.len(),
57        }
58    }
59}
60
61#[derive(Debug)]
62/// Enum representing a general abstraction for a collection of tracer payloads.
63pub enum TracerPayloadCollection {
64    /// Collection of TracerPayloads.
65    V07(Vec<pb::TracerPayload>),
66    /// Collection of TraceChunkSpan.
67    V04(Vec<Vec<v04::SpanBytes>>),
68    /// Collection of TraceChunkSpan with de-duplicated strings.
69    V05((SharedDictBytes, Vec<Vec<v05::Span>>)),
70    // /// V0.4-shaped spans that must be serialized as a V1 msgpack payload on send.
71    V1(Box<v1::TracerPayload<BytesData>>),
72}
73
74impl TracerPayloadCollection {
75    /// Appends `other` collection of the same type to the current collection.
76    ///
77    /// #Arguments
78    ///
79    /// * `other`: collection of the same type.
80    ///
81    /// # Examples:
82    ///
83    /// ```rust
84    /// use libdd_trace_protobuf::pb::TracerPayload;
85    /// use libdd_trace_utils::tracer_payload::TracerPayloadCollection;
86    /// let mut col1 = TracerPayloadCollection::V07(vec![TracerPayload::default()]);
87    /// let mut col2 = TracerPayloadCollection::V07(vec![TracerPayload::default()]);
88    /// col1.append(&mut col2);
89    /// ```
90    ///
91    /// # Returns
92    ///
93    /// `true` if `other`'s data was merged into `self`, `false` if the append was skipped (e.g.
94    /// diverging V1 tracer metadata). Callers that rely on `other` being fully drained must check
95    /// this return value rather than assuming success.
96    pub fn append(&mut self, other: &mut Self) -> bool {
97        match (self, other) {
98            (TracerPayloadCollection::V07(dest), TracerPayloadCollection::V07(src)) => {
99                dest.append(src);
100                true
101            }
102            (TracerPayloadCollection::V04(dest), TracerPayloadCollection::V04(src)) => {
103                dest.append(src);
104                true
105            }
106            (TracerPayloadCollection::V1(dest), TracerPayloadCollection::V1(src)) => {
107                // Same-target SendData entries are coalesced by
108                // trace_utils::coalesce_send_data, so both V1 payloads
109                // typically share tracer-level metadata. If all metadata
110                // fields match we append `src`'s chunks into `dest`; if any diverge we no-op
111                // (logging a warning) rather than silently dropping `src`'s metadata.
112                if metadata_matches_v1(dest, src) {
113                    dest.chunks.append(&mut src.chunks);
114                    true
115                } else {
116                    false
117                }
118            }
119            // TODO: Properly handle non-OK states to prevent possible panics (APMSP-18190).
120            #[allow(clippy::unimplemented)]
121            (TracerPayloadCollection::V05(_), _) => {
122                unimplemented!("Append for V05 not implemented")
123            }
124            _ => false,
125        }
126    }
127
128    /// Merges traces that came from the same origin together to reduce the payload size.
129    ///
130    /// # Examples:
131    ///
132    /// ```rust
133    /// use libdd_trace_protobuf::pb::TracerPayload;
134    /// use libdd_trace_utils::tracer_payload::TracerPayloadCollection;
135    /// let mut col1 =
136    ///     TracerPayloadCollection::V07(vec![TracerPayload::default(), TracerPayload::default()]);
137    /// col1.merge();
138    /// ```
139    pub fn merge(&mut self) {
140        if let TracerPayloadCollection::V07(collection) = self {
141            collection.sort_unstable_by(cmp_send_data_payloads);
142            collection.dedup_by(|a, b| {
143                if cmp_send_data_payloads(a, b) == Ordering::Equal {
144                    // Note: dedup_by drops a, and retains b.
145                    b.chunks.append(&mut a.chunks);
146                    return true;
147                }
148                false
149            })
150        }
151    }
152
153    /// Computes the size of the collection.
154    ///
155    /// # Returns
156    ///
157    /// The number of traces contained in the collection.
158    ///
159    /// # Examples:
160    ///
161    /// ```rust
162    /// use libdd_trace_protobuf::pb::TracerPayload;
163    /// use libdd_trace_utils::tracer_payload::TracerPayloadCollection;
164    /// let col1 = TracerPayloadCollection::V07(vec![TracerPayload::default()]);
165    /// col1.size();
166    /// ```
167    pub fn size(&self) -> usize {
168        match self {
169            TracerPayloadCollection::V07(collection) => {
170                collection.iter().map(|s| s.chunks.len()).sum()
171            }
172            TracerPayloadCollection::V04(collection) => collection.len(),
173            TracerPayloadCollection::V05((_, collection)) => collection.len(),
174            TracerPayloadCollection::V1(collection) => collection.chunks.len(),
175        }
176    }
177}
178
179/// A trait defining custom processing to be applied to `TraceChunks`.
180///
181/// TraceChunks are part of the v07 Trace payloads. Implementors of this trait can define specific
182/// logic to modify or enrich trace chunks and pass it to the `TracerPayloadCollection` via
183/// `TracerPayloadParams`.
184///
185/// # Examples
186///
187/// Implementing `TraceChunkProcessor` to add a custom tag to each span in a chunk:
188///
189/// ```rust
190/// use libdd_trace_protobuf::pb::{Span, TraceChunk};
191/// use libdd_trace_utils::tracer_payload::TraceChunkProcessor;
192/// use std::collections::HashMap;
193///
194/// struct CustomTagProcessor {
195///     tag_key: String,
196///     tag_value: String,
197/// }
198///
199/// impl TraceChunkProcessor for CustomTagProcessor {
200///     fn process(&mut self, chunk: &mut TraceChunk, index: usize) {
201///         for span in &mut chunk.spans {
202///             span.meta
203///                 .insert(self.tag_key.clone(), self.tag_value.clone());
204///         }
205///     }
206/// }
207/// ```
208pub trait TraceChunkProcessor {
209    fn process(&mut self, chunk: &mut pb::TraceChunk, index: usize);
210}
211
212#[derive(Default)]
213/// Default implementation of `TraceChunkProcessor` that does nothing.
214///
215/// If used, the compiler should optimize away calls to it.
216pub struct DefaultTraceChunkProcessor;
217
218impl TraceChunkProcessor for DefaultTraceChunkProcessor {
219    fn process(&mut self, _chunk: &mut pb::TraceChunk, _index: usize) {
220        // Default implementation does nothing.
221    }
222}
223
224/// This method processes the msgpack data contained within `data` based on
225/// the specified `encoding_type`, converting it into a collection of tracer payloads.
226///
227/// Note: Currently only the `TraceEncoding::V04` and `TraceEncoding::V05` encoding types are
228/// supported.
229///
230/// # Returns
231///
232/// A `Result` containing either the successfully converted `TraceChunks` and the length consummed
233/// from the data  or an error if the conversion fails. Possible errors include issues with
234/// deserializing the msgpack data or if the data does not conform to the expected format.
235///
236/// # Examples
237///
238/// ```rust
239/// use libdd_tinybytes;
240/// use libdd_trace_protobuf::pb;
241/// use libdd_trace_utils::trace_utils::TracerHeaderTags;
242/// use libdd_trace_utils::tracer_payload::{decode_to_trace_chunks, TraceEncoding};
243/// use std::convert::TryInto;
244/// // This will likely be a &[u8] slice in practice.
245/// let data: Vec<u8> = Vec::new();
246/// let data_as_bytes = libdd_tinybytes::Bytes::from(data);
247/// let result = decode_to_trace_chunks(data_as_bytes, TraceEncoding::V04)
248///     .map(|(chunks, _size)| chunks.into_tracer_payload_collection());
249///
250/// match result {
251///     Ok(collection) => println!("Successfully converted to TracerPayloadCollection."),
252///     Err(e) => println!("Failed to convert: {:?}", e),
253/// }
254/// ```
255pub fn decode_to_trace_chunks(
256    data: libdd_tinybytes::Bytes,
257    encoding_type: TraceEncoding,
258) -> Result<(TraceChunks<BytesData>, usize), anyhow::Error> {
259    match encoding_type {
260        TraceEncoding::V04 => {
261            let (data, size) = msgpack_decoder::v04::from_bytes(data).map_err(|e| {
262                anyhow::format_err!("Error deserializing trace from request body: {e}")
263            })?;
264            Ok((TraceChunks::V04(data), size))
265        }
266        TraceEncoding::V05 => {
267            let (data, size) = msgpack_decoder::v05::from_bytes(data).map_err(|e| {
268                anyhow::format_err!("Error deserializing trace from request body: {e}")
269            })?;
270            Ok((convert_trace_chunks_v04_to_v05(data)?, size))
271        }
272        TraceEncoding::V1 => {
273            let (data, size) = msgpack_decoder::v1::from_bytes(data).map_err(|e| {
274                anyhow::format_err!("Error deserializing trace from request body: {e}")
275            })?;
276            Ok((TraceChunks::V1(Box::new(data)), size))
277        }
278    }
279}
280
281/// Returns `true` if and only if every tracer-level metadata field (string fields and attributes)
282/// of `src` matches `dest`.
283///
284/// V1 payloads carry tracer metadata (env, hostname, language, …) inside the payload itself, so
285/// merging two payloads whose metadata diverges would silently drop one set of values. Callers
286/// use this to gate the merge: on a `false` return, append is skipped (no-op) and the two
287/// payloads stay separate. A warning is logged listing the diverging fields so the situation
288/// is observable rather than silent.
289fn metadata_matches_v1(
290    dest: &v1::TracerPayload<BytesData>,
291    src: &v1::TracerPayload<BytesData>,
292) -> bool {
293    let fields = [
294        ("container_id", dest.container_id == src.container_id),
295        ("language_name", dest.language_name == src.language_name),
296        (
297            "language_version",
298            dest.language_version == src.language_version,
299        ),
300        ("tracer_version", dest.tracer_version == src.tracer_version),
301        ("runtime_id", dest.runtime_id == src.runtime_id),
302        ("env", dest.env == src.env),
303        ("hostname", dest.hostname == src.hostname),
304        ("app_version", dest.app_version == src.app_version),
305        ("attributes", dest.attributes.slow_compare(&src.attributes)),
306    ];
307
308    if fields.iter().any(|(_, eq)| !eq) {
309        warn!(
310            "Skipping V1 TracerPayload append: diverging metadata fields {:?}",
311            fields
312                .iter()
313                .filter_map(|(label, eq)| (!eq).then_some(*label))
314                .collect::<Vec<_>>()
315        );
316        return false;
317    }
318    true
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324    use crate::span::v04::{SpanBytes, VecMap};
325    use crate::test_utils::create_test_no_alloc_span;
326    use libdd_tinybytes::BytesString;
327    use libdd_trace_protobuf::pb;
328    use serde_json::json;
329
330    fn create_dummy_collection_v07() -> TracerPayloadCollection {
331        TracerPayloadCollection::V07(vec![pb::TracerPayload {
332            container_id: "".to_string(),
333            language_name: "".to_string(),
334            language_version: "".to_string(),
335            tracer_version: "".to_string(),
336            runtime_id: "".to_string(),
337            chunks: vec![pb::TraceChunk {
338                priority: 0,
339                origin: "".to_string(),
340                spans: vec![],
341                tags: Default::default(),
342                dropped_trace: false,
343            }],
344            tags: Default::default(),
345            env: "".to_string(),
346            hostname: "".to_string(),
347            app_version: "".to_string(),
348            container_debug: None,
349        }])
350    }
351
352    fn create_trace() -> Vec<SpanBytes> {
353        vec![
354            // create a root span with metrics
355            create_test_no_alloc_span(1234, 12341, 0, 1, true),
356            create_test_no_alloc_span(1234, 12342, 12341, 1, false),
357            create_test_no_alloc_span(1234, 12343, 12342, 1, false),
358        ]
359    }
360
361    #[test]
362    fn test_append_traces_v07() {
363        let mut two_traces = create_dummy_collection_v07();
364        two_traces.append(&mut create_dummy_collection_v07());
365
366        let mut trace = create_dummy_collection_v07();
367
368        let mut empty = TracerPayloadCollection::V07(vec![]);
369
370        trace.append(&mut create_dummy_collection_v07());
371        assert_eq!(2, trace.size());
372
373        trace.append(&mut two_traces);
374        assert_eq!(4, trace.size());
375
376        trace.append(&mut empty);
377        assert_eq!(4, trace.size());
378    }
379
380    #[test]
381    fn test_append_traces_v04() {
382        fn create_trace() -> TracerPayloadCollection {
383            TracerPayloadCollection::V04(vec![vec![create_test_no_alloc_span(0, 1, 0, 2, true)]])
384        }
385
386        let mut two_traces = create_trace();
387        two_traces.append(&mut create_trace());
388
389        let mut trace = create_trace();
390
391        let mut empty = TracerPayloadCollection::V04(vec![]);
392
393        trace.append(&mut create_trace());
394        assert_eq!(2, trace.size());
395
396        trace.append(&mut two_traces);
397        assert_eq!(4, trace.size());
398
399        trace.append(&mut empty);
400        assert_eq!(4, trace.size());
401    }
402
403    #[test]
404    fn test_merge_traces() {
405        let mut trace = create_dummy_collection_v07();
406
407        trace.append(&mut create_dummy_collection_v07());
408        assert_eq!(2, trace.size());
409
410        trace.merge();
411        assert_eq!(2, trace.size());
412        if let TracerPayloadCollection::V07(collection) = trace {
413            assert_eq!(1, collection.len());
414        } else {
415            panic!("Unexpected type");
416        }
417    }
418
419    #[test]
420    fn test_try_into_success() {
421        let span_data1 = json!([{
422            "service": "test-service",
423            "name": "test-service-name",
424            "resource": "test-service-resource",
425            "trace_id": 111,
426            "span_id": 222,
427            "parent_id": 100,
428            "start": 1,
429            "duration": 5,
430            "error": 0,
431            "meta": {},
432            "metrics": {},
433            "type": "serverless",
434        }]);
435
436        let expected_serialized_span_data1 = vec![SpanBytes {
437            service: BytesString::from_slice("test-service".as_ref()).unwrap(),
438            name: BytesString::from_slice("test-service-name".as_ref()).unwrap(),
439            resource: BytesString::from_slice("test-service-resource".as_ref()).unwrap(),
440            trace_id: 111,
441            span_id: 222,
442            parent_id: 100,
443            start: 1,
444            duration: 5,
445            error: 0,
446            meta: VecMap::new(),
447            metrics: VecMap::new(),
448            meta_struct: VecMap::new(),
449            r#type: BytesString::from_slice("serverless".as_ref()).unwrap(),
450            span_links: vec![],
451            span_events: vec![],
452        }];
453
454        let span_data2 = json!([{
455            "service": "test-service",
456            "name": "test-service-name",
457            "resource": "test-service-resource",
458            "trace_id": 111,
459            "span_id": 333,
460            "parent_id": 100,
461            "start": 1,
462            "duration": 5,
463            "error": 1,
464            "meta": {},
465            "metrics": {},
466            "type": "",
467        }]);
468
469        let expected_serialized_span_data2 = vec![SpanBytes {
470            service: BytesString::from_slice("test-service".as_ref()).unwrap(),
471            name: BytesString::from_slice("test-service-name".as_ref()).unwrap(),
472            resource: BytesString::from_slice("test-service-resource".as_ref()).unwrap(),
473            trace_id: 111,
474            span_id: 333,
475            parent_id: 100,
476            start: 1,
477            duration: 5,
478            error: 1,
479            meta: VecMap::new(),
480            metrics: VecMap::new(),
481            meta_struct: VecMap::new(),
482            r#type: BytesString::default(),
483            span_links: vec![],
484            span_events: vec![],
485        }];
486
487        let data = rmp_serde::to_vec(&vec![span_data1, span_data2])
488            .expect("Failed to serialize test span.");
489        let data = libdd_tinybytes::Bytes::from(data);
490
491        let result = decode_to_trace_chunks(data, TraceEncoding::V04);
492
493        assert!(result.is_ok());
494
495        let (chunks, _) = result.unwrap();
496        assert_eq!(2, chunks.size());
497
498        if let TraceChunks::V04(traces) = chunks {
499            assert_eq!(expected_serialized_span_data1, traces[0]);
500            assert_eq!(expected_serialized_span_data2, traces[1]);
501        } else {
502            panic!("Invalid collection type returned for try_into");
503        }
504    }
505
506    #[cfg_attr(miri, ignore)]
507    #[test]
508    fn test_try_into_empty() {
509        let empty_data = vec![0x90];
510        let data = libdd_tinybytes::Bytes::from(empty_data);
511
512        let result = decode_to_trace_chunks(data, TraceEncoding::V04);
513
514        assert!(result.is_ok());
515
516        let (collection, _) = result.unwrap();
517        assert_eq!(0, collection.size());
518    }
519
520    #[test]
521    fn test_try_into_meta_metrics_success() {
522        let dummy_trace = create_trace();
523        let expected = vec![create_trace()];
524        let payload = rmp_serde::to_vec_named(&expected).unwrap();
525        let payload = libdd_tinybytes::Bytes::from(payload);
526
527        let result = decode_to_trace_chunks(payload, TraceEncoding::V04);
528
529        assert!(result.is_ok());
530
531        let (collection, _size) = result.unwrap();
532        assert_eq!(1, collection.size());
533        if let TraceChunks::V04(traces) = collection {
534            assert_eq!(dummy_trace, traces[0]);
535        } else {
536            panic!("Invalid collection type returned for try_into");
537        }
538    }
539}
libdd_trace_utils/tracer_payload.rs

libdd_trace_utils/
tracer_payload.rs
