dataprof_core/
memory_sampler.rs

1use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, System};
2
3/// Tracks the peak resident set size (RSS) of the current process across a
4/// profiling run.
5///
6/// Engines call [`sample`](Self::sample) at chunk/batch boundaries; the peak
7/// observed feeds `ExecutionMetadata::memory_peak_mb`. The reading is the
8/// whole-process RSS, so on embedded surfaces (e.g. the Python extension) it
9/// includes host-process memory — it is an upper bound on what profiling used,
10/// suitable for verifying the bounded-memory claim, not an exact attribution.
11pub struct PeakMemorySampler {
12    system: System,
13    pid: Option<Pid>,
14    peak_bytes: u64,
15}
16
17impl PeakMemorySampler {
18    /// Create a sampler and take an initial reading.
19    pub fn new() -> Self {
20        let mut sampler = Self {
21            system: System::new(),
22            pid: sysinfo::get_current_pid().ok(),
23            peak_bytes: 0,
24        };
25        sampler.sample();
26        sampler
27    }
28
29    /// Refresh the process RSS and fold it into the running peak.
30    pub fn sample(&mut self) {
31        let Some(pid) = self.pid else { return };
32        self.system.refresh_processes_specifics(
33            ProcessesToUpdate::Some(&[pid]),
34            false,
35            ProcessRefreshKind::nothing().with_memory(),
36        );
37        if let Some(process) = self.system.process(pid) {
38            self.peak_bytes = self.peak_bytes.max(process.memory());
39        }
40    }
41
42    /// Peak RSS observed so far, in megabytes.
43    ///
44    /// `None` when the platform provided no reading — callers should leave
45    /// `memory_peak_mb` absent (not analyzed) rather than report zero.
46    pub fn peak_mb(&self) -> Option<f64> {
47        (self.peak_bytes > 0).then(|| self.peak_bytes as f64 / (1024.0 * 1024.0))
48    }
49}
50
51impl Default for PeakMemorySampler {
52    fn default() -> Self {
53        Self::new()
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60
61    #[test]
62    fn test_sampler_reports_positive_peak() {
63        let sampler = PeakMemorySampler::new();
64        let peak = sampler
65            .peak_mb()
66            .expect("current process RSS should be readable on supported platforms");
67        assert!(peak > 0.0, "peak RSS must be positive, got {peak}");
68    }
69
70    #[test]
71    fn test_peak_is_monotonic_across_samples() {
72        let mut sampler = PeakMemorySampler::new();
73        let first = sampler.peak_mb().expect("initial reading");
74        // Allocate enough that RSS cannot legitimately shrink below the first peak.
75        let ballast = vec![1u8; 4 * 1024 * 1024];
76        sampler.sample();
77        let second = sampler.peak_mb().expect("second reading");
78        assert!(
79            second >= first,
80            "peak must never decrease: {first} -> {second}"
81        );
82        drop(ballast);
83    }
84}
dataprof_core/memory_sampler.rs

dataprof_core/
memory_sampler.rs
