lb_rs/service/
perf.rs

1//! Process RSS and CPU counters. Wall time is the parent span.
2
3use std::time::Duration;
4
5/// Emits `counter.rss_bytes` / `counter.cpu_us` on drop.
6pub struct Sample {
7    cpu: Option<Duration>,
8}
9
10impl Sample {
11    pub fn new() -> Option<Self> {
12        if !tracing::enabled!(tracing::Level::TRACE) {
13            return None;
14        }
15        Some(Self { cpu: process_cpu() })
16    }
17}
18
19impl Drop for Sample {
20    fn drop(&mut self) {
21        let cpu_us = process_cpu()
22            .and_then(|end| self.cpu.map(|start| end.saturating_sub(start)))
23            .map(|d| d.as_micros() as u64);
24        let rss = current_rss_bytes();
25        match (cpu_us, rss) {
26            (Some(cpu_us), Some(rss_bytes)) => {
27                tracing::trace!(counter.rss_bytes = rss_bytes, counter.cpu_us = cpu_us);
28            }
29            (Some(cpu_us), None) => {
30                tracing::trace!(counter.cpu_us = cpu_us);
31            }
32            (None, Some(rss_bytes)) => {
33                tracing::trace!(counter.rss_bytes = rss_bytes);
34            }
35            (None, None) => {}
36        }
37    }
38}
39
40#[cfg(not(any(unix, windows)))]
41fn process_cpu() -> Option<Duration> {
42    None
43}
44
45#[cfg(not(any(unix, windows)))]
46fn current_rss_bytes() -> Option<u64> {
47    None
48}
49
50#[cfg(unix)]
51fn process_cpu() -> Option<Duration> {
52    // SAFETY: `usg` is zeroed; getrusage writes it on success.
53    unsafe {
54        let mut usg = std::mem::zeroed();
55        if libc::getrusage(libc::RUSAGE_SELF, &mut usg) != 0 {
56            return None;
57        }
58        Some(timeval(usg.ru_utime) + timeval(usg.ru_stime))
59    }
60}
61
62#[cfg(unix)]
63fn timeval(tv: libc::timeval) -> Duration {
64    Duration::new(tv.tv_sec as u64, (tv.tv_usec as u32).saturating_mul(1_000))
65}
66
67#[cfg(target_os = "macos")]
68#[allow(deprecated)]
69fn current_rss_bytes() -> Option<u64> {
70    // SAFETY: `info`/`count` are valid out-params for MACH_TASK_BASIC_INFO.
71    unsafe {
72        let mut info: libc::mach_task_basic_info = std::mem::zeroed();
73        let mut count = libc::MACH_TASK_BASIC_INFO_COUNT;
74        let kr = libc::task_info(
75            libc::mach_task_self(),
76            libc::MACH_TASK_BASIC_INFO,
77            &mut info as *mut _ as libc::task_info_t,
78            &mut count,
79        );
80        if kr == libc::KERN_SUCCESS { Some(info.resident_size) } else { None }
81    }
82}
83
84#[cfg(all(unix, not(target_os = "macos")))]
85fn current_rss_bytes() -> Option<u64> {
86    let statm = std::fs::read_to_string("/proc/self/statm").ok()?;
87    let pages: u64 = statm.split_whitespace().nth(1)?.parse().ok()?;
88    // SAFETY: sysconf(_SC_PAGESIZE) takes no pointer args.
89    let page = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
90    if page <= 0 {
91        return None;
92    }
93    Some(pages.saturating_mul(page as u64))
94}
95
96#[cfg(windows)]
97fn process_cpu() -> Option<Duration> {
98    use windows_sys::Win32::Foundation::FILETIME;
99    use windows_sys::Win32::System::Threading::{GetCurrentProcess, GetProcessTimes};
100
101    // SAFETY: FILETIME out-params are zeroed; GetProcessTimes writes them on success.
102    unsafe {
103        let mut created = FILETIME { dwLowDateTime: 0, dwHighDateTime: 0 };
104        let mut exited = FILETIME { dwLowDateTime: 0, dwHighDateTime: 0 };
105        let mut kernel = FILETIME { dwLowDateTime: 0, dwHighDateTime: 0 };
106        let mut user = FILETIME { dwLowDateTime: 0, dwHighDateTime: 0 };
107        if GetProcessTimes(GetCurrentProcess(), &mut created, &mut exited, &mut kernel, &mut user)
108            == 0
109        {
110            return None;
111        }
112        Some(filetime(user) + filetime(kernel))
113    }
114}
115
116#[cfg(windows)]
117fn filetime(ft: windows_sys::Win32::Foundation::FILETIME) -> Duration {
118    let ticks = ((ft.dwHighDateTime as u64) << 32) | ft.dwLowDateTime as u64;
119    // FILETIME is 100 ns intervals.
120    Duration::from_nanos(ticks.saturating_mul(100))
121}
122
123#[cfg(windows)]
124fn current_rss_bytes() -> Option<u64> {
125    use windows_sys::Win32::System::ProcessStatus::{
126        GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS,
127    };
128    use windows_sys::Win32::System::Threading::GetCurrentProcess;
129
130    // SAFETY: pmc.cb is the struct size GetProcessMemoryInfo requires.
131    unsafe {
132        let mut pmc: PROCESS_MEMORY_COUNTERS = std::mem::zeroed();
133        pmc.cb = std::mem::size_of::<PROCESS_MEMORY_COUNTERS>() as u32;
134        if GetProcessMemoryInfo(GetCurrentProcess(), &mut pmc, pmc.cb) == 0 {
135            return None;
136        }
137        Some(pmc.WorkingSetSize as u64)
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn samples_do_not_panic() {
147        let _ = process_cpu();
148        let rss = current_rss_bytes();
149        assert!(rss.unwrap_or(0) > 0, "expected a current RSS reading, got {rss:?}");
150        let _s = Sample::new();
151    }
152}
lb_rs/service/perf.rs

lb_rs/service/
perf.rs
