rigtop-core 0.8.2

//! Data sources: anything that can produce [`Snapshot`]s.
//!
//! Front-ends only talk to the [`DataSource`] trait. The local machine is one
//! implementation ([`LocalSource`]); a remote host reached through
//! `ssh host rigtop --agent` will be another, so remote monitoring needs no
//! changes in the front-ends.

use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use crate::config::Config;
use crate::gpu::NvmlCollector;
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,
};

/// Shortest sensible interval between two polls. CPU usage is computed from
/// the difference between polls and is inaccurate below this.
pub const MIN_POLL_INTERVAL: Duration = sysinfo::MINIMUM_CPU_UPDATE_INTERVAL;

/// Errors a data source can report.
#[derive(Debug, thiserror::Error)]
pub enum SourceError {
    /// Reading from the underlying channel failed.
    #[error("I/O error: {0}")]
    Io(#[from] std::io::Error),
    /// The source is no longer available (e.g. the SSH connection closed).
    #[error("source disconnected: {0}")]
    Disconnected(String),
    /// The remote agent reported a problem or does not speak our protocol.
    #[error("remote: {0}")]
    Remote(String),
    /// Sending a signal to a process failed (no such process, permission denied, ...).
    #[error("cannot signal process {pid}: {source}")]
    Signal {
        /// Target process id.
        pid: u32,
        /// Underlying OS error.
        source: std::io::Error,
    },
}

/// Signals a front-end can send to a process.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum ProcessSignal {
    /// Ask the process to exit (SIGTERM).
    Terminate,
    /// Force the process to exit (SIGKILL).
    Kill,
}

/// A producer of snapshots, local or remote.
pub trait DataSource {
    /// Takes a new measurement.
    ///
    /// Rates (CPU %, throughput, power) cover the time since the previous
    /// call, so callers should poll at a steady interval of at least
    /// [`MIN_POLL_INTERVAL`].
    fn poll(&mut self) -> Result<Snapshot, SourceError>;

    /// Sends `signal` to process `pid` on the monitored host, with the
    /// privileges of the user the source runs as.
    fn signal_process(&mut self, pid: u32, signal: ProcessSignal) -> Result<(), SourceError>;

    /// Tells the source how often it will be polled. Remote sources forward
    /// this to the agent, which samples on its own schedule; local sources
    /// need nothing.
    fn set_interval(&mut self, _interval: Duration) {}

    /// Limits snapshots to the top `limit` processes by CPU (`None`: all).
    /// Remote sources use this to save bandwidth for hosts in an overview.
    fn set_process_limit(&mut self, _limit: Option<usize>) {}
}

/// Boxed sources, so front-ends can choose local or remote at run time.
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);
    }
}

/// Collects snapshots of the machine rigtop runs on.
pub struct LocalSource {
    system: SystemCollector,
    nvml: NvmlCollector,
    smart: SmartReader,
    last_poll: Option<Instant>,
    #[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,
    /// Directories searched for btrfs scrub status files, in order.
    #[cfg(target_os = "linux")]
    scrub_dirs: Vec<std::path::PathBuf>,
    /// AMD and Intel GPUs (NVIDIA ones come from NVML).
    #[cfg(target_os = "linux")]
    drm_gpus: linux::drm::DrmGpuReader,
    /// Temperatures, fans and CPU power from LibreHardwareMonitor.
    #[cfg(windows)]
    lhm: crate::lhm::LhmReader,
}

impl LocalSource {
    /// Creates a source with the default configuration. See [`LocalSource::with_config`].
    pub fn new() -> Self {
        Self::with_config(&Config::default())
    }

    /// Creates a source and takes a warm-up sample, so that the first
    /// [`DataSource::poll`] (after at least [`MIN_POLL_INTERVAL`]) already
    /// has rates.
    ///
    /// Only collection-related settings (the SMART directory) are used here;
    /// display settings such as sensor labels are applied by front-ends.
    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,
            #[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")]
            drm_gpus: linux::drm::DrmGpuReader::new(),
            // The copy made by the root collector first, then btrfs-progs'
            // own file in case rigtop runs as root.
            #[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(),
            ],
        };
        // The warm-up result itself is meaningless (no previous sample).
        let _ = source.collect();
        source
    }

    /// Why NVIDIA GPUs are unavailable, if NVML failed to initialize.
    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,
            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);

        let power = &mut snapshot.power;
        power.total_w = sum_present([power.cpu_or_soc_w().1, power.gpus_w]);
        snapshot
    }

    /// Adds what `sysinfo` offers beyond the basics on non-Linux platforms:
    /// per-volume I/O, temperatures and (on Windows) link speeds.
    #[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);
            // LibreHardwareMonitor, when running, knows far more than the
            // ACPI thermal zones sysinfo can see.
            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),
            }
        }
    }

    /// Adds sensors, power, per-device I/O and RAID state from `/sys` and `/proc`.
    #[cfg(target_os = "linux")]
    fn collect_linux(&mut self, snapshot: &mut Snapshot, now: Instant) {
        let chips = linux::hwmon::read_chips(&self.roots);

        snapshot.cpu.temperatures = chips
            .iter()
            .filter(|c| linux::hwmon::is_cpu_chip(&c.chip))
            .flat_map(|c| c.chip.temperatures.iter().cloned())
            .collect();
        snapshot.memory.cached_bytes = linux::meminfo::read_cached_bytes(&self.roots);
        snapshot.power.cpu_package_w = self.rapl.sample(&self.roots, now);
        // The APU's amdgpu power covers the whole SoC (see PowerStats).
        // Discrete AMD GPUs are not reported yet.
        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 = linux::battery::read_battery(&self.roots);

        // AMD and Intel GPUs after NVIDIA's; discrete ones add to GPU power.
        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);
        // sysinfo no longer lists threads (see SystemCollector::refresh).
        for process in &mut snapshot.processes {
            process.threads = linux::procstat::thread_count(&self.roots, process.pid);
        }

        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 })
    }

    /// With a limit, only the top processes are shown (overview), which do
    /// not include per-process GPU utilization: skip that NVML query.
    fn set_process_limit(&mut self, limit: Option<usize>) {
        self.nvml.set_process_utilization(limit.is_none());
    }
}

/// Sends a signal to a local process.
#[cfg(unix)]
fn send_signal(pid: u32, signal: ProcessSignal) -> std::io::Result<()> {
    // pid 0 and negative values address process groups in kill(2); never
    // let a bad pid turn into "signal everything".
    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,
    };
    // SAFETY: kill(2) has no memory-safety preconditions.
    if unsafe { libc::kill(pid, sig) } == 0 {
        Ok(())
    } else {
        Err(std::io::Error::last_os_error())
    }
}

#[cfg(windows)]
use crate::windows::send_signal;

/// Sends a signal to a local process (not implemented on this platform yet).
#[cfg(not(any(unix, windows)))]
fn send_signal(_pid: u32, _signal: ProcessSignal) -> std::io::Result<()> {
    Err(std::io::Error::from(std::io::ErrorKind::Unsupported))
}

/// Sums the values that are present; `None` if none are.
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))
}

/// Finds the composite temperature of the NVMe controller backing `disk`
/// (hwmon chip key `nvme:nvme0` for disk `nvme0n1`).
#[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(())?;
        // The first channel is the "Composite" temperature the drive reports.
        c.chip.temperatures.first().map(|t| t.celsius)
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[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
        );
        // A child that has been reaped no longer exists: ESRCH.
        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());
    }

    #[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);
    }
}