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