joularcore 0.2.0

/*
 * Copyright (c) 2025-2026, Adel Noureddine.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the
 * GNU Lesser General Public License v3.0 only (LGPL-3.0-only)
 * which accompanies this distribution, and is available at
 * https://www.gnu.org/licenses/lgpl-3.0.en.html
 *
 * Author : Adel Noureddine
 */

//! Discrete GPU power, discovered by running the vendor's command-line tool.
//!
//! Shared by the Linux and Windows backends, which read GPU power the same way.
//! macOS does not use this: `powermetrics` reports GPU power directly.
//!
//! Each vendor gets its own module below: `nvidia` runs `nvidia-smi`, `amd`
//! runs `amd-smi` or `rocm-smi`. Both read one power column out of CSV.

use crate::sensor::PowerSensor;
use crate::{Error, Result};
#[cfg(target_os = "windows")]
use std::os::windows::process::CommandExt;
use std::process::Command;

/// Suppresses the console window a child process would otherwise flash up when
/// spawned from a GUI program.
#[cfg(target_os = "windows")]
const CREATE_NO_WINDOW: u32 = 0x0800_0000;

/// Run a vendor tool and return its standard output as text.
///
/// The three outcomes are kept apart, because a GPU that is *missing* and one
/// that is *unreadable* are different answers:
///
/// * `Ok(Some(stdout))` — the tool ran and succeeded.
/// * `Ok(None)` — the tool could not be started at all, so this vendor has
///   nothing installed on this machine.
/// * `Err(..)` — the tool ran but failed, or answered in a way we cannot use.
///
/// `sensor` names the component in any error the caller surfaces, e.g.
/// `"NVIDIA GPU"`; `program` appears in the detail so a machine with several
/// tools installed shows which one failed.
///
/// Unlike the macOS backend, which insists on absolute paths because it runs
/// children as root, these are resolved through `PATH`: they run unprivileged,
/// and vendor tools install to different directories across distributions.
pub(super) fn run_tool(
    program: &str,
    args: &[&str],
    sensor: &'static str,
) -> Result<Option<String>> {
    let mut command = Command::new(program);
    command.args(args);

    #[cfg(target_os = "windows")]
    command.creation_flags(CREATE_NO_WINDOW);

    let output = match command.output() {
        Ok(output) => output,
        Err(e) => {
            log::debug!("vendor tool {program} could not be started: {e}");
            return Ok(None);
        }
    };

    if !output.status.success() {
        return Err(Error::sensor(
            sensor,
            format!("{program} exited with {}", output.status),
        ));
    }

    String::from_utf8(output.stdout)
        .map(Some)
        .map_err(|e| Error::sensor(sensor, format!("{program} output is not UTF-8: {e}")))
}

/// GPU power summed across every vendor tool present on the machine.
#[derive(Debug, Default)]
pub(super) struct VendorGpu;

impl PowerSensor for VendorGpu {
    fn power(&mut self) -> Result<f64> {
        // Both are read before either result is used, and a failure from either
        // propagates: letting a working AMD card mask a broken NVIDIA one would
        // report a total that is quietly missing a GPU.
        let nvidia = nvidia::power()?;
        let amd = amd::power()?;

        if nvidia.is_none() && amd.is_none() {
            return Err(Error::sensor(
                "GPU",
                "no nvidia-smi, amd-smi or rocm-smi reported GPU power",
            ));
        }

        Ok(nvidia.unwrap_or(0.0) + amd.unwrap_or(0.0))
    }
}

/// NVIDIA GPU power, read by running `nvidia-smi`.
mod nvidia {
    use super::run_tool;
    use crate::Result;
    use std::sync::OnceLock;

    /// How this sensor is named in errors.
    const SENSOR: &str = "NVIDIA GPU";

    /// The tool, and the query that makes it print one watt figure per GPU.
    const PROGRAM: &str = "nvidia-smi";
    const ARGS: &[&str] = &["--format=csv,noheader,nounits", "--query-gpu=power.draw"];

    /// Whether `nvidia-smi` can be started at all, once that is known.
    ///
    /// Spawning a program that does not exist is cheap but not free, and this runs
    /// on every sample, so a machine with no NVIDIA tooling stops being asked.
    static INSTALLED: OnceLock<bool> = OnceLock::new();

    /// Total power drawn by every NVIDIA GPU, in watts.
    ///
    /// `Ok(None)` means there is nothing here to measure: no driver, no card, or a
    /// card that does not report power. `Err` means the tool ran and failed, which
    /// must not be mistaken for an absent GPU — a machine's total would quietly
    /// lose a card.
    pub(super) fn power() -> Result<Option<f64>> {
        // `None` means "not decided yet", so the first sample always tries.
        if !INSTALLED.get().copied().unwrap_or(true) {
            return Ok(None);
        }

        let Some(stdout) = run_tool(PROGRAM, ARGS, SENSOR)? else {
            let _ = INSTALLED.set(false);
            return Ok(None);
        };
        let _ = INSTALLED.set(true);

        Ok(parse_power(&stdout))
    }

    /// Sum the per-GPU watt figures in `nvidia-smi` CSV output.
    ///
    /// Cards that do not report power print `[N/A]` and contribute nothing, so
    /// output holding no figure at all reports `None`.
    fn parse_power(output: &str) -> Option<f64> {
        let mut total = 0.0;
        let mut found = false;

        for line in output.lines() {
            if let Ok(watts) = line.trim().parse::<f64>() {
                total += watts;
                found = true;
            }
        }

        found.then_some(total)
    }

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

        #[test]
        fn sums_every_gpu() {
            assert_eq!(parse_power("120.5\n80.25\n"), Some(200.75));
        }

        #[test]
        fn cards_without_a_reading_contribute_nothing() {
            assert_eq!(parse_power("[N/A]\n42.0\n"), Some(42.0));
        }

        #[test]
        fn output_with_no_figure_at_all_reports_nothing() {
            // Distinct from 0 W: a card reporting `[N/A]` is not an idle card, and
            // an empty answer means the query found no GPU.
            assert_eq!(parse_power("[N/A]\n"), None);
            assert_eq!(parse_power(""), None);
        }
    }
}

/// AMD GPU power, read by running `amd-smi` or `rocm-smi`.
mod amd {
    use super::run_tool;
    use crate::{Error, Result};
    use std::sync::OnceLock;

    /// How this sensor is named in errors.
    const SENSOR: &str = "AMD GPU";

    /// Probing runs the tools, so the answer is cached for the process lifetime.
    static COMMAND: OnceLock<Option<AmdGpuCommand>> = OnceLock::new();

    /// The vendor tool available on this machine.
    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
    enum AmdGpuCommand {
        /// The current tool.
        AmdSmi,
        /// The older `ROCm` tool.
        RocmSmi,
    }

    impl AmdGpuCommand {
        fn program(self) -> &'static str {
            match self {
                AmdGpuCommand::AmdSmi => "amd-smi",
                AmdGpuCommand::RocmSmi => "rocm-smi",
            }
        }

        fn args(self) -> &'static [&'static str] {
            match self {
                AmdGpuCommand::AmdSmi => &["metric", "-p", "--csv"],
                AmdGpuCommand::RocmSmi => &["--showpower", "--csv"],
            }
        }
    }

    /// Column headers that hold the figure we want, lowercased.
    ///
    /// Two naming schemes, because the tools disagree: `rocm-smi` prints its
    /// human-readable labels verbatim, while `amd-smi` flattens its JSON keys to
    /// lowercase words joined by underscores. AMD has also renamed these across
    /// releases, so each spelling that has appeared is listed.
    ///
    /// Matching is on the whole column name rather than a substring, so a limit
    /// such as `power_cap` or `socket_power_limit` can never be mistaken for a
    /// reading. An unrecognised header reports no GPU rather than a wrong number.
    const POWER_COLUMNS: &[&str] = &[
        // rocm-smi
        "average graphics package power (w)",
        "current socket graphics package power (w)",
        "current graphics package power (w)",
        "power (w)",
        // amd-smi
        "power_socket_power_value",
        "socket_power_value",
        "current_socket_graphics_package_power",
        "current_graphics_package_power",
        "average_socket_power",
        "current_socket_power",
        "socket_power",
        "average_power",
    ];

    /// Total power drawn by every AMD GPU, in watts.
    ///
    /// `Ok(None)` means neither tool is installed, so there is no AMD GPU to
    /// account for. `Err` means a tool ran and failed, which must not be mistaken
    /// for an absent GPU — a machine's total would quietly lose a card.
    pub(super) fn power() -> Result<Option<f64>> {
        let Some(command) = detect_command() else {
            return Ok(None);
        };

        // The tool answered the probe, so it existed a moment ago; if it cannot be
        // started now it has been removed underneath us.
        let Some(stdout) = run_tool(command.program(), command.args(), SENSOR)? else {
            return Ok(None);
        };

        parse_power(&stdout).map(Some).ok_or_else(|| {
            Error::sensor(
                SENSOR,
                format!("no power column in {} CSV output", command.program()),
            )
        })
    }

    fn detect_command() -> Option<AmdGpuCommand> {
        *COMMAND.get_or_init(|| {
            [AmdGpuCommand::AmdSmi, AmdGpuCommand::RocmSmi]
                .into_iter()
                .find(|command| is_installed(command.program()))
        })
    }

    /// Whether `program` exists and runs. A tool that spawns but exits non-zero is
    /// broken or unsupported, so it does not count as installed.
    fn is_installed(program: &str) -> bool {
        matches!(run_tool(program, &["--version"], SENSOR), Ok(Some(_)))
    }

    /// Total watts across every GPU in a tool's CSV output.
    ///
    /// One header row names the columns, then one row per GPU; the power column is
    /// summed down the rows. Rows without a usable figure contribute nothing, so
    /// output holding no figure at all reports `None`.
    fn parse_power(csv: &str) -> Option<f64> {
        let mut rows = csv.lines().filter(|row| !row.trim().is_empty());
        let index = power_column(rows.next()?)?;

        let mut total = 0.0;
        let mut found = false;
        for row in rows {
            if let Some(watts) = row.split(',').nth(index).and_then(parse_watts) {
                total += watts;
                found = true;
            }
        }

        found.then_some(total)
    }

    /// Position of the power column in a header row.
    fn power_column(header: &str) -> Option<usize> {
        header.split(',').position(|column| {
            let column = column.trim().trim_matches('"').to_ascii_lowercase();
            POWER_COLUMNS.contains(&column.as_str())
        })
    }

    /// Parse one cell as watts, tolerating a trailing unit (`"75 W"`).
    fn parse_watts(cell: &str) -> Option<f64> {
        let cell = cell.trim().trim_matches('"');
        cell.parse()
            .ok()
            .or_else(|| cell.split_whitespace().next()?.parse().ok())
    }

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

        #[test]
        fn reads_rocm_smi_output() {
            let csv = "\
    device,Average Graphics Package Power (W)
    card0,142.0
    card1,38.0
    ";
            assert_eq!(parse_power(csv), Some(180.0));
        }

        #[test]
        fn reads_amd_smi_output() {
            let csv = "\
    gpu,power_socket_power_value,power_socket_power_unit
    0,100.5,W
    1,20.0,W
    ";
            assert_eq!(parse_power(csv), Some(120.5));
        }

        #[test]
        fn a_limit_column_is_never_mistaken_for_a_reading() {
            // Reporting a cap as a reading would be worse than reporting nothing:
            // the number looks entirely plausible.
            let csv = "gpu,socket_power_limit,power_cap\n0,550,550\n";
            assert_eq!(parse_power(csv), None);
        }

        #[test]
        fn cards_without_a_reading_contribute_nothing() {
            let csv = "\
    device,Average Graphics Package Power (W)
    card0,N/A (Secondary die)
    card1,42.0
    ";
            assert_eq!(parse_power(csv), Some(42.0));
        }

        #[test]
        fn numbers_may_carry_their_unit() {
            assert_eq!(parse_power("device,Power (W)\ncard0,75 W\n"), Some(75.0));
        }

        #[test]
        fn output_without_a_power_column_reports_nothing() {
            assert_eq!(parse_power("device,Temperature (C)\ncard0,55\n"), None);
            assert_eq!(parse_power("not csv"), None);
            assert_eq!(parse_power(""), None);
        }

        #[test]
        fn a_header_with_no_rows_reports_nothing() {
            // Distinct from 0 W: the query found no GPU, it did not find an idle
            // one.
            assert_eq!(parse_power("device,Power (W)\n"), None);
        }
    }
}