Lumotia/crates/core/src/power.rs

//! Power-state probe for inference thread tuning.
//!
//! Reports whether the machine is running on AC or battery so callers
//! can drop thread counts when energy matters more than throughput.
//! Linux uses the documented sysfs ABI under
//! `/sys/class/power_supply/`. macOS and Windows return `Unknown` for
//! now; native probes (IOPSGetProvidingPowerSourceType,
//! GetSystemPowerStatus) are deferred.

use std::fs;
use std::path::Path;
use std::sync::{Mutex, OnceLock};
use std::time::{Duration, Instant};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PowerState {
    OnAc,
    OnBattery,
    Unknown,
}

/// Parse a `/sys/class/power_supply/`-style directory and decide
/// whether the machine is on AC, on battery, or in an unknown state.
///
/// Rules (matches `Documentation/ABI/testing/sysfs-class-power`):
/// - Any entry with `type` in {`Mains`, `USB`} and `online == 1` → OnAc.
/// - Else any entry with `type == Battery` → OnBattery.
/// - Else → Unknown.
///
/// Top-level failures (missing dir, unreadable supply_dir) return
/// Unknown without panicking. Per-entry failures (unreadable
/// `type`/`online` file inside an individual supply, e.g. permission
/// denied or non-UTF-8 content) cause that entry to be silently
/// skipped via `read_trimmed().unwrap_or_default()` — which on a
/// stuck-AC machine could produce OnBattery if the Mains entry was
/// the unreadable one. In practice sysfs entries are world-readable,
/// so this is a theoretical hazard rather than a real one. `Unknown`
/// is treated as `OnAc` by the caller, preserving today's pre-clamp
/// behaviour on platforms or distros where the probe doesn't fire.
pub fn parse_power_state_from_dir(supply_dir: &Path) -> PowerState {
    let read_dir = match fs::read_dir(supply_dir) {
        Ok(rd) => rd,
        Err(_) => return PowerState::Unknown,
    };

    let mut saw_battery = false;
    for entry in read_dir.flatten() {
        let path = entry.path();
        let ty = read_trimmed(&path.join("type")).unwrap_or_default();
        let online = read_trimmed(&path.join("online")).unwrap_or_default();
        match ty.as_str() {
            "Mains" | "USB" => {
                if online == "1" {
                    return PowerState::OnAc;
                }
            }
            "Battery" => {
                saw_battery = true;
            }
            _ => {}
        }
    }

    if saw_battery {
        PowerState::OnBattery
    } else {
        PowerState::Unknown
    }
}

fn read_trimmed(path: &Path) -> Option<String> {
    fs::read_to_string(path).ok().map(|s| s.trim().to_owned())
}

const POWER_STATE_TTL: Duration = Duration::from_secs(10);

struct CachedState {
    state: PowerState,
    fetched_at: Instant,
}

fn cache_slot() -> &'static Mutex<Option<CachedState>> {
    static SLOT: OnceLock<Mutex<Option<CachedState>>> = OnceLock::new();
    SLOT.get_or_init(|| Mutex::new(None))
}

#[cfg(test)]
pub(crate) fn force_clear_cache() {
    *cache_slot().lock().expect("power cache mutex poisoned") = None;
}

#[cfg(test)]
pub(crate) fn force_set_cache(state: PowerState) {
    *cache_slot().lock().expect("power cache mutex poisoned") = Some(CachedState {
        state,
        fetched_at: Instant::now(),
    });
}

/// Top-level power-state probe.
///
/// Resolution order (highest to lowest priority):
/// 1. In-process test override (set via `with_override` from unit tests).
/// 2. `MAGNOTIA_POWER_STATE_OVERRIDE` env var (`ac` | `battery` | `unknown`,
///    case-insensitive). Used by `thread_sweep.rs` integration tests.
/// 3. Linux: `parse_power_state_from_dir("/sys/class/power_supply")`.
/// 4. macOS / Windows / other: `Unknown`.
///
/// `Unknown` is treated as `OnAc` by callers, which preserves today's
/// pre-clamp behaviour on platforms where the probe doesn't fire.
///
/// Results are cached for `POWER_STATE_TTL` (10 seconds). Override
/// paths (test and env-var) bypass the cache entirely so they always
/// take effect immediately.
pub fn probe_power_state() -> PowerState {
    #[cfg(test)]
    if let Some(state) = test_get_override() {
        return state;
    }
    if let Some(state) = env_override() {
        return state;
    }

    let mut slot = cache_slot().lock().expect("power cache mutex poisoned");
    if let Some(cached) = &*slot {
        if cached.fetched_at.elapsed() < POWER_STATE_TTL {
            return cached.state;
        }
    }
    let fresh = platform_probe();
    *slot = Some(CachedState {
        state: fresh,
        fetched_at: Instant::now(),
    });
    fresh
}

fn env_override() -> Option<PowerState> {
    let raw = std::env::var("MAGNOTIA_POWER_STATE_OVERRIDE").ok()?;
    match raw.trim().to_ascii_lowercase().as_str() {
        "ac" => Some(PowerState::OnAc),
        "battery" => Some(PowerState::OnBattery),
        "unknown" => Some(PowerState::Unknown),
        _ => None,
    }
}

#[cfg(target_os = "linux")]
fn platform_probe() -> PowerState {
    parse_power_state_from_dir(Path::new("/sys/class/power_supply"))
}

#[cfg(not(target_os = "linux"))]
fn platform_probe() -> PowerState {
    PowerState::Unknown
}

// In-process override slot. Tests must use `with_override` (below) to
// set it; production code never writes to this. Read on every probe.
#[cfg(test)]
static TEST_OVERRIDE: Mutex<Option<PowerState>> = Mutex::new(None);

#[cfg(test)]
fn test_get_override() -> Option<PowerState> {
    *TEST_OVERRIDE.lock().expect("power test override mutex poisoned")
}

/// Drop-guard that resets `TEST_OVERRIDE` to `None` on drop (including on unwind).
/// This prevents a panicking test body from leaking stale override state into
/// subsequent tests.
#[cfg(test)]
struct OverrideGuard;

#[cfg(test)]
impl Drop for OverrideGuard {
    fn drop(&mut self) {
        *TEST_OVERRIDE.lock().expect("power test override mutex poisoned") = None;
    }
}

/// Run `body` with the in-process override set to `state`. Restores
/// `TEST_OVERRIDE` to `None` on return, even if `body` panics.
///
/// Holds a dedicated `TEST_LOCK` mutex for the duration of the body,
/// so override-using unit tests run serially with respect to each
/// other even when cargo runs the test binary multi-threaded.
#[cfg(test)]
pub(crate) fn with_override<R>(state: Option<PowerState>, body: impl FnOnce() -> R) -> R {
    static TEST_LOCK: Mutex<()> = Mutex::new(());
    let _lock = TEST_LOCK.lock().expect("power TEST_LOCK poisoned");
    *TEST_OVERRIDE.lock().expect("power test override mutex poisoned") = state;
    let _guard = OverrideGuard;
    body()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::tempdir;

    #[test]
    fn power_state_variants_are_distinct() {
        assert_ne!(PowerState::OnAc, PowerState::OnBattery);
        assert_ne!(PowerState::OnAc, PowerState::Unknown);
        assert_ne!(PowerState::OnBattery, PowerState::Unknown);
    }

    fn write_supply(dir: &std::path::Path, name: &str, ty: &str, online: &str) {
        let entry = dir.join(name);
        fs::create_dir_all(&entry).unwrap();
        fs::write(entry.join("type"), format!("{ty}\n")).unwrap();
        fs::write(entry.join("online"), format!("{online}\n")).unwrap();
    }

    #[test]
    fn parses_mains_online_as_on_ac() {
        let dir = tempdir().unwrap();
        write_supply(dir.path(), "AC", "Mains", "1");
        write_supply(dir.path(), "BAT0", "Battery", "0");
        assert_eq!(parse_power_state_from_dir(dir.path()), PowerState::OnAc);
    }

    #[test]
    fn parses_battery_only_as_on_battery() {
        let dir = tempdir().unwrap();
        write_supply(dir.path(), "AC", "Mains", "0");
        write_supply(dir.path(), "BAT0", "Battery", "0");
        assert_eq!(parse_power_state_from_dir(dir.path()), PowerState::OnBattery);
    }

    #[test]
    fn parses_usb_pd_online_as_on_ac() {
        let dir = tempdir().unwrap();
        write_supply(dir.path(), "ucsi-source-psy-USBC000:001", "USB", "1");
        write_supply(dir.path(), "BAT0", "Battery", "0");
        assert_eq!(parse_power_state_from_dir(dir.path()), PowerState::OnAc);
    }

    #[test]
    fn parses_empty_dir_as_unknown() {
        let dir = tempdir().unwrap();
        assert_eq!(parse_power_state_from_dir(dir.path()), PowerState::Unknown);
    }

    #[test]
    fn parses_missing_dir_as_unknown() {
        let path = std::path::Path::new("/no/such/path/should/exist/at/this/inode/123456");
        assert_eq!(parse_power_state_from_dir(path), PowerState::Unknown);
    }

    #[test]
    fn parses_malformed_files_as_unknown_gracefully() {
        let dir = tempdir().unwrap();
        let entry = dir.path().join("garbage");
        fs::create_dir_all(&entry).unwrap();
        // No type, no online — should not panic.
        assert_eq!(parse_power_state_from_dir(dir.path()), PowerState::Unknown);
    }

    #[test]
    fn override_drives_battery() {
        with_override(Some(PowerState::OnBattery), || {
            assert_eq!(probe_power_state(), PowerState::OnBattery);
        });
    }

    #[test]
    fn override_drives_ac() {
        with_override(Some(PowerState::OnAc), || {
            assert_eq!(probe_power_state(), PowerState::OnAc);
        });
    }

    #[test]
    fn override_drives_unknown() {
        with_override(Some(PowerState::Unknown), || {
            assert_eq!(probe_power_state(), PowerState::Unknown);
        });
    }

    #[test]
    fn env_var_override_battery_via_set_var() {
        // env-var path tested in isolation under TEST_LOCK so it
        // doesn't race with the in-process override tests.
        with_override(None, || {
            std::env::set_var("MAGNOTIA_POWER_STATE_OVERRIDE", "battery");
            assert_eq!(probe_power_state(), PowerState::OnBattery);
            std::env::remove_var("MAGNOTIA_POWER_STATE_OVERRIDE");
        });
    }

    #[test]
    fn env_var_override_garbage_falls_through() {
        with_override(None, || {
            std::env::set_var("MAGNOTIA_POWER_STATE_OVERRIDE", "nonsense");
            // Garbage value falls through to the platform probe.
            // We can't assert the platform result so just assert it
            // doesn't panic.
            let _ = probe_power_state();
            std::env::remove_var("MAGNOTIA_POWER_STATE_OVERRIDE");
        });
    }

    #[test]
    fn ttl_cache_returns_cached_value_within_window() {
        with_override(None, || {
            force_clear_cache();
            force_set_cache(PowerState::OnBattery);
            // No override active, no env var; probe should hit cache.
            assert_eq!(probe_power_state(), PowerState::OnBattery);
        });
    }

    #[test]
    fn ttl_cache_clears_via_force_clear() {
        with_override(None, || {
            force_set_cache(PowerState::OnBattery);
            force_clear_cache();
            // Override flips to OnAc; with cleared cache, override
            // (which has higher priority) drives the result.
        });
        with_override(Some(PowerState::OnAc), || {
            force_clear_cache();
            assert_eq!(probe_power_state(), PowerState::OnAc);
        });
    }
}