dataprof_core/
memory_sampler.rs

1use std::time::{Duration, Instant};
2
3use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, System};
4
5/// Minimum spacing between two throttled samples.
6///
7/// One refresh costs about 14 ms on Windows, where sysinfo walks the process
8/// table even when asked for a single process, and the streaming engines
9/// sample once per chunk. Small chunks paid that cost thousands of times: a
10/// 300,000-row CSV took 345 s with 1 KB chunks against 1.6 s with 1 MB ones.
11/// Resident memory moves slowly next to that, so sampling at most four times
12/// a second loses nothing a peak reading needs, and callers take one
13/// unthrottled sample at the end.
14const MIN_SAMPLE_INTERVAL: Duration = Duration::from_millis(250);
15
16/// Tracks the peak resident set size (RSS) of the current process across a
17/// profiling run.
18///
19/// Engines call [`sample`](Self::sample) at chunk/batch boundaries and
20/// [`sample_now`](Self::sample_now) once at the end; the peak observed feeds
21/// `ExecutionMetadata::memory_peak_mb`. The reading is the
22/// whole-process RSS, so on embedded surfaces (e.g. the Python extension) it
23/// includes host-process memory — it is an upper bound on what profiling used,
24/// suitable for verifying the bounded-memory claim, not an exact attribution.
25pub struct PeakMemorySampler {
26    system: System,
27    pid: Option<Pid>,
28    peak_bytes: u64,
29    last_sample: Option<Instant>,
30    refreshes: u64,
31}
32
33impl PeakMemorySampler {
34    /// Create a sampler and take an initial reading.
35    pub fn new() -> Self {
36        let mut sampler = Self {
37            system: System::new(),
38            pid: sysinfo::get_current_pid().ok(),
39            peak_bytes: 0,
40            last_sample: None,
41            refreshes: 0,
42        };
43        sampler.sample_now();
44        sampler
45    }
46
47    /// Sample unless the previous sample is less than [`MIN_SAMPLE_INTERVAL`]
48    /// old. For per-chunk calls inside a loop.
49    pub fn sample(&mut self) {
50        if self
51            .last_sample
52            .is_some_and(|last| last.elapsed() < MIN_SAMPLE_INTERVAL)
53        {
54            return;
55        }
56        self.sample_now();
57    }
58
59    /// Refresh the process RSS and fold it into the running peak, whatever
60    /// the interval. For the final reading of a run, so the reported peak
61    /// covers the work done after the last throttled sample.
62    pub fn sample_now(&mut self) {
63        self.last_sample = Some(Instant::now());
64        let Some(pid) = self.pid else { return };
65        self.refreshes += 1;
66        self.system.refresh_processes_specifics(
67            ProcessesToUpdate::Some(&[pid]),
68            false,
69            ProcessRefreshKind::nothing().with_memory(),
70        );
71        if let Some(process) = self.system.process(pid) {
72            self.peak_bytes = self.peak_bytes.max(process.memory());
73        }
74    }
75
76    /// Peak RSS observed so far, in megabytes.
77    ///
78    /// `None` when the platform provided no reading — callers should leave
79    /// `memory_peak_mb` absent (not analyzed) rather than report zero.
80    pub fn peak_mb(&self) -> Option<f64> {
81        (self.peak_bytes > 0).then(|| self.peak_bytes as f64 / (1024.0 * 1024.0))
82    }
83}
84
85impl Default for PeakMemorySampler {
86    fn default() -> Self {
87        Self::new()
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn repeated_samples_within_the_interval_do_not_refresh() {
97        let mut sampler = PeakMemorySampler::new();
98        let before = sampler.refreshes;
99        for _ in 0..1_000 {
100            sampler.sample();
101        }
102        assert_eq!(sampler.refreshes, before, "a tight loop must not refresh");
103        sampler.sample_now();
104        assert_eq!(sampler.refreshes, before + 1, "sample_now always refreshes");
105    }
106
107    #[test]
108    fn test_sampler_reports_positive_peak() {
109        let sampler = PeakMemorySampler::new();
110        let peak = sampler
111            .peak_mb()
112            .expect("current process RSS should be readable on supported platforms");
113        assert!(peak > 0.0, "peak RSS must be positive, got {peak}");
114    }
115
116    #[test]
117    fn test_peak_is_monotonic_across_samples() {
118        let mut sampler = PeakMemorySampler::new();
119        let first = sampler.peak_mb().expect("initial reading");
120        // Allocate enough that RSS cannot legitimately shrink below the first peak.
121        let ballast = vec![1u8; 4 * 1024 * 1024];
122        sampler.sample_now();
123        let second = sampler.peak_mb().expect("second reading");
124        assert!(
125            second >= first,
126            "peak must never decrease: {first} -> {second}"
127        );
128        drop(ballast);
129    }
130}
dataprof_core/memory_sampler.rs

dataprof_core/
memory_sampler.rs
