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::v05::dict::SharedDict;
5use crate::span::{v04, v05, BytesData, SharedDictBytes, TraceData};
6use crate::trace_utils::collect_trace_chunks;
7use crate::{msgpack_decoder, trace_utils::cmp_send_data_payloads};
8use libdd_trace_protobuf::pb;
9use std::cmp::Ordering;
10use std::iter::Iterator;
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}
23
24#[derive(Debug)]
25pub enum TraceChunks<T: TraceData> {
26    /// Collection of TraceChunkSpan.
27    V04(Vec<Vec<v04::Span<T>>>),
28    /// Collection of TraceChunkSpan with de-duplicated strings.
29    V05((SharedDict<T::Text>, Vec<Vec<v05::Span>>)),
30    /// Collection of v0.4 spans to be serialized as a V1 msgpack payload.
31    V1(Vec<Vec<v04::Span<T>>>),
32}
33
34impl TraceChunks<BytesData> {
35    pub fn into_tracer_payload_collection(self) -> TracerPayloadCollection {
36        match self {
37            TraceChunks::V04(traces) => TracerPayloadCollection::V04(traces),
38            TraceChunks::V05(traces) => TracerPayloadCollection::V05(traces),
39            // V1 uses the same underlying span structure as V04.
40            TraceChunks::V1(traces) => TracerPayloadCollection::V04(traces),
41        }
42    }
43}
44
45impl<T: TraceData> TraceChunks<T> {
46    /// Returns the number of traces in the chunk
47    pub fn size(&self) -> usize {
48        match self {
49            TraceChunks::V04(traces) => traces.len(),
50            TraceChunks::V05((_, traces)) => traces.len(),
51            TraceChunks::V1(traces) => traces.len(),
52        }
53    }
54}
55
56#[derive(Debug)]
57/// Enum representing a general abstraction for a collection of tracer payloads.
58pub enum TracerPayloadCollection {
59    /// Collection of TracerPayloads.
60    V07(Vec<pb::TracerPayload>),
61    /// Collection of TraceChunkSpan.
62    V04(Vec<Vec<v04::SpanBytes>>),
63    /// Collection of TraceChunkSpan with de-duplicated strings.
64    V05((SharedDictBytes, Vec<Vec<v05::Span>>)),
65}
66
67impl TracerPayloadCollection {
68    /// Appends `other` collection of the same type to the current collection.
69    ///
70    /// #Arguments
71    ///
72    /// * `other`: collection of the same type.
73    ///
74    /// # Examples:
75    ///
76    /// ```rust
77    /// use libdd_trace_protobuf::pb::TracerPayload;
78    /// use libdd_trace_utils::tracer_payload::TracerPayloadCollection;
79    /// let mut col1 = TracerPayloadCollection::V07(vec![TracerPayload::default()]);
80    /// let mut col2 = TracerPayloadCollection::V07(vec![TracerPayload::default()]);
81    /// col1.append(&mut col2);
82    /// ```
83    pub fn append(&mut self, other: &mut Self) {
84        match self {
85            TracerPayloadCollection::V07(dest) => {
86                if let TracerPayloadCollection::V07(src) = other {
87                    dest.append(src)
88                }
89            }
90            TracerPayloadCollection::V04(dest) => {
91                if let TracerPayloadCollection::V04(src) = other {
92                    dest.append(src)
93                }
94            }
95            // TODO: Properly handle non-OK states to prevent possible panics (APMSP-18190).
96            #[allow(clippy::unimplemented)]
97            TracerPayloadCollection::V05(_) => unimplemented!("Append for V05 not implemented"),
98        }
99    }
100
101    /// Merges traces that came from the same origin together to reduce the payload size.
102    ///
103    /// # Examples:
104    ///
105    /// ```rust
106    /// use libdd_trace_protobuf::pb::TracerPayload;
107    /// use libdd_trace_utils::tracer_payload::TracerPayloadCollection;
108    /// let mut col1 =
109    ///     TracerPayloadCollection::V07(vec![TracerPayload::default(), TracerPayload::default()]);
110    /// col1.merge();
111    /// ```
112    pub fn merge(&mut self) {
113        if let TracerPayloadCollection::V07(collection) = self {
114            collection.sort_unstable_by(cmp_send_data_payloads);
115            collection.dedup_by(|a, b| {
116                if cmp_send_data_payloads(a, b) == Ordering::Equal {
117                    // Note: dedup_by drops a, and retains b.
118                    b.chunks.append(&mut a.chunks);
119                    return true;
120                }
121                false
122            })
123        }
124    }
125
126    /// Computes the size of the collection.
127    ///
128    /// # Returns
129    ///
130    /// The number of traces contained in the collection.
131    ///
132    /// # Examples:
133    ///
134    /// ```rust
135    /// use libdd_trace_protobuf::pb::TracerPayload;
136    /// use libdd_trace_utils::tracer_payload::TracerPayloadCollection;
137    /// let col1 = TracerPayloadCollection::V07(vec![TracerPayload::default()]);
138    /// col1.size();
139    /// ```
140    pub fn size(&self) -> usize {
141        match self {
142            TracerPayloadCollection::V07(collection) => {
143                collection.iter().map(|s| s.chunks.len()).sum()
144            }
145            TracerPayloadCollection::V04(collection) => collection.len(),
146            TracerPayloadCollection::V05((_, collection)) => collection.len(),
147        }
148    }
149}
150
151/// A trait defining custom processing to be applied to `TraceChunks`.
152///
153/// TraceChunks are part of the v07 Trace payloads. Implementors of this trait can define specific
154/// logic to modify or enrich trace chunks and pass it to the `TracerPayloadCollection` via
155/// `TracerPayloadParams`.
156///
157/// # Examples
158///
159/// Implementing `TraceChunkProcessor` to add a custom tag to each span in a chunk:
160///
161/// ```rust
162/// use libdd_trace_protobuf::pb::{Span, TraceChunk};
163/// use libdd_trace_utils::tracer_payload::TraceChunkProcessor;
164/// use std::collections::HashMap;
165///
166/// struct CustomTagProcessor {
167///     tag_key: String,
168///     tag_value: String,
169/// }
170///
171/// impl TraceChunkProcessor for CustomTagProcessor {
172///     fn process(&mut self, chunk: &mut TraceChunk, index: usize) {
173///         for span in &mut chunk.spans {
174///             span.meta
175///                 .insert(self.tag_key.clone(), self.tag_value.clone());
176///         }
177///     }
178/// }
179/// ```
180pub trait TraceChunkProcessor {
181    fn process(&mut self, chunk: &mut pb::TraceChunk, index: usize);
182}
183
184#[derive(Default)]
185/// Default implementation of `TraceChunkProcessor` that does nothing.
186///
187/// If used, the compiler should optimize away calls to it.
188pub struct DefaultTraceChunkProcessor;
189
190impl TraceChunkProcessor for DefaultTraceChunkProcessor {
191    fn process(&mut self, _chunk: &mut pb::TraceChunk, _index: usize) {
192        // Default implementation does nothing.
193    }
194}
195
196/// This method processes the msgpack data contained within `data` based on
197/// the specified `encoding_type`, converting it into a collection of tracer payloads.
198///
199/// Note: Currently only the `TraceEncoding::V04` and `TraceEncoding::V05` encoding types are
200/// supported.
201///
202/// # Returns
203///
204/// A `Result` containing either the successfully converted `TraceChunks` and the length consummed
205/// from the data  or an error if the conversion fails. Possible errors include issues with
206/// deserializing the msgpack data or if the data does not conform to the expected format.
207///
208/// # Examples
209///
210/// ```rust
211/// use libdd_tinybytes;
212/// use libdd_trace_protobuf::pb;
213/// use libdd_trace_utils::trace_utils::TracerHeaderTags;
214/// use libdd_trace_utils::tracer_payload::{decode_to_trace_chunks, TraceEncoding};
215/// use std::convert::TryInto;
216/// // This will likely be a &[u8] slice in practice.
217/// let data: Vec<u8> = Vec::new();
218/// let data_as_bytes = libdd_tinybytes::Bytes::from(data);
219/// let result = decode_to_trace_chunks(data_as_bytes, TraceEncoding::V04)
220///     .map(|(chunks, _size)| chunks.into_tracer_payload_collection());
221///
222/// match result {
223///     Ok(collection) => println!("Successfully converted to TracerPayloadCollection."),
224///     Err(e) => println!("Failed to convert: {:?}", e),
225/// }
226/// ```
227pub fn decode_to_trace_chunks(
228    data: libdd_tinybytes::Bytes,
229    encoding_type: TraceEncoding,
230) -> Result<(TraceChunks<BytesData>, usize), anyhow::Error> {
231    let (data, size) = match encoding_type {
232        TraceEncoding::V04 => msgpack_decoder::v04::from_bytes(data),
233        TraceEncoding::V05 => msgpack_decoder::v05::from_bytes(data),
234    }
235    .map_err(|e| anyhow::format_err!("Error deserializing trace from request body: {e}"))?;
236
237    Ok((collect_trace_chunks(data, encoding_type)?, size))
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243    use crate::span::v04::SpanBytes;
244    use crate::test_utils::create_test_no_alloc_span;
245    use libdd_tinybytes::BytesString;
246    use libdd_trace_protobuf::pb;
247    use serde_json::json;
248    use std::collections::HashMap;
249
250    fn create_dummy_collection_v07() -> TracerPayloadCollection {
251        TracerPayloadCollection::V07(vec![pb::TracerPayload {
252            container_id: "".to_string(),
253            language_name: "".to_string(),
254            language_version: "".to_string(),
255            tracer_version: "".to_string(),
256            runtime_id: "".to_string(),
257            chunks: vec![pb::TraceChunk {
258                priority: 0,
259                origin: "".to_string(),
260                spans: vec![],
261                tags: Default::default(),
262                dropped_trace: false,
263            }],
264            tags: Default::default(),
265            env: "".to_string(),
266            hostname: "".to_string(),
267            app_version: "".to_string(),
268        }])
269    }
270
271    fn create_trace() -> Vec<SpanBytes> {
272        vec![
273            // create a root span with metrics
274            create_test_no_alloc_span(1234, 12341, 0, 1, true),
275            create_test_no_alloc_span(1234, 12342, 12341, 1, false),
276            create_test_no_alloc_span(1234, 12343, 12342, 1, false),
277        ]
278    }
279
280    #[test]
281    fn test_append_traces_v07() {
282        let mut two_traces = create_dummy_collection_v07();
283        two_traces.append(&mut create_dummy_collection_v07());
284
285        let mut trace = create_dummy_collection_v07();
286
287        let mut empty = TracerPayloadCollection::V07(vec![]);
288
289        trace.append(&mut create_dummy_collection_v07());
290        assert_eq!(2, trace.size());
291
292        trace.append(&mut two_traces);
293        assert_eq!(4, trace.size());
294
295        trace.append(&mut empty);
296        assert_eq!(4, trace.size());
297    }
298
299    #[test]
300    fn test_append_traces_v04() {
301        fn create_trace() -> TracerPayloadCollection {
302            TracerPayloadCollection::V04(vec![vec![create_test_no_alloc_span(0, 1, 0, 2, true)]])
303        }
304
305        let mut two_traces = create_trace();
306        two_traces.append(&mut create_trace());
307
308        let mut trace = create_trace();
309
310        let mut empty = TracerPayloadCollection::V04(vec![]);
311
312        trace.append(&mut create_trace());
313        assert_eq!(2, trace.size());
314
315        trace.append(&mut two_traces);
316        assert_eq!(4, trace.size());
317
318        trace.append(&mut empty);
319        assert_eq!(4, trace.size());
320    }
321
322    #[test]
323    fn test_merge_traces() {
324        let mut trace = create_dummy_collection_v07();
325
326        trace.append(&mut create_dummy_collection_v07());
327        assert_eq!(2, trace.size());
328
329        trace.merge();
330        assert_eq!(2, trace.size());
331        if let TracerPayloadCollection::V07(collection) = trace {
332            assert_eq!(1, collection.len());
333        } else {
334            panic!("Unexpected type");
335        }
336    }
337
338    #[test]
339    fn test_try_into_success() {
340        let span_data1 = json!([{
341            "service": "test-service",
342            "name": "test-service-name",
343            "resource": "test-service-resource",
344            "trace_id": 111,
345            "span_id": 222,
346            "parent_id": 100,
347            "start": 1,
348            "duration": 5,
349            "error": 0,
350            "meta": {},
351            "metrics": {},
352            "type": "serverless",
353        }]);
354
355        let expected_serialized_span_data1 = vec![SpanBytes {
356            service: BytesString::from_slice("test-service".as_ref()).unwrap(),
357            name: BytesString::from_slice("test-service-name".as_ref()).unwrap(),
358            resource: BytesString::from_slice("test-service-resource".as_ref()).unwrap(),
359            trace_id: 111,
360            span_id: 222,
361            parent_id: 100,
362            start: 1,
363            duration: 5,
364            error: 0,
365            meta: HashMap::new(),
366            metrics: HashMap::new(),
367            meta_struct: HashMap::new(),
368            r#type: BytesString::from_slice("serverless".as_ref()).unwrap(),
369            span_links: vec![],
370            span_events: vec![],
371        }];
372
373        let span_data2 = json!([{
374            "service": "test-service",
375            "name": "test-service-name",
376            "resource": "test-service-resource",
377            "trace_id": 111,
378            "span_id": 333,
379            "parent_id": 100,
380            "start": 1,
381            "duration": 5,
382            "error": 1,
383            "meta": {},
384            "metrics": {},
385            "type": "",
386        }]);
387
388        let expected_serialized_span_data2 = vec![SpanBytes {
389            service: BytesString::from_slice("test-service".as_ref()).unwrap(),
390            name: BytesString::from_slice("test-service-name".as_ref()).unwrap(),
391            resource: BytesString::from_slice("test-service-resource".as_ref()).unwrap(),
392            trace_id: 111,
393            span_id: 333,
394            parent_id: 100,
395            start: 1,
396            duration: 5,
397            error: 1,
398            meta: HashMap::new(),
399            metrics: HashMap::new(),
400            meta_struct: HashMap::new(),
401            r#type: BytesString::default(),
402            span_links: vec![],
403            span_events: vec![],
404        }];
405
406        let data = rmp_serde::to_vec(&vec![span_data1, span_data2])
407            .expect("Failed to serialize test span.");
408        let data = libdd_tinybytes::Bytes::from(data);
409
410        let result = decode_to_trace_chunks(data, TraceEncoding::V04);
411
412        assert!(result.is_ok());
413
414        let (chunks, _) = result.unwrap();
415        assert_eq!(2, chunks.size());
416
417        if let TraceChunks::V04(traces) = chunks {
418            assert_eq!(expected_serialized_span_data1, traces[0]);
419            assert_eq!(expected_serialized_span_data2, traces[1]);
420        } else {
421            panic!("Invalid collection type returned for try_into");
422        }
423    }
424
425    #[cfg_attr(miri, ignore)]
426    #[test]
427    fn test_try_into_empty() {
428        let empty_data = vec![0x90];
429        let data = libdd_tinybytes::Bytes::from(empty_data);
430
431        let result = decode_to_trace_chunks(data, TraceEncoding::V04);
432
433        assert!(result.is_ok());
434
435        let (collection, _) = result.unwrap();
436        assert_eq!(0, collection.size());
437    }
438
439    #[test]
440    fn test_try_into_meta_metrics_success() {
441        let dummy_trace = create_trace();
442        let expected = vec![create_trace()];
443        let payload = rmp_serde::to_vec_named(&expected).unwrap();
444        let payload = libdd_tinybytes::Bytes::from(payload);
445
446        let result = decode_to_trace_chunks(payload, TraceEncoding::V04);
447
448        assert!(result.is_ok());
449
450        let (collection, _size) = result.unwrap();
451        assert_eq!(1, collection.size());
452        if let TraceChunks::V04(traces) = collection {
453            assert_eq!(dummy_trace, traces[0]);
454        } else {
455            panic!("Invalid collection type returned for try_into");
456        }
457    }
458}
libdd_trace_utils/tracer_payload.rs

libdd_trace_utils/
tracer_payload.rs
