use crate::{
format::{header_flags, SectionKind, FORMAT_VERSION, MAGIC},
id::EventId,
schema::{EventSchema, FieldRole, FieldType, Phase},
};
use alloc::{string::String, vec::Vec};
use bytes::Bytes;
const HEADER_LEN: usize = 8 + 2 + 2 + 2 + 2;
const SECTION_ENTRY_LEN: usize = 2 + 2 + 8 + 8;
mod ty_tag {
pub const U8: u8 = 0;
pub const U16: u8 = 1;
pub const U32: u8 = 2;
pub const U64: u8 = 3;
pub const I8: u8 = 4;
pub const I16: u8 = 5;
pub const I32: u8 = 6;
pub const I64: u8 = 7;
pub const BOOL: u8 = 8;
pub const BYTES: u8 = 9;
pub const ENUM: u8 = 10;
pub const INTERNED: u8 = 11;
}
mod role_tag {
pub const NONE: u8 = 0;
pub const KEY: u8 = 1;
pub const SPAN_ID: u8 = 2;
pub const PARENT_SPAN_ID: u8 = 3;
}
mod phase_tag {
pub const NONE: u8 = 0;
pub const ENTER: u8 = 1;
pub const EXIT: u8 = 2;
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Error {
UnexpectedEof,
BadMagic,
UnsupportedVersion(u16),
SectionOutOfBounds,
InvalidUtf8,
BadTag(u8),
}
impl core::fmt::Display for Error {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Error::UnexpectedEof => write!(f, "unexpected end of dump"),
Error::BadMagic => write!(f, "bad magic: not a backbeat dump"),
Error::UnsupportedVersion(v) => write!(f, "unsupported dump format version {v}"),
Error::SectionOutOfBounds => write!(f, "section offset/len out of bounds"),
Error::InvalidUtf8 => write!(f, "string field is not valid UTF-8"),
Error::BadTag(t) => write!(f, "unknown tag byte {t}"),
}
}
}
#[cfg(feature = "std")]
impl std::error::Error for Error {}
type Result<T> = core::result::Result<T, Error>;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OwnedEnumLabel {
pub value: u64,
pub label: String,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OwnedField {
pub name: String,
pub description: Option<String>,
pub ty: FieldType,
pub offset: u16,
pub width: u16,
pub role: FieldRole,
pub unit: Option<String>,
pub enum_labels: Vec<OwnedEnumLabel>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OwnedSchema {
pub id: EventId,
pub qualified_name: String,
pub description: Option<String>,
pub record_size: u16,
pub phase: Phase,
pub fields: Vec<OwnedField>,
}
impl OwnedSchema {
pub fn span_id(&self) -> Option<&OwnedField> {
self.fields.iter().find(|f| f.role == FieldRole::SpanId)
}
pub fn parent_span(&self) -> Option<&OwnedField> {
self.fields
.iter()
.find(|f| f.role == FieldRole::ParentSpanId)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OwnedMeta {
pub instance_id: u64,
pub host: String,
}
#[derive(Clone, Debug)]
pub struct ShardData {
pub shard_id: u32,
pub head: u64,
pub capacity: u64,
pub region: Bytes,
}
struct PendingSection {
kind: SectionKind,
body: Vec<u8>,
}
#[derive(Default)]
pub struct DumpWriter {
sections: Vec<PendingSection>,
}
impl DumpWriter {
pub fn new() -> Self {
Self::default()
}
pub fn schema_registry<'a>(&mut self, schemas: impl IntoIterator<Item = &'a EventSchema>) {
let mut body = Vec::new();
let schemas: Vec<&EventSchema> = schemas.into_iter().collect();
put_u32(&mut body, schemas.len() as u32);
for s in schemas {
put_schema(&mut body, s);
}
self.sections.push(PendingSection {
kind: SectionKind::Schema,
body,
});
}
pub fn intern_table<'a>(&mut self, entries: impl IntoIterator<Item = (u32, &'a [u8])>) {
let mut body = Vec::new();
let entries: Vec<(u32, &[u8])> = entries.into_iter().collect();
put_u32(&mut body, entries.len() as u32);
for (id, bytes) in entries {
put_u32(&mut body, id);
put_u32(&mut body, bytes.len() as u32);
body.extend_from_slice(bytes);
}
self.sections.push(PendingSection {
kind: SectionKind::Intern,
body,
});
}
pub fn meta(&mut self, instance_id: u64, host: &str) {
let mut body = Vec::new();
put_u64(&mut body, instance_id);
put_str(&mut body, host);
self.sections.push(PendingSection {
kind: SectionKind::Meta,
body,
});
}
pub fn shard(&mut self, shard_id: u32, head: u64, region: &[u8]) {
let mut body = Vec::with_capacity(20 + region.len());
put_u32(&mut body, shard_id);
put_u64(&mut body, head);
put_u64(&mut body, region.len() as u64);
body.extend_from_slice(region);
self.sections.push(PendingSection {
kind: SectionKind::Shard,
body,
});
}
pub fn finish(self) -> Vec<u8> {
let table_len = self.sections.len() * SECTION_ENTRY_LEN;
let mut body_offset = HEADER_LEN + table_len;
let mut out = Vec::new();
out.extend_from_slice(&MAGIC);
put_u16(&mut out, FORMAT_VERSION);
put_u16(&mut out, header_flags::LITTLE_ENDIAN);
put_u16(&mut out, self.sections.len() as u16);
put_u16(&mut out, 0);
for s in &self.sections {
put_u16(&mut out, s.kind as u16);
put_u16(&mut out, 0); put_u64(&mut out, body_offset as u64);
put_u64(&mut out, s.body.len() as u64);
body_offset += s.body.len();
}
for s in &self.sections {
out.extend_from_slice(&s.body);
}
out
}
}
fn put_schema(out: &mut Vec<u8>, s: &EventSchema) {
put_u64(out, s.id.get());
put_str(out, s.qualified_name);
put_opt_str(out, s.description);
put_u16(out, s.record_size);
out.push(encode_phase(s.phase));
put_u16(out, s.fields.len() as u16);
for f in s.fields {
put_str(out, f.name);
put_opt_str(out, f.description);
put_field_type(out, f.ty);
put_u16(out, f.offset);
put_u16(out, f.width);
out.push(encode_role(f.role));
put_opt_str(out, f.unit);
put_u16(out, f.enum_labels.len() as u16);
for l in f.enum_labels {
put_u64(out, l.value);
put_str(out, l.label);
}
}
}
fn encode_role(role: FieldRole) -> u8 {
match role {
FieldRole::None => role_tag::NONE,
FieldRole::Key => role_tag::KEY,
FieldRole::SpanId => role_tag::SPAN_ID,
FieldRole::ParentSpanId => role_tag::PARENT_SPAN_ID,
}
}
fn encode_phase(phase: Phase) -> u8 {
match phase {
Phase::None => phase_tag::NONE,
Phase::Enter => phase_tag::ENTER,
Phase::Exit => phase_tag::EXIT,
}
}
fn put_field_type(out: &mut Vec<u8>, ty: FieldType) {
match ty {
FieldType::U8 => out.push(ty_tag::U8),
FieldType::U16 => out.push(ty_tag::U16),
FieldType::U32 => out.push(ty_tag::U32),
FieldType::U64 => out.push(ty_tag::U64),
FieldType::I8 => out.push(ty_tag::I8),
FieldType::I16 => out.push(ty_tag::I16),
FieldType::I32 => out.push(ty_tag::I32),
FieldType::I64 => out.push(ty_tag::I64),
FieldType::Bool => out.push(ty_tag::BOOL),
FieldType::Bytes => out.push(ty_tag::BYTES),
FieldType::Enum { repr } => {
out.push(ty_tag::ENUM);
out.push(repr);
}
FieldType::Interned { dynamic } => {
out.push(ty_tag::INTERNED);
out.push(dynamic as u8);
}
}
}
fn put_u16(out: &mut Vec<u8>, v: u16) {
out.extend_from_slice(&v.to_le_bytes());
}
fn put_u32(out: &mut Vec<u8>, v: u32) {
out.extend_from_slice(&v.to_le_bytes());
}
fn put_u64(out: &mut Vec<u8>, v: u64) {
out.extend_from_slice(&v.to_le_bytes());
}
fn put_str(out: &mut Vec<u8>, s: &str) {
put_u16(out, s.len() as u16);
out.extend_from_slice(s.as_bytes());
}
fn put_opt_str(out: &mut Vec<u8>, s: Option<&str>) {
match s {
Some(s) => {
out.push(1);
put_str(out, s);
}
None => out.push(0),
}
}
pub struct DumpReader {
bytes: Bytes,
flags: u16,
sections: Vec<(u16, usize, usize)>,
}
impl DumpReader {
pub fn new(bytes: impl Into<Bytes>) -> Result<Self> {
let bytes = bytes.into();
let mut cur = Cursor::new(&bytes[..]);
let magic = cur.take(8)?;
if magic != MAGIC {
return Err(Error::BadMagic);
}
let format = cur.u16()?;
if format != FORMAT_VERSION {
return Err(Error::UnsupportedVersion(format));
}
let flags = cur.u16()?;
let section_count = cur.u16()? as usize;
let _reserved = cur.u16()?;
let mut sections = Vec::with_capacity(section_count);
for _ in 0..section_count {
let kind = cur.u16()?;
let _pad = cur.u16()?;
let offset = cur.u64()? as usize;
let len = cur.u64()? as usize;
let end = offset.checked_add(len).ok_or(Error::SectionOutOfBounds)?;
if end > bytes.len() {
return Err(Error::SectionOutOfBounds);
}
sections.push((kind, offset, len));
}
Ok(Self {
bytes,
flags,
sections,
})
}
pub fn flags(&self) -> u16 {
self.flags
}
pub fn section_count(&self) -> usize {
self.sections.len()
}
fn bodies_of(&self, kind: SectionKind) -> impl Iterator<Item = &[u8]> + '_ {
let bytes = &self.bytes;
self.sections
.iter()
.filter(move |(k, _, _)| *k == kind as u16)
.map(move |&(_, off, len)| &bytes[off..off + len])
}
pub fn schemas(&self) -> Result<Vec<OwnedSchema>> {
let Some(body) = self.bodies_of(SectionKind::Schema).next() else {
return Ok(Vec::new());
};
let mut cur = Cursor::new(body);
let count = cur.u32()? as usize;
let mut out = Vec::with_capacity(count.min(cur.remaining()));
for _ in 0..count {
out.push(get_schema(&mut cur)?);
}
Ok(out)
}
pub fn intern_table(&self) -> Result<Vec<(u32, Vec<u8>)>> {
let Some(body) = self.bodies_of(SectionKind::Intern).next() else {
return Ok(Vec::new());
};
let mut cur = Cursor::new(body);
let count = cur.u32()? as usize;
let mut out = Vec::with_capacity(count.min(cur.remaining()));
for _ in 0..count {
let id = cur.u32()?;
let len = cur.u32()? as usize;
out.push((id, cur.take(len)?.to_vec()));
}
Ok(out)
}
pub fn meta(&self) -> Result<Option<OwnedMeta>> {
let Some(body) = self.bodies_of(SectionKind::Meta).next() else {
return Ok(None);
};
let mut cur = Cursor::new(body);
let instance_id = cur.u64()?;
let host = cur.str()?;
Ok(Some(OwnedMeta { instance_id, host }))
}
pub fn shards(&self) -> Result<Vec<ShardData>> {
const PREFIX: usize = 4 + 8 + 8;
let mut out = Vec::new();
for &(kind, off, len) in &self.sections {
if kind != SectionKind::Shard as u16 {
continue;
}
let body = &self.bytes[off..off + len];
let mut cur = Cursor::new(body);
let shard_id = cur.u32()?;
let head = cur.u64()?;
let capacity = cur.u64()? as usize;
if PREFIX.checked_add(capacity).is_none_or(|end| end > len) {
return Err(Error::SectionOutOfBounds);
}
let region = self.bytes.slice(off + PREFIX..off + PREFIX + capacity);
out.push(ShardData {
shard_id,
head,
capacity: capacity as u64,
region,
});
}
Ok(out)
}
}
fn get_schema(cur: &mut Cursor) -> Result<OwnedSchema> {
let id = EventId(cur.u64()?);
let qualified_name = cur.str()?;
let description = cur.opt_str()?;
let record_size = cur.u16()?;
let phase = cur.phase()?;
let field_count = cur.u16()? as usize;
let mut fields = Vec::with_capacity(field_count.min(cur.remaining()));
for _ in 0..field_count {
let name = cur.str()?;
let description = cur.opt_str()?;
let ty = cur.field_type()?;
let offset = cur.u16()?;
let width = cur.u16()?;
let role = cur.role()?;
let unit = cur.opt_str()?;
let label_count = cur.u16()? as usize;
let mut enum_labels = Vec::with_capacity(label_count.min(cur.remaining()));
for _ in 0..label_count {
let value = cur.u64()?;
let label = cur.str()?;
enum_labels.push(OwnedEnumLabel { value, label });
}
fields.push(OwnedField {
name,
description,
ty,
offset,
width,
role,
unit,
enum_labels,
});
}
Ok(OwnedSchema {
id,
qualified_name,
description,
record_size,
phase,
fields,
})
}
struct Cursor<'a> {
bytes: &'a [u8],
pos: usize,
}
impl<'a> Cursor<'a> {
fn new(bytes: &'a [u8]) -> Self {
Self { bytes, pos: 0 }
}
fn take(&mut self, n: usize) -> Result<&'a [u8]> {
let end = self.pos.checked_add(n).ok_or(Error::UnexpectedEof)?;
let slice = self.bytes.get(self.pos..end).ok_or(Error::UnexpectedEof)?;
self.pos = end;
Ok(slice)
}
fn remaining(&self) -> usize {
self.bytes.len().saturating_sub(self.pos)
}
fn u8(&mut self) -> Result<u8> {
Ok(self.take(1)?[0])
}
fn u16(&mut self) -> Result<u16> {
Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap()))
}
fn u32(&mut self) -> Result<u32> {
Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
}
fn u64(&mut self) -> Result<u64> {
Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
}
fn str(&mut self) -> Result<String> {
let len = self.u16()? as usize;
let bytes = self.take(len)?;
core::str::from_utf8(bytes)
.map(String::from)
.map_err(|_| Error::InvalidUtf8)
}
fn opt_str(&mut self) -> Result<Option<String>> {
if self.u8()? == 0 {
Ok(None)
} else {
Ok(Some(self.str()?))
}
}
fn role(&mut self) -> Result<FieldRole> {
Ok(match self.u8()? {
role_tag::NONE => FieldRole::None,
role_tag::KEY => FieldRole::Key,
role_tag::SPAN_ID => FieldRole::SpanId,
role_tag::PARENT_SPAN_ID => FieldRole::ParentSpanId,
_ => FieldRole::None,
})
}
fn phase(&mut self) -> Result<Phase> {
Ok(match self.u8()? {
phase_tag::NONE => Phase::None,
phase_tag::ENTER => Phase::Enter,
phase_tag::EXIT => Phase::Exit,
_ => Phase::None,
})
}
fn field_type(&mut self) -> Result<FieldType> {
Ok(match self.u8()? {
ty_tag::U8 => FieldType::U8,
ty_tag::U16 => FieldType::U16,
ty_tag::U32 => FieldType::U32,
ty_tag::U64 => FieldType::U64,
ty_tag::I8 => FieldType::I8,
ty_tag::I16 => FieldType::I16,
ty_tag::I32 => FieldType::I32,
ty_tag::I64 => FieldType::I64,
ty_tag::BOOL => FieldType::Bool,
ty_tag::BYTES => FieldType::Bytes,
ty_tag::ENUM => FieldType::Enum { repr: self.u8()? },
ty_tag::INTERNED => FieldType::Interned {
dynamic: self.u8()? != 0,
},
other => return Err(Error::BadTag(other)),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::schema::{EnumLabel, FieldSchema};
const FIELDS: &[FieldSchema] = &[
FieldSchema {
name: "packet_number",
description: Some("the packet number"),
ty: FieldType::U64,
offset: 0,
width: 8,
role: FieldRole::Key,
unit: None,
enum_labels: &[],
},
FieldSchema {
name: "direction",
description: None,
ty: FieldType::Enum { repr: 1 },
offset: 8,
width: 1,
role: FieldRole::None,
unit: Some("dir"),
enum_labels: &[
EnumLabel {
value: 0,
label: "in",
},
EnumLabel {
value: 1,
label: "out",
},
],
},
];
const SCHEMA: EventSchema = EventSchema {
id: EventId::of("test::Demo"),
qualified_name: "test::Demo",
description: Some("a demo event"),
record_size: 9,
phase: Phase::None,
fields: FIELDS,
};
const SPAN_FIELDS: &[FieldSchema] = &[
FieldSchema {
name: "span",
description: None,
ty: FieldType::U64,
offset: 0,
width: 8,
role: FieldRole::SpanId,
unit: None,
enum_labels: &[],
},
FieldSchema {
name: "parent",
description: None,
ty: FieldType::U64,
offset: 8,
width: 8,
role: FieldRole::ParentSpanId,
unit: None,
enum_labels: &[],
},
];
const SPAN_SCHEMA: EventSchema = EventSchema {
id: EventId::of("test::Work"),
qualified_name: "test::Work",
description: None,
record_size: 16,
phase: Phase::Enter,
fields: SPAN_FIELDS,
};
#[test]
fn round_trips_a_full_dump() {
let mut w = DumpWriter::new();
w.schema_registry([&SCHEMA]);
w.intern_table([(7u32, b"hello".as_slice()), (9, b"world")]);
w.meta(0xDEADBEEF, "host-7");
w.shard(0, 32, &[0xAA; 64]);
w.shard(3, 8, &[0xBB; 16]);
let bytes = w.finish();
let r = DumpReader::new(bytes).unwrap();
assert_eq!(r.flags(), header_flags::LITTLE_ENDIAN);
assert_eq!(r.section_count(), 5);
let schemas = r.schemas().unwrap();
assert_eq!(schemas.len(), 1);
let s = &schemas[0];
assert_eq!(s.id, SCHEMA.id);
assert_eq!(s.qualified_name, "test::Demo");
assert_eq!(s.description.as_deref(), Some("a demo event"));
assert_eq!(s.record_size, 9);
assert_eq!(s.phase, Phase::None);
assert_eq!(s.fields.len(), 2);
assert_eq!(s.fields[0].name, "packet_number");
assert_eq!(
s.fields[0].description.as_deref(),
Some("the packet number")
);
assert_eq!(s.fields[0].role, FieldRole::Key);
assert_eq!(s.fields[1].ty, FieldType::Enum { repr: 1 });
assert_eq!(s.fields[1].unit.as_deref(), Some("dir"));
assert_eq!(s.fields[1].enum_labels[1].label, "out");
let intern = r.intern_table().unwrap();
assert_eq!(intern, [(7, b"hello".to_vec()), (9, b"world".to_vec())]);
let meta = r.meta().unwrap().unwrap();
assert_eq!(meta.instance_id, 0xDEADBEEF);
assert_eq!(meta.host, "host-7");
let shards = r.shards().unwrap();
assert_eq!(shards.len(), 2);
assert_eq!((shards[0].shard_id, shards[0].head), (0, 32));
assert_eq!(&shards[0].region[..], &[0xAA; 64]);
assert_eq!((shards[1].shard_id, shards[1].head), (3, 8));
assert_eq!(shards[1].region.len(), 16);
}
#[test]
fn round_trips_span_roles_and_phase() {
let mut w = DumpWriter::new();
w.schema_registry([&SPAN_SCHEMA]);
let bytes = w.finish();
let r = DumpReader::new(bytes).unwrap();
let s = &r.schemas().unwrap()[0];
assert_eq!(s.phase, Phase::Enter);
assert_eq!(s.fields[0].role, FieldRole::SpanId);
assert_eq!(s.fields[1].role, FieldRole::ParentSpanId);
}
#[test]
fn meta_absent_decodes_to_none() {
let mut w = DumpWriter::new();
w.schema_registry([&SCHEMA]);
let bytes = w.finish();
assert!(DumpReader::new(bytes).unwrap().meta().unwrap().is_none());
}
#[test]
fn rejects_bad_magic_and_truncation() {
assert_eq!(
DumpReader::new(b"nope".as_slice()).err(),
Some(Error::UnexpectedEof)
);
assert_eq!(
DumpReader::new(b"XXXXXXXX\x01\x00\x01\x00\x00\x00\x00\x00".as_slice()).err(),
Some(Error::BadMagic)
);
let mut w = DumpWriter::new();
w.schema_registry([&SCHEMA]);
let bytes = w.finish();
let truncated = bytes[..bytes.len() - 4].to_vec();
assert_eq!(
DumpReader::new(truncated).err(),
Some(Error::SectionOutOfBounds)
);
}
#[test]
fn empty_sections_decode_to_empty() {
let bytes = DumpWriter::new().finish();
let r = DumpReader::new(bytes).unwrap();
assert_eq!(r.section_count(), 0);
assert!(r.schemas().unwrap().is_empty());
assert!(r.intern_table().unwrap().is_empty());
assert!(r.shards().unwrap().is_empty());
}
}