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