use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use crate::config::Config;
use crate::gpu::NvmlCollector;
use crate::public_ip::{PublicIpChecker, PublicIpLookup};
use crate::smart::SmartReader;
#[cfg(target_os = "linux")]
use crate::snapshot::device_matches_disk;
use crate::snapshot::{PowerStats, Snapshot};
use crate::system::SystemCollector;
#[cfg(target_os = "linux")]
use crate::linux::{
self, SysRoots, diskstats::DiskStatsMeter, net::NetLinkReader, rapl::RaplMeter,
};
pub const MIN_POLL_INTERVAL: Duration = sysinfo::MINIMUM_CPU_UPDATE_INTERVAL;
#[derive(Debug, thiserror::Error)]
pub enum SourceError {
#[error("I/O error: {0}")]
Io(#[from] std::io::Error),
#[error("source disconnected: {0}")]
Disconnected(String),
#[error("remote: {0}")]
Remote(String),
#[error("cannot signal process {pid}: {source}")]
Signal {
pid: u32,
source: std::io::Error,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum ProcessSignal {
Terminate,
Kill,
Stop,
Continue,
Nice(i32),
}
impl ProcessSignal {
pub fn is_control(self) -> bool {
matches!(self, Self::Stop | Self::Continue | Self::Nice(_))
}
}
pub trait DataSource {
fn poll(&mut self) -> Result<Snapshot, SourceError>;
fn signal_process(&mut self, pid: u32, signal: ProcessSignal) -> Result<(), SourceError>;
fn set_interval(&mut self, _interval: Duration) {}
fn set_process_limit(&mut self, _limit: Option<usize>) {}
fn watch_processes(&mut self, _pids: Vec<u32>) {}
fn watch_process_names(&mut self, _names: Vec<String>) {}
fn set_public_ip(&mut self, _lookup: Option<PublicIpLookup>) {}
}
impl<T: DataSource + ?Sized> DataSource for Box<T> {
fn poll(&mut self) -> Result<Snapshot, SourceError> {
(**self).poll()
}
fn signal_process(&mut self, pid: u32, signal: ProcessSignal) -> Result<(), SourceError> {
(**self).signal_process(pid, signal)
}
fn set_interval(&mut self, interval: Duration) {
(**self).set_interval(interval);
}
fn set_process_limit(&mut self, limit: Option<usize>) {
(**self).set_process_limit(limit);
}
fn watch_processes(&mut self, pids: Vec<u32>) {
(**self).watch_processes(pids);
}
fn watch_process_names(&mut self, names: Vec<String>) {
(**self).watch_process_names(names);
}
fn set_public_ip(&mut self, lookup: Option<PublicIpLookup>) {
(**self).set_public_ip(lookup);
}
}
pub struct LocalSource {
system: SystemCollector,
nvml: NvmlCollector,
smart: SmartReader,
last_poll: Option<Instant>,
watched_pids: Vec<u32>,
watched_names: Vec<String>,
name_matches: Vec<u32>,
fill: crate::fill::FillTracker,
public_ip: Option<PublicIpChecker>,
#[cfg(target_os = "linux")]
roots: SysRoots,
#[cfg(target_os = "linux")]
rapl: RaplMeter,
#[cfg(target_os = "linux")]
diskstats: DiskStatsMeter,
#[cfg(target_os = "linux")]
net_links: NetLinkReader,
#[cfg(target_os = "linux")]
failed_units: linux::systemd::FailedUnitsReader,
#[cfg(target_os = "linux")]
zpool: linux::zpool::ZpoolReader,
#[cfg(target_os = "linux")]
containers: linux::containers::ContainerTracker,
#[cfg(target_os = "linux")]
oom: linux::oom::OomTracker,
#[cfg(target_os = "linux")]
battery: linux::battery::BatteryMeter,
#[cfg(target_os = "linux")]
core_locations: Vec<Option<linux::hwmon::CoreLocation>>,
#[cfg(target_os = "linux")]
scrub_dirs: Vec<std::path::PathBuf>,
#[cfg(target_os = "linux")]
drm_gpus: linux::drm::DrmGpuReader,
#[cfg(windows)]
lhm: crate::lhm::LhmReader,
}
impl LocalSource {
pub fn new() -> Self {
Self::with_config(&Config::default())
}
pub fn with_config(config: &Config) -> Self {
let mut source = Self {
system: SystemCollector::new(),
nvml: NvmlCollector::new(),
smart: SmartReader::new(&config.smart.dir),
last_poll: None,
watched_pids: Vec::new(),
watched_names: Vec::new(),
name_matches: Vec::new(),
fill: crate::fill::FillTracker::new(),
public_ip: None,
#[cfg(target_os = "linux")]
roots: SysRoots::default(),
#[cfg(target_os = "linux")]
rapl: RaplMeter::new(),
#[cfg(target_os = "linux")]
diskstats: DiskStatsMeter::new(),
#[cfg(target_os = "linux")]
net_links: NetLinkReader::new(),
#[cfg(target_os = "linux")]
failed_units: linux::systemd::FailedUnitsReader::new(),
#[cfg(target_os = "linux")]
zpool: linux::zpool::ZpoolReader::new(),
#[cfg(target_os = "linux")]
containers: linux::containers::ContainerTracker::new(),
#[cfg(target_os = "linux")]
oom: linux::oom::OomTracker::new(),
#[cfg(target_os = "linux")]
battery: linux::battery::BatteryMeter::default(),
#[cfg(target_os = "linux")]
core_locations: Vec::new(),
#[cfg(target_os = "linux")]
drm_gpus: linux::drm::DrmGpuReader::new(),
#[cfg(windows)]
lhm: crate::lhm::LhmReader::new(&config.lhm.url),
#[cfg(target_os = "linux")]
scrub_dirs: vec![
config.btrfs.scrub_dir.clone(),
linux::btrfs::BTRFS_PROGS_STATUS_DIR.into(),
],
};
let _ = source.collect();
source
}
pub fn gpu_unavailable_reason(&self) -> Option<&str> {
self.nvml.init_error()
}
fn collect(&mut self) -> Snapshot {
let now = Instant::now();
let elapsed = self
.last_poll
.map(|prev| now.saturating_duration_since(prev).as_secs_f64());
self.last_poll = Some(now);
self.system.refresh();
let processes = self.system.processes(elapsed);
let gpus = self.nvml.collect();
let gpus_w = sum_present(gpus.iter().map(|g| g.power_w));
let mut snapshot = Snapshot {
timestamp_ms: SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0),
host: self.system.host(),
cpu: self.system.cpu(),
memory: self.system.memory(),
networks: self.system.networks(elapsed),
filesystems: self.system.filesystems(),
process_count: processes.len(),
processes,
watched_pids: self.watched_pids.clone(),
gpus,
smart: self.smart.read(),
power: PowerStats {
gpus_w,
..Default::default()
},
..Default::default()
};
#[cfg(target_os = "linux")]
self.collect_linux(&mut snapshot, now);
#[cfg(not(target_os = "linux"))]
self.collect_portable(&mut snapshot, elapsed);
self.fill.annotate(&mut snapshot.filesystems, now);
if !(self.watched_names.is_empty() && self.name_matches.is_empty()) {
self.name_matches =
report_name_matches(&self.watched_names, &self.name_matches, &mut snapshot);
}
let power = &mut snapshot.power;
power.total_w = sum_present([power.cpu_or_soc_w().1, power.gpus_w]);
snapshot.public_ip = self
.public_ip
.as_mut()
.and_then(|checker| checker.poll().map(str::to_owned));
snapshot
}
#[cfg(not(target_os = "linux"))]
fn collect_portable(&mut self, snapshot: &mut Snapshot, elapsed: Option<f64>) {
snapshot.disks = self.system.volume_io(elapsed);
let (chips, cpu_temperatures) = self.system.components();
snapshot.sensors = chips;
snapshot.cpu.temperatures = cpu_temperatures;
#[cfg(windows)]
{
crate::windows::annotate_links(&mut snapshot.networks);
match self.lhm.read(Instant::now()) {
Some(lhm) => {
if !lhm.chips.is_empty() {
snapshot.sensors = lhm.chips;
}
if !lhm.cpu_temperatures.is_empty() {
snapshot.cpu.temperatures = lhm.cpu_temperatures;
}
snapshot.power.cpu_package_w = lhm.cpu_package_w;
if !lhm.smart.is_empty() {
snapshot.smart = lhm.smart;
}
}
None => snapshot.lhm_error = self.lhm.problem().map(str::to_owned),
}
}
}
#[cfg(target_os = "linux")]
fn collect_linux(&mut self, snapshot: &mut Snapshot, now: Instant) {
let chips = linux::hwmon::read_chips(&self.roots);
let cores = snapshot.cpu.per_core_percent.len();
if self.core_locations.len() != cores {
self.core_locations = linux::hwmon::read_core_locations(&self.roots, cores);
}
(snapshot.cpu.temperatures, snapshot.cpu.per_core_celsius) =
linux::hwmon::split_core_temperatures(&chips, &self.core_locations);
snapshot.memory.cached_bytes = linux::meminfo::read_cached_bytes(&self.roots);
snapshot.memory.oom_kills =
self.oom
.sample(&self.roots, &snapshot.processes, snapshot.timestamp_ms);
snapshot.power.cpu_package_w = self.rapl.sample(&self.roots, now);
snapshot.power.igpu_w = sum_present(
chips
.iter()
.filter(|c| c.is_apu)
.flat_map(|c| c.power_w.iter().map(|(_, w)| Some(*w))),
);
snapshot.power.battery = self.battery.read(&self.roots);
let first = snapshot.gpus.len() as u32;
snapshot.gpus.extend(self.drm_gpus.read(&self.roots, first));
snapshot.power.gpus_w = sum_present(snapshot.gpus.iter().map(|g| g.power_w));
let mut disks = self.diskstats.sample(&self.roots, now);
for disk in &mut disks {
disk.temperature_c = nvme_temperature(&chips, &disk.name);
}
snapshot.disks = disks;
let mounts = linux::mounts::read_mounts(&self.roots);
snapshot.raid.btrfs = linux::btrfs::read_filesystems(&self.roots, &mounts);
for fs in &mut snapshot.raid.btrfs {
fs.scrub = linux::btrfs::read_scrub(&self.scrub_dirs, &fs.fsid);
}
self.net_links
.annotate(&self.roots, &mut snapshot.networks, now);
snapshot.raid.md = linux::mdstat::read_arrays(&self.roots);
snapshot.raid.zfs = linux::zfs::read_pools(&self.roots);
self.zpool.annotate(&mut snapshot.raid.zfs, now);
snapshot.failed_units = self.failed_units.read(now);
snapshot.pressure = linux::pressure::read_pressure(&self.roots);
let release = std::fs::read_to_string(self.roots.proc.join("sys/kernel/osrelease"))
.unwrap_or_default();
snapshot.host.reboot_required = linux::reboot::reboot_required(
std::path::Path::new("/run"),
std::path::Path::new("/usr/lib/modules"),
&release,
);
for process in &mut snapshot.processes {
if let Some(stat) = linux::procstat::read(&self.roots, process.pid) {
process.threads = stat.threads;
process.nice = stat.nice;
} else {
process.threads = None;
}
}
self.containers
.annotate(&self.roots, &mut snapshot.processes, now);
snapshot.sensors = chips.into_iter().map(|c| c.chip).collect();
}
}
impl Default for LocalSource {
fn default() -> Self {
Self::new()
}
}
impl DataSource for LocalSource {
fn poll(&mut self) -> Result<Snapshot, SourceError> {
Ok(self.collect())
}
fn signal_process(&mut self, pid: u32, signal: ProcessSignal) -> Result<(), SourceError> {
send_signal(pid, signal).map_err(|source| SourceError::Signal { pid, source })
}
fn set_process_limit(&mut self, limit: Option<usize>) {
self.nvml.set_process_utilization(limit.is_none());
}
fn watch_processes(&mut self, pids: Vec<u32>) {
self.watched_pids = pids;
}
fn watch_process_names(&mut self, names: Vec<String>) {
self.watched_names = names;
}
fn set_public_ip(&mut self, lookup: Option<PublicIpLookup>) {
if self.public_ip.as_ref().map(PublicIpChecker::lookup) != lookup.as_ref() {
self.public_ip = lookup.map(PublicIpChecker::new);
}
}
}
fn report_name_matches(names: &[String], previous: &[u32], snapshot: &mut Snapshot) -> Vec<u32> {
let matches: Vec<u32> = snapshot
.processes
.iter()
.filter(|p| {
!p.kernel_thread
&& names
.iter()
.any(|name| crate::view::process_matches(p, name))
})
.map(|p| p.pid)
.collect();
let gone = previous
.iter()
.filter(|pid| !snapshot.processes.iter().any(|p| p.pid == **pid));
for pid in matches.iter().chain(gone) {
if !snapshot.watched_pids.contains(pid) {
snapshot.watched_pids.push(*pid);
}
}
matches
}
#[cfg(unix)]
fn send_signal(pid: u32, signal: ProcessSignal) -> std::io::Result<()> {
let pid = libc::pid_t::try_from(pid)
.ok()
.filter(|p| *p > 0)
.ok_or_else(|| std::io::Error::from(std::io::ErrorKind::InvalidInput))?;
let sig = match signal {
ProcessSignal::Terminate => libc::SIGTERM,
ProcessSignal::Kill => libc::SIGKILL,
ProcessSignal::Stop => libc::SIGSTOP,
ProcessSignal::Continue => libc::SIGCONT,
ProcessSignal::Nice(nice) => return set_nice(pid, nice),
};
if unsafe { libc::kill(pid, sig) } == 0 {
Ok(())
} else {
Err(std::io::Error::last_os_error())
}
}
#[cfg(unix)]
fn set_nice(pid: libc::pid_t, nice: i32) -> std::io::Result<()> {
let nice = nice.clamp(-20, 19);
if unsafe { libc::setpriority(libc::PRIO_PROCESS as _, pid as _, nice) } == 0 {
Ok(())
} else {
Err(std::io::Error::last_os_error())
}
}
#[cfg(windows)]
use crate::windows::send_signal;
#[cfg(not(any(unix, windows)))]
fn send_signal(_pid: u32, _signal: ProcessSignal) -> std::io::Result<()> {
Err(std::io::Error::from(std::io::ErrorKind::Unsupported))
}
fn sum_present(values: impl IntoIterator<Item = Option<f32>>) -> Option<f32> {
values
.into_iter()
.flatten()
.fold(None, |acc, v| Some(acc.unwrap_or(0.0) + v))
}
#[cfg(target_os = "linux")]
fn nvme_temperature(chips: &[linux::hwmon::HwmonChip], disk: &str) -> Option<f32> {
chips.iter().find_map(|c| {
let controller = c.chip.key.strip_prefix("nvme:")?;
device_matches_disk(controller, disk).then_some(())?;
c.chip.temperatures.first().map(|t| t.celsius)
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::snapshot::ProcessInfo;
#[test]
fn named_processes_are_reported_until_seen_gone() {
let process = |pid, name: &str, command: &str| ProcessInfo {
pid,
name: name.to_owned(),
command: command.to_owned(),
..Default::default()
};
let names = vec!["train.py".to_owned()];
let mut snapshot = Snapshot {
processes: vec![
process(10, "python3", "python3 jobs/train.py --epochs 3"),
process(11, "vim", "vim notes.txt"),
],
watched_pids: vec![99],
..Default::default()
};
let matches = report_name_matches(&names, &[], &mut snapshot);
assert_eq!(matches, [10]);
assert_eq!(snapshot.watched_pids, [99, 10]);
let mut later = Snapshot {
processes: vec![process(11, "vim", "vim notes.txt")],
..Default::default()
};
let matches = report_name_matches(&names, &matches, &mut later);
assert!(matches.is_empty());
assert_eq!(later.watched_pids, [10]);
let mut after = Snapshot::default();
assert!(report_name_matches(&names, &matches, &mut after).is_empty());
assert!(after.watched_pids.is_empty());
}
#[cfg(unix)]
#[test]
fn signal_rejects_group_pids_and_reports_missing_process() {
assert_eq!(
send_signal(0, ProcessSignal::Terminate).unwrap_err().kind(),
std::io::ErrorKind::InvalidInput
);
assert_eq!(
send_signal(u32::MAX, ProcessSignal::Kill)
.unwrap_err()
.kind(),
std::io::ErrorKind::InvalidInput
);
let mut child = std::process::Command::new("true").spawn().unwrap();
let pid = child.id();
child.wait().unwrap();
let err = send_signal(pid, ProcessSignal::Terminate).unwrap_err();
assert_eq!(err.raw_os_error(), Some(libc::ESRCH));
}
#[cfg(unix)]
#[test]
fn signal_terminates_a_child_process() {
let mut child = std::process::Command::new("sleep")
.arg("30")
.spawn()
.unwrap();
send_signal(child.id(), ProcessSignal::Kill).unwrap();
let status = child.wait().unwrap();
assert!(!status.success());
}
#[cfg(target_os = "linux")]
#[test]
fn pause_resume_and_renice_a_child_process() {
let mut child = std::process::Command::new("sleep")
.arg("30")
.spawn()
.unwrap();
let pid = child.id();
let roots = SysRoots::default();
let state = || {
let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")).unwrap();
let after = &stat[stat.rfind(')').unwrap() + 2..];
after.chars().next().unwrap()
};
send_signal(pid, ProcessSignal::Stop).unwrap();
for _ in 0..100 {
if state() == 'T' {
break;
}
std::thread::sleep(std::time::Duration::from_millis(10));
}
assert_eq!(state(), 'T');
send_signal(pid, ProcessSignal::Continue).unwrap();
send_signal(pid, ProcessSignal::Nice(7)).unwrap();
assert_eq!(linux::procstat::read(&roots, pid).unwrap().nice, Some(7));
child.kill().unwrap();
child.wait().unwrap();
}
#[test]
fn sum_present_ignores_missing_values() {
assert_eq!(sum_present([None, None]), None);
assert_eq!(sum_present([Some(1.5), None, Some(2.0)]), Some(3.5));
}
#[cfg(target_os = "linux")]
#[test]
fn nvme_temperature_matches_controller_exactly() {
use crate::linux::hwmon::HwmonChip;
use crate::snapshot::{SensorChip, Temperature};
let chip = |key: &str, celsius: f32| HwmonChip {
chip: SensorChip {
key: key.to_owned(),
name: "nvme".to_owned(),
temperatures: vec![Temperature {
label: "Composite".to_owned(),
celsius,
..Default::default()
}],
fans: Vec::new(),
},
power_w: Vec::new(),
is_apu: false,
};
let chips = vec![chip("nvme:nvme1", 40.0), chip("nvme:nvme10", 50.0)];
assert_eq!(nvme_temperature(&chips, "nvme10n1"), Some(50.0));
assert_eq!(nvme_temperature(&chips, "nvme1n1"), Some(40.0));
assert_eq!(nvme_temperature(&chips, "sda"), None);
}
}