Add span decoder for bytes slices

Author: VianneyRuhlmann

tinybytes/src/bytes_string.rs

@@ -82,6 +82,21 @@ impl BytesString {
         }
     }
 
+    /// Creates a `BytesString` from a string slice within the given buffer.
+    ///
+    /// # Arguments
+    ///
+    /// * `bytes` - A `tinybytes::Bytes` instance that will be converted into a `BytesString`.
+    /// * `slice` - The string slice pointing into the given bytes that will form the `BytesString`.
+    pub fn try_from_bytes_slice(bytes: &Bytes, slice: &str) -> Option<Self> {
+        // SAFETY: This is safe as a str slice is definitely a valid UTF-8 slice.
+        unsafe {
+            Some(Self::from_bytes_unchecked(
+                bytes.slice_ref(slice.as_bytes())?,
+            ))
+        }
+    }
+
     /// Creates a `BytesString` from a `tinybytes::Bytes` instance without validating the bytes.
     ///
     /// This function does not perform any validation on the provided bytes, and assumes that the

trace-utils/src/msgpack_decoder/v04/decoder/mod.rs

@@ -6,10 +6,11 @@ mod span_link;
 
 use self::span::decode_span;
 use super::error::DecodeError;
-use super::number::read_number_bytes;
-use crate::span_v04::Span;
+use super::number::{read_nullable_number_ref, read_number_bytes, read_number_ref};
+use crate::span_v04::{Span, SpanSlice};
 use rmp::decode::DecodeStringError;
 use rmp::{decode, decode::RmpRead, Marker};
+use span::decode_span_ref;
 use std::{collections::HashMap, f64};
 use tinybytes::{Bytes, BytesString};
 
@@ -58,11 +59,26 @@ const NULL_MARKER: &u8 = &0xc0;
 /// let decoded_span = &decoded_traces[0][0];
 /// assert_eq!("test-span", decoded_span.name.as_str());
 /// ```
-pub fn from_slice(mut data: tinybytes::Bytes) -> Result<(Vec<Vec<Span>>, usize), DecodeError> {
-    let trace_count =
-        rmp::decode::read_array_len(unsafe { data.as_mut_slice() }).map_err(|_| {
-            DecodeError::InvalidFormat("Unable to read array len for trace count".to_owned())
-        })?;
+pub fn from_slice(data: tinybytes::Bytes) -> Result<(Vec<Vec<Span>>, usize), DecodeError> {
+    let mut parsed_data = data.clone();
+    let (traces_ref, size) = from_slice_ref(unsafe { parsed_data.as_mut_slice() })?;
+    let traces_owned = traces_ref
+        .iter()
+        .map(|trace| {
+            trace
+                .iter()
+                // Safe to unwrap since the spans uses subslices of the `data` slice
+                .map(|span| span.try_to_bytes(&data).unwrap())
+                .collect()
+        })
+        .collect();
+    Ok((traces_owned, size))
+}
+
+pub fn from_slice_ref(mut data: &[u8]) -> Result<(Vec<Vec<SpanSlice>>, usize), DecodeError> {
+    let trace_count = rmp::decode::read_array_len(&mut data).map_err(|_| {
+        DecodeError::InvalidFormat("Unable to read array len for trace count".to_owned())
+    })?;
 
     let start_len = data.len();
 
@@ -74,12 +90,9 @@ pub fn from_slice(mut data: tinybytes::Bytes) -> Result<(Vec<Vec<Span>>, usize),
                     .expect("Unable to cast trace_count to usize"),
             ),
             |mut traces, _| {
-                let span_count = rmp::decode::read_array_len(unsafe { data.as_mut_slice() })
-                    .map_err(|_| {
-                        DecodeError::InvalidFormat(
-                            "Unable to read array len for span count".to_owned(),
-                        )
-                    })?;
+                let span_count = rmp::decode::read_array_len(&mut data).map_err(|_| {
+                    DecodeError::InvalidFormat("Unable to read array len for span count".to_owned())
+                })?;
 
                 let trace = (0..span_count).try_fold(
                     Vec::with_capacity(
@@ -88,7 +101,7 @@ pub fn from_slice(mut data: tinybytes::Bytes) -> Result<(Vec<Vec<Span>>, usize),
                             .expect("Unable to cast span_count to usize"),
                     ),
                     |mut trace, _| {
-                        let span = decode_span(&mut data)?;
+                        let span = decode_span_ref(&mut data)?;
                         trace.push(span);
                         Ok(trace)
                     },
@@ -143,6 +156,15 @@ fn read_nullable_string_bytes(buf: &mut Bytes) -> Result<BytesString, DecodeErro
     }
 }
 
+#[inline]
+fn read_nullable_string_ref<'a>(buf: &mut &'a [u8]) -> Result<&'a str, DecodeError> {
+    if is_null_marker(buf) {
+        Ok("")
+    } else {
+        read_string_ref(buf)
+    }
+}
+
 #[inline]
 // Safety: read_string_ref checks utf8 validity, so we don't do it again when creating the
 // BytesStrings.
@@ -172,6 +194,35 @@ fn read_nullable_str_map_to_bytes_strings(
     read_str_map_to_bytes_strings(buf)
 }
 
+#[inline]
+// Safety: read_string_ref checks utf8 validity, so we don't do it again when creating the
+// BytesStrings.
+fn read_str_map_to_str_ref<'a>(
+    buf: &mut &'a [u8],
+) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
+    let len = decode::read_map_len(buf)
+        .map_err(|_| DecodeError::InvalidFormat("Unable to get map len for str map".to_owned()))?;
+
+    let mut map = HashMap::with_capacity(len.try_into().expect("Unable to cast map len to usize"));
+    for _ in 0..len {
+        let key = read_string_ref(buf)?;
+        let value = read_string_ref(buf)?;
+        map.insert(key, value);
+    }
+    Ok(map)
+}
+
+#[inline]
+fn read_nullable_str_map_to_str_ref<'a>(
+    buf: &mut &'a [u8],
+) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
+    if is_null_marker(buf) {
+        return Ok(HashMap::default());
+    }
+
+    read_str_map_to_str_ref(buf)
+}
+
 #[inline]
 fn read_metric_pair(buf: &mut Bytes) -> Result<(BytesString, f64), DecodeError> {
     let key = read_string_bytes(buf)?;
@@ -215,6 +266,52 @@ fn read_meta_struct(buf: &mut Bytes) -> Result<HashMap<BytesString, Vec<u8>>, De
     read_map(len, buf, read_meta_struct_pair)
 }
 
+#[inline]
+fn read_metric_pair_ref<'a>(buf: &mut &'a [u8]) -> Result<(&'a str, f64), DecodeError> {
+    let key = read_string_ref(buf)?;
+    let v = read_number_ref(buf)?;
+
+    Ok((key, v))
+}
+
+#[inline]
+fn read_metrics_ref<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, f64>, DecodeError> {
+    if is_null_marker(buf) {
+        return Ok(HashMap::default());
+    }
+
+    let len = read_map_len(buf)?;
+
+    read_map(len, buf, read_metric_pair_ref)
+}
+
+#[inline]
+fn read_meta_struct_ref<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, Vec<u8>>, DecodeError> {
+    if is_null_marker(buf) {
+        return Ok(HashMap::default());
+    }
+
+    fn read_meta_struct_pair_ref<'a>(
+        buf: &mut &'a [u8],
+    ) -> Result<(&'a str, Vec<u8>), DecodeError> {
+        let key = read_string_ref(buf)?;
+        let array_len = decode::read_array_len(buf).map_err(|_| {
+            DecodeError::InvalidFormat("Unable to read array len for meta_struct".to_owned())
+        })?;
+
+        let mut v = Vec::with_capacity(array_len as usize);
+
+        for _ in 0..array_len {
+            let value = read_number_ref(buf)?;
+            v.push(value);
+        }
+        Ok((key, v))
+    }
+
+    let len = read_map_len(buf)?;
+    read_map(len, buf, read_meta_struct_pair_ref)
+}
+
 /// Reads a map from the buffer and returns it as a `HashMap`.
 ///
 /// This function is generic over the key and value types of the map, and it uses a provided
@@ -243,14 +340,10 @@ fn read_meta_struct(buf: &mut Bytes) -> Result<HashMap<BytesString, Vec<u8>>, De
 /// * `V` - The type of the values in the map.
 /// * `F` - The type of the function used to read key-value pairs from the buffer.
 #[inline]
-fn read_map<K, V, F>(
-    len: usize,
-    buf: &mut Bytes,
-    read_pair: F,
-) -> Result<HashMap<K, V>, DecodeError>
+fn read_map<K, B, V, F>(len: usize, buf: &mut B, read_pair: F) -> Result<HashMap<K, V>, DecodeError>
 where
     K: std::hash::Hash + Eq,
-    F: Fn(&mut Bytes) -> Result<(K, V), DecodeError>,
+    F: Fn(&mut B) -> Result<(K, V), DecodeError>,
 {
     let mut map = HashMap::with_capacity(len);
     for _ in 0..len {
@@ -297,6 +390,18 @@ where
     }
 }
 
+/// When you want to "peek" if the next value is a null marker, and only advance the buffer if it is
+/// null and return the default value. If it is not null, you can continue to decode as expected.
+#[inline]
+fn is_null_marker(buf: &mut &[u8]) -> bool {
+    if buf.first() == Some(NULL_MARKER) {
+        *buf = &buf[1..];
+        true
+    } else {
+        false
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

trace-utils/src/msgpack_decoder/v04/decoder/span.rs

@@ -2,12 +2,15 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::{
-    read_meta_struct, read_metrics, read_nullable_str_map_to_bytes_strings,
-    read_nullable_string_bytes, read_string_ref, span_link::read_span_links,
+    read_meta_struct, read_meta_struct_ref, read_metric_pair_ref, read_metrics, read_metrics_ref,
+    read_nullable_number_ref, read_nullable_str_map_to_bytes_strings,
+    read_nullable_str_map_to_str_ref, read_nullable_string_bytes, read_nullable_string_ref,
+    read_number_ref, read_str_map_to_str_ref, read_string_ref,
+    span_link::{read_span_links, read_span_links_ref},
 };
 use crate::msgpack_decoder::v04::error::DecodeError;
 use crate::msgpack_decoder::v04::number::read_nullable_number_bytes;
-use crate::span_v04::{Span, SpanKey};
+use crate::span_v04::{Span, SpanKey, SpanSlice};
 use tinybytes::Bytes;
 
 /// Decodes a slice of bytes into a `Span` object.
@@ -27,14 +30,24 @@ use tinybytes::Bytes;
 /// - The map length cannot be read.
 /// - Any key or value cannot be decoded.
 pub fn decode_span(buffer: &mut Bytes) -> Result<Span, DecodeError> {
-    let mut span = Span::default();
+    let span_ref = decode_span_ref(unsafe { buffer.as_mut_slice() });
+    println!("Span ref {:?}", span_ref);
 
-    let span_size = rmp::decode::read_map_len(unsafe { buffer.as_mut_slice() }).map_err(|_| {
+    span_ref
+        .unwrap()
+        .try_to_bytes(buffer)
+        .ok_or(DecodeError::IOError)
+}
+
+pub fn decode_span_ref<'a>(buffer: &mut &'a [u8]) -> Result<SpanSlice<'a>, DecodeError> {
+    let mut span = SpanSlice::default();
+
+    let span_size = rmp::decode::read_map_len(buffer).map_err(|_| {
         DecodeError::InvalidFormat("Unable to get map len for span size".to_owned())
     })?;
 
     for _ in 0..span_size {
-        fill_span(&mut span, buffer)?;
+        fill_span_ref(&mut span, buffer)?;
     }
 
     Ok(span)
@@ -66,6 +79,32 @@ fn fill_span(span: &mut Span, buf: &mut Bytes) -> Result<(), DecodeError> {
     Ok(())
 }
 
+// Safety: read_string_ref checks utf8 validity, so we don't do it again when creating the
+// BytesStrings
+fn fill_span_ref<'a>(span: &mut SpanSlice<'a>, buf: &mut &'a [u8]) -> Result<(), DecodeError> {
+    let key = read_string_ref(buf)?
+        .parse::<SpanKey>()
+        .map_err(|_| DecodeError::InvalidFormat("Invalid span key".to_owned()))?;
+
+    match key {
+        SpanKey::Service => span.service = read_nullable_string_ref(buf)?,
+        SpanKey::Name => span.name = read_nullable_string_ref(buf)?,
+        SpanKey::Resource => span.resource = read_nullable_string_ref(buf)?,
+        SpanKey::TraceId => span.trace_id = read_nullable_number_ref(buf)?,
+        SpanKey::SpanId => span.span_id = read_nullable_number_ref(buf)?,
+        SpanKey::ParentId => span.parent_id = read_nullable_number_ref(buf)?,
+        SpanKey::Start => span.start = read_nullable_number_ref(buf)?,
+        SpanKey::Duration => span.duration = read_nullable_number_ref(buf)?,
+        SpanKey::Error => span.error = read_nullable_number_ref(buf)?,
+        SpanKey::Type => span.r#type = read_nullable_string_ref(buf)?,
+        SpanKey::Meta => span.meta = read_nullable_str_map_to_str_ref(buf)?,
+        SpanKey::Metrics => span.metrics = read_metrics_ref(buf)?,
+        SpanKey::MetaStruct => span.meta_struct = read_meta_struct_ref(buf)?,
+        SpanKey::SpanLinks => span.span_links = read_span_links_ref(buf)?,
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::SpanKey;

trace-utils/src/msgpack_decoder/v04/decoder/span_link.rs

@@ -5,12 +5,14 @@ use crate::msgpack_decoder::v04::decoder::{
     handle_null_marker, read_str_map_to_bytes_strings, read_string_bytes, read_string_ref,
 };
 use crate::msgpack_decoder::v04::error::DecodeError;
-use crate::msgpack_decoder::v04::number::read_number_bytes;
-use crate::span_v04::SpanLink;
+use crate::msgpack_decoder::v04::number::{read_number_bytes, read_number_ref};
+use crate::span_v04::{SpanLink, SpanLinkSlice};
 use rmp::Marker;
 use std::str::FromStr;
 use tinybytes::Bytes;
 
+use super::{is_null_marker, read_str_map_to_str_ref};
+
 /// Reads a slice of bytes and decodes it into a vector of `SpanLink` objects.
 ///
 /// # Arguments
@@ -48,6 +50,30 @@ pub(crate) fn read_span_links(buf: &mut Bytes) -> Result<Vec<SpanLink>, DecodeEr
         )),
     }
 }
+
+pub(crate) fn read_span_links_ref<'a>(
+    buf: &mut &'a [u8],
+) -> Result<Vec<SpanLinkSlice<'a>>, DecodeError> {
+    if is_null_marker(buf) {
+        return Ok(Vec::default());
+    }
+
+    match rmp::decode::read_marker(buf).map_err(|_| {
+        DecodeError::InvalidFormat("Unable to read marker for span links".to_owned())
+    })? {
+        Marker::FixArray(len) => {
+            let mut vec: Vec<SpanLinkSlice<'a>> = Vec::with_capacity(len.into());
+            for _ in 0..len {
+                vec.push(decode_span_link_ref(buf)?);
+            }
+            Ok(vec)
+        }
+        _ => Err(DecodeError::InvalidType(
+            "Unable to read span link from buffer".to_owned(),
+        )),
+    }
+}
+
 #[derive(Debug, PartialEq)]
 enum SpanLinkKey {
     TraceId,
@@ -95,6 +121,25 @@ fn decode_span_link(buf: &mut Bytes) -> Result<SpanLink, DecodeError> {
     Ok(span)
 }
 
+fn decode_span_link_ref<'a>(buf: &mut &'a [u8]) -> Result<SpanLinkSlice<'a>, DecodeError> {
+    let mut span = SpanLinkSlice::default();
+    let span_size = rmp::decode::read_map_len(buf)
+        .map_err(|_| DecodeError::InvalidType("Unable to get map len for span size".to_owned()))?;
+
+    for _ in 0..span_size {
+        match read_string_ref(buf)?.parse::<SpanLinkKey>()? {
+            SpanLinkKey::TraceId => span.trace_id = read_number_ref(buf)?,
+            SpanLinkKey::TraceIdHigh => span.trace_id_high = read_number_ref(buf)?,
+            SpanLinkKey::SpanId => span.span_id = read_number_ref(buf)?,
+            SpanLinkKey::Attributes => span.attributes = read_str_map_to_str_ref(buf)?,
+            SpanLinkKey::Tracestate => span.tracestate = read_string_ref(buf)?,
+            SpanLinkKey::Flags => span.flags = read_number_ref(buf)?,
+        }
+    }
+
+    Ok(span)
+}
+
 #[cfg(test)]
 mod tests {
     use super::SpanLinkKey;