Skip to main content

vtcode_commons/
memory.rs

1//! Resident Set Size (RSS) sampling for memory diagnostics.
2//!
3//! Used by the allocator benchmark (`vtcode bench-allocator`) to measure whether
4//! the global allocator returns memory to the OS after bursty/sparse workloads.
5//! Unlike `performance_profiler::get_memory_usage_mb` (Linux `/proc` only, fake
6//! fallback on macOS), this returns a real value on every supported platform.
7use std::time::Duration;
8
9/// Returns the current process Resident Set Size in **megabytes**, or `None` if
10/// it cannot be determined on the current platform.
11#[cfg(target_os = "macos")]
12#[allow(deprecated, unsafe_code, unused_qualifications)] // libc::mach_task_self is deprecated; qualification is required here
13pub fn resident_set_size_mb() -> Option<f64> {
14    // SAFETY: `mach_task_basic_info` is a plain old data struct; zeroing it
15    // produces a valid (all-zero) starting value before `task_info` fills it.
16    let mut info: libc::mach_task_basic_info = unsafe { std::mem::zeroed() };
17    let mut count = (std::mem::size_of::<libc::mach_task_basic_info>()
18        / std::mem::size_of::<libc::integer_t>())
19        as libc::mach_msg_type_number_t;
20    // SAFETY: `mach_task_self()` returns a send-right to the current task with
21    // no preconditions; it cannot fail to produce a valid port name.
22    let task = unsafe { libc::mach_task_self() };
23    // SAFETY: `task` is our own task port; `info` and `count` are valid
24    // out-pointers of the expected size, and `task_info` only writes them on
25    // success.
26    let ret = unsafe {
27        libc::task_info(
28            task,
29            libc::MACH_TASK_BASIC_INFO,
30            &mut info as *mut _ as *mut libc::integer_t,
31            &mut count,
32        )
33    };
34    if ret != libc::KERN_SUCCESS {
35        return None;
36    }
37    Some(info.resident_size as f64 / (1024.0 * 1024.0))
38}
39
40/// Returns the current process Resident Set Size in **megabytes**, or `None` if
41/// it cannot be determined on the current platform.
42#[cfg(target_os = "linux")]
43pub fn resident_set_size_mb() -> Option<f64> {
44    let contents = std::fs::read_to_string("/proc/self/statm").ok()?;
45    let field = contents.split_whitespace().nth(1)?;
46    let pages: f64 = field.parse().ok()?;
47    let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) } as f64;
48    Some(pages * page_size / (1024.0 * 1024.0))
49}
50
51/// Fallback for unsupported platforms.
52#[cfg(not(any(target_os = "macos", target_os = "linux")))]
53pub fn resident_set_size_mb() -> Option<f64> {
54    None
55}
56
57/// Sample RSS once and return the value in MB (0.0 if unavailable).
58pub fn sample_rss_mb() -> f64 {
59    resident_set_size_mb().unwrap_or(0.0)
60}
61
62/// Sample RSS repeatedly, returning the maximum observed value in MB.
63/// Useful for capturing peak memory during a burst of activity.
64pub fn sample_peak_rss_mb(duration: Duration, poll_interval: Duration) -> f64 {
65    let start = std::time::Instant::now();
66    let mut peak = 0.0;
67    while start.elapsed() < duration {
68        let v = sample_rss_mb();
69        if v > peak {
70            peak = v;
71        }
72        std::thread::sleep(poll_interval);
73    }
74    peak
75}