libdd_trace_utils/msgpack_encoder/v04/
mod.rs

1// Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::span::v04::Span;
5use crate::span::TraceData;
6use libdd_common::ResultInfallibleExt;
7use rmp::encode::{write_array_len, ByteBuf, RmpWrite, ValueWriteError};
8
9mod span;
10
11#[inline(always)]
12fn to_writer<W: RmpWrite, T: TraceData, S: AsRef<[Span<T>]>>(
13    writer: &mut W,
14    traces: &[S],
15) -> Result<(), ValueWriteError<W::Error>> {
16    write_array_len(writer, traces.len() as u32)?;
17    for trace in traces {
18        write_array_len(writer, trace.as_ref().len() as u32)?;
19        for span in trace.as_ref() {
20            span::encode_span(writer, span)?;
21        }
22    }
23
24    Ok(())
25}
26
27/// Encodes a collection of traces into a slice of bytes.
28///
29/// # Arguments
30///
31/// * `slice` - A mutable reference to a byte slice.
32/// * `traces` - A reference to a slice of spans.
33///
34/// # Returns
35///
36/// * `Ok(())` - If encoding succeeds.
37/// * `Err(ValueWriteError)` - If encoding fails.
38///
39/// # Errors
40///
41/// This function will return an error if:
42/// - The array length for trace count or span count cannot be written.
43/// - Any span cannot be encoded.
44///
45/// # Examples
46///
47/// ```
48/// use libdd_trace_utils::msgpack_encoder::v04::write_to_slice;
49/// use libdd_trace_utils::span::v04::SpanSlice;
50///
51/// let mut buffer = vec![0u8; 1024];
52/// let span = SpanSlice {
53///     name: "test-span".into(),
54///     ..Default::default()
55/// };
56/// let traces = vec![vec![span]];
57///
58/// write_to_slice(&mut &mut buffer[..], &traces).expect("Encoding failed");
59/// ```
60pub fn write_to_slice<T: TraceData, S: AsRef<[Span<T>]>>(
61    slice: &mut &mut [u8],
62    traces: &[S],
63) -> Result<(), ValueWriteError> {
64    to_writer(slice, traces)
65}
66
67/// Serializes traces into a vector of bytes with a default capacity of 0.
68///
69/// # Arguments
70///
71/// * `traces` - A reference to a slice of spans.
72///
73/// # Returns
74///
75/// * `Vec<u8>` - A vector containing encoded traces.
76///
77/// # Examples
78///
79/// ```
80/// use libdd_trace_utils::msgpack_encoder::v04::to_vec;
81/// use libdd_trace_utils::span::v04::SpanSlice;
82///
83/// let span = SpanSlice {
84///     name: "test-span".into(),
85///     ..Default::default()
86/// };
87/// let traces = vec![vec![span]];
88/// let encoded = to_vec(&traces);
89///
90/// assert!(!encoded.is_empty());
91/// ```
92pub fn to_vec<T: TraceData, S: AsRef<[Span<T>]>>(traces: &[S]) -> Vec<u8> {
93    to_vec_with_capacity(traces, 0)
94}
95
96/// Serializes traces into a vector of bytes with specified capacity.
97///
98/// # Arguments
99///
100/// * `traces` - A reference to a slice of spans.
101/// * `capacity` - Desired initial capacity of the resulting vector.
102///
103/// # Returns
104///
105/// * `Vec<u8>` - A vector containing encoded traces.
106///
107/// # Examples
108///
109/// ```
110/// use libdd_trace_utils::msgpack_encoder::v04::to_vec_with_capacity;
111/// use libdd_trace_utils::span::v04::SpanSlice;
112///
113/// let span = SpanSlice {
114///     name: "test-span".into(),
115///     ..Default::default()
116/// };
117/// let traces = vec![vec![span]];
118/// let encoded = to_vec_with_capacity(&traces, 1024);
119///
120/// assert!(encoded.capacity() >= 1024);
121/// ```
122pub fn to_vec_with_capacity<T: TraceData, S: AsRef<[Span<T>]>>(
123    traces: &[S],
124    capacity: u32,
125) -> Vec<u8> {
126    let mut buf = ByteBuf::with_capacity(capacity as usize);
127    to_writer(&mut buf, traces)
128        .map_err(super::flatten_value_write_infallible)
129        .unwrap_infallible();
130    buf.into_vec()
131}
132
133/// Computes the number of bytes required to encode the given traces.
134///
135/// This does not allocate any actual buffer, but simulates writing in order to measure
136/// the encoded size of the traces.
137///
138/// # Arguments
139///
140/// * `traces` - A reference to a slice of spans.
141///
142/// # Returns
143///
144/// * `u32` - The number of bytes that would be written by the encoder.
145///
146/// # Examples
147///
148/// ```
149/// use libdd_trace_utils::msgpack_encoder::v04::to_encoded_byte_len;
150/// use libdd_trace_utils::span::v04::SpanSlice;
151///
152/// let span = SpanSlice {
153///     name: "test-span".into(),
154///     ..Default::default()
155/// };
156/// let traces = vec![vec![span]];
157/// let encoded_len = to_encoded_byte_len(&traces);
158///
159/// assert!(encoded_len > 0);
160/// ```
161pub fn to_encoded_byte_len<T: TraceData, S: AsRef<[Span<T>]>>(traces: &[S]) -> u32 {
162    let mut counter = super::CountLength(0);
163    // `CountLength` impls `std::io::Write` (whose error type is `std::io::Error`, not
164    // `Infallible`), so we can't statically prove infallibility via `unwrap_infallible`
165    // the way we do for `ByteBuf`. In practice `CountLength::write*` only ever return
166    // `Ok`, so the error path here is unreachable today; should `CountLength` ever grow
167    // a fallible code path, fuzz tests on the msgpack encoded length would catch it.
168    let _ = to_writer(&mut counter, traces);
169    counter.0
170}
libdd_trace_utils/msgpack_encoder/v04/mod.rs

libdd_trace_utils/msgpack_encoder/v04/
mod.rs
