refactor(trace-utils): replace use_v05_format bool and remove infallible expect (#1946)

# What does this PR do?        

Replaces the `use_v05_format`: bool parameter in `collect_trace_chunks` with the existing TraceEncoding enum, and removes `#[allow(clippy::expect_used)] + .expect("infallible...")` in the v04 msgpack encoder.

# Motivation                                              

Follow-up to [#1896](#1896), addressing review comments:                                                                                                                                   

- Passing a bare bool to `collect_trace_chunks` gave no context at the call site — you had to look up the function signature to understand what true/false meant. Using TraceEncoding makes the intent explicit.
- The `collect_and_process_traces` match had a catch-all arm that computed a bool from the format — replaced with exhaustive V04/V05 arms.
- The v04 encoder used `.expect("infallible...")` with a clippy allow to handle an error that can never occur. Replaced with `let _ = ...`, consistent with what was already done in the v1 encoder.

# Additional Notes

⚠️  Should be merged after [#1896](#1896).

# How to test the change?

Existing tests.

Co-authored-by: anais.raison <anais.raison@datadoghq.com>

Author: anais-raison

libdd-common/src/lib.rs

@@ -135,6 +135,40 @@ impl<T> RwLockExt<T> for RwLock<T> {
     }
 }
 
+/// Extension trait that extracts the value from a `Result` whose error type is uninhabited.
+///
+/// The signature constrains callers at compile time: the method is only available when the
+/// error type is [`core::convert::Infallible`]. No panics — the compiler proves the `Err`
+/// arm unreachable from the type.
+///
+/// # Examples
+///
+/// ```
+/// use libdd_common::ResultInfallibleExt;
+/// use std::convert::Infallible;
+///
+/// let result: Result<i32, Infallible> = Ok(42);
+/// assert_eq!(result.unwrap_infallible(), 42);
+/// ```
+pub trait ResultInfallibleExt<T>: sealed::Sealed {
+    fn unwrap_infallible(self) -> T;
+}
+
+impl<T> ResultInfallibleExt<T> for Result<T, core::convert::Infallible> {
+    #[inline(always)]
+    fn unwrap_infallible(self) -> T {
+        match self {
+            Ok(value) => value,
+            Err(never) => match never {},
+        }
+    }
+}
+
+mod sealed {
+    pub trait Sealed {}
+    impl<T> Sealed for Result<T, core::convert::Infallible> {}
+}
+
 pub mod header {
     #![allow(clippy::declare_interior_mutable_const)]
     use http::{header::HeaderName, HeaderValue};

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

@@ -17,7 +17,7 @@ use libdd_trace_utils::msgpack_encoder;
 use libdd_trace_utils::span::{v04::Span, TraceData};
 use libdd_trace_utils::trace_utils::{self, TracerHeaderTags};
 use libdd_trace_utils::tracer_metadata::TracerMetadata;
-use libdd_trace_utils::tracer_payload;
+use libdd_trace_utils::tracer_payload::{self, TraceEncoding};
 
 /// Minimal capacity of fresh buffers allocated to encode traces, in bytes.
 const MIN_BUFFER_CAPACITY: usize = 1024;
@@ -76,9 +76,13 @@ impl TraceSerializer {
     ) -> Result<tracer_payload::TraceChunks<T>, TraceExporterError> {
         match self.output_format {
             TraceExporterOutputFormat::V1 => Ok(tracer_payload::TraceChunks::V1(traces)),
-            format => {
-                let use_v05_format = matches!(format, TraceExporterOutputFormat::V05);
-                trace_utils::collect_trace_chunks(traces, use_v05_format).map_err(|e| {
+            TraceExporterOutputFormat::V04 => {
+                trace_utils::collect_trace_chunks(traces, TraceEncoding::V04).map_err(|e| {
+                    TraceExporterError::Deserialization(DecodeError::InvalidFormat(e.to_string()))
+                })
+            }
+            TraceExporterOutputFormat::V05 => {
+                trace_utils::collect_trace_chunks(traces, TraceEncoding::V05).map_err(|e| {
                     TraceExporterError::Deserialization(DecodeError::InvalidFormat(e.to_string()))
                 })
             }

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

@@ -4,6 +4,21 @@
 pub mod v04;
 pub mod v1;
 
+use rmp::encode::ValueWriteError;
+use std::convert::Infallible;
+
+/// Flatten `ValueWriteError<Infallible>` (uninhabited because both variants wrap
+/// `Infallible`) into the bare `Infallible` so callers can use
+/// [`libdd_common::ResultInfallibleExt`].
+#[inline(always)]
+pub(crate) fn flatten_value_write_infallible(err: ValueWriteError<Infallible>) -> Infallible {
+    match err {
+        ValueWriteError::InvalidMarkerWrite(never) | ValueWriteError::InvalidDataWrite(never) => {
+            never
+        }
+    }
+}
+
 /// A writer that counts bytes without storing them, used to compute encoded payload size.
 pub(crate) struct CountLength(u32);
 

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

@@ -3,6 +3,7 @@
 
 use crate::span::v04::Span;
 use crate::span::TraceData;
+use libdd_common::ResultInfallibleExt;
 use rmp::encode::{write_array_len, ByteBuf, RmpWrite, ValueWriteError};
 
 mod span;
@@ -123,8 +124,9 @@ pub fn to_vec_with_capacity<T: TraceData, S: AsRef<[Span<T>]>>(
     capacity: u32,
 ) -> Vec<u8> {
     let mut buf = ByteBuf::with_capacity(capacity as usize);
-    #[allow(clippy::expect_used)]
-    to_writer(&mut buf, traces).expect("infallible: the error is std::convert::Infallible");
+    to_writer(&mut buf, traces)
+        .map_err(super::flatten_value_write_infallible)
+        .unwrap_infallible();
     buf.into_vec()
 }
 
@@ -158,7 +160,11 @@ pub fn to_vec_with_capacity<T: TraceData, S: AsRef<[Span<T>]>>(
 /// ```
 pub fn to_encoded_byte_len<T: TraceData, S: AsRef<[Span<T>]>>(traces: &[S]) -> u32 {
     let mut counter = super::CountLength(0);
-    #[allow(clippy::expect_used)]
-    to_writer(&mut counter, traces).expect("infallible: CountLength never fails");
+    // `CountLength` impls `std::io::Write` (whose error type is `std::io::Error`, not
+    // `Infallible`), so we can't statically prove infallibility via `unwrap_infallible`
+    // the way we do for `ByteBuf`. In practice `CountLength::write*` only ever return
+    // `Ok`, so the error path here is unreachable today; should `CountLength` ever grow
+    // a fallible code path, fuzz tests on the msgpack encoded length would catch it.
+    let _ = to_writer(&mut counter, traces);
     counter.0
 }

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

@@ -6,6 +6,7 @@ mod span_v04;
 use crate::span::v04::Span;
 use crate::span::TraceData;
 use crate::tracer_metadata::TracerMetadata;
+use libdd_common::ResultInfallibleExt;
 use rmp::encode::{
     write_array_len, write_bin, write_map_len, write_sint, write_str, write_uint, write_uint8,
     ByteBuf, RmpWrite, ValueWriteError,
@@ -399,7 +400,9 @@ pub fn to_vec_with_capacity<T: TraceData, S: AsRef<[Span<T>]>>(
     metadata: &TracerMetadata,
 ) -> Vec<u8> {
     let mut buf = ByteBuf::with_capacity(capacity as usize);
-    let _ = encode_payload(&mut buf, traces, metadata); // infallible: ByteBuf write never fails
+    encode_payload(&mut buf, traces, metadata)
+        .map_err(super::flatten_value_write_infallible)
+        .unwrap_infallible();
     buf.into_vec()
 }
 
@@ -409,7 +412,12 @@ pub fn to_encoded_byte_len<T: TraceData, S: AsRef<[Span<T>]>>(
     metadata: &TracerMetadata,
 ) -> u32 {
     let mut counter = super::CountLength(0);
-    let _ = encode_payload(&mut counter, traces, metadata); // infallible: CountLength write never fails
+    // `CountLength` impls `std::io::Write` (whose error type is `std::io::Error`, not
+    // `Infallible`), so we can't statically prove infallibility via `unwrap_infallible`
+    // the way we do for `ByteBuf`. In practice `CountLength::write*` only ever return
+    // `Ok`, so the error path here is unreachable today; should `CountLength` ever grow
+    // a fallible code path, fuzz tests on the msgpack encoded length would catch it.
+    let _ = encode_payload(&mut counter, traces, metadata);
     counter.0
 }
 

libdd-trace-utils/src/trace_utils.rs

@@ -7,7 +7,7 @@ use crate::span::v05::dict::SharedDict;
 use crate::span::{v05, TraceData};
 pub use crate::tracer_header_tags::TracerHeaderTags;
 use crate::tracer_payload::TracerPayloadCollection;
-use crate::tracer_payload::{self, TraceChunks};
+use crate::tracer_payload::{self, TraceChunks, TraceEncoding};
 use anyhow::anyhow;
 use bytes::buf::Reader;
 use bytes::Buf;
@@ -592,26 +592,22 @@ pub fn enrich_span_with_azure_function_metadata(span: &mut pb::Span) {
 
 pub fn collect_trace_chunks<T: TraceData>(
     traces: Vec<Vec<crate::span::v04::Span<T>>>,
-    use_v05_format: bool,
+    format: TraceEncoding,
 ) -> 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 trace_len = trace.len();
-            let v05_trace = trace.into_iter().try_fold(
-                Vec::with_capacity(trace_len),
-                |mut acc, span| -> anyhow::Result<Vec<v05::Span>> {
-                    acc.push(v05::from_v04_span(span, &mut shared_dict)?);
-                    Ok(acc)
-                },
-            )?;
-
-            v05_traces.push(v05_trace);
+    match format {
+        TraceEncoding::V05 => {
+            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
+                    .into_iter()
+                    .map(|span| v05::from_v04_span(span, &mut shared_dict))
+                    .collect::<anyhow::Result<Vec<_>>>()?;
+                v05_traces.push(v05_trace);
+            }
+            Ok(TraceChunks::V05((shared_dict, v05_traces)))
         }
-        Ok(TraceChunks::V05((shared_dict, v05_traces)))
-    } else {
-        Ok(TraceChunks::V04(traces))
+        TraceEncoding::V04 => Ok(TraceChunks::V04(traces)),
     }
 }
 
@@ -1132,7 +1128,7 @@ mod tests {
     fn test_collect_trace_chunks_v05() {
         let chunk = vec![create_test_no_alloc_span(123, 456, 789, 1, true)];
 
-        let collection = collect_trace_chunks(vec![chunk], true).unwrap();
+        let collection = collect_trace_chunks(vec![chunk], TraceEncoding::V05).unwrap();
 
         let (dict, traces) = match collection {
             TraceChunks::V05(payload) => payload,
@@ -1203,6 +1199,27 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_collect_trace_chunks_v04() {
+        let chunk = vec![create_test_no_alloc_span(123, 456, 789, 1, true)];
+
+        let collection = collect_trace_chunks(vec![chunk], TraceEncoding::V04).unwrap();
+
+        let traces = match collection {
+            TraceChunks::V04(traces) => traces,
+            _ => panic!("Unexpected type"),
+        };
+
+        assert_eq!(traces.len(), 1);
+        assert_eq!(traces[0].len(), 1);
+        let span = &traces[0][0];
+        assert_eq!(span.trace_id, 123);
+        assert_eq!(span.span_id, 456);
+        assert_eq!(span.parent_id, 789);
+        assert_eq!(span.start, 1);
+        assert_eq!(span.error, 0);
+    }
+
     #[test]
     fn test_rmp_serde_deserialize_meta_with_null_values() {
         // Create a JSON representation with null value in meta

libdd-trace-utils/src/tracer_payload.rs

@@ -12,7 +12,7 @@ use std::iter::Iterator;
 pub type TracerPayloadV04 = Vec<v04::SpanBytes>;
 pub type TracerPayloadV05 = Vec<v05::Span>;
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 /// Enumerates the different encoding types.
 pub enum TraceEncoding {
     /// v0.4 encoding (TracerPayloadV04).
@@ -234,10 +234,7 @@ pub fn decode_to_trace_chunks(
     }
     .map_err(|e| anyhow::format_err!("Error deserializing trace from request body: {e}"))?;
 
-    Ok((
-        collect_trace_chunks(data, matches!(encoding_type, TraceEncoding::V05))?,
-        size,
-    ))
+    Ok((collect_trace_chunks(data, encoding_type)?, size))
 }
 
 #[cfg(test)]