1use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
9
10use crate::config::Config;
11use crate::gpu::NvmlCollector;
12use crate::smart::SmartReader;
13#[cfg(target_os = "linux")]
14use crate::snapshot::device_matches_disk;
15use crate::snapshot::{PowerStats, Snapshot};
16use crate::system::SystemCollector;
17
18#[cfg(target_os = "linux")]
19use crate::linux::{
20 self, SysRoots, diskstats::DiskStatsMeter, net::NetLinkReader, rapl::RaplMeter,
21};
22
23pub const MIN_POLL_INTERVAL: Duration = sysinfo::MINIMUM_CPU_UPDATE_INTERVAL;
26
27#[derive(Debug, thiserror::Error)]
29pub enum SourceError {
30 #[error("I/O error: {0}")]
32 Io(#[from] std::io::Error),
33 #[error("source disconnected: {0}")]
35 Disconnected(String),
36 #[error("remote: {0}")]
38 Remote(String),
39 #[error("cannot signal process {pid}: {source}")]
41 Signal {
42 pid: u32,
44 source: std::io::Error,
46 },
47}
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
51pub enum ProcessSignal {
52 Terminate,
54 Kill,
56 Stop,
59 Continue,
61 Nice(i32),
64}
65
66impl ProcessSignal {
67 pub fn is_control(self) -> bool {
69 matches!(self, Self::Stop | Self::Continue | Self::Nice(_))
70 }
71}
72
73pub trait DataSource {
75 fn poll(&mut self) -> Result<Snapshot, SourceError>;
81
82 fn signal_process(&mut self, pid: u32, signal: ProcessSignal) -> Result<(), SourceError>;
85
86 fn set_interval(&mut self, _interval: Duration) {}
90
91 fn set_process_limit(&mut self, _limit: Option<usize>) {}
94
95 fn watch_processes(&mut self, _pids: Vec<u32>) {}
99
100 fn watch_process_names(&mut self, _names: Vec<String>) {}
105}
106
107impl<T: DataSource + ?Sized> DataSource for Box<T> {
109 fn poll(&mut self) -> Result<Snapshot, SourceError> {
110 (**self).poll()
111 }
112
113 fn signal_process(&mut self, pid: u32, signal: ProcessSignal) -> Result<(), SourceError> {
114 (**self).signal_process(pid, signal)
115 }
116
117 fn set_interval(&mut self, interval: Duration) {
118 (**self).set_interval(interval);
119 }
120
121 fn set_process_limit(&mut self, limit: Option<usize>) {
122 (**self).set_process_limit(limit);
123 }
124
125 fn watch_processes(&mut self, pids: Vec<u32>) {
126 (**self).watch_processes(pids);
127 }
128
129 fn watch_process_names(&mut self, names: Vec<String>) {
130 (**self).watch_process_names(names);
131 }
132}
133
134pub struct LocalSource {
136 system: SystemCollector,
137 nvml: NvmlCollector,
138 smart: SmartReader,
139 last_poll: Option<Instant>,
140 watched_pids: Vec<u32>,
142 watched_names: Vec<String>,
145 name_matches: Vec<u32>,
146 fill: crate::fill::FillTracker,
147 #[cfg(target_os = "linux")]
148 roots: SysRoots,
149 #[cfg(target_os = "linux")]
150 rapl: RaplMeter,
151 #[cfg(target_os = "linux")]
152 diskstats: DiskStatsMeter,
153 #[cfg(target_os = "linux")]
154 net_links: NetLinkReader,
155 #[cfg(target_os = "linux")]
156 failed_units: linux::systemd::FailedUnitsReader,
157 #[cfg(target_os = "linux")]
158 zpool: linux::zpool::ZpoolReader,
159 #[cfg(target_os = "linux")]
160 containers: linux::containers::ContainerTracker,
161 #[cfg(target_os = "linux")]
162 oom: linux::oom::OomTracker,
163 #[cfg(target_os = "linux")]
164 battery: linux::battery::BatteryMeter,
165 #[cfg(target_os = "linux")]
168 core_locations: Vec<Option<linux::hwmon::CoreLocation>>,
169 #[cfg(target_os = "linux")]
171 scrub_dirs: Vec<std::path::PathBuf>,
172 #[cfg(target_os = "linux")]
174 drm_gpus: linux::drm::DrmGpuReader,
175 #[cfg(windows)]
177 lhm: crate::lhm::LhmReader,
178}
179
180impl LocalSource {
181 pub fn new() -> Self {
183 Self::with_config(&Config::default())
184 }
185
186 pub fn with_config(config: &Config) -> Self {
193 let mut source = Self {
194 system: SystemCollector::new(),
195 nvml: NvmlCollector::new(),
196 smart: SmartReader::new(&config.smart.dir),
197 last_poll: None,
198 watched_pids: Vec::new(),
199 watched_names: Vec::new(),
200 name_matches: Vec::new(),
201 fill: crate::fill::FillTracker::new(),
202 #[cfg(target_os = "linux")]
203 roots: SysRoots::default(),
204 #[cfg(target_os = "linux")]
205 rapl: RaplMeter::new(),
206 #[cfg(target_os = "linux")]
207 diskstats: DiskStatsMeter::new(),
208 #[cfg(target_os = "linux")]
209 net_links: NetLinkReader::new(),
210 #[cfg(target_os = "linux")]
211 failed_units: linux::systemd::FailedUnitsReader::new(),
212 #[cfg(target_os = "linux")]
213 zpool: linux::zpool::ZpoolReader::new(),
214 #[cfg(target_os = "linux")]
215 containers: linux::containers::ContainerTracker::new(),
216 #[cfg(target_os = "linux")]
217 oom: linux::oom::OomTracker::new(),
218 #[cfg(target_os = "linux")]
219 battery: linux::battery::BatteryMeter::default(),
220 #[cfg(target_os = "linux")]
221 core_locations: Vec::new(),
222 #[cfg(target_os = "linux")]
223 drm_gpus: linux::drm::DrmGpuReader::new(),
224 #[cfg(windows)]
227 lhm: crate::lhm::LhmReader::new(&config.lhm.url),
228 #[cfg(target_os = "linux")]
229 scrub_dirs: vec![
230 config.btrfs.scrub_dir.clone(),
231 linux::btrfs::BTRFS_PROGS_STATUS_DIR.into(),
232 ],
233 };
234 let _ = source.collect();
236 source
237 }
238
239 pub fn gpu_unavailable_reason(&self) -> Option<&str> {
241 self.nvml.init_error()
242 }
243
244 fn collect(&mut self) -> Snapshot {
245 let now = Instant::now();
246 let elapsed = self
247 .last_poll
248 .map(|prev| now.saturating_duration_since(prev).as_secs_f64());
249 self.last_poll = Some(now);
250
251 self.system.refresh();
252 let processes = self.system.processes(elapsed);
253 let gpus = self.nvml.collect();
254 let gpus_w = sum_present(gpus.iter().map(|g| g.power_w));
255
256 let mut snapshot = Snapshot {
257 timestamp_ms: SystemTime::now()
258 .duration_since(UNIX_EPOCH)
259 .map(|d| d.as_millis() as u64)
260 .unwrap_or(0),
261 host: self.system.host(),
262 cpu: self.system.cpu(),
263 memory: self.system.memory(),
264 networks: self.system.networks(elapsed),
265 filesystems: self.system.filesystems(),
266 process_count: processes.len(),
267 processes,
268 watched_pids: self.watched_pids.clone(),
269 gpus,
270 smart: self.smart.read(),
271 power: PowerStats {
272 gpus_w,
273 ..Default::default()
274 },
275 ..Default::default()
276 };
277
278 #[cfg(target_os = "linux")]
279 self.collect_linux(&mut snapshot, now);
280 #[cfg(not(target_os = "linux"))]
281 self.collect_portable(&mut snapshot, elapsed);
282
283 self.fill.annotate(&mut snapshot.filesystems, now);
284 if !(self.watched_names.is_empty() && self.name_matches.is_empty()) {
285 self.name_matches =
286 report_name_matches(&self.watched_names, &self.name_matches, &mut snapshot);
287 }
288 let power = &mut snapshot.power;
289 power.total_w = sum_present([power.cpu_or_soc_w().1, power.gpus_w]);
290 snapshot
291 }
292
293 #[cfg(not(target_os = "linux"))]
296 fn collect_portable(&mut self, snapshot: &mut Snapshot, elapsed: Option<f64>) {
297 snapshot.disks = self.system.volume_io(elapsed);
298 let (chips, cpu_temperatures) = self.system.components();
299 snapshot.sensors = chips;
300 snapshot.cpu.temperatures = cpu_temperatures;
301 #[cfg(windows)]
302 {
303 crate::windows::annotate_links(&mut snapshot.networks);
304 match self.lhm.read(Instant::now()) {
307 Some(lhm) => {
308 if !lhm.chips.is_empty() {
309 snapshot.sensors = lhm.chips;
310 }
311 if !lhm.cpu_temperatures.is_empty() {
312 snapshot.cpu.temperatures = lhm.cpu_temperatures;
313 }
314 snapshot.power.cpu_package_w = lhm.cpu_package_w;
315 if !lhm.smart.is_empty() {
316 snapshot.smart = lhm.smart;
317 }
318 }
319 None => snapshot.lhm_error = self.lhm.problem().map(str::to_owned),
320 }
321 }
322 }
323
324 #[cfg(target_os = "linux")]
326 fn collect_linux(&mut self, snapshot: &mut Snapshot, now: Instant) {
327 let chips = linux::hwmon::read_chips(&self.roots);
328
329 let cores = snapshot.cpu.per_core_percent.len();
330 if self.core_locations.len() != cores {
331 self.core_locations = linux::hwmon::read_core_locations(&self.roots, cores);
332 }
333 (snapshot.cpu.temperatures, snapshot.cpu.per_core_celsius) =
334 linux::hwmon::split_core_temperatures(&chips, &self.core_locations);
335 snapshot.memory.cached_bytes = linux::meminfo::read_cached_bytes(&self.roots);
336 snapshot.memory.oom_kills =
337 self.oom
338 .sample(&self.roots, &snapshot.processes, snapshot.timestamp_ms);
339 snapshot.power.cpu_package_w = self.rapl.sample(&self.roots, now);
340 snapshot.power.igpu_w = sum_present(
343 chips
344 .iter()
345 .filter(|c| c.is_apu)
346 .flat_map(|c| c.power_w.iter().map(|(_, w)| Some(*w))),
347 );
348 snapshot.power.battery = self.battery.read(&self.roots);
349
350 let first = snapshot.gpus.len() as u32;
352 snapshot.gpus.extend(self.drm_gpus.read(&self.roots, first));
353 snapshot.power.gpus_w = sum_present(snapshot.gpus.iter().map(|g| g.power_w));
354
355 let mut disks = self.diskstats.sample(&self.roots, now);
356 for disk in &mut disks {
357 disk.temperature_c = nvme_temperature(&chips, &disk.name);
358 }
359 snapshot.disks = disks;
360
361 let mounts = linux::mounts::read_mounts(&self.roots);
362 snapshot.raid.btrfs = linux::btrfs::read_filesystems(&self.roots, &mounts);
363 for fs in &mut snapshot.raid.btrfs {
364 fs.scrub = linux::btrfs::read_scrub(&self.scrub_dirs, &fs.fsid);
365 }
366 self.net_links
367 .annotate(&self.roots, &mut snapshot.networks, now);
368 snapshot.raid.md = linux::mdstat::read_arrays(&self.roots);
369 snapshot.raid.zfs = linux::zfs::read_pools(&self.roots);
370 self.zpool.annotate(&mut snapshot.raid.zfs, now);
371 snapshot.failed_units = self.failed_units.read(now);
372 snapshot.pressure = linux::pressure::read_pressure(&self.roots);
373 let release = std::fs::read_to_string(self.roots.proc.join("sys/kernel/osrelease"))
374 .unwrap_or_default();
375 snapshot.host.reboot_required = linux::reboot::reboot_required(
376 std::path::Path::new("/run"),
377 std::path::Path::new("/usr/lib/modules"),
378 &release,
379 );
380 for process in &mut snapshot.processes {
382 if let Some(stat) = linux::procstat::read(&self.roots, process.pid) {
383 process.threads = stat.threads;
384 process.nice = stat.nice;
385 } else {
386 process.threads = None;
387 }
388 }
389 self.containers
390 .annotate(&self.roots, &mut snapshot.processes, now);
391
392 snapshot.sensors = chips.into_iter().map(|c| c.chip).collect();
393 }
394}
395
396impl Default for LocalSource {
397 fn default() -> Self {
398 Self::new()
399 }
400}
401
402impl DataSource for LocalSource {
403 fn poll(&mut self) -> Result<Snapshot, SourceError> {
404 Ok(self.collect())
405 }
406
407 fn signal_process(&mut self, pid: u32, signal: ProcessSignal) -> Result<(), SourceError> {
408 send_signal(pid, signal).map_err(|source| SourceError::Signal { pid, source })
409 }
410
411 fn set_process_limit(&mut self, limit: Option<usize>) {
414 self.nvml.set_process_utilization(limit.is_none());
415 }
416
417 fn watch_processes(&mut self, pids: Vec<u32>) {
418 self.watched_pids = pids;
419 }
420
421 fn watch_process_names(&mut self, names: Vec<String>) {
422 self.watched_names = names;
423 }
424}
425
426fn report_name_matches(names: &[String], previous: &[u32], snapshot: &mut Snapshot) -> Vec<u32> {
431 let matches: Vec<u32> = snapshot
432 .processes
433 .iter()
434 .filter(|p| {
435 !p.kernel_thread
436 && names
437 .iter()
438 .any(|name| crate::view::process_matches(p, name))
439 })
440 .map(|p| p.pid)
441 .collect();
442 let gone = previous
443 .iter()
444 .filter(|pid| !snapshot.processes.iter().any(|p| p.pid == **pid));
445 for pid in matches.iter().chain(gone) {
446 if !snapshot.watched_pids.contains(pid) {
447 snapshot.watched_pids.push(*pid);
448 }
449 }
450 matches
451}
452
453#[cfg(unix)]
455fn send_signal(pid: u32, signal: ProcessSignal) -> std::io::Result<()> {
456 let pid = libc::pid_t::try_from(pid)
459 .ok()
460 .filter(|p| *p > 0)
461 .ok_or_else(|| std::io::Error::from(std::io::ErrorKind::InvalidInput))?;
462 let sig = match signal {
463 ProcessSignal::Terminate => libc::SIGTERM,
464 ProcessSignal::Kill => libc::SIGKILL,
465 ProcessSignal::Stop => libc::SIGSTOP,
466 ProcessSignal::Continue => libc::SIGCONT,
467 ProcessSignal::Nice(nice) => return set_nice(pid, nice),
468 };
469 if unsafe { libc::kill(pid, sig) } == 0 {
471 Ok(())
472 } else {
473 Err(std::io::Error::last_os_error())
474 }
475}
476
477#[cfg(unix)]
479fn set_nice(pid: libc::pid_t, nice: i32) -> std::io::Result<()> {
480 let nice = nice.clamp(-20, 19);
481 if unsafe { libc::setpriority(libc::PRIO_PROCESS as _, pid as _, nice) } == 0 {
484 Ok(())
485 } else {
486 Err(std::io::Error::last_os_error())
487 }
488}
489
490#[cfg(windows)]
491use crate::windows::send_signal;
492
493#[cfg(not(any(unix, windows)))]
495fn send_signal(_pid: u32, _signal: ProcessSignal) -> std::io::Result<()> {
496 Err(std::io::Error::from(std::io::ErrorKind::Unsupported))
497}
498
499fn sum_present(values: impl IntoIterator<Item = Option<f32>>) -> Option<f32> {
501 values
502 .into_iter()
503 .flatten()
504 .fold(None, |acc, v| Some(acc.unwrap_or(0.0) + v))
505}
506
507#[cfg(target_os = "linux")]
510fn nvme_temperature(chips: &[linux::hwmon::HwmonChip], disk: &str) -> Option<f32> {
511 chips.iter().find_map(|c| {
512 let controller = c.chip.key.strip_prefix("nvme:")?;
513 device_matches_disk(controller, disk).then_some(())?;
514 c.chip.temperatures.first().map(|t| t.celsius)
516 })
517}
518
519#[cfg(test)]
520mod tests {
521 use super::*;
522 use crate::snapshot::ProcessInfo;
523
524 #[test]
525 fn named_processes_are_reported_until_seen_gone() {
526 let process = |pid, name: &str, command: &str| ProcessInfo {
527 pid,
528 name: name.to_owned(),
529 command: command.to_owned(),
530 ..Default::default()
531 };
532 let names = vec!["train.py".to_owned()];
533 let mut snapshot = Snapshot {
534 processes: vec![
535 process(10, "python3", "python3 jobs/train.py --epochs 3"),
536 process(11, "vim", "vim notes.txt"),
537 ],
538 watched_pids: vec![99],
539 ..Default::default()
540 };
541 let matches = report_name_matches(&names, &[], &mut snapshot);
542 assert_eq!(matches, [10]);
543 assert_eq!(snapshot.watched_pids, [99, 10]);
544
545 let mut later = Snapshot {
547 processes: vec![process(11, "vim", "vim notes.txt")],
548 ..Default::default()
549 };
550 let matches = report_name_matches(&names, &matches, &mut later);
551 assert!(matches.is_empty());
552 assert_eq!(later.watched_pids, [10]);
553 let mut after = Snapshot::default();
554 assert!(report_name_matches(&names, &matches, &mut after).is_empty());
555 assert!(after.watched_pids.is_empty());
556 }
557
558 #[cfg(unix)]
559 #[test]
560 fn signal_rejects_group_pids_and_reports_missing_process() {
561 assert_eq!(
562 send_signal(0, ProcessSignal::Terminate).unwrap_err().kind(),
563 std::io::ErrorKind::InvalidInput
564 );
565 assert_eq!(
566 send_signal(u32::MAX, ProcessSignal::Kill)
567 .unwrap_err()
568 .kind(),
569 std::io::ErrorKind::InvalidInput
570 );
571 let mut child = std::process::Command::new("true").spawn().unwrap();
573 let pid = child.id();
574 child.wait().unwrap();
575 let err = send_signal(pid, ProcessSignal::Terminate).unwrap_err();
576 assert_eq!(err.raw_os_error(), Some(libc::ESRCH));
577 }
578
579 #[cfg(unix)]
580 #[test]
581 fn signal_terminates_a_child_process() {
582 let mut child = std::process::Command::new("sleep")
583 .arg("30")
584 .spawn()
585 .unwrap();
586 send_signal(child.id(), ProcessSignal::Kill).unwrap();
587 let status = child.wait().unwrap();
588 assert!(!status.success());
589 }
590
591 #[cfg(target_os = "linux")]
592 #[test]
593 fn pause_resume_and_renice_a_child_process() {
594 let mut child = std::process::Command::new("sleep")
595 .arg("30")
596 .spawn()
597 .unwrap();
598 let pid = child.id();
599 let roots = SysRoots::default();
600 let state = || {
601 let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")).unwrap();
602 let after = &stat[stat.rfind(')').unwrap() + 2..];
603 after.chars().next().unwrap()
604 };
605 send_signal(pid, ProcessSignal::Stop).unwrap();
606 for _ in 0..100 {
608 if state() == 'T' {
609 break;
610 }
611 std::thread::sleep(std::time::Duration::from_millis(10));
612 }
613 assert_eq!(state(), 'T');
614 send_signal(pid, ProcessSignal::Continue).unwrap();
615 send_signal(pid, ProcessSignal::Nice(7)).unwrap();
617 assert_eq!(linux::procstat::read(&roots, pid).unwrap().nice, Some(7));
618 child.kill().unwrap();
619 child.wait().unwrap();
620 }
621
622 #[test]
623 fn sum_present_ignores_missing_values() {
624 assert_eq!(sum_present([None, None]), None);
625 assert_eq!(sum_present([Some(1.5), None, Some(2.0)]), Some(3.5));
626 }
627
628 #[cfg(target_os = "linux")]
629 #[test]
630 fn nvme_temperature_matches_controller_exactly() {
631 use crate::linux::hwmon::HwmonChip;
632 use crate::snapshot::{SensorChip, Temperature};
633
634 let chip = |key: &str, celsius: f32| HwmonChip {
635 chip: SensorChip {
636 key: key.to_owned(),
637 name: "nvme".to_owned(),
638 temperatures: vec![Temperature {
639 label: "Composite".to_owned(),
640 celsius,
641 ..Default::default()
642 }],
643 fans: Vec::new(),
644 },
645 power_w: Vec::new(),
646 is_apu: false,
647 };
648 let chips = vec![chip("nvme:nvme1", 40.0), chip("nvme:nvme10", 50.0)];
649 assert_eq!(nvme_temperature(&chips, "nvme10n1"), Some(50.0));
650 assert_eq!(nvme_temperature(&chips, "nvme1n1"), Some(40.0));
651 assert_eq!(nvme_temperature(&chips, "sda"), None);
652 }
653}