statsig_rust/interned_values/interned_store/mmap_reader/
memory.rs

1use std::fs::File;
2
3use memmap2::Mmap;
4
5use crate::StatsigErr;
6
7/// Process-local memory accounting for the currently loaded interned mmap
8/// reader generation.
9///
10/// The snapshot contains no SDK key, filesystem path, inode, device, or virtual
11/// address. Linux residency fields are derived from the exact address range of
12/// the retained mapping. Platforms without an equivalent implementation leave
13/// optional residency fields unset while still reporting format and mapped
14/// bytes. A failed Linux probe likewise leaves only its optional fields unset.
15#[non_exhaustive]
16#[derive(Clone, Debug, Eq, PartialEq)]
17pub struct MmapReaderMemorySnapshot {
18    pub format_version: u32,
19    pub mapped_bytes: u64,
20    pub resident_bytes: Option<u64>,
21    pub proportional_set_bytes: Option<u64>,
22    pub private_dirty_bytes: Option<u64>,
23    pub deleted_mapped_bytes: Option<u64>,
24    /// Number of reader generations retained by this process. The current
25    /// single-generation reader reports one when a snapshot is present.
26    pub loaded_generation_count: u64,
27    pub vma_segment_count: Option<u64>,
28}
29
30pub(super) fn snapshot(
31    file: &File,
32    mmap: &Mmap,
33    format_version: u32,
34) -> Result<MmapReaderMemorySnapshot, StatsigErr> {
35    let mapped_bytes = u64::try_from(mmap.len()).map_err(|_| {
36        StatsigErr::InvalidOperation("Loaded mmap length does not fit in u64".to_string())
37    })?;
38
39    #[cfg(target_os = "linux")]
40    {
41        use std::os::unix::fs::MetadataExt;
42
43        let residency = linux_residency(mmap.as_ptr() as usize, mmap.len()).ok();
44        let link_count = file.metadata().ok().map(|metadata| metadata.nlink());
45
46        Ok(linux_snapshot(
47            format_version,
48            mapped_bytes,
49            residency,
50            link_count,
51        ))
52    }
53
54    #[cfg(not(target_os = "linux"))]
55    {
56        let _ = file;
57        Ok(MmapReaderMemorySnapshot {
58            format_version,
59            mapped_bytes,
60            resident_bytes: None,
61            proportional_set_bytes: None,
62            private_dirty_bytes: None,
63            deleted_mapped_bytes: None,
64            loaded_generation_count: 1,
65            vma_segment_count: None,
66        })
67    }
68}
69
70#[cfg(target_os = "linux")]
71fn linux_snapshot(
72    format_version: u32,
73    mapped_bytes: u64,
74    residency: Option<LinuxMmapResidency>,
75    link_count: Option<u64>,
76) -> MmapReaderMemorySnapshot {
77    let (resident_bytes, proportional_set_bytes, private_dirty_bytes, vma_segment_count) =
78        match residency {
79            Some(residency) => (
80                Some(residency.resident_bytes),
81                Some(residency.proportional_set_bytes),
82                Some(residency.private_dirty_bytes),
83                Some(residency.vma_segment_count),
84            ),
85            None => (None, None, None, None),
86        };
87
88    MmapReaderMemorySnapshot {
89        format_version,
90        mapped_bytes,
91        resident_bytes,
92        proportional_set_bytes,
93        private_dirty_bytes,
94        deleted_mapped_bytes: link_count.map(|count| if count == 0 { mapped_bytes } else { 0 }),
95        loaded_generation_count: 1,
96        vma_segment_count,
97    }
98}
99
100#[cfg(target_os = "linux")]
101#[derive(Debug, Default, Eq, PartialEq)]
102struct LinuxMmapResidency {
103    resident_bytes: u64,
104    proportional_set_bytes: u64,
105    private_dirty_bytes: u64,
106    vma_segment_count: u64,
107}
108
109#[cfg(target_os = "linux")]
110fn linux_residency(
111    mapping_start: usize,
112    mapping_len: usize,
113) -> Result<LinuxMmapResidency, StatsigErr> {
114    use std::io::BufReader;
115
116    let smaps =
117        File::open("/proc/self/smaps").map_err(|error| StatsigErr::FileError(error.to_string()))?;
118    parse_smaps(BufReader::new(smaps), mapping_start, mapping_len)
119}
120
121#[cfg(target_os = "linux")]
122fn parse_smaps(
123    mut reader: impl std::io::BufRead,
124    mapping_start: usize,
125    mapping_len: usize,
126) -> Result<LinuxMmapResidency, StatsigErr> {
127    let mapping_end = mapping_start.checked_add(mapping_len).ok_or_else(|| {
128        StatsigErr::InvalidOperation("Loaded mmap address range overflowed".to_string())
129    })?;
130    if mapping_start == mapping_end {
131        return Err(StatsigErr::InvalidOperation(
132            "Loaded mmap address range is empty".to_string(),
133        ));
134    }
135
136    let mut totals = LinuxMmapResidency::default();
137    let mut overlaps_mapping = false;
138    let mut line = String::new();
139    loop {
140        line.clear();
141        let bytes_read = reader
142            .read_line(&mut line)
143            .map_err(|error| StatsigErr::FileError(error.to_string()))?;
144        if bytes_read == 0 {
145            break;
146        }
147
148        if let Some((vma_start, vma_end)) = parse_vma_header(&line)? {
149            if vma_start >= mapping_end {
150                break;
151            }
152            overlaps_mapping = vma_start < mapping_end && mapping_start < vma_end;
153            if overlaps_mapping {
154                totals.vma_segment_count =
155                    checked_add(totals.vma_segment_count, 1, "VMA segment count overflowed")?;
156            }
157            continue;
158        }
159        if !overlaps_mapping {
160            continue;
161        }
162
163        for (field, destination) in [
164            ("Rss:", &mut totals.resident_bytes),
165            ("Pss:", &mut totals.proportional_set_bytes),
166            ("Private_Dirty:", &mut totals.private_dirty_bytes),
167        ] {
168            if let Some(value) = parse_kibibyte_field(&line, field)? {
169                *destination = checked_add(*destination, value, "smaps byte count overflowed")?;
170                break;
171            }
172        }
173    }
174
175    if totals.vma_segment_count == 0 {
176        return Err(StatsigErr::FileError(
177            "Loaded mmap range was not present in process smaps".to_string(),
178        ));
179    }
180    Ok(totals)
181}
182
183#[cfg(target_os = "linux")]
184fn parse_vma_header(line: &str) -> Result<Option<(usize, usize)>, StatsigErr> {
185    let mut fields = line.split_ascii_whitespace();
186    let Some(range) = fields.next() else {
187        return Ok(None);
188    };
189    let Some(permissions) = fields.next() else {
190        return Ok(None);
191    };
192    let permissions = permissions.as_bytes();
193    if permissions.len() != 4
194        || !matches!(permissions[0], b'r' | b'-')
195        || !matches!(permissions[1], b'w' | b'-')
196        || !matches!(permissions[2], b'x' | b'-')
197        || !matches!(permissions[3], b'p' | b's')
198    {
199        return Ok(None);
200    }
201    let Some((start, end)) = range.split_once('-') else {
202        return Err(StatsigErr::FileError(
203            "Invalid VMA range in process smaps".to_string(),
204        ));
205    };
206    let parse_address = |value: &str| {
207        usize::from_str_radix(value, 16)
208            .map_err(|_| StatsigErr::FileError("Invalid VMA range in process smaps".to_string()))
209    };
210    let start = parse_address(start)?;
211    let end = parse_address(end)?;
212    if start >= end {
213        return Err(StatsigErr::FileError(
214            "Invalid VMA range in process smaps".to_string(),
215        ));
216    }
217    Ok(Some((start, end)))
218}
219
220#[cfg(target_os = "linux")]
221fn parse_kibibyte_field(line: &str, expected_field: &str) -> Result<Option<u64>, StatsigErr> {
222    const KIBIBYTE: u64 = 1024;
223
224    let mut fields = line.split_ascii_whitespace();
225    if fields.next() != Some(expected_field) {
226        return Ok(None);
227    }
228    let value = fields
229        .next()
230        .ok_or_else(|| invalid_smaps_field(expected_field))?
231        .parse::<u64>()
232        .map_err(|_| invalid_smaps_field(expected_field))?;
233    if fields.next() != Some("kB") {
234        return Err(invalid_smaps_field(expected_field));
235    }
236    value
237        .checked_mul(KIBIBYTE)
238        .map(Some)
239        .ok_or_else(|| invalid_smaps_field(expected_field))
240}
241
242#[cfg(target_os = "linux")]
243fn checked_add(left: u64, right: u64, message: &str) -> Result<u64, StatsigErr> {
244    left.checked_add(right)
245        .ok_or_else(|| StatsigErr::FileError(message.to_string()))
246}
247
248#[cfg(target_os = "linux")]
249fn invalid_smaps_field(field: &str) -> StatsigErr {
250    StatsigErr::FileError(format!("Invalid {field} field in process smaps"))
251}
252
253#[cfg(all(test, target_os = "linux"))]
254mod tests {
255    use std::io::Cursor;
256
257    use super::*;
258
259    const SMAPS: &str = "\
26000000000-00001000 r--p 00000000 00:00 0\n\
261Rss:                 100 kB\n\
262Pss:                  90 kB\n\
263Private_Dirty:        80 kB\n\
26400001000-00002000 r--s 00000000 00:00 0\n\
265Rss:                   4 kB\n\
266Pss:                   2 kB\n\
267Private_Dirty:         1 kB\n\
26800002000-00003000 r--s 00001000 00:00 0\n\
269Rss:                   8 kB\n\
270Pss:                   5 kB\n\
271Private_Dirty:         3 kB\n\
27200003000-00004000 r--p 00000000 00:00 0\n\
273Rss:                 200 kB\n\
274Pss:                 190 kB\n\
275Private_Dirty:       180 kB\n";
276
277    #[test]
278    fn smaps_parser_sums_only_overlapping_vmas() {
279        let result = parse_smaps(Cursor::new(SMAPS), 0x1800, 0x1000).unwrap();
280        assert_eq!(
281            result,
282            LinuxMmapResidency {
283                resident_bytes: 12 * 1024,
284                proportional_set_bytes: 7 * 1024,
285                private_dirty_bytes: 4 * 1024,
286                vma_segment_count: 2,
287            }
288        );
289    }
290
291    #[test]
292    fn smaps_parser_treats_vma_boundaries_as_non_overlapping() {
293        let result = parse_smaps(Cursor::new(SMAPS), 0x2000, 0x1000).unwrap();
294        assert_eq!(result.resident_bytes, 8 * 1024);
295        assert_eq!(result.vma_segment_count, 1);
296    }
297
298    #[test]
299    fn smaps_parser_stops_after_the_target_range() {
300        let smaps = "\
30100001000-00002000 r--p 00000000 00:00 0\n\
302Rss:                   4 kB\n\
30300002000-00003000 r--p 00000000 00:00 0\n\
304not-a-range r--p 00000000 00:00 0\n";
305        let result = parse_smaps(Cursor::new(smaps), 0x1000, 0x1000).unwrap();
306        assert_eq!(result.resident_bytes, 4 * 1024);
307        assert_eq!(result.vma_segment_count, 1);
308    }
309
310    #[test]
311    fn smaps_parser_rejects_malformed_metrics_and_missing_mapping() {
312        let malformed = "1000-2000 r--s 00000000 00:00 0\nRss: nope kB\n";
313        assert!(matches!(
314            parse_smaps(Cursor::new(malformed), 0x1000, 0x1000),
315            Err(StatsigErr::FileError(message)) if message.contains("Rss:")
316        ));
317        assert!(matches!(
318            parse_smaps(Cursor::new(SMAPS), 0x5000, 0x1000),
319            Err(StatsigErr::FileError(message))
320                if message == "Loaded mmap range was not present in process smaps"
321        ));
322    }
323
324    #[test]
325    fn snapshot_retains_baseline_when_smaps_probe_fails() {
326        let snapshot = linux_snapshot(2, 4096, None, Some(0));
327
328        assert_eq!(snapshot.format_version, 2);
329        assert_eq!(snapshot.mapped_bytes, 4096);
330        assert_eq!(snapshot.loaded_generation_count, 1);
331        assert_eq!(snapshot.resident_bytes, None);
332        assert_eq!(snapshot.proportional_set_bytes, None);
333        assert_eq!(snapshot.private_dirty_bytes, None);
334        assert_eq!(snapshot.vma_segment_count, None);
335        assert_eq!(snapshot.deleted_mapped_bytes, Some(4096));
336    }
337
338    #[test]
339    fn snapshot_retains_residency_when_fstat_probe_fails() {
340        let snapshot = linux_snapshot(
341            2,
342            4096,
343            Some(LinuxMmapResidency {
344                resident_bytes: 2048,
345                proportional_set_bytes: 1024,
346                private_dirty_bytes: 512,
347                vma_segment_count: 1,
348            }),
349            None,
350        );
351
352        assert_eq!(snapshot.format_version, 2);
353        assert_eq!(snapshot.mapped_bytes, 4096);
354        assert_eq!(snapshot.loaded_generation_count, 1);
355        assert_eq!(snapshot.resident_bytes, Some(2048));
356        assert_eq!(snapshot.proportional_set_bytes, Some(1024));
357        assert_eq!(snapshot.private_dirty_bytes, Some(512));
358        assert_eq!(snapshot.vma_segment_count, Some(1));
359        assert_eq!(snapshot.deleted_mapped_bytes, None);
360    }
361}
statsig_rust/interned_values/interned_store/mmap_reader/memory.rs

statsig_rust/interned_values/interned_store/mmap_reader/
memory.rs
