resource-meter 0.2.0

//! # resource-meter
//!
//! A lightweight resource usage measurement library for Rust, providing scoped measurement of wall-clock time, user CPU time, and system CPU time.
//!
//! 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.
//!
//! ## Features
//!
//! - Scoped measurement using a stack-based API
//! - Hierarchical (tree) and flat reporting
//! - Tracks wall time, user CPU time, and system CPU time
//! - Simple API for integration
//! - Platform support for Unix-like systems (uses `libc::getrusage`)
//!
//! ## Example
//!
//! ```rust
//! use resource_meter::ResourceMeterStack;
//! use std::thread::sleep;
//! use std::time::Duration;
//!
//! fn main() {
//!     // Create a stack to manage resource measurement scopes
//!     let mut stack = ResourceMeterStack::new();
//!
//!     // Start a measurement scope named "outer"
//!     stack.push("outer");
//!     sleep(Duration::from_millis(100));
//!
//!     // Start a nested measurement scope named "outer/inner"
//!     stack.push("outer/inner");
//!     sleep(Duration::from_millis(200));
//!     stack.pop(); // End "outer/inner"
//!
//!     sleep(Duration::from_millis(50));
//!     stack.pop(); // End "outer"
//!
//!     // Generate and print a hierarchical report
//!     let report = stack.finish();
//!     println!("{}", report);
//! }
//! ```
//!
//! ## Usage
//!
//! - Use [`ResourceMeterStack`] to manage measurement scopes with `push`/`pop`.
//! - Generate a report with [`ResourceMeterStack::finish`] (tree report) or [`ResourceMeterStack::into_report`] (choose flat or tree).
//!
//! ## Main Types
//!
//! - [`ResourceMeterStack`]: Main stack for managing resource measurement scopes.
//! - [`TreeReport`]: Hierarchical report of resource usage.
//! - [`UsageMeasurement`]: Stores individual usage statistics.
//!
//! ## Platform Support
//!
//! This crate uses `libc::getrusage` and is intended for Unix-like platforms.
//!
//! ## License
//!
//! Licensed under MIT or Apache-2.0.
//!
//! ## See Also
//!
//! - [`std::time::Instant`]
//! - [`libc::getrusage`]
//!
use std::{
    collections::{BTreeMap, HashMap},
    ops::{AddAssign, Sub},
    time::{Duration, Instant},
};

#[inline(always)]
fn tv_to_duration(tv: libc::timeval) -> Duration {
    Duration::new(tv.tv_sec as u64, (tv.tv_usec * 1000) as u32)
}

#[derive(Clone, Copy)]
struct ResourceUsage(libc::rusage);

impl ResourceUsage {
    pub fn now() -> Self {
        unsafe {
            let mut rusage = std::mem::MaybeUninit::<libc::rusage>::uninit();
            assert!(0 == libc::getrusage(libc::RUSAGE_SELF, rusage.as_mut_ptr()));
            Self(rusage.assume_init())
        }
    }

    #[inline(always)]
    pub fn utime(&self) -> Duration {
        tv_to_duration(self.0.ru_utime)
    }

    #[inline(always)]
    pub fn stime(&self) -> Duration {
        tv_to_duration(self.0.ru_stime)
    }

    #[inline(always)]
    pub fn cpu_time(&self) -> Duration {
        self.utime() + self.stime()
    }
}

impl std::fmt::Debug for ResourceUsage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ResourceUsage")
            .field("ru_utime.tv_sec", &self.0.ru_utime.tv_sec)
            .field("ru_utime.tv_usec", &self.0.ru_utime.tv_usec)
            .field("ru_stime.tv_sec", &self.0.ru_stime.tv_sec)
            .field("ru_stime.tv_usec", &self.0.ru_stime.tv_usec)
            .finish()
    }
}

#[derive(Debug, Clone, Copy)]
struct UsageSnapshot {
    pub time: Instant,
    pub usage: ResourceUsage,
}

impl UsageSnapshot {
    pub fn now() -> Self {
        Self {
            time: Instant::now(),
            usage: ResourceUsage::now(),
        }
    }
}

impl Sub for UsageSnapshot {
    type Output = UsageMeasurement;

    fn sub(self, rhs: Self) -> Self::Output {
        UsageMeasurement {
            utime: self.usage.utime() - rhs.usage.utime(),
            stime: self.usage.stime() - rhs.usage.stime(),
            wtime: self.time.duration_since(rhs.time),
        }
    }
}

#[derive(Debug)]
struct ScopeFrame {
    path: String,
    start: UsageSnapshot,
}

#[derive(Debug, Default, Clone, Copy)]
pub struct UsageMeasurement {
    /* User time in this process.  */
    utime: Duration,
    /* System time (if applicable for this host platform) in this
    process.  */
    stime: Duration,
    /* Wall clock time.  */
    wtime: Duration,
}

impl UsageMeasurement {}

impl AddAssign for UsageMeasurement {
    fn add_assign(&mut self, rhs: Self) {
        self.utime += rhs.utime;
        self.stime += rhs.stime;
        self.wtime += rhs.wtime;
    }
}

#[derive(Debug)]
pub struct ResourceMeterStack {
    store: HashMap<String, UsageMeasurement>,
    stack: Vec<ScopeFrame>,
    last_snapshot: UsageSnapshot,
}

impl ResourceMeterStack {
    pub fn new() -> Self {
        Self {
            store: HashMap::new(),
            stack: vec![ScopeFrame {
                path: String::new(),
                start: UsageSnapshot::now(),
            }],
            last_snapshot: UsageSnapshot::now(),
        }
    }

    pub fn push<T: Into<String>>(&mut self, key: T) {
        self.stack.push(ScopeFrame {
            path: key.into(),
            start: UsageSnapshot::now(),
        });
    }

    pub fn pop(&mut self) {
        let now = UsageSnapshot::now();
        let curr = self.stack.pop().unwrap();
        let slot = self.store.entry(curr.path).or_default();
        *slot += now - curr.start;
        self.last_snapshot = now;
    }

    pub fn into_report(mut self, format: ReportFormat) -> Report {
        while !self.stack.is_empty() {
            self.pop();
        }
        match format {
            ReportFormat::Flat => Report::Flat(self.store),
            ReportFormat::Tree => {
                let mut root = TreeNode::default();
                for (path, usage) in self.store {
                    root.insert(path, usage);
                }
                Report::Tree(root)
            }
        }
    }

    pub fn finish(mut self) -> TreeReport {
        while !self.stack.is_empty() {
            self.pop();
        }
        TreeReport::new(self.store)
    }
}

pub enum ReportFormat {
    Flat,
    Tree,
}

pub enum Report {
    Flat(HashMap<String, UsageMeasurement>),
    Tree(TreeNode),
}

#[derive(Debug)]
pub struct TreeReport(TreeNode);

impl TreeReport {
    pub fn new(measurements: HashMap<String, UsageMeasurement>) -> Self {
        let mut root = TreeNode::default();
        for (path, usage) in measurements {
            root.insert(path, usage);
        }
        Self(root)
    }
}

#[derive(Debug, Default)]
struct TreeNode {
    usage: Option<UsageMeasurement>,
    children: BTreeMap<String, TreeNode>,
}

impl TreeNode {
    fn insert(&mut self, path: String, usage: UsageMeasurement) {
        let mut current = self;
        for part in path.split('/') {
            current = current.children.entry(part.to_string()).or_default();
        }
        current.usage = Some(usage);
    }

    fn fmt_with_indent(&self, f: &mut std::fmt::Formatter<'_>, depth: usize) -> std::fmt::Result {
        let width = 20;
        for (name, child) in &self.children {
            let indent = " ".repeat(depth);

            if let Some(usage) = child.usage {
                writeln!(
                    f,
                    "{indent}{:<} real: {:>6?}, user: {:>6?}, sys: {:?}",
                    name,
                    usage.wtime.as_millis(),
                    usage.utime.as_millis(),
                    usage.stime.as_millis(),
                )?;
            } else {
                writeln!(f, "{indent}{name}")?;
            }
            child.fmt_with_indent(f, depth + 1)?;
        }
        Ok(())
    }
}

/* Summarize timing variables to FP.  The timing variable TV_TOTAL has
a special meaning -- it's considered to be the total elapsed time,
for normalizing the others, and is displayed last.  */
impl std::fmt::Display for TreeReport {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "--- Resource Usage Report ---")?;
        self.0.fmt_with_indent(f, 0)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_simple() {
        let mut stk = ResourceMeterStack::new();
        let now = Instant::now();
        std::thread::sleep(Duration::from_millis(200));
        stk.push("sleep600");
        stk.push("sleep600/300");
        std::thread::sleep(Duration::from_millis(300));
        stk.pop();
        std::thread::sleep(Duration::from_millis(100));
        stk.push("sleep600/200");
        std::thread::sleep(Duration::from_millis(200));
        stk.pop();
        stk.pop();
        std::thread::sleep(Duration::from_millis(200));
        let report = stk.finish();
        eprintln!("{}, took: {:?}", report, now.elapsed());
    }

    #[test]
    fn test_resource_usage() {
        let rusage = ResourceUsage::now();
        eprintln!("{:#?}", rusage);
    }
}