Skip to content

Commit 2beef4c

Browse files
committed
Fix meta struct serialization.
* Now 'meta_struct' is defined as HashMap<T, Bytes> so it holds a a slice from the buffer rather than copying the bytes in anther container. * Bytes implements serialization so it's always serialized by serde as a byte array.
1 parent 9393721 commit 2beef4c

File tree

4 files changed

+67
-69
lines changed

4 files changed

+67
-69
lines changed

tinybytes/src/lib.rs

+13
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@ use std::{
77
sync::Arc,
88
};
99

10+
#[cfg(feature = "serde")]
11+
use serde::Serialize;
12+
1013
/// Immutable bytes type with zero copy cloning and slicing.
1114
#[derive(Clone)]
1215
pub struct Bytes {
@@ -270,6 +273,16 @@ impl fmt::Debug for Bytes {
270273
}
271274
}
272275

276+
#[cfg(feature = "serde")]
277+
impl Serialize for Bytes {
278+
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
279+
where
280+
S: serde::Serializer,
281+
{
282+
serializer.serialize_bytes(self.as_slice())
283+
}
284+
}
285+
273286
#[cfg(feature = "bytes_string")]
274287
mod bytes_string;
275288
#[cfg(feature = "bytes_string")]

trace-utils/src/msgpack_decoder/decode/meta_struct.rs

+41-12
Original file line numberDiff line numberDiff line change
@@ -3,33 +3,62 @@
33

44
use crate::msgpack_decoder::decode::error::DecodeError;
55
use crate::msgpack_decoder::decode::map::{read_map, read_map_len};
6-
use crate::msgpack_decoder::decode::number::read_number_bytes;
76
use crate::msgpack_decoder::decode::string::{handle_null_marker, read_string_bytes};
87
use rmp::decode;
98
use std::collections::HashMap;
109
use tinybytes::{Bytes, BytesString};
1110

11+
fn read_byte_array_len(buf: &mut &[u8]) -> Result<u32, DecodeError> {
12+
decode::read_bin_len(buf).map_err(|_| {
13+
DecodeError::InvalidFormat("Unable to read binary len for meta_struct".to_owned())
14+
})
15+
}
16+
1217
#[inline]
13-
pub fn read_meta_struct(buf: &mut Bytes) -> Result<HashMap<BytesString, Vec<u8>>, DecodeError> {
18+
pub fn read_meta_struct(buf: &mut Bytes) -> Result<HashMap<BytesString, Bytes>, DecodeError> {
1419
if let Some(empty_map) = handle_null_marker(buf, HashMap::default) {
1520
return Ok(empty_map);
1621
}
1722

18-
fn read_meta_struct_pair(buf: &mut Bytes) -> Result<(BytesString, Vec<u8>), DecodeError> {
23+
fn read_meta_struct_pair(buf: &mut Bytes) -> Result<(BytesString, Bytes), DecodeError> {
1924
let key = read_string_bytes(buf)?;
20-
let array_len = decode::read_array_len(unsafe { buf.as_mut_slice() }).map_err(|_| {
21-
DecodeError::InvalidFormat("Unable to read array len for meta_struct".to_owned())
22-
})?;
25+
let byte_array_len = read_byte_array_len(unsafe { buf.as_mut_slice() })? as usize;
2326

24-
let mut v = Vec::with_capacity(array_len as usize);
25-
26-
for _ in 0..array_len {
27-
let value = read_number_bytes(buf)?;
28-
v.push(value);
27+
let data = buf.slice_ref(&buf[0..byte_array_len]).unwrap();
28+
unsafe {
29+
// SAFETY: forwarding the buffer requires that buf is borrowed from static.
30+
*buf.as_mut_slice() = &buf.as_mut_slice()[byte_array_len..];
2931
}
30-
Ok((key, v))
32+
33+
Ok((key, data))
3134
}
3235

3336
let len = read_map_len(unsafe { buf.as_mut_slice() })?;
3437
read_map(len, buf, read_meta_struct_pair)
3538
}
39+
40+
#[cfg(test)]
41+
mod tests {
42+
use super::*;
43+
44+
#[test]
45+
fn read_meta_test() {
46+
let meta = HashMap::from([("key".to_string(), Bytes::from(vec![1, 2, 3, 4]))]);
47+
48+
let mut bytes = Bytes::from(rmp_serde::to_vec_named(&meta).unwrap());
49+
let res = read_meta_struct(&mut bytes).unwrap();
50+
51+
assert_eq!(res.get("key").unwrap().to_vec(), vec![1, 2, 3, 4]);
52+
}
53+
54+
#[test]
55+
fn read_meta_wrong_family_test() {
56+
let meta = HashMap::from([("key".to_string(), vec![1, 2, 3, 4])]);
57+
58+
let mut bytes = Bytes::from(rmp_serde::to_vec_named(&meta).unwrap());
59+
let res = read_meta_struct(&mut bytes);
60+
61+
assert!(res.is_err());
62+
matches!(res.unwrap_err(), DecodeError::InvalidFormat(_));
63+
}
64+
}

trace-utils/src/msgpack_decoder/v04/mod.rs

+11-55
Original file line numberDiff line numberDiff line change
@@ -96,29 +96,14 @@ pub fn from_slice(mut data: tinybytes::Bytes) -> Result<(Vec<Vec<SpanBytes>>, us
9696
#[cfg(test)]
9797
mod tests {
9898
use super::*;
99-
use crate::test_utils::create_test_json_span;
99+
use crate::test_utils::{create_test_json_span, create_test_no_alloc_span};
100100
use bolero::check;
101101
use rmp_serde;
102102
use rmp_serde::to_vec_named;
103103
use serde_json::json;
104104
use std::collections::HashMap;
105-
use tinybytes::BytesString;
105+
use tinybytes::{Bytes, BytesString};
106106

107-
fn generate_meta_struct_element(i: u8) -> (String, Vec<u8>) {
108-
let map = HashMap::from([
109-
(
110-
format!("meta_struct_map_key {}", i + 1),
111-
format!("meta_struct_map_val {}", i + 1),
112-
),
113-
(
114-
format!("meta_struct_map_key {}", i + 2),
115-
format!("meta_struct_map_val {}", i + 2),
116-
),
117-
]);
118-
let key = format!("key {}", i).to_owned();
119-
120-
(key, rmp_serde::to_vec_named(&map).unwrap())
121-
}
122107
#[test]
123108
fn test_empty_array() {
124109
let encoded_data = vec![0x90];
@@ -224,38 +209,11 @@ mod tests {
224209
}
225210

226211
#[test]
227-
fn test_decoder_meta_struct_fixed_map_success() {
228-
let expected_meta_struct = HashMap::from([
229-
generate_meta_struct_element(0),
230-
generate_meta_struct_element(1),
231-
]);
232-
233-
let mut span = create_test_json_span(1, 2, 0, 0, false);
234-
span["meta_struct"] = json!(expected_meta_struct.clone());
235-
236-
let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap();
237-
let (decoded_traces, _) =
238-
from_slice(tinybytes::Bytes::from(encoded_data)).expect("Decoding failed");
239-
240-
assert_eq!(1, decoded_traces.len());
241-
assert_eq!(1, decoded_traces[0].len());
242-
let decoded_span = &decoded_traces[0][0];
243-
244-
for (key, value) in expected_meta_struct.iter() {
245-
assert_eq!(
246-
value,
247-
&decoded_span.meta_struct[&BytesString::from_slice(key.as_ref()).unwrap()]
248-
);
249-
}
250-
}
251-
252-
#[test]
253-
fn test_decoder_meta_struct_map_16_success() {
254-
let expected_meta_struct: HashMap<String, Vec<u8>> =
255-
(0..20).map(generate_meta_struct_element).collect();
256-
257-
let mut span = create_test_json_span(1, 2, 0, 0, false);
258-
span["meta_struct"] = json!(expected_meta_struct.clone());
212+
fn test_decoder_meta_struct_success() {
213+
let data = vec![1, 2, 3, 4];
214+
let mut span = create_test_no_alloc_span(1, 2, 0, 0, false);
215+
span.meta_struct =
216+
HashMap::from([(BytesString::from("meta_key"), Bytes::from(data.clone()))]);
259217

260218
let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap();
261219
let (decoded_traces, _) =
@@ -265,12 +223,10 @@ mod tests {
265223
assert_eq!(1, decoded_traces[0].len());
266224
let decoded_span = &decoded_traces[0][0];
267225

268-
for (key, value) in expected_meta_struct.iter() {
269-
assert_eq!(
270-
value,
271-
&decoded_span.meta_struct[&BytesString::from_slice(key.as_ref()).unwrap()]
272-
);
273-
}
226+
assert_eq!(
227+
decoded_span.meta_struct.get("meta_key").unwrap().to_vec(),
228+
data
229+
);
274230
}
275231

276232
#[test]

trace-utils/src/span/mod.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ use std::collections::HashMap;
1010
use std::fmt;
1111
use std::hash::Hash;
1212
use std::str::FromStr;
13-
use tinybytes::BytesString;
13+
use tinybytes::{Bytes, BytesString};
1414

1515
#[derive(Debug, PartialEq)]
1616
pub enum SpanKey {
@@ -101,7 +101,7 @@ where
101101
#[serde(skip_serializing_if = "HashMap::is_empty")]
102102
pub metrics: HashMap<T, f64>,
103103
#[serde(skip_serializing_if = "HashMap::is_empty")]
104-
pub meta_struct: HashMap<T, Vec<u8>>,
104+
pub meta_struct: HashMap<T, Bytes>,
105105
#[serde(skip_serializing_if = "Vec::is_empty")]
106106
pub span_links: Vec<SpanLink<T>>,
107107
}

0 commit comments

Comments
 (0)