1use std::{
2 borrow::Cow,
3 convert::Infallible,
4 fmt::{Display, Formatter},
5};
6
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(transparent)]
15pub struct DeltaEncodedTimestamp(pub i64);
16
17pub const TRACE_HEADER: &[u8] = b"TRACEv1";
21
22pub const TRACE_HEADER_PREFIX: &[u8] = b"TRACEv";
24
25pub fn check_trace_header(data: &[u8]) -> anyhow::Result<()> {
28 if data.starts_with(TRACE_HEADER) {
29 return Ok(());
30 }
31 let expected = String::from_utf8_lossy(TRACE_HEADER);
32 if data.starts_with(TRACE_HEADER_PREFIX) {
33 let found = &data[..data.len().min(TRACE_HEADER.len())];
34 anyhow::bail!(
35 "Unsupported trace file version: expected {expected}, found {}",
36 String::from_utf8_lossy(found)
37 );
38 }
39 anyhow::bail!("Not a trace file: missing {expected} header");
40}
41
42#[derive(Debug, Serialize, Deserialize)]
62pub enum TraceRow<'a, T = DeltaEncodedTimestamp> {
63 Start {
65 ts: T,
67 id: u64,
69 parent: Option<u64>,
71 #[serde(borrow)]
73 name: Cow<'a, str>,
74 #[serde(borrow)]
76 target: Cow<'a, str>,
77 #[serde(borrow)]
79 values: Vec<(Cow<'a, str>, TraceValue<'a>)>,
80 },
81 End {
83 ts: T,
85 id: u64,
87 },
88 Enter {
90 ts: T,
92 id: u64,
94 thread_id: u64,
96 allocations: Option<Allocations>,
99 },
100 Exit {
102 ts: T,
104 id: u64,
106 thread_id: u64,
108 allocations: Option<Allocations>,
111 },
112 Event {
114 ts: T,
116 parent: Option<u64>,
118 #[serde(borrow)]
120 values: Vec<(Cow<'a, str>, TraceValue<'a>)>,
121 },
122 Record {
124 id: u64,
126 #[serde(borrow)]
128 values: Vec<(Cow<'a, str>, TraceValue<'a>)>,
129 },
130 MemorySample {
132 ts: T,
134 memory: u64,
136 memory_pressure: u8,
140 memory_footprint: u64,
146 active_worker_threads: u64,
148 },
149 TimestampBase {
152 ts: u64,
154 },
155}
156
157impl<'a, T> TraceRow<'a, T> {
158 pub fn timestamp(&self) -> Option<&T> {
161 match self {
162 TraceRow::Start { ts, .. }
163 | TraceRow::End { ts, .. }
164 | TraceRow::Enter { ts, .. }
165 | TraceRow::Exit { ts, .. }
166 | TraceRow::Event { ts, .. }
167 | TraceRow::MemorySample { ts, .. } => Some(ts),
168 TraceRow::Record { .. } | TraceRow::TimestampBase { .. } => None,
169 }
170 }
171
172 pub fn map_timestamp<U>(self, f: impl FnOnce(T) -> U) -> TraceRow<'a, U> {
174 let Ok(row) = self.try_map_timestamp(|ts| Ok::<_, Infallible>(f(ts)));
175 row
176 }
177
178 pub fn try_map_timestamp<U, E>(
180 self,
181 f: impl FnOnce(T) -> Result<U, E>,
182 ) -> Result<TraceRow<'a, U>, E> {
183 Ok(match self {
184 TraceRow::Start {
185 ts,
186 id,
187 parent,
188 name,
189 target,
190 values,
191 } => TraceRow::Start {
192 ts: f(ts)?,
193 id,
194 parent,
195 name,
196 target,
197 values,
198 },
199 TraceRow::End { ts, id } => TraceRow::End { ts: f(ts)?, id },
200 TraceRow::Enter {
201 ts,
202 id,
203 thread_id,
204 allocations,
205 } => TraceRow::Enter {
206 ts: f(ts)?,
207 id,
208 thread_id,
209 allocations,
210 },
211 TraceRow::Exit {
212 ts,
213 id,
214 thread_id,
215 allocations,
216 } => TraceRow::Exit {
217 ts: f(ts)?,
218 id,
219 thread_id,
220 allocations,
221 },
222 TraceRow::Event { ts, parent, values } => TraceRow::Event {
223 ts: f(ts)?,
224 parent,
225 values,
226 },
227 TraceRow::Record { id, values } => TraceRow::Record { id, values },
228 TraceRow::MemorySample {
229 ts,
230 memory,
231 memory_pressure,
232 memory_footprint,
233 active_worker_threads,
234 } => TraceRow::MemorySample {
235 ts: f(ts)?,
236 memory,
237 memory_pressure,
238 memory_footprint,
239 active_worker_threads,
240 },
241 TraceRow::TimestampBase { ts } => TraceRow::TimestampBase { ts },
242 })
243 }
244}
245
246#[derive(Debug, Default, Clone, Copy)]
251pub struct TimestampEncoder {
252 last: Option<u64>,
254}
255
256impl TimestampEncoder {
257 pub fn encode(&mut self, ts: u64) -> (Option<u64>, DeltaEncodedTimestamp) {
263 let delta = self
264 .last
265 .and_then(|last| i64::try_from(i128::from(ts) - i128::from(last)).ok());
266 self.last = Some(ts);
267 match delta {
268 Some(delta) => (None, DeltaEncodedTimestamp(delta)),
269 None => (Some(ts), DeltaEncodedTimestamp(0)),
270 }
271 }
272
273 pub fn encode_row<'a>(
276 &mut self,
277 row: TraceRow<'a, u64>,
278 ) -> (Option<TraceRow<'static>>, TraceRow<'a>) {
279 let mut base = None;
280 let row = row.map_timestamp(|ts| {
281 let (new_base, encoded) = self.encode(ts);
282 base = new_base.map(|ts| TraceRow::TimestampBase { ts });
283 encoded
284 });
285 (base, row)
286 }
287}
288
289#[derive(Debug, Clone, Copy, PartialEq, Eq)]
291pub enum TimestampDecodeError {
292 MissingBase,
294 Overflow,
296}
297
298impl Display for TimestampDecodeError {
299 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
300 match self {
301 TimestampDecodeError::MissingBase => {
302 write!(
303 f,
304 "invalid trace: timestamp before the first timestamp base"
305 )
306 }
307 TimestampDecodeError::Overflow => {
308 write!(f, "invalid trace: timestamp out of range")
309 }
310 }
311 }
312}
313
314impl std::error::Error for TimestampDecodeError {}
315
316#[derive(Debug, Default, Clone, Copy)]
320pub struct TimestampDecoder {
321 last: Option<u64>,
323}
324
325impl TimestampDecoder {
326 pub fn decode<'a>(
329 &mut self,
330 row: TraceRow<'a>,
331 ) -> Result<TraceRow<'a, u64>, TimestampDecodeError> {
332 if let TraceRow::TimestampBase { ts } = row {
333 self.last = Some(ts);
334 }
335 row.try_map_timestamp(|DeltaEncodedTimestamp(delta)| {
336 let last = self.last.ok_or(TimestampDecodeError::MissingBase)?;
337 let decoded = last
338 .checked_add_signed(delta)
339 .ok_or(TimestampDecodeError::Overflow)?;
340 self.last = Some(decoded);
341 Ok(decoded)
342 })
343 }
344}
345
346#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
356pub struct Allocations {
357 pub allocations: u64,
359 pub allocation_count: u64,
361 pub deallocations: u64,
363 pub deallocation_count: u64,
365}
366
367#[derive(Debug, Serialize, Deserialize)]
368pub enum TraceValue<'a> {
369 String(#[serde(borrow)] Cow<'a, str>),
370 Bool(bool),
371 UInt(u64),
372 Int(i64),
373 Float(f64),
374}
375
376impl Display for TraceValue<'_> {
377 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
378 match self {
379 TraceValue::String(s) => write!(f, "{s}"),
380 TraceValue::Bool(b) => write!(f, "{b}"),
381 TraceValue::UInt(u) => write!(f, "{u}"),
382 TraceValue::Int(i) => write!(f, "{i}"),
383 TraceValue::Float(fl) => write!(f, "{fl}"),
384 }
385 }
386}
387
388impl TraceValue<'_> {
389 pub fn as_u64(&self) -> Option<u64> {
390 match self {
391 TraceValue::UInt(u) => Some(*u),
392 _ => None,
393 }
394 }
395
396 pub fn as_str(&self) -> Option<&str> {
397 match self {
398 TraceValue::String(s) => Some(s),
399 _ => None,
400 }
401 }
402
403 pub fn into_static(self) -> TraceValue<'static> {
404 match self {
405 TraceValue::String(s) => TraceValue::String(s.into_owned().into()),
406 TraceValue::Bool(b) => TraceValue::Bool(b),
407 TraceValue::UInt(u) => TraceValue::UInt(u),
408 TraceValue::Int(i) => TraceValue::Int(i),
409 TraceValue::Float(fl) => TraceValue::Float(fl),
410 }
411 }
412}
413
414#[cfg(test)]
415mod tests {
416 use crate::tracing::{
417 DeltaEncodedTimestamp, TimestampDecodeError, TimestampDecoder, TimestampEncoder, TraceRow,
418 check_trace_header,
419 };
420
421 fn end(ts: u64) -> TraceRow<'static, u64> {
422 TraceRow::End { ts, id: 1 }
423 }
424
425 fn encoded_end(delta: i64) -> TraceRow<'static> {
426 TraceRow::End {
427 ts: DeltaEncodedTimestamp(delta),
428 id: 1,
429 }
430 }
431
432 fn base(ts: u64) -> TraceRow<'static> {
433 TraceRow::TimestampBase { ts }
434 }
435
436 fn delta_of(row: &TraceRow<'_>) -> i64 {
438 match row {
439 TraceRow::End { ts, .. } => ts.0,
440 _ => unreachable!(),
441 }
442 }
443
444 fn ts_of(row: &TraceRow<'_, u64>) -> u64 {
446 match row {
447 TraceRow::End { ts, .. } => *ts,
448 _ => unreachable!(),
449 }
450 }
451
452 fn round_trip(timestamps: &[u64]) -> (Vec<TraceRow<'static>>, Vec<u64>) {
454 let mut encoder = TimestampEncoder::default();
455 let mut encoded = Vec::new();
456 for &ts in timestamps {
457 let (base, row) = encoder.encode_row(end(ts));
458 encoded.extend(base);
459 encoded.push(row);
460 }
461 let bytes: Vec<Vec<u8>> = encoded
463 .iter()
464 .map(|row| postcard::to_stdvec(row).unwrap())
465 .collect();
466 let mut decoder = TimestampDecoder::default();
467 let mut decoded = Vec::new();
468 for bytes in &bytes {
469 let row: TraceRow<'_> = postcard::from_bytes(bytes).unwrap();
470 let row = decoder.decode(row).unwrap();
471 if !matches!(row, TraceRow::TimestampBase { .. }) {
472 decoded.push(ts_of(&row));
473 }
474 }
475 (encoded, decoded)
476 }
477
478 #[test]
479 fn timestamps_round_trip_with_negative_deltas() {
480 let timestamps = [1000, 1005, 1003, 1003, 2000, 0, 7];
481 let (encoded, decoded) = round_trip(×tamps);
482 assert_eq!(decoded, timestamps);
483 assert!(matches!(encoded[0], TraceRow::TimestampBase { ts: 1000 }));
485 let deltas: Vec<i64> = encoded[1..].iter().map(delta_of).collect();
486 assert_eq!(deltas, [0, 5, -2, 0, 997, -2000, 7]);
487 }
488
489 #[test]
492 fn deltas_are_serialized_as_zigzag_varints() {
493 fn zigzag(value: i64) -> u64 {
494 ((value << 1) ^ (value >> 63)) as u64
495 }
496 for delta in [0, 1, -1, 5, -2, 1994, -3999, i64::MIN, i64::MAX] {
497 let mut expected = postcard::to_stdvec(&TraceRow::End {
498 ts: DeltaEncodedTimestamp(0),
499 id: 7,
500 })
501 .unwrap();
502 expected.splice(1..2, postcard::to_stdvec(&zigzag(delta)).unwrap());
504 assert_eq!(
505 postcard::to_stdvec(&TraceRow::End {
506 ts: DeltaEncodedTimestamp(delta),
507 id: 7
508 })
509 .unwrap(),
510 expected,
511 "delta {delta}"
512 );
513 }
514 }
515
516 #[test]
517 fn out_of_range_delta_starts_a_new_base() {
518 let mut encoder = TimestampEncoder::default();
519 let encoded: Vec<_> = [0, u64::MAX, u64::MAX - 1, 0]
520 .into_iter()
521 .map(|ts| encoder.encode(ts))
522 .collect();
523 assert_eq!(
524 encoded,
525 [
526 (Some(0), DeltaEncodedTimestamp(0)),
527 (Some(u64::MAX), DeltaEncodedTimestamp(0)),
528 (None, DeltaEncodedTimestamp(-1)),
529 (Some(0), DeltaEncodedTimestamp(0))
530 ]
531 );
532 }
533
534 #[test]
535 fn reset_encoder_writes_a_new_base() {
536 let mut encoder = TimestampEncoder::default();
537 assert!(encoder.encode_row(end(10)).0.is_some());
538 assert!(encoder.encode_row(end(12)).0.is_none());
539 encoder = TimestampEncoder::default();
540 assert!(matches!(
541 encoder.encode_row(end(13)).0,
542 Some(TraceRow::TimestampBase { ts: 13 })
543 ));
544 }
545
546 #[test]
547 fn rows_without_timestamp_are_unchanged() {
548 let mut encoder = TimestampEncoder::default();
549 let mut decoder = TimestampDecoder::default();
550 let (base, record) = encoder.encode_row(TraceRow::<u64>::Record {
551 id: 5,
552 values: Vec::new(),
553 });
554 assert!(base.is_none());
555 let record = decoder.decode(record).unwrap();
556 assert!(matches!(record, TraceRow::Record { id: 5, .. }));
557 assert!(encoder.encode_row(end(1)).0.is_some());
559 }
560
561 #[test]
562 fn memory_sample_keeps_all_fields() {
563 let mut encoder = TimestampEncoder::default();
564 let (base, row) = encoder.encode_row(TraceRow::<u64>::MemorySample {
565 ts: 1234,
566 memory: 1 << 30,
567 memory_pressure: 42,
568 memory_footprint: 3 << 30,
569 active_worker_threads: 7,
570 });
571 let bytes: Vec<Vec<u8>> = base
573 .into_iter()
574 .chain([row])
575 .map(|row| postcard::to_stdvec(&row).unwrap())
576 .collect();
577 let mut decoder = TimestampDecoder::default();
578 let mut decoded = None;
579 for bytes in &bytes {
580 let row: TraceRow<'_> = postcard::from_bytes(bytes).unwrap();
581 decoded = Some(decoder.decode(row).unwrap());
582 }
583 match decoded.unwrap() {
584 TraceRow::MemorySample {
585 ts,
586 memory,
587 memory_pressure,
588 memory_footprint,
589 active_worker_threads,
590 } => {
591 assert_eq!(ts, 1234);
592 assert_eq!(memory, 1 << 30);
593 assert_eq!(memory_pressure, 42);
594 assert_eq!(memory_footprint, 3 << 30);
595 assert_eq!(active_worker_threads, 7);
596 }
597 row => panic!("expected a MemorySample, got {row:?}"),
598 }
599 }
600
601 #[test]
602 fn decoder_keeps_timestamp_bases() {
603 let mut decoder = TimestampDecoder::default();
604 assert!(matches!(
605 decoder.decode(base(42)),
606 Ok(TraceRow::TimestampBase { ts: 42 })
607 ));
608 assert_eq!(ts_of(&decoder.decode(encoded_end(-2)).unwrap()), 40);
609 }
610
611 #[test]
612 fn decoder_requires_a_base() {
613 let mut decoder = TimestampDecoder::default();
614 assert_eq!(
615 decoder.decode(encoded_end(4)).unwrap_err(),
616 TimestampDecodeError::MissingBase
617 );
618 }
619
620 #[test]
621 fn decoder_detects_overflow() {
622 let mut decoder = TimestampDecoder::default();
623 decoder.decode(base(u64::MAX - 1)).unwrap();
624 assert_eq!(
625 decoder.decode(encoded_end(5)).unwrap_err(),
626 TimestampDecodeError::Overflow
627 );
628 let mut decoder = TimestampDecoder::default();
629 decoder.decode(base(1)).unwrap();
630 assert_eq!(
631 decoder.decode(encoded_end(-2)).unwrap_err(),
632 TimestampDecodeError::Overflow
633 );
634 }
635
636 #[test]
638 fn decoder_accepts_timestamps_beyond_i64() {
639 let beyond = i64::MAX as u64 + 5;
640 let (_, decoded) = round_trip(&[beyond, beyond + 1, u64::MAX]);
641 assert_eq!(decoded, [beyond, beyond + 1, u64::MAX]);
642 let mut decoder = TimestampDecoder::default();
643 decoder.decode(base(i64::MAX as u64)).unwrap();
644 assert_eq!(ts_of(&decoder.decode(encoded_end(5)).unwrap()), beyond);
645 }
646
647 #[test]
648 fn checks_trace_header() {
649 assert!(check_trace_header(b"TRACEv1").is_ok());
650 assert!(check_trace_header(b"TRACEv1\x00\x01").is_ok());
651 assert_eq!(
652 check_trace_header(b"TRACEv0\x00").unwrap_err().to_string(),
653 "Unsupported trace file version: expected TRACEv1, found TRACEv0"
654 );
655 assert_eq!(
656 check_trace_header(b"TRACEv2").unwrap_err().to_string(),
657 "Unsupported trace file version: expected TRACEv1, found TRACEv2"
658 );
659 assert_eq!(
660 check_trace_header(b"\x00\x01\x02").unwrap_err().to_string(),
661 "Not a trace file: missing TRACEv1 header"
662 );
663 assert!(check_trace_header(b"").is_err());
664 }
665}