resource_meter/
lib.rs

1//! # resource-meter
2//!
3//! A lightweight resource usage measurement library for Rust, providing scoped measurement of wall-clock time, user CPU time, and system CPU time.
4//!
5//! This crate allows you to track resource usage (CPU and wall time) in a hierarchical or flat manner, making it easy to profile sections of your code. It leverages `libc::getrusage` for accurate process resource usage statistics.
6//!
7//! ## Features
8//!
9//! - Scoped measurement using a stack-based API
10//! - Hierarchical (tree) and flat reporting
11//! - Tracks wall time, user CPU time, and system CPU time
12//! - Simple API for integration
13//! - Platform support for Unix-like systems (uses `libc::getrusage`)
14//!
15//! ## Example
16//!
17//! ```rust
18//! use resource_meter::ResourceMeterStack;
19//! use std::thread::sleep;
20//! use std::time::Duration;
21//!
22//! fn main() {
23//!     // Create a stack to manage resource measurement scopes
24//!     let mut stack = ResourceMeterStack::new();
25//!
26//!     // Start a measurement scope named "outer"
27//!     stack.push("outer");
28//!     sleep(Duration::from_millis(100));
29//!
30//!     // Start a nested measurement scope named "outer/inner"
31//!     stack.push("outer/inner");
32//!     sleep(Duration::from_millis(200));
33//!     stack.pop(); // End "outer/inner"
34//!
35//!     sleep(Duration::from_millis(50));
36//!     stack.pop(); // End "outer"
37//!
38//!     // Generate and print a hierarchical report
39//!     let report = stack.finish();
40//!     println!("{}", report);
41//! }
42//! ```
43//!
44//! ## Usage
45//!
46//! - Use [`ResourceMeterStack`] to manage measurement scopes with `push`/`pop`.
47//! - Generate a report with [`ResourceMeterStack::finish`] (tree report) or [`ResourceMeterStack::into_report`] (choose flat or tree).
48//!
49//! ## Main Types
50//!
51//! - [`ResourceMeterStack`]: Main stack for managing resource measurement scopes.
52//! - [`TreeReport`]: Hierarchical report of resource usage.
53//! - [`UsageMeasurement`]: Stores individual usage statistics.
54//!
55//! ## Platform Support
56//!
57//! This crate uses `libc::getrusage` and is intended for Unix-like platforms.
58//!
59//! ## License
60//!
61//! Licensed under MIT or Apache-2.0.
62//!
63//! ## See Also
64//!
65//! - [`std::time::Instant`]
66//! - [`libc::getrusage`]
67//!
68use std::{
69    collections::{BTreeMap, HashMap},
70    ops::{AddAssign, Sub},
71    time::{Duration, Instant},
72};
73
74#[inline(always)]
75fn tv_to_duration(tv: libc::timeval) -> Duration {
76    Duration::new(tv.tv_sec as u64, (tv.tv_usec * 1000) as u32)
77}
78
79#[derive(Clone, Copy)]
80struct ResourceUsage(libc::rusage);
81
82impl ResourceUsage {
83    pub fn now() -> Self {
84        unsafe {
85            let mut rusage = std::mem::MaybeUninit::<libc::rusage>::uninit();
86            assert!(0 == libc::getrusage(libc::RUSAGE_SELF, rusage.as_mut_ptr()));
87            Self(rusage.assume_init())
88        }
89    }
90
91    #[inline(always)]
92    pub fn utime(&self) -> Duration {
93        tv_to_duration(self.0.ru_utime)
94    }
95
96    #[inline(always)]
97    pub fn stime(&self) -> Duration {
98        tv_to_duration(self.0.ru_stime)
99    }
100
101    #[inline(always)]
102    pub fn cpu_time(&self) -> Duration {
103        self.utime() + self.stime()
104    }
105}
106
107impl std::fmt::Debug for ResourceUsage {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        f.debug_struct("ResourceUsage")
110            .field("ru_utime.tv_sec", &self.0.ru_utime.tv_sec)
111            .field("ru_utime.tv_usec", &self.0.ru_utime.tv_usec)
112            .field("ru_stime.tv_sec", &self.0.ru_stime.tv_sec)
113            .field("ru_stime.tv_usec", &self.0.ru_stime.tv_usec)
114            .finish()
115    }
116}
117
118#[derive(Debug, Clone, Copy)]
119struct UsageSnapshot {
120    pub time: Instant,
121    pub usage: ResourceUsage,
122}
123
124impl UsageSnapshot {
125    pub fn now() -> Self {
126        Self {
127            time: Instant::now(),
128            usage: ResourceUsage::now(),
129        }
130    }
131}
132
133impl Sub for UsageSnapshot {
134    type Output = UsageMeasurement;
135
136    fn sub(self, rhs: Self) -> Self::Output {
137        UsageMeasurement {
138            utime: self.usage.utime() - rhs.usage.utime(),
139            stime: self.usage.stime() - rhs.usage.stime(),
140            wtime: self.time.duration_since(rhs.time),
141        }
142    }
143}
144
145#[derive(Debug)]
146struct ScopeFrame {
147    path: String,
148    start: UsageSnapshot,
149}
150
151#[derive(Debug, Default, Clone, Copy)]
152pub struct UsageMeasurement {
153    /* User time in this process.  */
154    utime: Duration,
155    /* System time (if applicable for this host platform) in this
156    process.  */
157    stime: Duration,
158    /* Wall clock time.  */
159    wtime: Duration,
160}
161
162impl UsageMeasurement {}
163
164impl AddAssign for UsageMeasurement {
165    fn add_assign(&mut self, rhs: Self) {
166        self.utime += rhs.utime;
167        self.stime += rhs.stime;
168        self.wtime += rhs.wtime;
169    }
170}
171
172#[derive(Debug)]
173pub struct ResourceMeterStack {
174    store: HashMap<String, UsageMeasurement>,
175    stack: Vec<ScopeFrame>,
176    last_snapshot: UsageSnapshot,
177}
178
179impl ResourceMeterStack {
180    pub fn new() -> Self {
181        Self {
182            store: HashMap::new(),
183            stack: vec![ScopeFrame {
184                path: String::new(),
185                start: UsageSnapshot::now(),
186            }],
187            last_snapshot: UsageSnapshot::now(),
188        }
189    }
190
191    pub fn push<T: Into<String>>(&mut self, key: T) {
192        self.stack.push(ScopeFrame {
193            path: key.into(),
194            start: UsageSnapshot::now(),
195        });
196    }
197
198    pub fn pop(&mut self) {
199        let now = UsageSnapshot::now();
200        let curr = self.stack.pop().unwrap();
201        let slot = self.store.entry(curr.path).or_default();
202        *slot += now - curr.start;
203        self.last_snapshot = now;
204    }
205
206    pub fn into_report(mut self, format: ReportFormat) -> Report {
207        while !self.stack.is_empty() {
208            self.pop();
209        }
210        match format {
211            ReportFormat::Flat => Report::Flat(self.store),
212            ReportFormat::Tree => {
213                let mut root = TreeNode::default();
214                for (path, usage) in self.store {
215                    root.insert(path, usage);
216                }
217                Report::Tree(root)
218            }
219        }
220    }
221
222    pub fn finish(mut self) -> TreeReport {
223        while !self.stack.is_empty() {
224            self.pop();
225        }
226        TreeReport::new(self.store)
227    }
228}
229
230pub enum ReportFormat {
231    Flat,
232    Tree,
233}
234
235pub enum Report {
236    Flat(HashMap<String, UsageMeasurement>),
237    Tree(TreeNode),
238}
239
240#[derive(Debug)]
241pub struct TreeReport(TreeNode);
242
243impl TreeReport {
244    pub fn new(measurements: HashMap<String, UsageMeasurement>) -> Self {
245        let mut root = TreeNode::default();
246        for (path, usage) in measurements {
247            root.insert(path, usage);
248        }
249        Self(root)
250    }
251}
252
253#[derive(Debug, Default)]
254struct TreeNode {
255    usage: Option<UsageMeasurement>,
256    children: BTreeMap<String, TreeNode>,
257}
258
259impl TreeNode {
260    fn insert(&mut self, path: String, usage: UsageMeasurement) {
261        let mut current = self;
262        for part in path.split('/') {
263            current = current.children.entry(part.to_string()).or_default();
264        }
265        current.usage = Some(usage);
266    }
267
268    fn fmt_with_indent(&self, f: &mut std::fmt::Formatter<'_>, depth: usize) -> std::fmt::Result {
269        let width = 20;
270        for (name, child) in &self.children {
271            let indent = " ".repeat(depth);
272
273            if let Some(usage) = child.usage {
274                writeln!(
275                    f,
276                    "{indent}{:<} real: {:>6?}, user: {:>6?}, sys: {:?}",
277                    name,
278                    usage.wtime.as_millis(),
279                    usage.utime.as_millis(),
280                    usage.stime.as_millis(),
281                )?;
282            } else {
283                writeln!(f, "{indent}{name}")?;
284            }
285            child.fmt_with_indent(f, depth + 1)?;
286        }
287        Ok(())
288    }
289}
290
291/* Summarize timing variables to FP.  The timing variable TV_TOTAL has
292a special meaning -- it's considered to be the total elapsed time,
293for normalizing the others, and is displayed last.  */
294impl std::fmt::Display for TreeReport {
295    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
296        writeln!(f, "--- Resource Usage Report ---")?;
297        self.0.fmt_with_indent(f, 0)?;
298        Ok(())
299    }
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305
306    #[test]
307    fn test_simple() {
308        let mut stk = ResourceMeterStack::new();
309        let now = Instant::now();
310        std::thread::sleep(Duration::from_millis(200));
311        stk.push("sleep600");
312        stk.push("sleep600/300");
313        std::thread::sleep(Duration::from_millis(300));
314        stk.pop();
315        std::thread::sleep(Duration::from_millis(100));
316        stk.push("sleep600/200");
317        std::thread::sleep(Duration::from_millis(200));
318        stk.pop();
319        stk.pop();
320        std::thread::sleep(Duration::from_millis(200));
321        let report = stk.finish();
322        eprintln!("{}, took: {:?}", report, now.elapsed());
323    }
324
325    #[test]
326    fn test_resource_usage() {
327        let rusage = ResourceUsage::now();
328        eprintln!("{:#?}", rusage);
329    }
330}