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