use std::collections::{HashSet, VecDeque};
use log::trace;
use procfs::{Current, FromRead, Meminfo, process::Process};
use crate::{
Result,
snapshot::{GlobalSnapshot, ProcSnapshot},
utils::read_child_processes,
};
#[cfg_attr(test, mockall::automock)]
pub trait Backend: Default + Send + Sync + 'static {
fn read_proc(&self, pid: i32) -> Result<ProcSnapshot>;
fn measure_global(&self) -> Result<GlobalSnapshot>;
fn collect_children(&self, pid: i32) -> HashSet<i32>;
fn mem_total(&self) -> Result<u64>;
}
#[derive(Default, Clone, Copy)]
pub struct ProcfsBackend;
impl Backend for ProcfsBackend {
fn read_proc(&self, pid: i32) -> Result<ProcSnapshot> {
let process = Process::new(pid)?;
let mut snapshot = ProcSnapshot::default();
trace!("Querying process {} stat.", process.pid);
snapshot.vm_size += process.stat()?.vsize;
trace!("Querying process {} smaps_rollup.", process.pid);
let smaps = process.smaps_rollup()?;
if let Some(entry) = smaps.memory_map_rollup.0.first() {
let map = &entry.extension.map;
snapshot.rss += map.get("Rss").copied().unwrap_or(0);
snapshot.pss += map.get("Pss").copied().unwrap_or(0);
snapshot.anon += map.get("Anonymous").copied().unwrap_or(0);
snapshot.shared += map.get("Shared_Clean").copied().unwrap_or(0)
+ map.get("Shared_Dirty").copied().unwrap_or(0);
}
trace!("Querying process {} io.", process.pid);
match process.io() {
Ok(io) => {
snapshot.read_bytes += io.rchar;
snapshot.write_bytes += io.wchar;
}
Err(procfs::ProcError::PermissionDenied(_)) => {}
Err(err) => return Err(err.into()),
}
Ok(snapshot)
}
fn measure_global(&self) -> Result<GlobalSnapshot> {
let meminfo = Meminfo::from_file(Meminfo::PATH)?;
trace!("Querying global meminfo from {}.", procfs::Meminfo::PATH);
Ok(GlobalSnapshot {
mem_available: meminfo.mem_available,
mem_free: meminfo.mem_free,
cached: meminfo.cached,
anon: meminfo.anon_pages,
swap_free: meminfo.swap_free,
})
}
fn collect_children(&self, pid: i32) -> HashSet<i32> {
let mut pids: HashSet<i32> = HashSet::new();
pids.insert(pid);
let mut queue: VecDeque<i32> = VecDeque::from([pid]);
while let Some(pid) = queue.pop_front() {
for child_pid in read_child_processes(pid) {
if !pids.contains(&child_pid) {
pids.insert(child_pid);
queue.push_back(child_pid);
}
}
}
pids
}
fn mem_total(&self) -> Result<u64> {
trace!("Retrieving Meminfo from {}.", Meminfo::PATH);
Ok(Meminfo::from_file(Meminfo::PATH)?.mem_total)
}
}