Skip to main content

libdd_trace_utils/span/
mod.rs

1// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4pub mod trace_utils;
5pub mod v04;
6pub mod v05;
7pub mod vec_map;
8
9use crate::msgpack_decoder::decode::buffer::read_string_ref_nomut;
10use crate::msgpack_decoder::decode::error::DecodeError;
11use crate::span::v05::dict::SharedDict;
12use libdd_tinybytes::{Bytes, BytesString};
13use serde::Serialize;
14use std::borrow::{Borrow, Cow};
15use std::fmt::Debug;
16use std::hash::Hash;
17use std::marker::PhantomData;
18use std::ptr::NonNull;
19use std::{fmt, ptr};
20
21/// Trait representing the requirements for a type to be used as a Span "string" type.
22/// Note: Borrow<str> is not required by the derived traits, but allows to access HashMap elements
23/// from a static str and check if the string is empty.
24pub trait SpanText: Debug + Eq + Hash + Borrow<str> + Serialize + Default + From<String> {
25    fn from_static_str(value: &'static str) -> Self;
26}
27
28impl SpanText for Cow<'_, str> {
29    fn from_static_str(value: &'static str) -> Self {
30        Cow::Borrowed(value)
31    }
32}
33
34impl SpanText for BytesString {
35    fn from_static_str(value: &'static str) -> Self {
36        BytesString::from_static(value)
37    }
38}
39
40pub trait SpanBytes: Debug + Eq + Hash + Borrow<[u8]> + Serialize + Default {
41    fn from_static_bytes(value: &'static [u8]) -> Self;
42}
43
44impl SpanBytes for &[u8] {
45    fn from_static_bytes(value: &'static [u8]) -> Self {
46        value
47    }
48}
49
50impl SpanBytes for Bytes {
51    fn from_static_bytes(value: &'static [u8]) -> Self {
52        Bytes::from_static(value)
53    }
54}
55
56/// Trait representing a tuple of (Text, Bytes) types used for different underlying data structures.
57/// Note: The functions are internal to the msgpack decoder and should not be used directly: they're
58/// only exposed here due to the unavailability of min_specialization in stable Rust.
59/// Also note that the Clone and PartialEq bounds are only present for tests.
60pub trait TraceData: Default + Clone + Debug + PartialEq {
61    type Text: SpanText;
62    type Bytes: SpanBytes;
63}
64
65pub trait DeserializableTraceData: TraceData {
66    fn get_mut_slice(buf: &mut Self::Bytes) -> &mut &'static [u8];
67
68    fn try_slice_and_advance(buf: &mut Self::Bytes, bytes: usize) -> Option<Self::Bytes>;
69
70    fn read_string(buf: &mut Self::Bytes) -> Result<Self::Text, DecodeError>;
71}
72
73/// TraceData implementation using `Bytes` and `BytesString`.
74#[derive(Clone, Default, Debug, PartialEq, Serialize)]
75pub struct BytesData;
76impl TraceData for BytesData {
77    type Text = BytesString;
78    type Bytes = Bytes;
79}
80
81impl DeserializableTraceData for BytesData {
82    #[inline]
83    fn get_mut_slice(buf: &mut Bytes) -> &mut &'static [u8] {
84        // SAFETY: Bytes has the same layout
85        unsafe { std::mem::transmute::<&mut Bytes, &mut &[u8]>(buf) }
86    }
87
88    #[inline]
89    fn try_slice_and_advance(buf: &mut Bytes, bytes: usize) -> Option<Bytes> {
90        let data = buf.slice_ref(&buf[0..bytes])?;
91        unsafe {
92            // SAFETY: forwarding the buffer requires that buf is borrowed from static.
93            let (ptr, len, underlying) = ptr::read(buf).into_raw();
94            ptr::write(
95                buf,
96                Bytes::from_raw(ptr.add(bytes), len - bytes, underlying),
97            );
98        }
99        Some(data)
100    }
101
102    #[inline]
103    fn read_string(buf: &mut Bytes) -> Result<BytesString, DecodeError> {
104        // Note: we need to pass a &'static lifetime here, otherwise it'll complain
105        let (str, newbuf) = read_string_ref_nomut(buf.as_ref())?;
106        let string = BytesString::from_bytes_slice(buf, str);
107        unsafe {
108            // SAFETY: forwarding the buffer requires that buf is borrowed from static.
109            let (_, _, underlying) = ptr::read(buf).into_raw();
110            let new = Bytes::from_raw(
111                NonNull::new_unchecked(newbuf.as_ptr() as *mut _),
112                newbuf.len(),
113                underlying,
114            );
115            ptr::write(buf, new);
116        }
117        Ok(string)
118    }
119}
120
121/// TraceData implementation using `Cow<'a, str>` and `&[u8]`.
122#[derive(Clone, Default, Debug, PartialEq, Serialize)]
123pub struct SliceData<'a>(PhantomData<&'a u8>);
124impl<'a> TraceData for SliceData<'a> {
125    type Text = Cow<'a, str>;
126    type Bytes = &'a [u8];
127}
128
129impl<'a> DeserializableTraceData for SliceData<'a> {
130    #[inline]
131    fn get_mut_slice<'b>(buf: &'b mut Self::Bytes) -> &'b mut &'static [u8] {
132        unsafe { std::mem::transmute::<&'b mut &[u8], &'b mut &'static [u8]>(buf) }
133    }
134
135    #[inline]
136    fn try_slice_and_advance(buf: &mut &'a [u8], bytes: usize) -> Option<&'a [u8]> {
137        let slice = buf.get(0..bytes)?;
138        *buf = &buf[bytes..];
139        Some(slice)
140    }
141
142    #[inline]
143    fn read_string(buf: &mut &'a [u8]) -> Result<Cow<'a, str>, DecodeError> {
144        read_string_ref_nomut(buf).map(|(str, newbuf)| {
145            *buf = newbuf;
146            Cow::Borrowed(str)
147        })
148    }
149}
150
151#[derive(Debug)]
152pub struct SpanKeyParseError {
153    pub message: String,
154}
155
156impl SpanKeyParseError {
157    pub fn new(message: impl Into<String>) -> Self {
158        SpanKeyParseError {
159            message: message.into(),
160        }
161    }
162}
163impl fmt::Display for SpanKeyParseError {
164    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
165        write!(f, "SpanKeyParseError: {}", self.message)
166    }
167}
168impl std::error::Error for SpanKeyParseError {}
169
170pub type SharedDictBytes = SharedDict<BytesString>;