1use 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
13pub struct TraceWriter {
16 inner: BufWriter<File>,
17 count: u64,
18}
19
20impl TraceWriter {
21 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 pub fn write(&mut self, record: &TraceRecord) -> Result<()> {
61 self.write_raw(record)?;
62 self.count += 1;
63 Ok(())
64 }
65
66 pub fn write_event(&mut self, event: &katra_core::TraceEvent) -> Result<()> {
68 self.write(&TraceRecord::Event(event.clone()))
69 }
70
71 pub fn count(&self) -> u64 {
73 self.count
74 }
75
76 pub fn flush_buffered(&mut self) -> Result<()> {
78 self.inner.flush().map_err(KatraError::from)
79 }
80
81 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 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
97pub 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}