Skip to content

Commit 50d3a77

Browse files
VianneyRuhlmannVianneyRuhlmann
committed
Move number, string and map decoding to separate modules
1 parent 84f8853 commit 50d3a77

File tree

9 files changed

+265
-241
lines changed

9 files changed

+265
-241
lines changed

Diff for: ‎trace-utils/src/msgpack_decoder/mod.rs

+1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
22
// SPDX-License-Identifier: Apache-2.0
33

4+
/// Decoding logic for V04 encoded trace payload
45
pub mod v04;

Diff for: ‎trace-utils/src/msgpack_decoder/v04/decoder/mod.rs

+6-185
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,12 @@
11
// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
22
// SPDX-License-Identifier: Apache-2.0
33

4+
/// Read spans from msgpack
45
mod span;
5-
mod span_link;
66

77
use super::error::DecodeError;
8-
use super::number::{read_nullable_number_ref, read_number_ref};
98
use crate::span_v04::{Span, SpanSlice};
10-
use rmp::decode::DecodeStringError;
11-
use rmp::{decode, decode::RmpRead, Marker};
129
use span::decode_span;
13-
use std::{collections::HashMap, f64};
14-
15-
// https://docs.rs/rmp/latest/rmp/enum.Marker.html#variant.Null (0xc0 == 192)
16-
const NULL_MARKER: &u8 = &0xc0;
1710

1811
/// Decodes a Bytes buffer into a vector of `TracerPayloadV04` objects.
1912
///
@@ -152,181 +145,6 @@ pub fn from_slice(mut data: &[u8]) -> Result<(Vec<Vec<SpanSlice>>, usize), Decod
152145
))
153146
}
154147

155-
#[inline]
156-
fn read_string_nomut(buf: &[u8]) -> Result<(&str, &[u8]), DecodeError> {
157-
decode::read_str_from_slice(buf).map_err(|e| match e {
158-
DecodeStringError::InvalidMarkerRead(e) => DecodeError::InvalidFormat(e.to_string()),
159-
DecodeStringError::InvalidDataRead(e) => DecodeError::InvalidConversion(e.to_string()),
160-
DecodeStringError::TypeMismatch(marker) => {
161-
DecodeError::InvalidType(format!("Type mismatch at marker {:?}", marker))
162-
}
163-
DecodeStringError::InvalidUtf8(_, e) => DecodeError::Utf8Error(e.to_string()),
164-
_ => DecodeError::IOError,
165-
})
166-
}
167-
168-
#[inline]
169-
fn read_string<'a>(buf: &mut &'a [u8]) -> Result<&'a str, DecodeError> {
170-
read_string_nomut(buf).map(|(str, newbuf)| {
171-
*buf = newbuf;
172-
str
173-
})
174-
}
175-
176-
#[inline]
177-
fn read_nullable_string<'a>(buf: &mut &'a [u8]) -> Result<&'a str, DecodeError> {
178-
if is_null_marker(buf) {
179-
Ok("")
180-
} else {
181-
read_string(buf)
182-
}
183-
}
184-
185-
#[inline]
186-
fn read_str_map_to_str<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
187-
let len = decode::read_map_len(buf)
188-
.map_err(|_| DecodeError::InvalidFormat("Unable to get map len for str map".to_owned()))?;
189-
190-
let mut map = HashMap::with_capacity(len.try_into().expect("Unable to cast map len to usize"));
191-
for _ in 0..len {
192-
let key = read_string(buf)?;
193-
let value = read_string(buf)?;
194-
map.insert(key, value);
195-
}
196-
Ok(map)
197-
}
198-
199-
#[inline]
200-
fn read_nullable_str_map_to_str<'a>(
201-
buf: &mut &'a [u8],
202-
) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
203-
if is_null_marker(buf) {
204-
return Ok(HashMap::default());
205-
}
206-
207-
read_str_map_to_str(buf)
208-
}
209-
210-
#[inline]
211-
fn read_metrics<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, f64>, DecodeError> {
212-
if is_null_marker(buf) {
213-
return Ok(HashMap::default());
214-
}
215-
216-
fn read_metric_pair<'a>(buf: &mut &'a [u8]) -> Result<(&'a str, f64), DecodeError> {
217-
let key = read_string(buf)?;
218-
let v = read_number_ref(buf)?;
219-
220-
Ok((key, v))
221-
}
222-
223-
let len = read_map_len(buf)?;
224-
225-
read_map(len, buf, read_metric_pair)
226-
}
227-
228-
#[inline]
229-
fn read_meta_struct<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, Vec<u8>>, DecodeError> {
230-
if is_null_marker(buf) {
231-
return Ok(HashMap::default());
232-
}
233-
234-
fn read_meta_struct_pair<'a>(buf: &mut &'a [u8]) -> Result<(&'a str, Vec<u8>), DecodeError> {
235-
let key = read_string(buf)?;
236-
let array_len = decode::read_array_len(buf).map_err(|_| {
237-
DecodeError::InvalidFormat("Unable to read array len for meta_struct".to_owned())
238-
})?;
239-
240-
let mut v = Vec::with_capacity(array_len as usize);
241-
242-
for _ in 0..array_len {
243-
let value = read_number_ref(buf)?;
244-
v.push(value);
245-
}
246-
Ok((key, v))
247-
}
248-
249-
let len = read_map_len(buf)?;
250-
read_map(len, buf, read_meta_struct_pair)
251-
}
252-
253-
/// Reads a map from the buffer and returns it as a `HashMap`.
254-
///
255-
/// This function is generic over the key and value types of the map, and it uses a provided
256-
/// function to read key-value pairs from the buffer.
257-
///
258-
/// # Arguments
259-
///
260-
/// * `len` - The number of key-value pairs to read from the buffer.
261-
/// * `buf` - A reference to the slice containing the encoded map data.
262-
/// * `read_pair` - A function that reads a key-value pair from the buffer and returns it as a
263-
/// `Result<(K, V), DecodeError>`.
264-
///
265-
/// # Returns
266-
///
267-
/// * `Ok(HashMap<K, V>)` - A `HashMap` containing the decoded key-value pairs if successful.
268-
/// * `Err(DecodeError)` - An error if the decoding process fails.
269-
///
270-
/// # Errors
271-
///
272-
/// This function will return an error if:
273-
/// - The `read_pair` function returns an error while reading a key-value pair.
274-
///
275-
/// # Type Parameters
276-
///
277-
/// * `K` - The type of the keys in the map. Must implement `std::hash::Hash` and `Eq`.
278-
/// * `V` - The type of the values in the map.
279-
/// * `F` - The type of the function used to read key-value pairs from the buffer.
280-
#[inline]
281-
fn read_map<'a, K, V, F>(
282-
len: usize,
283-
buf: &mut &'a [u8],
284-
read_pair: F,
285-
) -> Result<HashMap<K, V>, DecodeError>
286-
where
287-
K: std::hash::Hash + Eq,
288-
F: Fn(&mut &'a [u8]) -> Result<(K, V), DecodeError>,
289-
{
290-
let mut map = HashMap::with_capacity(len);
291-
for _ in 0..len {
292-
let (k, v) = read_pair(buf)?;
293-
map.insert(k, v);
294-
}
295-
Ok(map)
296-
}
297-
298-
#[inline]
299-
fn read_map_len(buf: &mut &[u8]) -> Result<usize, DecodeError> {
300-
match decode::read_marker(buf)
301-
.map_err(|_| DecodeError::InvalidFormat("Unable to read marker for map".to_owned()))?
302-
{
303-
Marker::FixMap(len) => Ok(len as usize),
304-
Marker::Map16 => buf
305-
.read_data_u16()
306-
.map_err(|_| DecodeError::IOError)
307-
.map(|len| len as usize),
308-
Marker::Map32 => buf
309-
.read_data_u32()
310-
.map_err(|_| DecodeError::IOError)
311-
.map(|len| len as usize),
312-
_ => Err(DecodeError::InvalidType(
313-
"Unable to read map from buffer".to_owned(),
314-
)),
315-
}
316-
}
317-
318-
/// When you want to "peek" if the next value is a null marker, and only advance the buffer if it is
319-
/// null. If it is not null, you can continue to decode as expected.
320-
#[inline]
321-
fn is_null_marker(buf: &mut &[u8]) -> bool {
322-
if buf.first() == Some(NULL_MARKER) {
323-
*buf = &buf[1..];
324-
true
325-
} else {
326-
false
327-
}
328-
}
329-
330148
#[cfg(test)]
331149
mod tests {
332150
use super::*;
@@ -335,6 +153,7 @@ mod tests {
335153
use rmp_serde;
336154
use rmp_serde::to_vec_named;
337155
use serde_json::json;
156+
use std::collections::HashMap;
338157
use tinybytes::BytesString;
339158

340159
fn generate_meta_struct_element(i: u8) -> (String, Vec<u8>) {
@@ -352,6 +171,7 @@ mod tests {
352171

353172
(key, rmp_serde::to_vec_named(&map).unwrap())
354173
}
174+
355175
#[test]
356176
fn test_empty_array() {
357177
let encoded_data = vec![0x90];
@@ -824,7 +644,7 @@ mod tests {
824644

825645
#[test]
826646
#[cfg_attr(miri, ignore)]
827-
fn fuzz_from_slice() {
647+
fn fuzz_from_bytes() {
828648
check!()
829649
.with_type::<(
830650
String,
@@ -875,10 +695,11 @@ mod tests {
875695
start,
876696
..Default::default()
877697
};
878-
let encoded_data = to_vec_named(&vec![vec![span]]).unwrap();
698+
let encoded_data = to_vec_named(&vec![vec![span.clone()]]).unwrap();
879699
let result = from_bytes(tinybytes::Bytes::from(encoded_data));
880700

881701
assert!(result.is_ok());
702+
assert_eq!(result.unwrap().0, vec![vec![span]])
882703
},
883704
);
884705
}

Diff for: ‎trace-utils/src/msgpack_decoder/v04/decoder/span/map.rs

+149
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
use super::is_null_marker;
2+
use super::number::read_number;
3+
use super::string::read_string;
4+
use crate::msgpack_decoder::v04::error::DecodeError;
5+
use rmp::{decode, decode::RmpRead, Marker};
6+
use std::collections::HashMap;
7+
8+
/// Read a map of string to string from `buf`.
9+
#[inline]
10+
pub fn read_str_map_to_str<'a>(
11+
buf: &mut &'a [u8],
12+
) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
13+
let len = decode::read_map_len(buf)
14+
.map_err(|_| DecodeError::InvalidFormat("Unable to get map len for str map".to_owned()))?;
15+
16+
let mut map = HashMap::with_capacity(len.try_into().expect("Unable to cast map len to usize"));
17+
for _ in 0..len {
18+
let key = read_string(buf)?;
19+
let value = read_string(buf)?;
20+
map.insert(key, value);
21+
}
22+
Ok(map)
23+
}
24+
25+
/// Read a nullable map of string to string from `buf`.
26+
#[inline]
27+
pub fn read_nullable_str_map_to_str<'a>(
28+
buf: &mut &'a [u8],
29+
) -> Result<HashMap<&'a str, &'a str>, DecodeError> {
30+
if is_null_marker(buf) {
31+
return Ok(HashMap::default());
32+
}
33+
34+
read_str_map_to_str(buf)
35+
}
36+
37+
/// Read a map of string to f64 from `buf`.
38+
#[inline]
39+
pub fn read_metrics<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, f64>, DecodeError> {
40+
if is_null_marker(buf) {
41+
return Ok(HashMap::default());
42+
}
43+
44+
fn read_metric_pair<'a>(buf: &mut &'a [u8]) -> Result<(&'a str, f64), DecodeError> {
45+
let key = read_string(buf)?;
46+
let v = read_number(buf)?;
47+
48+
Ok((key, v))
49+
}
50+
51+
let len = read_map_len(buf)?;
52+
53+
read_map(len, buf, read_metric_pair)
54+
}
55+
56+
/// Read a map of string to u8 array from `buf`.
57+
///
58+
/// The struct can't be a u8 slice since it is encoded as a msgpack array and not as a raw bytes
59+
/// buffer.
60+
#[inline]
61+
pub fn read_meta_struct<'a>(buf: &mut &'a [u8]) -> Result<HashMap<&'a str, Vec<u8>>, DecodeError> {
62+
if is_null_marker(buf) {
63+
return Ok(HashMap::default());
64+
}
65+
66+
fn read_meta_struct_pair<'a>(buf: &mut &'a [u8]) -> Result<(&'a str, Vec<u8>), DecodeError> {
67+
let key = read_string(buf)?;
68+
let array_len = decode::read_array_len(buf).map_err(|_| {
69+
DecodeError::InvalidFormat("Unable to read array len for meta_struct".to_owned())
70+
})?;
71+
72+
let mut v = Vec::with_capacity(array_len as usize);
73+
74+
for _ in 0..array_len {
75+
let value = read_number(buf)?;
76+
v.push(value);
77+
}
78+
Ok((key, v))
79+
}
80+
81+
let len = read_map_len(buf)?;
82+
read_map(len, buf, read_meta_struct_pair)
83+
}
84+
85+
/// Reads a map from the buffer and returns it as a `HashMap`.
86+
///
87+
/// This function is generic over the key and value types of the map, and it uses a provided
88+
/// function to read key-value pairs from the buffer.
89+
///
90+
/// # Arguments
91+
///
92+
/// * `len` - The number of key-value pairs to read from the buffer.
93+
/// * `buf` - A reference to the slice containing the encoded map data.
94+
/// * `read_pair` - A function that reads a key-value pair from the buffer and returns it as a
95+
/// `Result<(K, V), DecodeError>`.
96+
///
97+
/// # Returns
98+
///
99+
/// * `Ok(HashMap<K, V>)` - A `HashMap` containing the decoded key-value pairs if successful.
100+
/// * `Err(DecodeError)` - An error if the decoding process fails.
101+
///
102+
/// # Errors
103+
///
104+
/// This function will return an error if:
105+
/// - The `read_pair` function returns an error while reading a key-value pair.
106+
///
107+
/// # Type Parameters
108+
///
109+
/// * `K` - The type of the keys in the map. Must implement `std::hash::Hash` and `Eq`.
110+
/// * `V` - The type of the values in the map.
111+
/// * `F` - The type of the function used to read key-value pairs from the buffer.
112+
#[inline]
113+
fn read_map<'a, K, V, F>(
114+
len: usize,
115+
buf: &mut &'a [u8],
116+
read_pair: F,
117+
) -> Result<HashMap<K, V>, DecodeError>
118+
where
119+
K: std::hash::Hash + Eq,
120+
F: Fn(&mut &'a [u8]) -> Result<(K, V), DecodeError>,
121+
{
122+
let mut map = HashMap::with_capacity(len);
123+
for _ in 0..len {
124+
let (k, v) = read_pair(buf)?;
125+
map.insert(k, v);
126+
}
127+
Ok(map)
128+
}
129+
130+
/// Read the length of a msgpack map.
131+
#[inline]
132+
fn read_map_len(buf: &mut &[u8]) -> Result<usize, DecodeError> {
133+
match decode::read_marker(buf)
134+
.map_err(|_| DecodeError::InvalidFormat("Unable to read marker for map".to_owned()))?
135+
{
136+
Marker::FixMap(len) => Ok(len as usize),
137+
Marker::Map16 => buf
138+
.read_data_u16()
139+
.map_err(|_| DecodeError::IOError)
140+
.map(|len| len as usize),
141+
Marker::Map32 => buf
142+
.read_data_u32()
143+
.map_err(|_| DecodeError::IOError)
144+
.map(|len| len as usize),
145+
_ => Err(DecodeError::InvalidType(
146+
"Unable to read map from buffer".to_owned(),
147+
)),
148+
}
149+
}

0 commit comments

Comments
 (0)