Make TraceChunk generic on the span text for easier migration

Author: paullegranddc

data-pipeline/src/trace_exporter/mod.rs

@@ -624,10 +624,9 @@ impl TraceExporter {
                 unreachable!("Codepath invalid for proxy mode",)
             }
         };
-        let payload: tracer_payload::TraceChunks =
-            trace_utils::collect_trace_chunks(traces, use_v05_format).map_err(|e| {
-                TraceExporterError::Deserialization(DecodeError::InvalidFormat(e.to_string()))
-            })?;
+        let payload = trace_utils::collect_trace_chunks(traces, use_v05_format).map_err(|e| {
+            TraceExporterError::Deserialization(DecodeError::InvalidFormat(e.to_string()))
+        })?;
 
         let chunks = payload.size();
         let endpoint = Endpoint {

trace-utils/src/span/mod.rs

@@ -12,6 +12,9 @@ use std::fmt;
 use std::hash::Hash;
 use std::str::FromStr;
 use tinybytes::{Bytes, BytesString};
+use v05::dict::SharedDict;
+
+use crate::tracer_payload::TraceChunks;
 
 #[derive(Debug, PartialEq)]
 pub enum SpanKey {
@@ -60,9 +63,9 @@ impl FromStr for SpanKey {
 /// 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: Eq + Hash + Borrow<str> + Serialize {}
+pub trait SpanText: Eq + Hash + Borrow<str> + Serialize + Default + Clone {}
 /// Implement the SpanText trait for any type which satisfies the sub traits.
-impl<T: Eq + Hash + Borrow<str> + Serialize> SpanText for T {}
+impl<T: Eq + Hash + Borrow<str> + Serialize + Default + Clone> SpanText for T {}
 
 /// Checks if the `value` represents an empty string. Used to skip serializing empty strings
 /// with serde.
@@ -259,6 +262,10 @@ pub type SpanEventSlice<'a> = SpanEvent<&'a str>;
 pub type AttributeAnyValueSlice<'a> = AttributeAnyValue<&'a str>;
 pub type AttributeArrayValueSlice<'a> = AttributeArrayValue<&'a str>;
 
+pub type TraceChunksBytes = TraceChunks<BytesString>;
+
+pub type SharedDictBytes = SharedDict<BytesString>;
+
 impl SpanSlice<'_> {
     /// Converts a borrowed `SpanSlice` into an owned `SpanBytes`, by resolving all internal
     /// references into slices of the provided `Bytes` buffer. Returns `None` if any slice is

trace-utils/src/span/v05/dict.rs

@@ -2,58 +2,60 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use std::collections::HashMap;
-use tinybytes::{Bytes, BytesString};
+
+use crate::span::SpanText;
 
 /// This struct represents the shared dictionary used for interning all the strings belonging to a
 /// v05 trace chunk.
-pub struct SharedDict {
+pub struct SharedDict<T> {
     /// Map strings with their index (O(1) retrieval complexity).
-    string_map: HashMap<BytesString, usize>,
+    string_map: HashMap<T, usize>,
     /// Since the collection needs to be ordered an additional vector to keep the insertion order.
-    dict: Vec<BytesString>,
+    dict: Vec<T>,
 }
 
-impl SharedDict {
+impl<T: SpanText> SharedDict<T> {
     /// Gets the index of the interned string. If the string is not part of the dictionary it is
     /// added and its corresponding index returned.
     ///
     /// # Arguments:
     ///
     /// * `str`: string to look up in the dictionary.
-    pub fn get_or_insert(&mut self, str: &BytesString) -> Result<u32, std::num::TryFromIntError> {
-        if let Some(index) = self.string_map.get(str) {
+    pub fn get_or_insert(&mut self, s: &T) -> Result<u32, std::num::TryFromIntError> {
+        if let Some(index) = self.string_map.get(s.borrow()) {
             (*index).try_into()
         } else {
             let index = self.dict.len();
-            self.dict.push(str.clone());
-            self.string_map.insert(str.clone(), index);
+            self.dict.push(s.clone());
+            self.string_map.insert(s.clone(), index);
             index.try_into()
         }
     }
 
     /// Returns the dictionary. This method consumes the structure.
-    pub fn dict(mut self) -> Vec<BytesString> {
+    pub fn dict(mut self) -> Vec<T> {
         std::mem::take(&mut self.dict)
     }
 }
 
-impl Default for SharedDict {
+impl<T: SpanText> Default for SharedDict<T> {
     fn default() -> Self {
-        let empty_str = unsafe { BytesString::from_bytes_unchecked(Bytes::from_static(b"")) };
         Self {
-            string_map: HashMap::from([(empty_str.clone(), 0)]),
-            dict: vec![empty_str.clone()],
+            string_map: HashMap::from([(T::default(), 0)]),
+            dict: vec![T::default()],
         }
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use tinybytes::{Bytes, BytesString};
+
     use super::*;
 
     #[test]
     fn default_test() {
-        let dict = SharedDict::default();
+        let dict: SharedDict<BytesString> = SharedDict::default();
 
         assert_eq!(dict.string_map.len(), 1);
         assert_eq!(dict.dict.len(), 1);

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

@@ -3,8 +3,7 @@
 
 pub mod dict;
 
-use crate::span::v05::dict::SharedDict;
-use crate::span::SpanBytes;
+use crate::span::{v05::dict::SharedDict, SpanText};
 use anyhow::Result;
 use serde::Serialize;
 use std::collections::HashMap;
@@ -29,7 +28,10 @@ pub struct Span {
     pub r#type: u32,
 }
 
-pub fn from_span_bytes(span: &SpanBytes, dict: &mut SharedDict) -> Result<Span> {
+pub fn from_span<T: SpanText>(
+    span: &crate::span::Span<T>,
+    dict: &mut SharedDict<T>,
+) -> Result<Span> {
     Ok(Span {
         service: dict.get_or_insert(&span.service)?,
         name: dict.get_or_insert(&span.name)?,
@@ -61,6 +63,7 @@ pub fn from_span_bytes(span: &SpanBytes, dict: &mut SharedDict) -> Result<Span>
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::span::SpanBytes;
     use tinybytes::BytesString;
 
     #[test]
@@ -87,7 +90,7 @@ mod tests {
         };
 
         let mut dict = SharedDict::default();
-        let v05_span = from_span_bytes(&span, &mut dict).unwrap();
+        let v05_span = from_span(&span, &mut dict).unwrap();
 
         let dict = dict.dict();
 

trace-utils/src/test_utils/mod.rs

@@ -15,9 +15,8 @@ use std::collections::HashMap;
 use std::time::Duration;
 
 use crate::send_data::SendData;
-use crate::span::v05;
-use crate::span::v05::dict::SharedDict;
 use crate::span::SpanBytes;
+use crate::span::{v05, SharedDictBytes};
 use crate::trace_utils::TracerHeaderTags;
 use crate::tracer_payload::TracerPayloadCollection;
 use datadog_trace_protobuf::pb;
@@ -230,7 +229,7 @@ pub fn create_test_v05_span(
     parent_id: u64,
     start: i64,
     is_top_level: bool,
-    dict: &mut SharedDict,
+    dict: &mut SharedDictBytes,
     metrics: Option<Vec<(String, f64)>>,
 ) -> v05::Span {
     let mut meta = HashMap::from([

trace-utils/src/trace_utils.rs

@@ -3,8 +3,8 @@
 
 pub use crate::send_data::send_data_result::SendDataResult;
 pub use crate::send_data::SendData;
-use crate::span::v05;
 use crate::span::v05::dict::SharedDict;
+use crate::span::{v05, SpanText};
 pub use crate::tracer_header_tags::TracerHeaderTags;
 use crate::tracer_payload::TracerPayloadCollection;
 use crate::tracer_payload::{self, TraceChunks};
@@ -597,18 +597,18 @@ macro_rules! parse_root_span_tags {
     }
 }
 
-pub fn collect_trace_chunks(
-    traces: Vec<Vec<crate::span::SpanBytes>>,
+pub fn collect_trace_chunks<T: SpanText>(
+    traces: Vec<Vec<crate::span::Span<T>>>,
     use_v05_format: bool,
-) -> anyhow::Result<TraceChunks> {
+) -> anyhow::Result<TraceChunks<T>> {
     if use_v05_format {
         let mut shared_dict = SharedDict::default();
         let mut v05_traces: Vec<Vec<v05::Span>> = Vec::with_capacity(traces.len());
         for trace in traces {
             let v05_trace = trace.iter().try_fold(
                 Vec::with_capacity(trace.len()),
                 |mut acc, span| -> anyhow::Result<Vec<v05::Span>> {
-                    acc.push(v05::from_span_bytes(span, &mut shared_dict)?);
+                    acc.push(v05::from_span(span, &mut shared_dict)?);
                     Ok(acc)
                 },
             )?;

trace-utils/src/tracer_payload.rs

@@ -1,13 +1,13 @@
 // Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::span::{v05, SpanBytes};
+use crate::span::{v05, Span, SpanBytes, SpanText};
 use crate::trace_utils::collect_trace_chunks;
 use crate::{msgpack_decoder, trace_utils::cmp_send_data_payloads};
 use datadog_trace_protobuf::pb;
 use std::cmp::Ordering;
 use std::iter::Iterator;
-use tinybytes;
+use tinybytes::{self, BytesString};
 
 pub type TracerPayloadV04 = Vec<SpanBytes>;
 pub type TracerPayloadV05 = Vec<v05::Span>;
@@ -22,14 +22,14 @@ pub enum TraceEncoding {
 }
 
 #[derive(Debug, Clone)]
-pub enum TraceChunks {
+pub enum TraceChunks<T: SpanText> {
     /// Collection of TraceChunkSpan.
-    V04(Vec<Vec<SpanBytes>>),
+    V04(Vec<Vec<Span<T>>>),
     /// Collection of TraceChunkSpan with de-duplicated strings.
-    V05((Vec<tinybytes::BytesString>, Vec<Vec<v05::Span>>)),
+    V05((Vec<T>, Vec<Vec<v05::Span>>)),
 }
 
-impl TraceChunks {
+impl TraceChunks<BytesString> {
     pub fn into_tracer_payload_collection(self) -> TracerPayloadCollection {
         match self {
             TraceChunks::V04(traces) => TracerPayloadCollection::V04(traces),
@@ -38,7 +38,7 @@ impl TraceChunks {
     }
 }
 
-impl TraceChunks {
+impl<T: SpanText> TraceChunks<T> {
     /// Returns the number of traces in the chunk
     pub fn size(&self) -> usize {
         match self {
@@ -222,7 +222,7 @@ impl TraceChunkProcessor for DefaultTraceChunkProcessor {
 pub fn decode_to_trace_chunks(
     data: tinybytes::Bytes,
     encoding_type: TraceEncoding,
-) -> Result<(TraceChunks, usize), anyhow::Error> {
+) -> Result<(TraceChunks<BytesString>, usize), anyhow::Error> {
     let (data, size) = match encoding_type {
         TraceEncoding::V04 => msgpack_decoder::v04::from_bytes(data),
         TraceEncoding::V05 => msgpack_decoder::v05::from_bytes(data),