libdd-trace-utils 6.0.0

// Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/
// SPDX-License-Identifier: Apache-2.0

use crate::span::v04::Span;
use crate::span::TraceData;
use libdd_common::ResultInfallibleExt;
use rmp::encode::{write_array_len, ByteBuf, RmpWrite, ValueWriteError};

mod span;

#[inline(always)]
fn to_writer<W: RmpWrite, T: TraceData, S: AsRef<[Span<T>]>>(
    writer: &mut W,
    traces: &[S],
) -> Result<(), ValueWriteError<W::Error>> {
    write_array_len(writer, traces.len() as u32)?;
    for trace in traces {
        write_array_len(writer, trace.as_ref().len() as u32)?;
        for span in trace.as_ref() {
            span::encode_span(writer, span)?;
        }
    }

    Ok(())
}

/// Encodes a collection of traces into a slice of bytes.
///
/// # Arguments
///
/// * `slice` - A mutable reference to a byte slice.
/// * `traces` - A reference to a slice of spans.
///
/// # Returns
///
/// * `Ok(())` - If encoding succeeds.
/// * `Err(ValueWriteError)` - If encoding fails.
///
/// # Errors
///
/// This function will return an error if:
/// - The array length for trace count or span count cannot be written.
/// - Any span cannot be encoded.
///
/// # Examples
///
/// ```
/// use libdd_trace_utils::msgpack_encoder::v04::write_to_slice;
/// use libdd_trace_utils::span::v04::SpanSlice;
///
/// let mut buffer = vec![0u8; 1024];
/// let span = SpanSlice {
///     name: "test-span".into(),
///     ..Default::default()
/// };
/// let traces = vec![vec![span]];
///
/// write_to_slice(&mut &mut buffer[..], &traces).expect("Encoding failed");
/// ```
pub fn write_to_slice<T: TraceData, S: AsRef<[Span<T>]>>(
    slice: &mut &mut [u8],
    traces: &[S],
) -> Result<(), ValueWriteError> {
    to_writer(slice, traces)
}

/// Serializes traces into a vector of bytes with a default capacity of 0.
///
/// # Arguments
///
/// * `traces` - A reference to a slice of spans.
///
/// # Returns
///
/// * `Vec<u8>` - A vector containing encoded traces.
///
/// # Examples
///
/// ```
/// use libdd_trace_utils::msgpack_encoder::v04::to_vec;
/// use libdd_trace_utils::span::v04::SpanSlice;
///
/// let span = SpanSlice {
///     name: "test-span".into(),
///     ..Default::default()
/// };
/// let traces = vec![vec![span]];
/// let encoded = to_vec(&traces);
///
/// assert!(!encoded.is_empty());
/// ```
pub fn to_vec<T: TraceData, S: AsRef<[Span<T>]>>(traces: &[S]) -> Vec<u8> {
    to_vec_with_capacity(traces, 0)
}

/// Serializes traces into a vector of bytes with specified capacity.
///
/// # Arguments
///
/// * `traces` - A reference to a slice of spans.
/// * `capacity` - Desired initial capacity of the resulting vector.
///
/// # Returns
///
/// * `Vec<u8>` - A vector containing encoded traces.
///
/// # Examples
///
/// ```
/// use libdd_trace_utils::msgpack_encoder::v04::to_vec_with_capacity;
/// use libdd_trace_utils::span::v04::SpanSlice;
///
/// let span = SpanSlice {
///     name: "test-span".into(),
///     ..Default::default()
/// };
/// let traces = vec![vec![span]];
/// let encoded = to_vec_with_capacity(&traces, 1024);
///
/// assert!(encoded.capacity() >= 1024);
/// ```
pub fn to_vec_with_capacity<T: TraceData, S: AsRef<[Span<T>]>>(
    traces: &[S],
    capacity: u32,
) -> Vec<u8> {
    let mut buf = ByteBuf::with_capacity(capacity as usize);
    to_writer(&mut buf, traces)
        .map_err(super::flatten_value_write_infallible)
        .unwrap_infallible();
    buf.into_vec()
}

/// Computes the number of bytes required to encode the given traces.
///
/// This does not allocate any actual buffer, but simulates writing in order to measure
/// the encoded size of the traces.
///
/// # Arguments
///
/// * `traces` - A reference to a slice of spans.
///
/// # Returns
///
/// * `u32` - The number of bytes that would be written by the encoder.
///
/// # Examples
///
/// ```
/// use libdd_trace_utils::msgpack_encoder::v04::to_encoded_byte_len;
/// use libdd_trace_utils::span::v04::SpanSlice;
///
/// let span = SpanSlice {
///     name: "test-span".into(),
///     ..Default::default()
/// };
/// let traces = vec![vec![span]];
/// let encoded_len = to_encoded_byte_len(&traces);
///
/// assert!(encoded_len > 0);
/// ```
pub fn to_encoded_byte_len<T: TraceData, S: AsRef<[Span<T>]>>(traces: &[S]) -> u32 {
    let mut counter = super::CountLength(0);
    // `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
}