grafos_profile/
recorder.rs

1//! Live profile recorder — wraps observe event sink to capture spans as they complete.
2
3use grafos_observe::span::ResourceSpan;
4
5use crate::recording::ProfileRecording;
6
7/// Records [`ResourceSpan`] data from a running program.
8///
9/// Usage:
10/// ```rust,no_run
11/// use grafos_profile::ProfileRecorder;
12/// use grafos_observe::span::ResourceSpan;
13/// use grafos_observe::trace::TraceContext;
14///
15/// let mut recorder = ProfileRecorder::start();
16/// // ... instrument program, record spans ...
17/// let recording = recorder.stop();
18/// recording.save("/tmp/profile.ndjson").unwrap();
19/// ```
20pub struct ProfileRecorder {
21    spans: Vec<ResourceSpan>,
22    program_name: Option<String>,
23    start_time_us: Option<u64>,
24}
25
26impl ProfileRecorder {
27    /// Begin recording. Call [`stop`](Self::stop) to finish and get the recording.
28    pub fn start() -> Self {
29        ProfileRecorder {
30            spans: Vec::new(),
31            program_name: None,
32            start_time_us: None,
33        }
34    }
35
36    /// Set an optional program name for the recording header.
37    pub fn with_program_name(mut self, name: &str) -> Self {
38        self.program_name = Some(name.to_string());
39        self
40    }
41
42    /// Record a completed span.
43    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    /// Number of spans captured so far.
51    pub fn span_count(&self) -> usize {
52        self.spans.len()
53    }
54
55    /// Stop recording and return the captured data.
56    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}