Skip to main content

ferrum_types/
process_memory.rs

1use crate::MemorySnapshot;
2
3#[derive(Clone, Debug, PartialEq, Eq)]
4pub struct ProcessMemorySample {
5    pub current_bytes: u64,
6    pub high_water_bytes: u64,
7    pub source: &'static str,
8}
9
10#[derive(Clone, Debug, PartialEq, Eq)]
11pub struct ProcessMemoryObservation {
12    pub before_bytes: u64,
13    pub after_bytes: u64,
14    pub current_bytes: u64,
15    pub high_water_bytes: u64,
16    pub source: &'static str,
17}
18
19#[derive(Clone, Copy, Debug, Default)]
20pub struct ProcessMemorySampler;
21
22impl ProcessMemorySampler {
23    pub fn sample(&self) -> Option<ProcessMemorySample> {
24        sample_process_memory()
25    }
26
27    pub fn observe(&self, before: Option<ProcessMemorySample>) -> Option<ProcessMemoryObservation> {
28        let after = self.sample()?;
29        Some(ProcessMemoryObservation::from_samples(before, after))
30    }
31}
32
33impl ProcessMemoryObservation {
34    pub fn from_samples(before: Option<ProcessMemorySample>, after: ProcessMemorySample) -> Self {
35        let before_bytes = before
36            .as_ref()
37            .map(|sample| sample.current_bytes)
38            .unwrap_or(after.current_bytes);
39        let high_water_bytes = before
40            .as_ref()
41            .map(|sample| sample.high_water_bytes)
42            .unwrap_or(0)
43            .max(after.high_water_bytes)
44            .max(before_bytes)
45            .max(after.current_bytes);
46        Self {
47            before_bytes,
48            after_bytes: after.current_bytes,
49            current_bytes: after.current_bytes,
50            high_water_bytes,
51            source: after.source,
52        }
53    }
54
55    pub fn from_sample(sample: ProcessMemorySample) -> Self {
56        Self::from_samples(Some(sample.clone()), sample)
57    }
58
59    pub fn to_snapshot(&self, scope: impl Into<String>, backend: Option<&str>) -> MemorySnapshot {
60        MemorySnapshot {
61            scope: scope.into(),
62            backend: backend.map(str::to_string),
63            before_bytes: Some(self.before_bytes),
64            after_bytes: Some(self.after_bytes),
65            current_bytes: Some(self.current_bytes),
66            high_water_bytes: Some(self.high_water_bytes),
67            available_bytes: None,
68        }
69    }
70}
71
72#[cfg(target_os = "linux")]
73fn current_resident_bytes() -> Option<u64> {
74    let statm = std::fs::read_to_string("/proc/self/statm").ok()?;
75    let resident_pages = statm.split_whitespace().nth(1)?.parse::<u64>().ok()?;
76    let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
77    if page_size <= 0 {
78        return None;
79    }
80    resident_pages.checked_mul(page_size as u64)
81}
82
83#[cfg(not(target_os = "linux"))]
84fn current_resident_bytes() -> Option<u64> {
85    None
86}
87
88#[cfg(unix)]
89fn high_water_bytes() -> Option<u64> {
90    let mut usage = std::mem::MaybeUninit::<libc::rusage>::zeroed();
91    let rc = unsafe { libc::getrusage(libc::RUSAGE_SELF, usage.as_mut_ptr()) };
92    if rc != 0 {
93        return None;
94    }
95    let max_rss = unsafe { usage.assume_init() }.ru_maxrss;
96    if max_rss <= 0 {
97        return None;
98    }
99    #[cfg(target_os = "macos")]
100    {
101        Some(max_rss as u64)
102    }
103    #[cfg(not(target_os = "macos"))]
104    {
105        (max_rss as u64).checked_mul(1024)
106    }
107}
108
109#[cfg(not(unix))]
110fn high_water_bytes() -> Option<u64> {
111    None
112}
113
114pub fn sample_process_memory() -> Option<ProcessMemorySample> {
115    let high_water = high_water_bytes()?;
116    let current = current_resident_bytes().unwrap_or(high_water);
117    Some(ProcessMemorySample {
118        current_bytes: current,
119        high_water_bytes: high_water.max(current),
120        source: process_memory_source(),
121    })
122}
123
124#[cfg(target_os = "linux")]
125fn process_memory_source() -> &'static str {
126    "procfs_statm_plus_getrusage"
127}
128
129#[cfg(all(unix, not(target_os = "linux")))]
130fn process_memory_source() -> &'static str {
131    "getrusage_maxrss"
132}
133
134#[cfg(not(unix))]
135fn process_memory_source() -> &'static str {
136    "unsupported"
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn process_memory_observation_builds_valid_snapshot() {
145        let before = ProcessMemorySample {
146            current_bytes: 100,
147            high_water_bytes: 110,
148            source: "test",
149        };
150        let after = ProcessMemorySample {
151            current_bytes: 150,
152            high_water_bytes: 160,
153            source: "test",
154        };
155        let observation = ProcessMemoryObservation::from_samples(Some(before), after);
156        let snapshot = observation.to_snapshot("process", Some("actual"));
157        assert_eq!(snapshot.before_bytes, Some(100));
158        assert_eq!(snapshot.after_bytes, Some(150));
159        assert_eq!(snapshot.current_bytes, Some(150));
160        assert_eq!(snapshot.high_water_bytes, Some(160));
161        snapshot.validate().unwrap();
162    }
163
164    #[cfg(unix)]
165    #[test]
166    fn process_memory_sampler_returns_non_zero_on_unix() {
167        let sample = ProcessMemorySampler.sample();
168        assert!(sample
169            .as_ref()
170            .is_some_and(|sample| sample.current_bytes > 0));
171        assert!(sample
172            .as_ref()
173            .is_some_and(|sample| sample.high_water_bytes > 0));
174    }
175}