Skip to main content

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