Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix meta struct serialization. #927

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix meta struct serialization. #927

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions tinybytes/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@ use std::{
sync::Arc,
};

#[cfg(feature = "serde")]
use serde::Serialize;

/// Immutable bytes type with zero copy cloning and slicing.
#[derive(Clone)]
pub struct Bytes {
Expand Down Expand Up @@ -270,6 +273,16 @@ impl fmt::Debug for Bytes {
}
}

#[cfg(feature = "serde")]
impl Serialize for Bytes {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I know it's covered in trace-utils, but Is this something we should explicitly test in tinybytes as well?

fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serializer.serialize_bytes(self.as_slice())
}
}

#[cfg(feature = "bytes_string")]
mod bytes_string;
#[cfg(feature = "bytes_string")]
Expand Down
53 changes: 41 additions & 12 deletions trace-utils/src/msgpack_decoder/decode/meta_struct.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,33 +3,62 @@

use crate::msgpack_decoder::decode::error::DecodeError;
use crate::msgpack_decoder::decode::map::{read_map, read_map_len};
use crate::msgpack_decoder::decode::number::read_number_bytes;
use crate::msgpack_decoder::decode::string::{handle_null_marker, read_string_bytes};
use rmp::decode;
use std::collections::HashMap;
use tinybytes::{Bytes, BytesString};

fn read_byte_array_len(buf: &mut &[u8]) -> Result<u32, DecodeError> {
decode::read_bin_len(buf).map_err(|_| {
DecodeError::InvalidFormat("Unable to read binary len for meta_struct".to_owned())
})
}

#[inline]
pub fn read_meta_struct(buf: &mut Bytes) -> Result<HashMap<BytesString, Vec<u8>>, DecodeError> {
pub fn read_meta_struct(buf: &mut Bytes) -> Result<HashMap<BytesString, Bytes>, DecodeError> {
if let Some(empty_map) = handle_null_marker(buf, HashMap::default) {
return Ok(empty_map);
}

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

let mut v = Vec::with_capacity(array_len as usize);

for _ in 0..array_len {
let value = read_number_bytes(buf)?;
v.push(value);
let data = buf.slice_ref(&buf[0..byte_array_len]).unwrap();
unsafe {
// SAFETY: forwarding the buffer requires that buf is borrowed from static.
*buf.as_mut_slice() = &buf.as_mut_slice()[byte_array_len..];
}
Ok((key, v))

Ok((key, data))
}

let len = read_map_len(unsafe { buf.as_mut_slice() })?;
read_map(len, buf, read_meta_struct_pair)
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn read_meta_test() {
let meta = HashMap::from([("key".to_string(), Bytes::from(vec![1, 2, 3, 4]))]);

let mut bytes = Bytes::from(rmp_serde::to_vec_named(&meta).unwrap());
let res = read_meta_struct(&mut bytes).unwrap();

assert_eq!(res.get("key").unwrap().to_vec(), vec![1, 2, 3, 4]);
}

#[test]
fn read_meta_wrong_family_test() {
let meta = HashMap::from([("key".to_string(), vec![1, 2, 3, 4])]);

let mut bytes = Bytes::from(rmp_serde::to_vec_named(&meta).unwrap());
let res = read_meta_struct(&mut bytes);

assert!(res.is_err());
matches!(res.unwrap_err(), DecodeError::InvalidFormat(_));
}
}
66 changes: 11 additions & 55 deletions trace-utils/src/msgpack_decoder/v04/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,29 +97,14 @@ pub fn from_slice(mut data: tinybytes::Bytes) -> Result<(Vec<Vec<SpanBytes>>, us
#[cfg(test)]
mod tests {
use super::*;
use crate::test_utils::create_test_json_span;
use crate::test_utils::{create_test_json_span, create_test_no_alloc_span};
use bolero::check;
use rmp_serde;
use rmp_serde::to_vec_named;
use serde_json::json;
use std::collections::HashMap;
use tinybytes::BytesString;
use tinybytes::{Bytes, BytesString};

fn generate_meta_struct_element(i: u8) -> (String, Vec<u8>) {
let map = HashMap::from([
(
format!("meta_struct_map_key {}", i + 1),
format!("meta_struct_map_val {}", i + 1),
),
(
format!("meta_struct_map_key {}", i + 2),
format!("meta_struct_map_val {}", i + 2),
),
]);
let key = format!("key {}", i).to_owned();

(key, rmp_serde::to_vec_named(&map).unwrap())
}
#[test]
fn test_empty_array() {
let encoded_data = vec![0x90];
Expand Down Expand Up @@ -225,38 +210,11 @@ mod tests {
}

#[test]
fn test_decoder_meta_struct_fixed_map_success() {
let expected_meta_struct = HashMap::from([
generate_meta_struct_element(0),
generate_meta_struct_element(1),
]);

let mut span = create_test_json_span(1, 2, 0, 0, false);
span["meta_struct"] = json!(expected_meta_struct.clone());

let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap();
let (decoded_traces, _) =
from_slice(tinybytes::Bytes::from(encoded_data)).expect("Decoding failed");

assert_eq!(1, decoded_traces.len());
assert_eq!(1, decoded_traces[0].len());
let decoded_span = &decoded_traces[0][0];

for (key, value) in expected_meta_struct.iter() {
assert_eq!(
value,
&decoded_span.meta_struct[&BytesString::from_slice(key.as_ref()).unwrap()]
);
}
}

#[test]
fn test_decoder_meta_struct_map_16_success() {
let expected_meta_struct: HashMap<String, Vec<u8>> =
(0..20).map(generate_meta_struct_element).collect();

let mut span = create_test_json_span(1, 2, 0, 0, false);
span["meta_struct"] = json!(expected_meta_struct.clone());
fn test_decoder_meta_struct_success() {
let data = vec![1, 2, 3, 4];
let mut span = create_test_no_alloc_span(1, 2, 0, 0, false);
span.meta_struct =
HashMap::from([(BytesString::from("meta_key"), Bytes::from(data.clone()))]);

let encoded_data = rmp_serde::to_vec_named(&vec![vec![span]]).unwrap();
let (decoded_traces, _) =
Expand All @@ -266,12 +224,10 @@ mod tests {
assert_eq!(1, decoded_traces[0].len());
let decoded_span = &decoded_traces[0][0];

for (key, value) in expected_meta_struct.iter() {
assert_eq!(
value,
&decoded_span.meta_struct[&BytesString::from_slice(key.as_ref()).unwrap()]
);
}
assert_eq!(
decoded_span.meta_struct.get("meta_key").unwrap().to_vec(),
data
);
}

#[test]
Expand Down
4 changes: 2 additions & 2 deletions trace-utils/src/span/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use std::collections::HashMap;
use std::fmt;
use std::hash::Hash;
use std::str::FromStr;
use tinybytes::BytesString;
use tinybytes::{Bytes, BytesString};

#[derive(Debug, PartialEq)]
pub enum SpanKey {
Expand Down Expand Up @@ -104,7 +104,7 @@ where
#[serde(skip_serializing_if = "HashMap::is_empty")]
pub metrics: HashMap<T, f64>,
#[serde(skip_serializing_if = "HashMap::is_empty")]
pub meta_struct: HashMap<T, Vec<u8>>,
pub meta_struct: HashMap<T, Bytes>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub span_links: Vec<SpanLink<T>>,
#[serde(skip_serializing_if = "Vec::is_empty")]
Expand Down
Loading