fix: dedup before msgpack span encoding

Author: yannham

libdd-data-pipeline/src/trace_buffer/mod.rs

@@ -22,7 +22,7 @@ use crate::trace_exporter::{
 
 /// Trait for types stored in a [`TraceBuffer`] that can report their approximate byte size.
 pub trait BufferSize {
-    fn byte_size(&self) -> usize;
+    fn byte_size(&mut self) -> usize;
 }
 
 impl<T> BufferSize for libdd_trace_utils::span::v04::Span<T>
@@ -31,9 +31,11 @@ where
     T::Text: AsRef<str>,
     T::Bytes: AsRef<[u8]>,
 {
-    fn byte_size(&self) -> usize {
+    fn byte_size(&mut self) -> usize {
         use libdd_trace_utils::span::v04::AttributeAnyValue;
 
+        self.dedup();
+
         // trace_id(16) + span_id(8) + parent_id(8) + start(8) + duration(8) + error(4)
         let mut size: usize = 52;
 
@@ -235,7 +237,7 @@ impl<T> Batch<T> {
     /// batch. So the batch can be over the maximum size after this call.
     /// This is because we don't want to always drop traces that contain more bytes than the maximum
     /// size.
-    fn add_trace_chunk(&mut self, chunk: Vec<T>) -> Result<(), BatchFullError>
+    fn add_trace_chunk(&mut self, mut chunk: Vec<T>) -> Result<(), BatchFullError>
     where
         T: BufferSize,
     {
@@ -248,7 +250,7 @@ impl<T> Batch<T> {
             return Ok(());
         }
 
-        self.byte_count += chunk.iter().map(|s| s.byte_size()).sum::<usize>();
+        self.byte_count += chunk.iter_mut().map(|s| s.byte_size()).sum::<usize>();
         self.chunks.push(chunk);
         Ok(())
     }
@@ -788,7 +790,7 @@ mod tests {
 
     // Used for tests, 1 byte per item so size computations are easier
     impl BufferSize for () {
-        fn byte_size(&self) -> usize {
+        fn byte_size(&mut self) -> usize {
             1
         }
     }

libdd-data-pipeline/src/trace_exporter/mod.rs

@@ -584,6 +584,12 @@ impl<C: HttpClientCapability + SleepCapability + MaybeSend + Sync + 'static> Tra
             self.client_computed_top_level,
         );
 
+        for chunk in &mut traces {
+            for span in chunk.iter_mut() {
+                span.dedup();
+            }
+        }
+
         // OTLP path: send sampled traces via OTLP when an OTLP endpoint is configured.
         // Unlike the agent path, there is no downstream agent to drop unsampled traces,
         // so drop_chunks is always called here regardless of whether stats are enabled.

libdd-trace-utils/src/msgpack_encoder/v04/span.rs

@@ -288,17 +288,19 @@ pub fn encode_span<W: RmpWrite, T: TraceData>(
 
     if !span.meta.is_empty() {
         write_const_msg_pack_str!(writer, "meta")?;
-        rmp::encode::write_map_len(writer, span.meta.len() as u32)?;
-        for (k, v) in span.meta.iter() {
+        let meta_dd = span.meta.defensive_dedup();
+        rmp::encode::write_map_len(writer, meta_dd.len() as u32)?;
+        for (k, v) in meta_dd.iter() {
             write_str(writer, k.borrow())?;
             write_str(writer, v.borrow())?;
         }
     }
 
     if !span.metrics.is_empty() {
         write_const_msg_pack_str!(writer, "metrics")?;
-        rmp::encode::write_map_len(writer, span.metrics.len() as u32)?;
-        for (k, v) in span.metrics.iter() {
+        let metrics_dd = span.metrics.defensive_dedup();
+        rmp::encode::write_map_len(writer, metrics_dd.len() as u32)?;
+        for (k, v) in metrics_dd.iter() {
             write_str(writer, k.borrow())?;
             write_f64(writer, *v)?;
         }
@@ -311,8 +313,9 @@ pub fn encode_span<W: RmpWrite, T: TraceData>(
 
     if !span.meta_struct.is_empty() {
         write_const_msg_pack_str!(writer, "meta_struct")?;
-        rmp::encode::write_map_len(writer, span.meta_struct.len() as u32)?;
-        for (k, v) in span.meta_struct.iter() {
+        let meta_struct_dd = span.meta_struct.defensive_dedup();
+        rmp::encode::write_map_len(writer, meta_struct_dd.len() as u32)?;
+        for (k, v) in meta_struct_dd.iter() {
             write_str(writer, k.borrow())?;
             write_bin(writer, v.borrow())?;
         }

libdd-trace-utils/src/msgpack_encoder/v1/span_v04.rs

@@ -255,8 +255,14 @@ pub fn encode_span<W: RmpWrite, T: TraceData>(
             "env" | "version" | "component" | "span.kind" | "_dd.p.tid"
         )
     };
-    let non_promoted_meta = span.meta.iter().filter(|(k, _)| !is_promoted(k)).count() as u32;
-    let attr_count = non_promoted_meta + span.metrics.len() as u32 + span.meta_struct.len() as u32;
+    let meta_dd = span.meta.defensive_dedup();
+    let metrics_dd = span.metrics.defensive_dedup();
+    let meta_struct_dd = span.meta_struct.defensive_dedup();
+
+    let non_promoted_meta = meta_dd.iter().filter(|(k, _)| !is_promoted(k)).count() as u32;
+    let metrics_len = metrics_dd.len() as u32;
+    let meta_struct_len = meta_struct_dd.len() as u32;
+    let attr_count = non_promoted_meta + metrics_len + meta_struct_len;
     let has_attributes = attr_count > 0;
 
     let env = span.meta.get("env").map(|v| v.borrow());
@@ -344,25 +350,25 @@ pub fn encode_span<W: RmpWrite, T: TraceData>(
         write_uint8(writer, SpanKey::Attributes as u8)?;
         rmp::encode::write_array_len(writer, attr_count * 3)?;
 
-        for (k, v) in span.meta.iter() {
+        for (k, v) in meta_dd.iter() {
             if is_promoted(k) {
                 continue;
             }
-            table.write_interned(writer, k.borrow())?;
+            table.write_interned(writer, (*k).borrow())?;
             write_uint8(writer, AnyValueKey::String as u8)?;
-            table.write_interned(writer, v.borrow())?;
+            table.write_interned(writer, (*v).borrow())?;
         }
 
-        for (k, v) in span.metrics.iter() {
-            table.write_interned(writer, k.borrow())?;
+        for (k, v) in metrics_dd.iter() {
+            table.write_interned(writer, (*k).borrow())?;
             write_uint8(writer, AnyValueKey::Double as u8)?;
             write_f64(writer, *v)?;
         }
 
-        for (k, v) in span.meta_struct.iter() {
-            table.write_interned(writer, k.borrow())?;
+        for (k, v) in meta_struct_dd.iter() {
+            table.write_interned(writer, (*k).borrow())?;
             write_uint8(writer, AnyValueKey::Bytes as u8)?;
-            write_bin(writer, v.borrow())?;
+            write_bin(writer, (*v).borrow())?;
         }
     }
 

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

@@ -21,7 +21,9 @@ use std::{fmt, ptr};
 /// 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 + From<String> {
+pub trait SpanText:
+    Debug + Eq + Hash + Clone + Borrow<str> + Serialize + Default + From<String>
+{
     fn from_static_str(value: &'static str) -> Self;
 }
 

libdd-trace-utils/src/span/v04/mod.rs

@@ -310,6 +310,14 @@ impl<T: TraceData> From<&AttributeArrayValue<T>> for u8 {
     }
 }
 
+impl<T: TraceData> Span<T> {
+    pub fn dedup(&mut self) {
+        self.meta.dedup();
+        self.metrics.dedup();
+        self.meta_struct.dedup();
+    }
+}
+
 fn is_default<T: Default + PartialEq>(t: &T) -> bool {
     t == &T::default()
 }

libdd-trace-utils/src/span/vec_map.rs

@@ -194,13 +194,6 @@ impl<K: Eq + Hash, V> VecMap<K, V> {
         if self.deduped {
             DedupedVecMap::Borrowed(self)
         } else {
-            static WARNED: AtomicBool = AtomicBool::new(false);
-            if !WARNED.swap(true, Ordering::Relaxed) {
-                tracing::warn!(
-                    "VecMap not deduped before encoding. Performing defensive on-the-fly dedup"
-                );
-            }
-
             DedupedVecMap::Owned(
                 self.data
                     .iter()
@@ -215,7 +208,6 @@ impl<K: Eq + Hash, V> VecMap<K, V> {
     /// This is a convenience wrapper around [Self::as_deduped_map] used in the msgpack encoder,
     /// where we expect the map to be deduped, but call `as_deduped_map` as a defensive measure. If
     /// the latter had to deduplicate and allocate a new vec, we log a warning (at most once).
-    #[allow(unused)]
     pub(crate) fn defensive_dedup(&self) -> DedupedVecMap<'_, K, V> {
         if !self.is_deduped() {
             static WARNED: AtomicBool = AtomicBool::new(false);