libdd_trace_utils/msgpack_decoder/decode/
span_link.rs

1// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::msgpack_decoder::decode::buffer::Buffer;
5use crate::msgpack_decoder::decode::error::DecodeError;
6use crate::msgpack_decoder::decode::number::read_number;
7use crate::msgpack_decoder::decode::string::{handle_null_marker, read_str_map_to_strings};
8use crate::span::v04::SpanLink;
9use crate::span::DeserializableTraceData;
10use std::borrow::Borrow;
11use std::str::FromStr;
12
13/// Reads a slice of bytes and decodes it into a vector of `SpanLink` objects.
14///
15/// # Arguments
16///
17/// * `buf` - A mutable reference to a slice of bytes containing the encoded data.
18///
19/// # Returns
20///
21/// * `Ok(Vec<SpanLink>)` - A vector of decoded `SpanLink` objects if successful.
22/// * `Err(DecodeError)` - An error if the decoding process fails.
23///
24/// # Errors
25///
26/// This function will return an error if:
27/// - The marker for the array length cannot be read.
28/// - Any `SpanLink` cannot be decoded.
29/// ```
30pub(crate) fn read_span_links<T: DeserializableTraceData>(
31    buf: &mut Buffer<T>,
32) -> Result<Vec<SpanLink<T>>, DecodeError> {
33    if handle_null_marker(buf) {
34        return Ok(Vec::default());
35    }
36
37    let len = rmp::decode::read_array_len(buf.as_mut_slice()).map_err(|_| {
38        DecodeError::InvalidType("Unable to get array len for span links".to_owned())
39    })?;
40
41    let mut vec: Vec<SpanLink<T>> = Vec::with_capacity(len as usize);
42    for _ in 0..len {
43        vec.push(decode_span_link(buf)?);
44    }
45    Ok(vec)
46}
47#[derive(Debug, PartialEq)]
48enum SpanLinkKey {
49    TraceId,
50    TraceIdHigh,
51    SpanId,
52    Attributes,
53    Tracestate,
54    Flags,
55}
56
57impl FromStr for SpanLinkKey {
58    type Err = DecodeError;
59
60    fn from_str(s: &str) -> Result<Self, Self::Err> {
61        match s {
62            "trace_id" => Ok(SpanLinkKey::TraceId),
63            "trace_id_high" => Ok(SpanLinkKey::TraceIdHigh),
64            "span_id" => Ok(SpanLinkKey::SpanId),
65            "attributes" => Ok(SpanLinkKey::Attributes),
66            "tracestate" => Ok(SpanLinkKey::Tracestate),
67            "flags" => Ok(SpanLinkKey::Flags),
68            _ => Err(DecodeError::InvalidFormat(
69                format!("Invalid span link key: {s}").to_owned(),
70            )),
71        }
72    }
73}
74
75fn decode_span_link<T: DeserializableTraceData>(
76    buf: &mut Buffer<T>,
77) -> Result<SpanLink<T>, DecodeError> {
78    let mut span = SpanLink::default();
79    let span_size = rmp::decode::read_map_len(buf.as_mut_slice())
80        .map_err(|_| DecodeError::InvalidType("Unable to get map len for span size".to_owned()))?;
81
82    for _ in 0..span_size {
83        match buf.read_string()?.borrow().parse::<SpanLinkKey>()? {
84            SpanLinkKey::TraceId => span.trace_id = read_number(buf)?,
85            SpanLinkKey::TraceIdHigh => span.trace_id_high = read_number(buf)?,
86            SpanLinkKey::SpanId => span.span_id = read_number(buf)?,
87            SpanLinkKey::Attributes => span.attributes = read_str_map_to_strings(buf)?,
88            SpanLinkKey::Tracestate => span.tracestate = buf.read_string()?,
89            SpanLinkKey::Flags => span.flags = read_number(buf)?,
90        }
91    }
92
93    Ok(span)
94}
95
96#[cfg(test)]
97mod tests {
98    use super::SpanLinkKey;
99    use crate::msgpack_decoder::decode::error::DecodeError;
100    use std::str::FromStr;
101
102    #[test]
103    fn test_span_link_key_from_str() {
104        // Valid cases
105        assert_eq!(
106            SpanLinkKey::from_str("trace_id").unwrap(),
107            SpanLinkKey::TraceId
108        );
109        assert_eq!(
110            SpanLinkKey::from_str("trace_id_high").unwrap(),
111            SpanLinkKey::TraceIdHigh
112        );
113        assert_eq!(
114            SpanLinkKey::from_str("span_id").unwrap(),
115            SpanLinkKey::SpanId
116        );
117        assert_eq!(
118            SpanLinkKey::from_str("attributes").unwrap(),
119            SpanLinkKey::Attributes
120        );
121        assert_eq!(
122            SpanLinkKey::from_str("tracestate").unwrap(),
123            SpanLinkKey::Tracestate
124        );
125        assert_eq!(SpanLinkKey::from_str("flags").unwrap(), SpanLinkKey::Flags);
126
127        // Invalid case
128        assert!(matches!(
129            SpanLinkKey::from_str("invalid_key"),
130            Err(DecodeError::InvalidFormat(_))
131        ));
132    }
133}
libdd_trace_utils/msgpack_decoder/decode/span_link.rs

libdd_trace_utils/msgpack_decoder/decode/
span_link.rs
