1use crate::TraceEvent;
8use crate::codec::{self, PoolEntry, StackPoolEntry, WireTypeId};
9use crate::schema::{SchemaEntry, SchemaRegistry};
10use crate::types::{
11 CountingWriter, EncodeState, EventEncoder, InternedStackFrames, InternedString,
12};
13use std::any::TypeId;
14use std::collections::{HashMap, HashSet};
15use std::hash::{BuildHasherDefault, Hasher};
16use std::io::{self, Write};
17use std::sync::Arc;
18
19#[doc(hidden)]
24#[derive(Default)]
25pub struct FxHasher(u64);
26
27impl FxHasher {
28 #[inline]
29 fn hash_word(&mut self, word: u64) {
30 self.0 = (self.0.rotate_left(5) ^ word).wrapping_mul(0x517cc1b727220a95);
31 }
32}
33
34impl Hasher for FxHasher {
35 #[inline]
36 fn write(&mut self, mut bytes: &[u8]) {
37 while bytes.len() >= 8 {
38 self.hash_word(u64::from_ne_bytes(bytes[..8].try_into().unwrap()));
39 bytes = &bytes[8..];
40 }
41 if bytes.len() >= 4 {
42 self.hash_word(u32::from_ne_bytes(bytes[..4].try_into().unwrap()) as u64);
43 bytes = &bytes[4..];
44 }
45 for &b in bytes {
46 self.hash_word(b as u64);
47 }
48 }
49
50 #[inline]
51 fn write_u8(&mut self, i: u8) {
52 self.hash_word(i as u64);
53 }
54
55 #[inline]
56 fn write_u16(&mut self, i: u16) {
57 self.hash_word(i as u64);
58 }
59
60 #[inline]
61 fn write_u32(&mut self, i: u32) {
62 self.hash_word(i as u64);
63 }
64
65 #[inline]
66 fn write_u64(&mut self, i: u64) {
67 self.hash_word(i);
68 }
69
70 #[inline]
71 fn write_usize(&mut self, i: usize) {
72 self.hash_word(i as u64);
73 }
74
75 #[inline]
76 fn write_u128(&mut self, i: u128) {
77 self.hash_word(i as u64);
78 self.hash_word((i >> 64) as u64);
79 }
80
81 #[inline]
82 fn finish(&self) -> u64 {
83 self.0
84 }
85}
86
87#[doc(hidden)]
88pub type FxBuildHasher = BuildHasherDefault<FxHasher>;
89#[doc(hidden)]
90pub type FxHashMap<K, V> = HashMap<K, V, FxBuildHasher>;
91#[doc(hidden)]
92pub type FxHashSet<T> = HashSet<T, FxBuildHasher>;
93
94#[derive(Clone, Debug)]
104pub struct Schema {
105 pub(crate) entry: Arc<SchemaEntry>,
106 name_key: Arc<str>,
109}
110
111impl Schema {
112 pub fn new(name: &str, fields: Vec<crate::schema::FieldDef>) -> Self {
117 let name_key: Arc<str> = Arc::from(name);
118 Self {
119 entry: Arc::new(SchemaEntry {
120 name: name.to_string(),
121 has_timestamp: true,
122 fields,
123 annotations: Vec::new(),
124 }),
125 name_key,
126 }
127 }
128
129 pub fn from_entry(entry: SchemaEntry) -> Self {
131 let name_key: Arc<str> = Arc::from(entry.name.as_str());
132 Self {
133 entry: Arc::new(entry),
134 name_key,
135 }
136 }
137
138 pub fn name(&self) -> &str {
140 &self.entry.name
141 }
142
143 pub fn fields(&self) -> &[crate::schema::FieldDef] {
145 &self.entry.fields
146 }
147}
148
149#[derive(Clone, PartialEq, Eq, Hash)]
152enum SchemaKey {
153 Name(Arc<str>),
154 RustType(TypeId),
155}
156
157const DYNAMIC_SCHEMA_CACHE_LIMIT: usize = 1024;
168
169pub struct Encoder<W: Write = Vec<u8>> {
170 state: EncodeState<W>,
171 registry: SchemaRegistry,
172 string_pool: FxHashMap<String, u32>,
173 next_pool_id: u32,
174 stack_pool: FxHashMap<Box<[u64]>, u32>,
175 next_stack_pool_id: u32,
176 schema_ids: FxHashMap<SchemaKey, WireTypeId>,
177 dynamic_schema_cache: FxHashMap<usize, (Arc<SchemaEntry>, WireTypeId)>,
189 slot_cache: Vec<u32>,
192 registered_ids: [u64; (crate::STATIC_WIRE_ID_LIMIT as usize) / 64],
196}
197
198impl Default for Encoder<Vec<u8>> {
199 fn default() -> Self {
200 Self::new()
201 }
202}
203
204impl Encoder<Vec<u8>> {
205 pub fn new() -> Self {
206 let mut buf = Vec::new();
207 codec::encode_header(&mut buf).expect("Vec::write_all cannot fail");
208 Self {
209 state: EncodeState::new(buf),
210 registry: SchemaRegistry::new(),
211 string_pool: FxHashMap::default(),
212 next_pool_id: 0,
213 stack_pool: FxHashMap::default(),
214 next_stack_pool_id: 0,
215 schema_ids: FxHashMap::default(),
216 dynamic_schema_cache: FxHashMap::default(),
217 slot_cache: Vec::new(),
218 registered_ids: [0; (crate::STATIC_WIRE_ID_LIMIT as usize) / 64],
219 }
220 }
221
222 pub fn finish(self) -> Vec<u8> {
224 self.state.writer.into_inner()
225 }
226}
227
228impl<W: Write> Encoder<W> {
229 pub fn new_to(mut writer: W) -> io::Result<Self> {
232 codec::encode_header(&mut writer)?;
233 Ok(Self {
234 state: EncodeState::new(writer),
235 registry: SchemaRegistry::new(),
236 string_pool: FxHashMap::default(),
237 next_pool_id: 0,
238 stack_pool: FxHashMap::default(),
239 next_stack_pool_id: 0,
240 schema_ids: FxHashMap::default(),
241 dynamic_schema_cache: FxHashMap::default(),
242 slot_cache: Vec::new(),
243 registered_ids: [0; (crate::STATIC_WIRE_ID_LIMIT as usize) / 64],
244 })
245 }
246
247 pub(crate) fn from_decoder(
250 mut registry: SchemaRegistry,
251 string_pool: crate::decoder::StringPool,
252 stack_pool: crate::decoder::StackPool,
253 timestamp_base_ns: u64,
254 writer: W,
255 ) -> Self {
256 let mut pool = FxHashMap::default();
257 let mut next_pool_id: u32 = 0;
258 for (id, value) in string_pool.0.into_iter() {
259 pool.insert(value, id.raw_id());
260 if id.raw_id() >= next_pool_id {
261 next_pool_id = id.raw_id() + 1;
262 }
263 }
264
265 let mut new_stack_pool: FxHashMap<Box<[u64]>, u32> = FxHashMap::default();
266 let mut next_stack_pool_id: u32 = 0;
267 for (id, frames) in stack_pool.0.into_iter() {
268 new_stack_pool.insert(frames.into_boxed_slice(), id.raw_id());
269 if id.raw_id() >= next_stack_pool_id {
270 next_stack_pool_id = id.raw_id() + 1;
271 }
272 }
273
274 let mut schema_ids = FxHashMap::default();
275 for (wire_id, entry) in registry.entries() {
276 schema_ids.insert(SchemaKey::Name(Arc::from(entry.name.as_str())), wire_id);
277 }
278 registry.sync_next_id();
279
280 let mut state = EncodeState::new(writer);
281 state.set_ts_base_unchecked(timestamp_base_ns);
282
283 Self {
284 state,
285 registry,
286 string_pool: pool,
287 next_pool_id,
288 stack_pool: new_stack_pool,
289 next_stack_pool_id,
290 schema_ids,
291 dynamic_schema_cache: FxHashMap::default(),
292 slot_cache: Vec::new(),
293 registered_ids: [0; (crate::STATIC_WIRE_ID_LIMIT as usize) / 64],
294 }
295 }
296
297 pub fn into_inner(self) -> W {
299 self.state.writer.into_inner()
300 }
301
302 pub fn as_inner(&self) -> &W {
304 self.state.writer.inner()
305 }
306
307 pub fn bytes_written(&self) -> u64 {
309 self.state.writer.bytes_written()
310 }
311
312 pub fn reset_to(&mut self, mut new_writer: W) -> io::Result<W> {
315 codec::encode_header(&mut new_writer)?;
316 self.string_pool.clear();
317 self.next_pool_id = 0;
318 self.stack_pool.clear();
319 self.next_stack_pool_id = 0;
320 self.registry.clear();
321 self.schema_ids.clear();
322 self.dynamic_schema_cache.clear();
323 self.slot_cache.fill(0);
324 self.registered_ids.fill(0);
325 let old_state = std::mem::replace(&mut self.state, EncodeState::new(new_writer));
327 Ok(old_state.writer.into_inner())
328 }
329
330 fn ensure_registered(&mut self, schema: &Schema) -> io::Result<WireTypeId> {
336 let identity = Arc::as_ptr(&schema.entry) as usize;
337 if let Some((_, wire_id)) = self.dynamic_schema_cache.get(&identity) {
338 return Ok(*wire_id);
339 }
340 let wire_id = self.ensure_registered_slow(schema)?;
341 if self.dynamic_schema_cache.len() >= DYNAMIC_SCHEMA_CACHE_LIMIT {
342 self.dynamic_schema_cache.clear();
346 }
347 self.dynamic_schema_cache
348 .insert(identity, (Arc::clone(&schema.entry), wire_id));
349 Ok(wire_id)
350 }
351
352 fn ensure_registered_slow(&mut self, schema: &Schema) -> io::Result<WireTypeId> {
355 let key = SchemaKey::Name(Arc::clone(&schema.name_key));
356 if let Some(&wire_id) = self.schema_ids.get(&key) {
357 let Some(existing) = self.registry.get(wire_id) else {
359 return Err(io::Error::other(format!(
360 "corrupted internal state. {wire_id:?} in schema_ids but not in registry."
361 )));
362 };
363 if *existing == *schema.entry {
364 return Ok(wire_id);
365 }
366 return Err(io::Error::new(
367 io::ErrorKind::InvalidInput,
368 format!(
369 "schema already registered with different definition: {}",
370 schema.name()
371 ),
372 ));
373 }
374 let id = self.registry.next_type_id();
375 codec::encode_schema(id, &schema.entry, &mut self.state.writer)?;
376 if !schema.entry.annotations.is_empty() {
377 codec::encode_schema_annotations(
378 id,
379 &schema.entry.annotations,
380 &mut self.state.writer,
381 )?;
382 }
383 self.registry
384 .register(id, (*schema.entry).clone())
385 .expect("schema registration failed");
386 self.schema_ids.insert(key, id);
387 Ok(id)
388 }
389
390 pub fn register_schema(
400 &mut self,
401 name: &str,
402 fields: Vec<crate::schema::FieldDef>,
403 ) -> io::Result<Schema> {
404 let schema = Schema::new(name, fields);
405 self.ensure_registered(&schema)?;
406 Ok(schema)
407 }
408
409 pub fn register_existing(&mut self, schema: &Schema) -> io::Result<WireTypeId> {
414 self.ensure_registered(schema)
415 }
416
417 pub fn write_event(
432 &mut self,
433 schema: &Schema,
434 values: &[crate::types::FieldValue],
435 ) -> io::Result<()> {
436 use crate::types::FieldValue;
437
438 let type_id = self.ensure_registered(schema)?;
439 let expected_fields = schema.entry.fields.len();
440
441 let ts_ns = match values.first() {
442 Some(FieldValue::Varint(ns)) => *ns,
443 _ => {
444 return Err(io::Error::new(
445 io::ErrorKind::InvalidInput,
446 "first value must be FieldValue::Varint(timestamp_ns)",
447 ));
448 }
449 };
450 let field_values = &values[1..];
451
452 if field_values.len() != expected_fields {
453 return Err(io::Error::new(
454 io::ErrorKind::InvalidInput,
455 format!(
456 "value count ({}) does not match schema field count ({}) for schema '{}'",
457 field_values.len(),
458 expected_fields,
459 schema.name(),
460 ),
461 ));
462 }
463
464 let ts_delta = self.state.encode_timestamp_delta(ts_ns)?;
465 self.state.writer.write_all(&[codec::TAG_EVENT])?;
466 self.state.writer.write_all(&type_id.0.to_le_bytes())?;
467 codec::encode_u24_le(ts_delta, &mut self.state.writer)?;
468 let mut enc = EventEncoder::new(&mut self.state);
469 for (i, v) in field_values.iter().enumerate() {
470 enc.write_field_value(v, schema.entry.fields[i].field_type)?;
471 }
472 Ok(())
473 }
474
475 pub fn write<T: TraceEvent + 'static>(&mut self, event: &T) -> io::Result<()> {
478 let slot = T::type_slot();
479 let tid = if slot != 0 && slot < crate::STATIC_WIRE_ID_LIMIT {
480 let word = (slot >> 6) as usize;
481 let bit = 1u64 << (slot & 63);
482 if self.registered_ids[word] & bit == 0 {
483 self.register_fast_id::<T>(slot)?;
484 }
485 WireTypeId(slot)
486 } else {
487 let s = slot as usize;
488 let cached = self.slot_cache.get(s).copied().unwrap_or(0);
489 if cached != 0 {
490 WireTypeId((cached - 1) as u16)
491 } else {
492 self.resolve_dynamic_wire_id::<T>(s)?
493 }
494 };
495 let ts_ns = event.timestamp();
496 let ts_delta = self.state.encode_timestamp_delta(ts_ns)?;
497 self.state.writer.write_all(&[codec::TAG_EVENT])?;
498 self.state.writer.write_all(&tid.0.to_le_bytes())?;
499 codec::encode_u24_le(ts_delta, &mut self.state.writer)?;
500 let mut enc = EventEncoder::new(&mut self.state);
501 event.encode_fields(&mut enc)
502 }
503
504 #[cold]
508 fn resolve_dynamic_wire_id<T: TraceEvent + 'static>(
509 &mut self,
510 slot: usize,
511 ) -> io::Result<WireTypeId> {
512 let key = SchemaKey::RustType(TypeId::of::<T>());
513 let tid = if let Some(&existing) = self.schema_ids.get(&key) {
514 existing
515 } else {
516 let schema = Schema::from_entry(T::schema_entry());
517 let id = self.ensure_registered(&schema)?;
518 self.schema_ids.insert(key, id);
519 id
520 };
521 if slot != 0 {
522 if self.slot_cache.len() <= slot {
523 self.slot_cache.resize(slot + 1, 0);
524 }
525 self.slot_cache[slot] = (tid.0 as u32) + 1;
526 }
527 Ok(tid)
528 }
529
530 #[cold]
533 fn register_fast_id<T: TraceEvent + 'static>(&mut self, id: u16) -> io::Result<()> {
534 let entry = T::schema_entry();
535 let wire = WireTypeId(id);
536 codec::encode_schema(wire, &entry, &mut self.state.writer)?;
537 if !entry.annotations.is_empty() {
538 codec::encode_schema_annotations(wire, &entry.annotations, &mut self.state.writer)?;
539 }
540 self.registry.register(wire, entry).map_err(|e| {
541 io::Error::new(
542 io::ErrorKind::InvalidInput,
543 format!("wire id {id} collision: {e}"),
544 )
545 })?;
546 self.registered_ids[(id >> 6) as usize] |= 1u64 << (id & 63);
548 Ok(())
549 }
550
551 pub fn intern_string(&mut self, s: &str) -> io::Result<InternedString> {
553 if let Some(&id) = self.string_pool.get(s) {
554 return Ok(InternedString(id));
555 }
556 let id = self.next_pool_id;
557 self.next_pool_id += 1;
558 self.string_pool.insert(s.to_string(), id);
559 codec::encode_string_pool(
560 &[PoolEntry {
561 pool_id: id,
562 data: s.as_bytes().to_vec(),
563 }],
564 &mut self.state.writer,
565 )?;
566 Ok(InternedString(id))
567 }
568
569 pub fn write_string_pool(&mut self, entries: &[PoolEntry]) -> io::Result<()> {
570 codec::encode_string_pool(entries, &mut self.state.writer)
571 }
572
573 pub fn intern_stack_frames(&mut self, frames: &[u64]) -> io::Result<InternedStackFrames> {
576 if let Some(&id) = self.stack_pool.get(frames) {
577 return Ok(InternedStackFrames(id));
578 }
579 let id = self.next_stack_pool_id;
580 self.next_stack_pool_id += 1;
581 self.stack_pool.insert(frames.into(), id);
582 codec::encode_stack_pool(
583 &[StackPoolEntry {
584 pool_id: id,
585 frames: frames.to_vec(),
588 }],
589 &mut self.state.writer,
590 )?;
591 Ok(InternedStackFrames(id))
592 }
593
594 pub fn write_stack_pool(&mut self, entries: &[StackPoolEntry]) -> io::Result<()> {
595 codec::encode_stack_pool(entries, &mut self.state.writer)
596 }
597
598 pub fn flush(&mut self) -> io::Result<()> {
600 self.state.writer.flush()
601 }
602
603 pub fn into_raw_encoder(self) -> RawEncoder<W> {
610 RawEncoder {
611 writer: self.state.writer,
612 }
613 }
614}
615
616pub struct RawEncoder<W> {
623 writer: CountingWriter<W>,
624}
625
626impl<W: Write> RawEncoder<W> {
627 pub fn write_raw(&mut self, bytes: &[u8]) -> io::Result<()> {
629 self.writer.write_all(bytes)
630 }
631
632 pub fn bytes_written(&self) -> u64 {
635 self.writer.bytes_written()
636 }
637
638 pub fn flush(&mut self) -> io::Result<()> {
640 self.writer.flush()
641 }
642
643 pub fn into_inner(self) -> W {
645 self.writer.into_inner()
646 }
647}
648
649impl Encoder<Vec<u8>> {
650 pub fn write_infallible<T: TraceEvent + 'static>(&mut self, event: &T) {
651 self.write(event).expect("writing to Vec<u8> is infallible")
652 }
653
654 pub fn intern_string_infallible(&mut self, s: &str) -> InternedString {
655 self.intern_string(s)
656 .expect("interning into Vec<u8> is infallible")
657 }
658
659 pub fn intern_stack_frames_infallible(&mut self, frames: &[u64]) -> InternedStackFrames {
660 self.intern_stack_frames(frames)
661 .expect("interning into Vec<u8> is infallible")
662 }
663
664 pub fn reset_to_infallible(&mut self, data: Vec<u8>) -> Vec<u8> {
666 self.reset_to(data)
667 .expect("writing to Vec<u8> is infallible")
668 }
669}
670
671#[cfg(test)]
672mod tests {
673 use super::*;
674 use crate::schema::FieldDef;
675 use crate::types::{FieldType, FieldValue};
676
677 #[test]
678 fn encoder_writes_header() {
679 let enc = Encoder::new();
680 let data = enc.finish();
681 assert_eq!(&data[..5], &[0x54, 0x52, 0x43, 0x00, 1]);
682 }
683
684 #[test]
685 fn encoder_register_and_write_event() {
686 let mut enc = Encoder::new();
687 let schema = enc
688 .register_schema(
689 "Ev",
690 vec![FieldDef {
691 name: "v".into(),
692 field_type: FieldType::Varint,
693 }],
694 )
695 .unwrap();
696 enc.write_event(
697 &schema,
698 &[FieldValue::Varint(1_000), FieldValue::Varint(42)],
699 )
700 .unwrap();
701 let data = enc.finish();
702 assert!(data.len() > 5);
703 }
704
705 #[test]
706 fn idempotent_re_registration() {
707 let mut enc = Encoder::new();
708 let fields = vec![FieldDef {
709 name: "v".into(),
710 field_type: FieldType::Varint,
711 }];
712 let _s1 = enc.register_schema("Ev", fields.clone()).unwrap();
713 let _s2 = enc.register_schema("Ev", fields).unwrap();
714 }
716
717 #[test]
718 fn re_registration_different_schema_errors() {
719 let mut enc = Encoder::new();
720 enc.register_schema(
721 "Ev",
722 vec![FieldDef {
723 name: "v".into(),
724 field_type: FieldType::Varint,
725 }],
726 )
727 .unwrap();
728 let result = enc.register_schema(
729 "Ev",
730 vec![FieldDef {
731 name: "different".into(),
732 field_type: FieldType::Bool,
733 }],
734 );
735 assert!(result.is_err());
736 }
737
738 #[test]
739 fn schema_auto_registers_on_write() {
740 use crate::decoder::{DecodedFrame, Decoder};
741
742 let schema = Schema::new(
744 "Lazy",
745 vec![FieldDef {
746 name: "v".into(),
747 field_type: FieldType::Varint,
748 }],
749 );
750
751 let mut enc = Encoder::new();
753 enc.write_event(
754 &schema,
755 &[FieldValue::Varint(1_000), FieldValue::Varint(42)],
756 )
757 .unwrap();
758
759 let bytes = enc.finish();
760 let mut dec = Decoder::new(&bytes).unwrap();
761 let frames = dec.decode_all();
762 assert!(matches!(&frames[0], DecodedFrame::Schema(s) if s.name == "Lazy"));
763 if let DecodedFrame::Event { values, .. } = &frames[1] {
764 assert_eq!(*values, vec![FieldValue::Varint(42)]);
765 } else {
766 panic!("expected event");
767 }
768 }
769
770 #[test]
771 fn schema_portable_across_encoders() {
772 use crate::decoder::{DecodedFrame, Decoder};
773
774 let mut enc1 = Encoder::new();
775 let schema = enc1
776 .register_schema(
777 "Shared",
778 vec![FieldDef {
779 name: "v".into(),
780 field_type: FieldType::Varint,
781 }],
782 )
783 .unwrap();
784 enc1.write_event(&schema, &[FieldValue::Varint(1_000), FieldValue::Varint(1)])
785 .unwrap();
786
787 let mut enc2 = Encoder::new();
789 enc2.write_event(&schema, &[FieldValue::Varint(2_000), FieldValue::Varint(2)])
790 .unwrap();
791
792 for (enc, expected_val) in [(enc1, 1u64), (enc2, 2u64)] {
794 let bytes = enc.finish();
795 let mut dec = Decoder::new(&bytes).unwrap();
796 let frames = dec.decode_all();
797 let event = frames
798 .iter()
799 .find(|f| matches!(f, DecodedFrame::Event { .. }))
800 .unwrap();
801 if let DecodedFrame::Event { values, .. } = event {
802 assert_eq!(values[0], FieldValue::Varint(expected_val));
803 }
804 }
805 }
806
807 #[test]
808 fn encoder_intern_string_deduplicates() {
809 let mut enc = Encoder::new();
810 let id1 = enc.intern_string("hello").unwrap();
811 let id2 = enc.intern_string("hello").unwrap();
812 let id3 = enc.intern_string("world").unwrap();
813 assert_eq!(id1, id2);
814 assert_ne!(id1, id3);
815 }
816
817 #[test]
818 fn encoder_intern_stack_frames_deduplicates() {
819 let mut enc = Encoder::new();
820 let stack_a: &[u64] = &[0x1000, 0x2000, 0x3000];
821 let stack_b: &[u64] = &[0x4000, 0x5000];
822 let id1 = enc.intern_stack_frames(stack_a).unwrap();
823 let id2 = enc.intern_stack_frames(stack_a).unwrap();
824 let id3 = enc.intern_stack_frames(stack_b).unwrap();
825 assert_eq!(id1, id2);
826 assert_ne!(id1, id3);
827 }
828
829 #[test]
830 fn stack_pool_round_trip_via_decoder() {
831 use crate::decoder::Decoder;
832 use crate::types::InternedStackFrames;
833
834 let mut enc = Encoder::new();
835 let stack_a: &[u64] = &[0xdead, 0xbeef, 0xcafe];
836 let stack_b: &[u64] = &[0x1, 0x2];
837 let id_a = enc.intern_stack_frames(stack_a).unwrap();
838 let id_b = enc.intern_stack_frames(stack_b).unwrap();
839 let bytes = enc.finish();
840
841 let mut dec = Decoder::new(&bytes).unwrap();
842 let _ = dec.decode_all();
843 assert_eq!(
844 dec.stack_pool().get(InternedStackFrames(id_a.raw_id())),
845 Some(stack_a)
846 );
847 assert_eq!(
848 dec.stack_pool().get(InternedStackFrames(id_b.raw_id())),
849 Some(stack_b)
850 );
851 }
852
853 #[test]
854 fn for_each_event_populates_stack_pool() {
855 use crate::decoder::Decoder;
856 use crate::schema::FieldDef;
857 use crate::types::{FieldType, FieldValue, InternedStackFrames};
858
859 let mut enc = Encoder::new();
860 let schema = enc
861 .register_schema(
862 "CpuSampleEvent",
863 vec![FieldDef {
864 name: "callchain".into(),
865 field_type: FieldType::PooledStackFrames,
866 }],
867 )
868 .unwrap();
869 let stack: &[u64] = &[0x1234, 0x5678, 0x9abc];
870 let id = enc.intern_stack_frames(stack).unwrap();
871 enc.write_event(
872 &schema,
873 &[
874 FieldValue::Varint(1_000_000),
875 FieldValue::PooledStackFrames(id),
876 ],
877 )
878 .unwrap();
879 let bytes = enc.finish();
880
881 let mut dec = Decoder::new(&bytes).unwrap();
882 let mut event_count = 0;
883 dec.for_each_event(|_ev| {
884 event_count += 1;
885 })
886 .unwrap();
887 assert_eq!(event_count, 1);
888 assert_eq!(
889 dec.stack_pool().get(InternedStackFrames(id.raw_id())),
890 Some(stack),
891 );
892 }
893
894 #[test]
895 fn encoder_intern_empty_stack_frames() {
896 use crate::decoder::Decoder;
897 use crate::types::InternedStackFrames;
898
899 let mut enc = Encoder::new();
900 let id1 = enc.intern_stack_frames(&[]).unwrap();
901 let id2 = enc.intern_stack_frames(&[]).unwrap();
902 assert_eq!(id1, id2);
903 let bytes = enc.finish();
904
905 let mut dec = Decoder::new(&bytes).unwrap();
906 let _ = dec.decode_all();
907 assert_eq!(
908 dec.stack_pool().get(InternedStackFrames(id1.raw_id())),
909 Some(&[][..])
910 );
911 }
912
913 #[test]
914 fn write_stack_pool_multi_entry_round_trip() {
915 use crate::decoder::Decoder;
916 use crate::types::InternedStackFrames;
917
918 let mut enc = Encoder::new();
919 let entries = vec![
920 StackPoolEntry {
921 pool_id: 0,
922 frames: vec![0xaaaa, 0xbbbb, 0xcccc],
923 },
924 StackPoolEntry {
925 pool_id: 1,
926 frames: vec![0x1111],
927 },
928 StackPoolEntry {
929 pool_id: 2,
930 frames: vec![],
931 },
932 ];
933 enc.write_stack_pool(&entries).unwrap();
934 let bytes = enc.finish();
935
936 let mut dec = Decoder::new(&bytes).unwrap();
937 let _ = dec.decode_all();
938 assert_eq!(
939 dec.stack_pool().get(InternedStackFrames(0)),
940 Some(&[0xaaaa, 0xbbbb, 0xcccc][..])
941 );
942 assert_eq!(
943 dec.stack_pool().get(InternedStackFrames(1)),
944 Some(&[0x1111][..])
945 );
946 assert_eq!(dec.stack_pool().get(InternedStackFrames(2)), Some(&[][..]));
947 }
948
949 #[test]
950 fn decoder_into_encoder_deduplicates_interned_stack_frames() {
951 use crate::decoder::Decoder;
952
953 let mut enc = Encoder::new();
954 let id1 = enc.intern_stack_frames(&[0x10, 0x20]).unwrap();
955 let base = enc.finish();
956
957 let mut decoder = Decoder::new(&base).unwrap();
958 while decoder.next_frame_ref().ok().flatten().is_some() {}
959 let mut output = Vec::new();
960 let mut ext = decoder.into_encoder(&mut output);
961 let id2 = ext.intern_stack_frames(&[0x10, 0x20]).unwrap();
962 let id3 = ext.intern_stack_frames(&[0x30]).unwrap();
963 assert_eq!(id1.raw_id(), id2.raw_id());
964 assert_ne!(id2.raw_id(), id3.raw_id());
965 }
966
967 #[test]
968 fn timestamp_round_trip() {
969 use crate::decoder::{DecodedFrame, Decoder};
970
971 let mut enc = Encoder::new();
972 let schema = enc
973 .register_schema(
974 "TS",
975 vec![FieldDef {
976 name: "v".into(),
977 field_type: FieldType::Varint,
978 }],
979 )
980 .unwrap();
981
982 let ts1 = 100_000u64;
983 let ts2 = 50_000u64;
984 let ts3 = 200_000_000u64;
985 let ts4 = 100_000_000u64;
986 enc.write_event(&schema, &[FieldValue::Varint(ts1), FieldValue::Varint(1)])
987 .unwrap();
988 enc.write_event(&schema, &[FieldValue::Varint(ts2), FieldValue::Varint(2)])
989 .unwrap();
990 enc.write_event(&schema, &[FieldValue::Varint(ts3), FieldValue::Varint(3)])
991 .unwrap();
992 enc.write_event(&schema, &[FieldValue::Varint(ts4), FieldValue::Varint(4)])
993 .unwrap();
994
995 let bytes = enc.finish();
996 let mut dec = Decoder::new(&bytes).unwrap();
997 let events: Vec<_> = dec
998 .decode_all()
999 .into_iter()
1000 .filter_map(|f| match f {
1001 DecodedFrame::Event {
1002 timestamp_ns,
1003 values,
1004 ..
1005 } => Some((timestamp_ns, values)),
1006 _ => None,
1007 })
1008 .collect();
1009
1010 assert_eq!(events.len(), 4);
1011 assert_eq!(events[0].0, Some(ts1));
1012 assert_eq!(events[0].1, vec![FieldValue::Varint(1)]);
1013 assert_eq!(events[1].0, Some(ts2));
1014 assert_eq!(events[1].1, vec![FieldValue::Varint(2)]);
1015 assert_eq!(events[2].0, Some(ts3));
1016 assert_eq!(events[2].1, vec![FieldValue::Varint(3)]);
1017 assert_eq!(events[3].0, Some(ts4));
1018 assert_eq!(events[3].1, vec![FieldValue::Varint(4)]);
1019 }
1020
1021 #[test]
1022 fn encoder_new_to_writer() {
1023 let mut buf = Vec::new();
1024 let enc = Encoder::new_to(&mut buf).unwrap();
1025 drop(enc);
1026 assert!(buf.len() >= 5);
1027 assert_eq!(&buf[..5], &[0x54, 0x52, 0x43, 0x00, 1]);
1028 }
1029
1030 #[test]
1031 fn decoder_into_encoder_appends_without_header() {
1032 use crate::decoder::{DecodedFrame, Decoder};
1033
1034 let mut enc = Encoder::new();
1036 let schema = enc
1037 .register_schema(
1038 "Ev",
1039 vec![FieldDef {
1040 name: "v".into(),
1041 field_type: FieldType::Varint,
1042 }],
1043 )
1044 .unwrap();
1045 enc.write_event(&schema, &[FieldValue::Varint(1_000), FieldValue::Varint(1)])
1046 .unwrap();
1047 let base = enc.finish();
1048
1049 let mut decoder = Decoder::new(&base).unwrap();
1051 while decoder.next_frame_ref().ok().flatten().is_some() {}
1052 let mut output = Vec::new();
1053 let mut ext = decoder.into_encoder(&mut output);
1054 ext.write_event(&schema, &[FieldValue::Varint(2_000), FieldValue::Varint(2)])
1056 .unwrap();
1057 drop(ext);
1058
1059 let mut combined = base.clone();
1061 combined.extend_from_slice(&output);
1062 let mut dec = Decoder::new(&combined).unwrap();
1063 let events: Vec<_> = dec
1064 .decode_all()
1065 .into_iter()
1066 .filter_map(|f| match f {
1067 DecodedFrame::Event {
1068 timestamp_ns,
1069 values,
1070 ..
1071 } => Some((timestamp_ns, values)),
1072 _ => None,
1073 })
1074 .collect();
1075 assert_eq!(events.len(), 2);
1076 assert_eq!(events[0].0, Some(1_000));
1077 assert_eq!(events[1].0, Some(2_000));
1078 }
1079
1080 #[test]
1081 fn decoder_into_encoder_deduplicates_interned_strings() {
1082 use crate::decoder::{DecodedFrame, Decoder};
1083
1084 let mut enc = Encoder::new();
1086 let id1 = enc.intern_string("hello").unwrap();
1087 let base = enc.finish();
1088
1089 let mut decoder = Decoder::new(&base).unwrap();
1091 while decoder.next_frame_ref().ok().flatten().is_some() {}
1092 let mut output = Vec::new();
1093 let mut ext = decoder.into_encoder(&mut output);
1094 let id2 = ext.intern_string("hello").unwrap();
1096 let id3 = ext.intern_string("world").unwrap();
1097 drop(ext);
1098
1099 assert_eq!(id1, id2, "existing string should reuse pool ID");
1100 assert_ne!(id2, id3);
1101
1102 let mut combined = base.clone();
1104 combined.extend_from_slice(&output);
1105 let mut dec = Decoder::new(&combined).unwrap();
1106 let frames = dec.decode_all();
1107 let pool_frames: Vec<_> = frames
1108 .iter()
1109 .filter(|f| matches!(f, DecodedFrame::StringPool(_)))
1110 .collect();
1111 assert_eq!(pool_frames.len(), 2);
1113 }
1114
1115 struct FastSlot {
1118 ts: u64,
1119 }
1120 impl TraceEvent for FastSlot {
1121 fn type_slot() -> u16 {
1122 5
1123 }
1124 fn event_name() -> &'static str {
1125 "FastSlot"
1126 }
1127 fn field_defs() -> Vec<FieldDef> {
1128 Vec::new()
1129 }
1130 fn timestamp(&self) -> u64 {
1131 self.ts
1132 }
1133 fn encode_fields<W: Write>(&self, _enc: &mut EventEncoder<'_, W>) -> io::Result<()> {
1134 Ok(())
1135 }
1136 }
1137
1138 #[test]
1139 fn fast_slot_registers_at_slot_id() {
1140 use crate::decoder::Decoder;
1141
1142 let mut enc = Encoder::new();
1143 enc.write(&FastSlot { ts: 1_000_000 }).unwrap();
1144 let dynamic = enc
1146 .register_schema(
1147 "Dyn",
1148 vec![FieldDef {
1149 name: "v".into(),
1150 field_type: FieldType::Varint,
1151 }],
1152 )
1153 .unwrap();
1154 enc.write_event(
1155 &dynamic,
1156 &[FieldValue::Varint(2_000), FieldValue::Varint(1)],
1157 )
1158 .unwrap();
1159 let bytes = enc.finish();
1160
1161 let mut dec = Decoder::new(&bytes).unwrap();
1162 let _ = dec.decode_all();
1163 assert_eq!(
1165 dec.registry().get(WireTypeId(5)).unwrap().name(),
1166 "FastSlot"
1167 );
1168 assert_eq!(
1170 dec.registry()
1171 .get(WireTypeId(crate::STATIC_WIRE_ID_LIMIT))
1172 .unwrap()
1173 .name(),
1174 "Dyn"
1175 );
1176 }
1177
1178 #[test]
1179 fn register_and_write() {
1180 use crate::decoder::{DecodedFrame, Decoder};
1181
1182 let mut enc = Encoder::new();
1183 let schema = enc
1184 .register_schema(
1185 "MyEvent",
1186 vec![
1187 FieldDef {
1188 name: "count".into(),
1189 field_type: FieldType::Varint,
1190 },
1191 FieldDef {
1192 name: "name".into(),
1193 field_type: FieldType::String,
1194 },
1195 ],
1196 )
1197 .unwrap();
1198
1199 enc.write_event(
1200 &schema,
1201 &[
1202 FieldValue::Varint(1_000_000),
1203 FieldValue::Varint(42),
1204 FieldValue::String("hello".into()),
1205 ],
1206 )
1207 .unwrap();
1208
1209 let bytes = enc.finish();
1210 let mut dec = Decoder::new(&bytes).unwrap();
1211 let frames = dec.decode_all();
1212 let events: Vec<_> = frames
1213 .into_iter()
1214 .filter_map(|f| match f {
1215 DecodedFrame::Event {
1216 timestamp_ns,
1217 values,
1218 ..
1219 } => Some((timestamp_ns, values)),
1220 _ => None,
1221 })
1222 .collect();
1223 assert_eq!(events.len(), 1);
1224 assert_eq!(events[0].0, Some(1_000_000));
1225 assert_eq!(events[0].1[0], FieldValue::Varint(42));
1226 assert_eq!(events[0].1[1], FieldValue::String("hello".into()));
1227 }
1228
1229 #[test]
1230 fn register_conflict_errors() {
1231 let mut enc = Encoder::new();
1232 enc.register_schema(
1233 "Ev",
1234 vec![FieldDef {
1235 name: "v".into(),
1236 field_type: FieldType::Varint,
1237 }],
1238 )
1239 .unwrap();
1240 let result = enc.register_schema(
1241 "Ev",
1242 vec![FieldDef {
1243 name: "other".into(),
1244 field_type: FieldType::Bool,
1245 }],
1246 );
1247 assert!(result.is_err());
1248 }
1249
1250 #[test]
1251 fn write_wrong_field_count_errors() {
1252 let mut enc = Encoder::new();
1253 let schema = enc
1254 .register_schema(
1255 "Ev",
1256 vec![FieldDef {
1257 name: "v".into(),
1258 field_type: FieldType::Varint,
1259 }],
1260 )
1261 .unwrap();
1262 let result = enc.write_event(
1264 &schema,
1265 &[
1266 FieldValue::Varint(0),
1267 FieldValue::Varint(1),
1268 FieldValue::Varint(2),
1269 ],
1270 );
1271 assert!(result.is_err());
1272 }
1273
1274 #[test]
1277 fn timestamp_base_advances_per_event() {
1278 use crate::decoder::{DecodedFrame, Decoder};
1279
1280 let mut enc = Encoder::new();
1281 let schema = enc
1282 .register_schema(
1283 "Ev",
1284 vec![FieldDef {
1285 name: "v".into(),
1286 field_type: FieldType::Varint,
1287 }],
1288 )
1289 .unwrap();
1290
1291 let ts1 = 12_000_000u64;
1292 let ts2 = 24_000_000u64;
1293 enc.write_event(&schema, &[FieldValue::Varint(ts1), FieldValue::Varint(1)])
1294 .unwrap();
1295 enc.write_event(&schema, &[FieldValue::Varint(ts2), FieldValue::Varint(2)])
1296 .unwrap();
1297
1298 let bytes = enc.finish();
1299
1300 let reset_count = bytes.iter().filter(|&&b| b == 0x05).count();
1301 assert_eq!(
1302 reset_count, 0,
1303 "base should advance per event, avoiding unnecessary resets"
1304 );
1305
1306 let mut dec = Decoder::new(&bytes).unwrap();
1307 let events: Vec<_> = dec
1308 .decode_all()
1309 .into_iter()
1310 .filter_map(|f| match f {
1311 DecodedFrame::Event { timestamp_ns, .. } => timestamp_ns,
1312 _ => None,
1313 })
1314 .collect();
1315 assert_eq!(events, vec![ts1, ts2]);
1316 }
1317
1318 #[test]
1319 fn reset_to_preserves_capacity() {
1320 let mut enc = Encoder::new();
1321 for i in 0..100 {
1322 enc.intern_string(&format!("string_{}", i)).unwrap();
1323 }
1324 let cap_before = enc.string_pool.capacity();
1325 let _bytes = enc.reset_to(Vec::new());
1326 let cap_after = enc.string_pool.capacity();
1327 assert_eq!(
1328 cap_before, cap_after,
1329 "string_pool capacity should be preserved after reset_to"
1330 );
1331 }
1332
1333 #[test]
1334 fn reset_to_returns_old_data_and_clears_state() {
1335 use crate::decoder::{DecodedFrame, Decoder};
1336
1337 let mut enc = Encoder::new();
1338 let schema = enc
1339 .register_schema(
1340 "Ev",
1341 vec![FieldDef {
1342 name: "v".into(),
1343 field_type: FieldType::Varint,
1344 }],
1345 )
1346 .unwrap();
1347 enc.write_event(
1348 &schema,
1349 &[FieldValue::Varint(1_000), FieldValue::Varint(42)],
1350 )
1351 .unwrap();
1352 let _s = enc.intern_string("hello").unwrap();
1353
1354 let old_bytes_written = enc.bytes_written();
1355 assert!(old_bytes_written > 0);
1356
1357 let old = enc.reset_to_infallible(Vec::new());
1359
1360 let mut dec = Decoder::new(&old).unwrap();
1362 let frames = dec.decode_all();
1363 assert!(frames.iter().any(|f| matches!(f, DecodedFrame::Schema(_))));
1364 assert!(
1365 frames
1366 .iter()
1367 .any(|f| matches!(f, DecodedFrame::Event { .. }))
1368 );
1369 assert!(
1370 frames
1371 .iter()
1372 .any(|f| matches!(f, DecodedFrame::StringPool(_)))
1373 );
1374
1375 assert!(
1377 enc.bytes_written() < old_bytes_written,
1378 "bytes_written should reset (got {} vs old {})",
1379 enc.bytes_written(),
1380 old_bytes_written
1381 );
1382
1383 enc.write_event(
1386 &schema,
1387 &[FieldValue::Varint(2_000), FieldValue::Varint(99)],
1388 )
1389 .unwrap();
1390
1391 let _s2 = enc.intern_string("hello").unwrap();
1393
1394 let new_bytes = enc.reset_to_infallible(Vec::new());
1396 let mut dec2 = Decoder::new(&new_bytes).unwrap();
1397 let new_frames = dec2.decode_all();
1398 assert!(
1400 new_frames
1401 .iter()
1402 .any(|f| matches!(f, DecodedFrame::Schema(s) if s.name == "Ev")),
1403 "new trace must contain schema definition"
1404 );
1405 assert!(
1407 new_frames
1408 .iter()
1409 .any(|f| matches!(f, DecodedFrame::StringPool(_))),
1410 "new trace must contain string pool"
1411 );
1412 let event = new_frames
1414 .iter()
1415 .find_map(|f| match f {
1416 DecodedFrame::Event {
1417 timestamp_ns,
1418 values,
1419 ..
1420 } => Some((timestamp_ns, values)),
1421 _ => None,
1422 })
1423 .expect("new trace must contain event");
1424 assert_eq!(*event.0, Some(2_000));
1425 assert_eq!(event.1[0], FieldValue::Varint(99));
1426 }
1427
1428 #[test]
1429 fn into_raw_encoder_preserves_byte_count() {
1430 let mut enc = Encoder::new();
1431 let schema = enc
1432 .register_schema(
1433 "Ev",
1434 vec![FieldDef {
1435 name: "v".into(),
1436 field_type: FieldType::Varint,
1437 }],
1438 )
1439 .unwrap();
1440 enc.write_event(
1441 &schema,
1442 &[FieldValue::Varint(1_000), FieldValue::Varint(42)],
1443 )
1444 .unwrap();
1445
1446 let bytes_before = enc.bytes_written();
1447 assert!(bytes_before > 0);
1448
1449 let raw = enc.into_raw_encoder();
1450 assert_eq!(
1451 raw.bytes_written(),
1452 bytes_before,
1453 "byte count must be preserved across conversion"
1454 );
1455 }
1456
1457 #[test]
1458 fn raw_encoder_write_raw_and_bytes_written() {
1459 let enc = Encoder::new();
1460 let initial = enc.bytes_written();
1461 let mut raw = enc.into_raw_encoder();
1462
1463 let payload = [0xAA; 100];
1464 raw.write_raw(&payload).unwrap();
1465
1466 assert_eq!(
1467 raw.bytes_written(),
1468 initial + payload.len() as u64,
1469 "bytes_written must include raw payload"
1470 );
1471 }
1472
1473 #[test]
1474 fn raw_encoder_into_inner_returns_all_data() {
1475 use crate::decoder::{DecodedFrame, Decoder};
1476
1477 let mut enc = Encoder::new();
1480 let schema = enc
1481 .register_schema(
1482 "Ev",
1483 vec![FieldDef {
1484 name: "v".into(),
1485 field_type: FieldType::Varint,
1486 }],
1487 )
1488 .unwrap();
1489 enc.write_event(&schema, &[FieldValue::Varint(1_000), FieldValue::Varint(1)])
1490 .unwrap();
1491
1492 let raw_batch = {
1494 let mut batch_enc = Encoder::new();
1495 batch_enc
1496 .write_event(&schema, &[FieldValue::Varint(2_000), FieldValue::Varint(2)])
1497 .unwrap();
1498 batch_enc.finish()
1499 };
1500
1501 let mut raw = enc.into_raw_encoder();
1502 raw.write_raw(&raw_batch).unwrap();
1503 let combined = raw.into_inner();
1504
1505 let mut dec = Decoder::new(&combined).unwrap();
1506 let events: Vec<_> = dec
1507 .decode_all()
1508 .into_iter()
1509 .filter_map(|f| match f {
1510 DecodedFrame::Event {
1511 timestamp_ns,
1512 values,
1513 ..
1514 } => Some((timestamp_ns, values)),
1515 _ => None,
1516 })
1517 .collect();
1518
1519 assert_eq!(events.len(), 2);
1520 assert_eq!(events[0].0, Some(1_000));
1521 assert_eq!(events[0].1, vec![FieldValue::Varint(1)]);
1522 assert_eq!(events[1].0, Some(2_000));
1523 assert_eq!(events[1].1, vec![FieldValue::Varint(2)]);
1524 }
1525}
1526
1527#[cfg(test)]
1528mod dynamic_schema_cache_tests {
1529 use super::*;
1530 use crate::schema::{FieldDef, SchemaEntry};
1531 use crate::types::{FieldType, FieldValue};
1532
1533 fn schema(name: &str) -> Schema {
1534 Schema::from_entry(SchemaEntry::new(
1535 name,
1536 true,
1537 vec![FieldDef::new("v", FieldType::Varint)],
1538 ))
1539 }
1540
1541 #[test]
1545 fn identity_cache_agrees_with_name_registration() {
1546 let mut enc = Encoder::new();
1547 let a = schema("Ev");
1548 let id1 = enc.ensure_registered(&a).unwrap();
1549 let id2 = enc.ensure_registered(&a).unwrap();
1550 assert_eq!(id1, id2, "same handle must reuse its wire id");
1551
1552 let b = schema("Ev"); let id3 = enc.ensure_registered(&b).unwrap();
1554 assert_eq!(id1, id3, "same name must resolve to the same wire id");
1555 assert_eq!(enc.dynamic_schema_cache.len(), 2);
1556 }
1557
1558 #[test]
1561 fn identity_cache_is_bounded() {
1562 let mut enc = Encoder::new();
1563 for i in 0..(DYNAMIC_SCHEMA_CACHE_LIMIT * 2 + 7) {
1564 let s = schema(&format!("Ev{}", i % 3));
1567 enc.ensure_registered(&s).unwrap();
1568 assert!(
1569 enc.dynamic_schema_cache.len() <= DYNAMIC_SCHEMA_CACHE_LIMIT,
1570 "cache exceeded its bound at iteration {i}"
1571 );
1572 }
1573 }
1574
1575 #[test]
1578 fn events_across_cache_clears_decode() {
1579 let mut enc = Encoder::new();
1580 for i in 0..(DYNAMIC_SCHEMA_CACHE_LIMIT + 3) {
1581 let s = schema("Ev");
1582 enc.write_event(
1584 &s,
1585 &[FieldValue::Varint(i as u64), FieldValue::Varint(i as u64)],
1586 )
1587 .unwrap();
1588 }
1589 let data = enc.finish();
1590 let mut decoder = crate::decoder::Decoder::new(&data).unwrap();
1591 let mut count = 0u64;
1592 decoder
1593 .for_each_event(|ev| {
1594 assert_eq!(ev.name, "Ev");
1595 count += 1;
1596 })
1597 .unwrap();
1598 assert_eq!(count, (DYNAMIC_SCHEMA_CACHE_LIMIT + 3) as u64);
1599 }
1600}