rust_zero_core/
profile.rs

1use std::{
2    collections::BTreeMap,
3    fmt::Write,
4    sync::{
5        atomic::{AtomicBool, Ordering},
6        Mutex,
7    },
8    time::{Duration, Instant},
9};
10
11/// A timing point returned by [`Profiler::start`].
12#[derive(Debug)]
13pub struct ProfilePoint {
14    started_at: Option<Instant>,
15}
16
17impl ProfilePoint {
18    /// Returns the elapsed time when profiling was enabled when this point was created.
19    pub fn elapsed(&self) -> Option<Duration> {
20        self.started_at.map(|started_at| started_at.elapsed())
21    }
22}
23
24#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
25struct ProfileSlot {
26    lifetime_count: u64,
27    lifetime_duration: Duration,
28    interval_count: u64,
29    interval_duration: Duration,
30}
31
32/// An immutable aggregate for one named profiling operation.
33#[derive(Debug, Clone, PartialEq, Eq)]
34pub struct ProfileSnapshot {
35    pub name: String,
36    pub lifetime_count: u64,
37    pub lifetime_duration: Duration,
38    pub interval_count: u64,
39    pub interval_duration: Duration,
40}
41
42impl ProfileSnapshot {
43    pub fn lifetime_average(&self) -> Option<Duration> {
44        average(self.lifetime_duration, self.lifetime_count)
45    }
46
47    pub fn interval_average(&self) -> Option<Duration> {
48        average(self.interval_duration, self.interval_count)
49    }
50}
51
52/// Low-overhead, named duration profiling.
53///
54/// Profiling starts disabled, matching go-zero's opt-in profiler. Calling
55/// [`Profiler::snapshot`] returns lifetime aggregates and resets only the interval aggregates.
56#[derive(Debug, Default)]
57pub struct Profiler {
58    enabled: AtomicBool,
59    slots: Mutex<BTreeMap<String, ProfileSlot>>,
60}
61
62impl Profiler {
63    pub fn new() -> Self {
64        Self::default()
65    }
66
67    pub fn enable(&self) {
68        self.enabled.store(true, Ordering::Release);
69    }
70
71    pub fn disable(&self) {
72        self.enabled.store(false, Ordering::Release);
73    }
74
75    pub fn is_enabled(&self) -> bool {
76        self.enabled.load(Ordering::Acquire)
77    }
78
79    pub fn start(&self) -> ProfilePoint {
80        ProfilePoint {
81            started_at: self.is_enabled().then(Instant::now),
82        }
83    }
84
85    pub fn report(&self, name: impl Into<String>, point: ProfilePoint) {
86        let Some(duration) = point.elapsed() else {
87            return;
88        };
89        self.record(name, duration);
90    }
91
92    pub fn record(&self, name: impl Into<String>, duration: Duration) {
93        if !self.is_enabled() {
94            return;
95        }
96
97        let mut slots = self.slots.lock().expect("profiler mutex poisoned");
98        let slot = slots.entry(name.into()).or_default();
99        slot.lifetime_count = slot.lifetime_count.saturating_add(1);
100        slot.lifetime_duration = slot.lifetime_duration.saturating_add(duration);
101        slot.interval_count = slot.interval_count.saturating_add(1);
102        slot.interval_duration = slot.interval_duration.saturating_add(duration);
103    }
104
105    /// Takes a report snapshot and begins a fresh reporting interval.
106    pub fn snapshot(&self) -> Vec<ProfileSnapshot> {
107        let mut slots = self.slots.lock().expect("profiler mutex poisoned");
108        slots
109            .iter_mut()
110            .map(|(name, slot)| {
111                let snapshot = ProfileSnapshot {
112                    name: name.clone(),
113                    lifetime_count: slot.lifetime_count,
114                    lifetime_duration: slot.lifetime_duration,
115                    interval_count: slot.interval_count,
116                    interval_duration: slot.interval_duration,
117                };
118                slot.interval_count = 0;
119                slot.interval_duration = Duration::ZERO;
120                snapshot
121            })
122            .collect()
123    }
124
125    /// Renders the same lifetime/last-interval view as go-zero's profiler report.
126    pub fn render_report(&self) -> String {
127        let mut output = String::from(
128            "Profiling report\nOPERATION,LIFETIME_COUNT,LIFETIME_AVERAGE,INTERVAL_COUNT,INTERVAL_AVERAGE\n",
129        );
130        for snapshot in self.snapshot() {
131            let _ = writeln!(
132                output,
133                "{},{},{},{},{}",
134                snapshot.name,
135                snapshot.lifetime_count,
136                display_average(snapshot.lifetime_average()),
137                snapshot.interval_count,
138                display_average(snapshot.interval_average()),
139            );
140        }
141        output
142    }
143}
144
145fn average(total: Duration, count: u64) -> Option<Duration> {
146    (count > 0).then(|| Duration::from_secs_f64(total.as_secs_f64() / count as f64))
147}
148
149fn display_average(duration: Option<Duration>) -> String {
150    duration
151        .map(|duration| format!("{:.6}s", duration.as_secs_f64()))
152        .unwrap_or_else(|| "-".to_owned())
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    #[test]
160    fn disabled_profiler_has_negligible_noop_points() {
161        let profiler = Profiler::new();
162        let point = profiler.start();
163        assert_eq!(point.elapsed(), None);
164        profiler.report("ignored", point);
165        assert!(profiler.snapshot().is_empty());
166    }
167
168    #[test]
169    fn snapshots_keep_lifetime_and_reset_interval_values() {
170        let profiler = Profiler::new();
171        profiler.enable();
172        profiler.record("database", Duration::from_millis(10));
173        profiler.record("database", Duration::from_millis(30));
174
175        let first = profiler.snapshot();
176        assert_eq!(first[0].lifetime_count, 2);
177        assert_eq!(first[0].lifetime_average(), Some(Duration::from_millis(20)));
178        assert_eq!(first[0].interval_count, 2);
179
180        let second = profiler.snapshot();
181        assert_eq!(second[0].lifetime_count, 2);
182        assert_eq!(second[0].interval_count, 0);
183        assert_eq!(second[0].interval_average(), None);
184    }
185
186    #[test]
187    fn report_is_stable_and_human_readable() {
188        let profiler = Profiler::new();
189        profiler.enable();
190        profiler.record("cache", Duration::from_millis(5));
191
192        let report = profiler.render_report();
193        assert!(report.contains("OPERATION,LIFETIME_COUNT"));
194        assert!(report.contains("cache,1,0.005000s,1,0.005000s"));
195    }
196}
rust_zero_core/profile.rs

rust_zero_core/
profile.rs
