grafos_profile/
recorder.rs1use grafos_observe::span::ResourceSpan;
4
5use crate::recording::ProfileRecording;
6
7pub struct ProfileRecorder {
21 spans: Vec<ResourceSpan>,
22 program_name: Option<String>,
23 start_time_us: Option<u64>,
24}
25
26impl ProfileRecorder {
27 pub fn start() -> Self {
29 ProfileRecorder {
30 spans: Vec::new(),
31 program_name: None,
32 start_time_us: None,
33 }
34 }
35
36 pub fn with_program_name(mut self, name: &str) -> Self {
38 self.program_name = Some(name.to_string());
39 self
40 }
41
42 pub fn record_span(&mut self, span: ResourceSpan) {
44 if self.start_time_us.is_none() {
45 self.start_time_us = Some(span.start_time_unix_us);
46 }
47 self.spans.push(span);
48 }
49
50 pub fn span_count(&self) -> usize {
52 self.spans.len()
53 }
54
55 pub fn stop(self) -> ProfileRecording {
57 let mut recording = ProfileRecording::from_spans(self.spans);
58 if let Some(name) = self.program_name {
59 recording.header.program = Some(name);
60 }
61 recording
62 }
63}
64
65#[cfg(test)]
66mod tests {
67 use super::*;
68 use grafos_observe::trace::TraceContext;
69
70 fn test_ctx() -> TraceContext {
71 let mut bytes = [0u8; 24];
72 for (i, b) in bytes.iter_mut().enumerate() {
73 *b = (i as u8).wrapping_add(0x42);
74 }
75 TraceContext::new_root(&bytes)
76 }
77
78 #[test]
79 fn recorder_start_stop() {
80 let mut recorder = ProfileRecorder::start();
81 assert_eq!(recorder.span_count(), 0);
82
83 let mut span = ResourceSpan::new("test", test_ctx());
84 span.start_time_unix_us = 100;
85 span.end_time_unix_us = 200;
86 recorder.record_span(span);
87 assert_eq!(recorder.span_count(), 1);
88
89 let recording = recorder.stop();
90 assert_eq!(recording.spans.len(), 1);
91 assert_eq!(recording.header.format, "grafos-profile-v1");
92 }
93
94 #[test]
95 fn recorder_with_program_name() {
96 let recorder = ProfileRecorder::start().with_program_name("my_app");
97 let recording = recorder.stop();
98 assert_eq!(recording.header.program.as_deref(), Some("my_app"));
99 }
100
101 #[test]
102 fn recorder_save_load_roundtrip() {
103 let mut recorder = ProfileRecorder::start().with_program_name("roundtrip_test");
104 let mut span = ResourceSpan::new("op1", test_ctx());
105 span.start_time_unix_us = 1000;
106 span.end_time_unix_us = 2000;
107 span.add_lease_id(42);
108 span.bytes_read = 4096;
109 recorder.record_span(span);
110
111 let recording = recorder.stop();
112
113 let dir = std::env::temp_dir();
114 let path = dir.join("grafos_recorder_test.ndjson");
115 let path_str = path.to_str().unwrap();
116
117 recording.save(path_str).unwrap();
118 let loaded = ProfileRecording::load(path_str).unwrap();
119
120 assert_eq!(loaded.spans.len(), 1);
121 assert_eq!(loaded.spans[0].name, "op1");
122 assert_eq!(loaded.spans[0].bytes_read, 4096);
123 assert_eq!(loaded.header.program.as_deref(), Some("roundtrip_test"));
124
125 let _ = std::fs::remove_file(path_str);
126 }
127}