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",
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",
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",
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
131struct CountLength(u32);
132
133impl std::io::Write for CountLength {
134    #[inline]
135    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
136        self.write_all(buf)?;
137        Ok(buf.len())
138    }
139
140    #[inline]
141    fn flush(&mut self) -> std::io::Result<()> {
142        Ok(())
143    }
144
145    #[inline]
146    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
147        self.0 += buf.len() as u32;
148        Ok(())
149    }
150}
151
152/// Computes the number of bytes required to encode the given traces.
153///
154/// This does not allocate any actual buffer, but simulates writing in order to measure
155/// the encoded size of the traces.
156///
157/// # Arguments
158///
159/// * `traces` - A reference to a slice of spans.
160///
161/// # Returns
162///
163/// * `u32` - The number of bytes that would be written by the encoder.
164///
165/// # Examples
166///
167/// ```
168/// use libdd_trace_utils::msgpack_encoder::v04::to_len;
169/// use libdd_trace_utils::span::v04::SpanSlice;
170///
171/// let span = SpanSlice {
172///     name: "test-span",
173///     ..Default::default()
174/// };
175/// let traces = vec![vec![span]];
176/// let encoded_len = to_len(&traces);
177///
178/// assert!(encoded_len > 0);
179/// ```
180pub fn to_len<T: TraceData, S: AsRef<[Span<T>]>>(traces: &[S]) -> u32 {
181    let mut counter = CountLength(0);
182    #[allow(clippy::expect_used)]
183    to_writer(&mut counter, traces).expect("infallible: CountLength never fails");
184    counter.0
185}
libdd_trace_utils/msgpack_encoder/v04/mod.rs

libdd_trace_utils/msgpack_encoder/v04/
mod.rs
