Skip to main content

katra_trace/
writer.rs

1//! Trace writer: appends length-prefixed bincode records to a file.
2
3use std::fs::{File, OpenOptions};
4use std::io::{BufWriter, Write};
5use std::path::Path;
6
7use katra_core::{KatraError, Result};
8use serde::Serialize;
9
10use crate::header::{FORMAT_VERSION, SCHEMA_HASH, TRACE_MAGIC, TraceHeader};
11use crate::record::{TraceRecord, TraceSummary};
12
13/// Appends records to a trace file. Not `Sync`; a profiler must serialize
14/// writes (the profiler's sink mutex does this).
15pub struct TraceWriter {
16    inner: BufWriter<File>,
17    count: u64,
18}
19
20impl TraceWriter {
21    /// Create (truncate) a trace file and write the header.
22    pub fn create(path: &Path, header: &TraceHeader) -> Result<Self> {
23        let file = OpenOptions::new()
24            .create(true)
25            .truncate(true)
26            .write(true)
27            .open(path)
28            .map_err(KatraError::from)?;
29        let mut w = TraceWriter { inner: BufWriter::new(file), count: 0 };
30        w.write_header(header)?;
31        Ok(w)
32    }
33
34    fn write_raw<T: Serialize>(&mut self, value: &T) -> Result<()> {
35        let bytes =
36            bincode::serialize(value).map_err(|e| KatraError::Protocol(format!("bincode: {e}")))?;
37        let len = bytes.len() as u32;
38        self.inner
39            .write_all(&len.to_le_bytes())
40            .and_then(|_| self.inner.write_all(&bytes))
41            .map_err(KatraError::from)?;
42        Ok(())
43    }
44
45    fn write_header(&mut self, header: &TraceHeader) -> Result<()> {
46        debug_assert_eq!(header.format_version, FORMAT_VERSION);
47        debug_assert_eq!(header.schema_hash, SCHEMA_HASH);
48        let bytes = bincode::serialize(header)
49            .map_err(|e| KatraError::Protocol(format!("bincode header: {e}")))?;
50        self.inner
51            .write_all(&TRACE_MAGIC)
52            .and_then(|_| self.inner.write_all(&FORMAT_VERSION.to_le_bytes()))
53            .and_then(|_| self.inner.write_all(&(bytes.len() as u32).to_le_bytes()))
54            .and_then(|_| self.inner.write_all(&bytes))
55            .map_err(KatraError::from)?;
56        Ok(())
57    }
58
59    /// Append one record.
60    pub fn write(&mut self, record: &TraceRecord) -> Result<()> {
61        self.write_raw(record)?;
62        self.count += 1;
63        Ok(())
64    }
65
66    /// Append one event.
67    pub fn write_event(&mut self, event: &katra_core::TraceEvent) -> Result<()> {
68        self.write(&TraceRecord::Event(event.clone()))
69    }
70
71    /// Number of records written.
72    pub fn count(&self) -> u64 {
73        self.count
74    }
75
76    /// Flush buffered bytes to the file without finishing the session.
77    pub fn flush_buffered(&mut self) -> Result<()> {
78        self.inner.flush().map_err(KatraError::from)
79    }
80
81    /// Write the session summary and flush. The file is left open for
82    /// inspection; call `close` to release it.
83    pub fn finish(&mut self, summary: &TraceSummary) -> Result<()> {
84        self.write(&TraceRecord::SessionSummary(summary.clone()))?;
85        self.inner.flush().map_err(KatraError::from)
86    }
87
88    /// Close the file (flush and drop).
89    pub fn close(self) -> Result<()> {
90        let mut w = self;
91        w.inner.flush().map_err(KatraError::from)?;
92        drop(w.inner);
93        Ok(())
94    }
95}
96
97/// Helper to assemble a header with sensible defaults for a given workload.
98pub fn default_header(
99    workload: &str,
100    capture: &katra_core::CaptureOptions,
101    notes: Vec<String>,
102) -> TraceHeader {
103    let now_wall = katra_core::wall_now_ns();
104    let now_mono = katra_core::monotonic_now_ns();
105    TraceHeader {
106        format_version: FORMAT_VERSION,
107        schema_hash: SCHEMA_HASH,
108        tool_version: format!("katra3d {}", katra_core::KATRA_VERSION),
109        capture_options: capture.clone(),
110        start_wall_ns: now_wall,
111        start_mono_ns: now_mono,
112        process_id: std::process::id() as u64,
113        workload: workload.to_string(),
114        host: crate::reader::host_info(),
115        notes,
116    }
117}