libdd_trace_utils/span/

mod.rs

// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
// SPDX-License-Identifier: Apache-2.0

pub mod trace_utils;
pub mod trace_utils_v1;
pub mod v04;
pub mod v05;
pub mod v1;
pub mod vec_map;

use crate::msgpack_decoder::decode::buffer::read_string_ref_nomut;
use crate::msgpack_decoder::decode::error::DecodeError;
use crate::span::v05::dict::SharedDict;
use libdd_tinybytes::{Bytes, BytesString};
use serde::Serialize;
use std::borrow::Borrow;
use std::fmt::Debug;
use std::hash::Hash;
use std::marker::PhantomData;
use std::ptr::NonNull;
use std::{fmt, ptr};

/// A `SpanLink`'s `flags` field reserves bit 31 to mean "a value was explicitly set", separate
/// from the sampling decision carried in the low bits. The sentinel bit distinguishes
/// `flags == 0` (never set) from an explicit decision of `0`, for example a dropped context.
/// Without the sentinel, both cases look identical on the wire.
///
/// Every non-JSON wire format keeps this bit raw, except OTLP. The native v0.4 msgpack format
/// (`msgpack_encoder::v04::span_v04`), the v1 msgpack format, and the native protobuf format
/// (`libdd_trace_protobuf::pb::SpanLink`) all keep the sentinel raw in `flags`. Tracers already
/// send the bit set in these formats. JSON formats and OTLP protobuf must mask this bit before
/// they emit `flags`, because those consumers treat `flags` as the real W3C trace-flags value.
/// The JSON formats are the v0.5 `_dd.span_links` dictionary, agentless JSON, and structured
/// JSON logging.
pub(crate) const SPAN_LINK_FLAGS_SET_SENTINEL: u32 = 1 << 31;

/// Trait representing the requirements for a type to be used as a Span "string" type.
/// Note: Borrow<str> is not required by the derived traits, but allows to access HashMap elements
/// from a static str and check if the string is empty.
pub trait SpanText: Debug + Eq + Hash + Borrow<str> + Serialize + Default {
    fn from_static_str(value: &'static str) -> Self;

    /// Copies this text into an owned [`BytesString`].
    ///
    /// Used by the v0.5 conversion, whose shared dictionary always owns its strings so it
    /// can hold both interned span text and dynamically-built JSON (span links / events).
    /// The default copies the bytes; owned text types (e.g. `BytesString`) should override
    /// with a cheaper reference-counted clone.
    fn to_bytes_string(&self) -> BytesString {
        BytesString::from(<Self as Borrow<str>>::borrow(self).to_string())
    }
}

impl SpanText for &str {
    fn from_static_str(value: &'static str) -> Self {
        value
    }
}

impl SpanText for BytesString {
    fn from_static_str(value: &'static str) -> Self {
        BytesString::from_static(value)
    }

    fn to_bytes_string(&self) -> BytesString {
        self.clone()
    }
}

pub trait SpanBytes: Debug + Eq + Hash + Borrow<[u8]> + Serialize + Default + Clone {
    fn from_static_bytes(value: &'static [u8]) -> Self;
}

impl SpanBytes for &[u8] {
    fn from_static_bytes(value: &'static [u8]) -> Self {
        value
    }
}

impl SpanBytes for Bytes {
    fn from_static_bytes(value: &'static [u8]) -> Self {
        Bytes::from_static(value)
    }
}

/// Trait representing a tuple of (Text, Bytes) types used for different underlying data structures.
/// Note: The functions are internal to the msgpack decoder and should not be used directly: they're
/// only exposed here due to the unavailability of min_specialization in stable Rust.
/// Also note that the Clone and PartialEq bounds are only present for tests.
pub trait TraceData: Default + Clone + Debug + PartialEq {
    type Text: SpanText;
    type Bytes: SpanBytes;
}

pub trait DeserializableTraceData: TraceData {
    fn get_mut_slice(buf: &mut Self::Bytes) -> &mut &'static [u8];

    fn try_slice_and_advance(buf: &mut Self::Bytes, bytes: usize) -> Option<Self::Bytes>;

    fn read_string(buf: &mut Self::Bytes) -> Result<Self::Text, DecodeError>;

    /// Interns a string found while walking a value through `get_mut_slice`'s lied `'static`
    /// view (e.g. skipping an unrecognized V1 field for forward compatibility). `s` really
    /// borrows from `owner`'s memory, not `'static`: implementations must derive `Self::Text`
    /// from `owner` itself rather than trusting that lifetime, so a refcounted backing
    /// allocation isn't freed out from under the interned string.
    fn intern_skipped_str(owner: &Self::Bytes, s: &'static str) -> Self::Text;
}

/// TraceData implementation using `Bytes` and `BytesString`.
#[derive(Clone, Default, Debug, PartialEq, Serialize)]
pub struct BytesData;
impl TraceData for BytesData {
    type Text = BytesString;
    type Bytes = Bytes;
}

impl DeserializableTraceData for BytesData {
    #[inline]
    fn get_mut_slice(buf: &mut Bytes) -> &mut &'static [u8] {
        // SAFETY: Bytes has the same layout
        unsafe { std::mem::transmute::<&mut Bytes, &mut &[u8]>(buf) }
    }

    #[inline]
    fn try_slice_and_advance(buf: &mut Bytes, bytes: usize) -> Option<Bytes> {
        if bytes > buf.len() {
            return None;
        }
        let data = buf.slice_ref(&buf[0..bytes])?;
        unsafe {
            // SAFETY: forwarding the buffer requires that buf is borrowed from static.
            let (ptr, len, underlying) = ptr::read(buf).into_raw();
            ptr::write(
                buf,
                Bytes::from_raw(ptr.add(bytes), len - bytes, underlying),
            );
        }
        Some(data)
    }

    #[inline]
    fn read_string(buf: &mut Bytes) -> Result<BytesString, DecodeError> {
        // Note: we need to pass a &'static lifetime here, otherwise it'll complain
        let (str, newbuf) = read_string_ref_nomut(buf.as_ref())?;
        let string = BytesString::from_bytes_slice(buf, str);
        unsafe {
            // SAFETY: forwarding the buffer requires that buf is borrowed from static.
            let (_, _, underlying) = ptr::read(buf).into_raw();
            let new = Bytes::from_raw(
                NonNull::new_unchecked(newbuf.as_ptr() as *mut _),
                newbuf.len(),
                underlying,
            );
            ptr::write(buf, new);
        }
        Ok(string)
    }

    #[inline]
    fn intern_skipped_str(owner: &Bytes, s: &'static str) -> BytesString {
        BytesString::from_bytes_slice(owner, s)
    }
}

/// TraceData implementation using `&str` and `&[u8]`.
#[derive(Clone, Default, Debug, PartialEq, Serialize)]
pub struct SliceData<'a>(PhantomData<&'a u8>);
impl<'a> TraceData for SliceData<'a> {
    type Text = &'a str;
    type Bytes = &'a [u8];
}

impl<'a> DeserializableTraceData for SliceData<'a> {
    #[inline]
    fn get_mut_slice<'b>(buf: &'b mut Self::Bytes) -> &'b mut &'static [u8] {
        unsafe { std::mem::transmute::<&'b mut &[u8], &'b mut &'static [u8]>(buf) }
    }

    #[inline]
    fn try_slice_and_advance(buf: &mut &'a [u8], bytes: usize) -> Option<&'a [u8]> {
        let slice = buf.get(0..bytes)?;
        *buf = &buf[bytes..];
        Some(slice)
    }

    #[inline]
    fn read_string(buf: &mut &'a [u8]) -> Result<&'a str, DecodeError> {
        read_string_ref_nomut(buf).map(|(str, newbuf)| {
            *buf = newbuf;
            str
        })
    }

    #[inline]
    fn intern_skipped_str(_owner: &&'a [u8], s: &'static str) -> &'a str {
        // No refcounted allocation to preserve here: `s` borrows from a plain slice the
        // caller owns for `'a`, and a `'static` reference is always a valid `'a` reference.
        s
    }
}

#[derive(Debug)]
pub struct SpanKeyParseError {
    pub message: String,
}

impl SpanKeyParseError {
    pub fn new(message: impl Into<String>) -> Self {
        SpanKeyParseError {
            message: message.into(),
        }
    }
}
impl fmt::Display for SpanKeyParseError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "SpanKeyParseError: {}", self.message)
    }
}
impl std::error::Error for SpanKeyParseError {}

pub type SharedDictBytes = SharedDict<BytesString>;