use std::collections::HashMap;
use procfs::{process, process::Process};
use super::ProcessSampler;
use crate::sampler::ProcessSnapshot;
pub struct LinuxSampler {
ticks_per_second: f64,
}
impl LinuxSampler {
#[must_use]
pub fn new() -> Self {
Self { ticks_per_second: procfs::ticks_per_second() as f64 }
}
}
impl Default for LinuxSampler {
fn default() -> Self {
Self::new()
}
}
impl ProcessSampler for LinuxSampler {
fn sample_tree(&mut self, root_pid: i32) -> Vec<ProcessSnapshot> {
let pids = collect_descendants(root_pid);
pids.into_iter().filter_map(|pid| sample_process(pid, self.ticks_per_second)).collect()
}
}
fn collect_descendants(root: i32) -> Vec<i32> {
let mut by_parent: HashMap<i32, Vec<i32>> = HashMap::new();
if let Ok(iter) = process::all_processes() {
for proc in iter.flatten() {
if let Ok(stat) = proc.stat() {
by_parent.entry(stat.ppid).or_default().push(stat.pid);
}
}
}
let mut out = vec![root];
let mut stack = vec![root];
while let Some(parent) = stack.pop() {
if let Some(children) = by_parent.get(&parent) {
for &child in children {
out.push(child);
stack.push(child);
}
}
}
out
}
fn sample_process(pid: i32, ticks_per_second: f64) -> Option<ProcessSnapshot> {
let proc = Process::new(pid).ok()?;
let stat = proc.stat().ok()?;
let status = proc.status().ok()?;
let rss_bytes = status.vmrss.unwrap_or(0).saturating_mul(1024);
let vms_bytes = status.vmsize.unwrap_or(0).saturating_mul(1024);
let (uss_bytes, pss_bytes) = read_smaps_rollup(&proc);
let (io_read_bytes, io_write_bytes) = match proc.io() {
Ok(io) => (Some(io.read_bytes), Some(io.write_bytes)),
Err(_) => (None, None),
};
let cpu_ticks = stat.utime.saturating_add(stat.stime);
let cpu_time_seconds = (cpu_ticks as f64) / ticks_per_second;
Some(ProcessSnapshot {
pid,
rss_bytes,
vms_bytes,
uss_bytes,
pss_bytes,
io_read_bytes,
io_write_bytes,
cpu_time_seconds,
})
}
fn read_smaps_rollup(proc: &Process) -> (Option<u64>, Option<u64>) {
match proc.smaps_rollup() {
Ok(rollup) => {
let Some(map) = rollup.memory_map_rollup.0.first() else {
return (None, None);
};
let extension = &map.extension.map;
let private_clean = extension.get("Private_Clean").copied().unwrap_or(0);
let private_dirty = extension.get("Private_Dirty").copied().unwrap_or(0);
let uss = private_clean.saturating_add(private_dirty);
let pss = extension.get("Pss").copied();
(Some(uss), pss)
}
Err(_) => (None, None),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn self_pid_produces_a_snapshot() {
let mut sampler = LinuxSampler::new();
let pid = std::process::id().cast_signed();
let snaps = sampler.sample_tree(pid);
let me = snaps.iter().find(|s| s.pid == pid).expect("self snapshot");
assert!(me.rss_bytes > 0, "expected non-zero RSS");
assert!(me.vms_bytes > 0, "expected non-zero VMS");
}
}