Try TraceProjection (incomplete)

Signed-off-by: Bob Weinand <[email protected]>

Author: bwoebi

datadog-trace-utils/src/msgpack_decoder/mod.rs

@@ -4,3 +4,4 @@
 pub mod decode;
 pub mod v04;
 pub mod v05;
+pub mod v1;

datadog-trace-utils/src/msgpack_decoder/v1/mod.rs

@@ -0,0 +1,120 @@
+// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+pub(crate) mod span;
+
+use self::span::decode_span;
+use crate::msgpack_decoder::decode::buffer::Buffer;
+use crate::msgpack_decoder::decode::error::DecodeError;
+use crate::span::v04::{Span, SpanBytes, SpanSlice};
+use crate::span::{TraceData, Traces, TracesBytes, TracesSlice};
+
+/// Decodes a Bytes buffer into a `Vec<Vec<SpanBytes>>` object, also represented as a vector of
+/// `TracerPayloadV04` objects.
+///
+/// # Arguments
+///
+/// * `data` - A tinybytes Bytes buffer containing the encoded data. Bytes are expected to be
+///   encoded msgpack data containing a list of a list of v04 spans.
+///
+/// # Returns
+///
+/// * `Ok(Vec<TracerPayloadV04>, usize)` - A vector of decoded `Vec<SpanSlice>` objects if
+///   successful. and the number of bytes in the slice used by the decoder.
+/// * `Err(DecodeError)` - An error if the decoding process fails.
+///
+/// # Errors
+///
+/// This function will return an error if:
+/// - The array length for trace count or span count cannot be read.
+/// - Any span cannot be decoded.
+///
+/// # Examples
+///
+/// ```
+/// use datadog_trace_protobuf::pb::Span;
+/// use datadog_trace_utils::msgpack_decoder::v04::from_bytes;
+/// use rmp_serde::to_vec_named;
+/// use tinybytes;
+///
+/// let span = Span {
+///     name: "test-span".to_owned(),
+///     ..Default::default()
+/// };
+/// let encoded_data = to_vec_named(&vec![vec![span]]).unwrap();
+/// let encoded_data_as_tinybytes = tinybytes::Bytes::from(encoded_data);
+/// let (decoded_traces, _payload_size) =
+///     from_bytes(encoded_data_as_tinybytes).expect("Decoding failed");
+///
+/// assert_eq!(1, decoded_traces.len());
+/// assert_eq!(1, decoded_traces[0].len());
+/// let decoded_span = &decoded_traces[0][0];
+/// assert_eq!("test-span", decoded_span.name.as_str());
+/// ```
+pub fn from_bytes(data: tinybytes::Bytes) -> Result<(TracesBytes, usize), DecodeError> {
+    from_buffer(&mut Buffer::new(data))
+}
+
+/// Decodes a slice of bytes into a `Vec<Vec<SpanSlice>>` object.
+/// The resulting spans have the same lifetime as the initial buffer.
+///
+/// # Arguments
+///
+/// * `data` - A slice of bytes containing the encoded data. Bytes are expected to be encoded
+///   msgpack data containing a list of a list of v04 spans.
+///
+/// # Returns
+///
+/// * `Ok(Vec<TracerPayloadV04>, usize)` - A vector of decoded `Vec<SpanSlice>` objects if
+///   successful. and the number of bytes in the slice used by the decoder.
+/// * `Err(DecodeError)` - An error if the decoding process fails.
+///
+/// # Errors
+///
+/// This function will return an error if:
+/// - The array length for trace count or span count cannot be read.
+/// - Any span cannot be decoded.
+///
+/// # Examples
+///
+/// ```
+/// use datadog_trace_protobuf::pb::Span;
+/// use datadog_trace_utils::msgpack_decoder::v04::from_slice;
+/// use rmp_serde::to_vec_named;
+/// use tinybytes;
+///
+/// let span = Span {
+///     name: "test-span".to_owned(),
+///     ..Default::default()
+/// };
+/// let encoded_data = to_vec_named(&vec![vec![span]]).unwrap();
+/// let encoded_data_as_tinybytes = tinybytes::Bytes::from(encoded_data);
+/// let (decoded_traces, _payload_size) =
+///     from_slice(&encoded_data_as_tinybytes).expect("Decoding failed");
+///
+/// assert_eq!(1, decoded_traces.len());
+/// assert_eq!(1, decoded_traces[0].len());
+/// let decoded_span = &decoded_traces[0][0];
+/// assert_eq!("test-span", decoded_span.name);
+/// ```
+pub fn from_slice(data: &[u8]) -> Result<(TracesSlice<'_>, usize), DecodeError> {
+    from_buffer(&mut Buffer::new(data))
+}
+
+#[allow(clippy::type_complexity)]
+pub fn from_buffer<T: TraceData>(
+    data: &mut Buffer<T>,
+) -> Result<(Traces<T>, usize), DecodeError> {
+    let trace_count = rmp::decode::read_array_len(data.as_mut_slice()).map_err(|_| {
+        DecodeError::InvalidFormat("Unable to read array len for trace count".to_owned())
+    })?;
+
+    let traces = Traces::default();
+
+    // Intentionally skip the size of the array (as it will be recomputed after coalescing).
+    let start_len = data.len();
+
+
+
+    Ok((traces, start_len - data.len()))
+}

datadog-trace-utils/src/msgpack_decoder/v1/span.rs

@@ -0,0 +1,110 @@
+// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::msgpack_decoder::decode::buffer::Buffer;
+use crate::msgpack_decoder::decode::error::DecodeError;
+use crate::msgpack_decoder::decode::number::read_nullable_number;
+use crate::msgpack_decoder::decode::span_event::read_span_events;
+use crate::msgpack_decoder::decode::span_link::read_span_links;
+use crate::msgpack_decoder::decode::string::{
+    read_nullable_str_map_to_strings, read_nullable_string,
+};
+use crate::msgpack_decoder::decode::{meta_struct::read_meta_struct, metrics::read_metrics};
+use crate::span::{v04::Span, SpanKey, TraceData};
+use std::borrow::Borrow;
+
+/// Decodes a slice of bytes into a `Span` object.
+///
+/// # Arguments
+///
+/// * `buf` - A mutable reference to a slice of bytes containing the encoded data.
+///
+/// # Returns
+///
+/// * `Ok(Span)` - A decoded `Span` object if successful.
+/// * `Err(DecodeError)` - An error if the decoding process fails.
+///
+/// # Errors
+///
+/// This function will return an error if:
+/// - The map length cannot be read.
+/// - Any key or value cannot be decoded.
+pub fn decode_span<T: TraceData>(buffer: &mut Buffer<T>) -> Result<Span<T>, DecodeError> {
+    let mut span = Span::<T>::default();
+
+    let span_size = rmp::decode::read_map_len(buffer.as_mut_slice()).map_err(|_| {
+        DecodeError::InvalidFormat("Unable to get map len for span size".to_owned())
+    })?;
+
+    for _ in 0..span_size {
+        fill_span(&mut span, buffer)?;
+    }
+
+    Ok(span)
+}
+
+// Safety: read_string_ref checks utf8 validity, so we don't do it again when creating the
+// BytesStrings
+fn fill_span<T: TraceData>(span: &mut Span<T>, buf: &mut Buffer<T>) -> Result<(), DecodeError> {
+    let key = buf
+        .read_string()?
+        .borrow()
+        .parse::<SpanKey>()
+        .map_err(|e| DecodeError::InvalidFormat(e.message))?;
+
+    match key {
+        SpanKey::Service => span.service = read_nullable_string(buf)?,
+        SpanKey::Name => span.name = read_nullable_string(buf)?,
+        SpanKey::Resource => span.resource = read_nullable_string(buf)?,
+        SpanKey::TraceId => span.trace_id = read_nullable_number(buf)?,
+        SpanKey::SpanId => span.span_id = read_nullable_number(buf)?,
+        SpanKey::ParentId => span.parent_id = read_nullable_number(buf)?,
+        SpanKey::Start => span.start = read_nullable_number(buf)?,
+        SpanKey::Duration => span.duration = read_nullable_number(buf)?,
+        SpanKey::Error => span.error = read_nullable_number(buf)?,
+        SpanKey::Type => span.r#type = read_nullable_string(buf)?,
+        SpanKey::Meta => span.meta = read_nullable_str_map_to_strings(buf)?,
+        SpanKey::Metrics => span.metrics = read_metrics(buf)?,
+        SpanKey::MetaStruct => span.meta_struct = read_meta_struct(buf)?,
+        SpanKey::SpanLinks => span.span_links = read_span_links(buf)?,
+        SpanKey::SpanEvents => span.span_events = read_span_events(buf)?,
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::SpanKey;
+    use crate::span::SpanKeyParseError;
+    use std::str::FromStr;
+
+    #[test]
+    fn test_span_key_from_str() {
+        assert_eq!(SpanKey::from_str("service").unwrap(), SpanKey::Service);
+        assert_eq!(SpanKey::from_str("name").unwrap(), SpanKey::Name);
+        assert_eq!(SpanKey::from_str("resource").unwrap(), SpanKey::Resource);
+        assert_eq!(SpanKey::from_str("trace_id").unwrap(), SpanKey::TraceId);
+        assert_eq!(SpanKey::from_str("span_id").unwrap(), SpanKey::SpanId);
+        assert_eq!(SpanKey::from_str("parent_id").unwrap(), SpanKey::ParentId);
+        assert_eq!(SpanKey::from_str("start").unwrap(), SpanKey::Start);
+        assert_eq!(SpanKey::from_str("duration").unwrap(), SpanKey::Duration);
+        assert_eq!(SpanKey::from_str("error").unwrap(), SpanKey::Error);
+        assert_eq!(SpanKey::from_str("meta").unwrap(), SpanKey::Meta);
+        assert_eq!(SpanKey::from_str("metrics").unwrap(), SpanKey::Metrics);
+        assert_eq!(SpanKey::from_str("type").unwrap(), SpanKey::Type);
+        assert_eq!(
+            SpanKey::from_str("meta_struct").unwrap(),
+            SpanKey::MetaStruct
+        );
+        assert_eq!(SpanKey::from_str("span_links").unwrap(), SpanKey::SpanLinks);
+        assert_eq!(
+            SpanKey::from_str("span_events").unwrap(),
+            SpanKey::SpanEvents
+        );
+
+        let invalid_result = SpanKey::from_str("invalid_key");
+        let msg = format!("SpanKeyParseError: Invalid span key: {}", "invalid_key");
+        assert!(matches!(invalid_result, Err(SpanKeyParseError { .. })));
+        assert_eq!(invalid_result.unwrap_err().to_string(), msg);
+    }
+}

datadog-trace-utils/src/span/mod.rs

@@ -1,9 +1,14 @@
 // Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+pub mod table;
 pub mod trace_utils;
 pub mod v04;
 pub mod v05;
+pub mod v1;
+
+mod trace;
+pub use trace::*;
 
 use crate::msgpack_decoder::decode::buffer::read_string_ref_nomut;
 use crate::msgpack_decoder::decode::error::DecodeError;
@@ -14,6 +19,7 @@ use std::fmt;
 use std::fmt::Debug;
 use std::hash::Hash;
 use std::marker::PhantomData;
+use std::ops::Deref;
 use tinybytes::{Bytes, BytesString};
 
 /// Trait representing the requirements for a type to be used as a Span "string" type.
@@ -126,6 +132,7 @@ impl<'a> TraceData for SliceData<'a> {
     }
 }
 
+
 #[derive(Debug)]
 pub struct SpanKeyParseError {
     pub message: String,

datadog-trace-utils/src/span/table.rs

@@ -0,0 +1,91 @@
+use std::collections::HashMap;
+use std::marker::PhantomData;
+use crate::span::TraceData;
+
+trait TraceDataType {
+    type Data<T: TraceData>;
+}
+#[derive(Debug, Default, Eq, PartialEq, Hash)]
+pub struct TraceDataBytes;
+impl TraceDataType for TraceDataBytes {
+    type Data<T: TraceData> = T::Bytes;
+}
+#[derive(Debug, Default, Eq, PartialEq, Hash)]
+pub struct TraceDataText;
+impl TraceDataType for TraceDataText {
+    type Data<T: TraceData> = T::Text;
+}
+
+#[derive(Copy, Debug, Default, Eq, PartialEq, Hash)]
+#[repr(transparent)]
+pub struct TraceDataRef<T: TraceDataType> {
+    index: u32,
+    _phantom: PhantomData<T>,
+}
+
+impl<T: TraceDataType> TraceDataRef<T> {
+    fn new(r#ref: u32) -> Self {
+        Self {
+            index: r#ref,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+pub type TraceStringRef = TraceDataRef<TraceDataText>;
+pub type TraceBytesRef = TraceDataRef<TraceDataBytes>;
+
+struct StaticDataValue<T> {
+    value: T,
+    rc: u32,
+}
+
+pub struct StaticDataVec<T: TraceData, D: TraceDataType> {
+    vec: Vec<StaticDataValue<D::Data::<T>>>,
+    // This HashMap is probably the bottleneck. However we are required to ensure every string only exists once.
+    table: HashMap<D::Data::<T>, TraceDataRef<D>>,
+}
+
+impl<T: TraceData, D: TraceDataType> Default for StaticDataVec<T, D> {
+    fn default() -> Self {
+        Self {
+            vec: vec![StaticDataValue {
+                value: D::Data::<T>::default(),
+                rc: 1 << 30, // so that we can just have TraceDataRef::new(0) as default without the rc ever reaching 0
+            }],
+            table: HashMap::from([(D::Data::<T>::default(), TraceDataRef::new(0))]),
+        }
+    }
+}
+
+impl<T: TraceData, D: TraceDataType> StaticDataVec<T, D> {
+    pub fn get(&self, r#ref: TraceDataRef<D>) -> &D::Data::<T> {
+        &self.vec[r#ref.index as usize].value
+    }
+
+    pub fn add(&mut self, value: D::Data::<T>) -> TraceDataRef<D> {
+        if let Some(r#ref) = self.table.get(&value) {
+            self.vec[r#ref.index as usize].rc += 1;
+            return *r#ref;
+        }
+        let index = self.vec.len() as u32;
+        self.table.insert(value.clone(), TraceDataRef::new(index));
+        self.vec.push(StaticDataValue {
+            value,
+            rc: 1,
+        });
+        TraceDataRef::new(index)
+    }
+
+    pub fn update(&mut self, r#ref: &mut TraceDataRef<D>, value: D::Data::<T>) {
+        let entry = &mut self.vec[r#ref.index as usize];
+        if entry.rc == 1 {
+            self.table.remove(&entry.value);
+            self.table.insert(value.clone(), *r#ref);
+            entry.value = value;
+        } else {
+            entry.rc -= 1;
+            *r#ref = self.add(entry.value);
+        }
+    }
+}