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

libdd_trace_utils/msgpack_encoder/v04/
mod.rs
