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

libdd_trace_utils/msgpack_decoder/decode/
span_link.rs
