hyperhive/hive-c0re/src/stats/warnings.rs

//! Process-global server-warning registry backing the dashboard banner.
//!
//! The dashboard shows a top-of-page banner for hive-level degradation
//! (disk pressure, a failed converge sweep, agents that can't log in, …).
//! Historically each warning was *pulled*: [`host_stats::server_warnings`]
//! recomputed the whole list on every `/api/state` by re-probing (a
//! `statvfs`, a scan of the container snapshot). Background sweeps that
//! failed had nowhere to record it, so they just `warn!`'d to the journal
//! and vanished from anywhere an operator looks.
//!
//! This module replaces the pull model with a **push** one: a single
//! `Mutex<BTreeMap<kind, …>>` any subsystem writes into. [`snapshot`] is a
//! cheap read of that map. The default API is RAII: [`set_warning`] returns
//! a [`WarningGuard`] whose `Drop` removes the key, so a producer holds the
//! guard exactly as long as the condition is true and lets it drop to clear
//! — no manual dismiss, no stale banner. A failing periodic sweep does
//! `let _w = set_warning(...)` in its failure branch and drops it on the
//! next success.
//!
//! Concurrency: all mutation is behind the one mutex. Each guard carries a
//! unique id and only touches its own map entry (`Drop` / [`update`]
//! no-op once a later `set_warning` for the same kind has superseded it),
//! so a re-`set_warning` for the same kind is never clobbered when the
//! stale guard later drops.
//!
//! [`host_stats::server_warnings`]: crate::host_stats::server_warnings
//! [`update`]: WarningGuard::update

use std::collections::BTreeMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Mutex, OnceLock};

use crate::host_stats::ServerWarning;

/// Monotonic guard-id source. Each [`set_warning`] mints a fresh id so a
/// guard's `Drop` can tell "my entry" from "an entry a later guard replaced
/// under the same kind".
static NEXT_ID: AtomicU64 = AtomicU64::new(1);

/// One registry slot: the warning payload plus the id of the guard that
/// currently owns the `kind`.
struct Slot {
    level: &'static str,
    message: String,
    owner: u64,
}

/// The one process-global registry. `BTreeMap` (not `HashMap`) so
/// [`snapshot`] yields a deterministic, kind-sorted banner order across
/// renders.
fn registry() -> &'static Mutex<BTreeMap<&'static str, Slot>> {
    static REG: OnceLock<Mutex<BTreeMap<&'static str, Slot>>> = OnceLock::new();
    REG.get_or_init(|| Mutex::new(BTreeMap::new()))
}

/// Lock the registry, recovering the guard even if a previous holder
/// panicked mid-mutation — a poisoned lock must not wedge every future
/// warning update.
fn lock() -> std::sync::MutexGuard<'static, BTreeMap<&'static str, Slot>> {
    registry()
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
}

/// Raise (or replace) the warning under `kind` and return a guard that
/// clears it on drop. This is the default way to surface a hive-level
/// degradation: hold the returned guard while the condition is true, and
/// let it drop (or call `drop(guard)`) the moment it clears.
///
/// Re-calling `set_warning` for a `kind` that already has a live guard
/// replaces the payload and takes ownership of the entry (fresh id): the
/// superseded guard's later `Drop`/`update` become no-ops, so a producer
/// can freely re-arm without racing its own previous guard.
#[must_use = "dropping the guard immediately clears the warning; bind it for as long as the condition holds"]
pub fn set_warning(
    kind: &'static str,
    level: &'static str,
    message: impl Into<String>,
) -> WarningGuard {
    let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
    lock().insert(
        kind,
        Slot {
            level,
            message: message.into(),
            owner: id,
        },
    );
    WarningGuard { kind, id }
}

/// Snapshot of the currently-active warnings, kind-sorted. Cheap read
/// behind the registry mutex — safe to call on every `/api/state`.
#[must_use]
pub fn snapshot() -> Vec<ServerWarning> {
    lock()
        .iter()
        .map(|(kind, slot)| ServerWarning {
            kind,
            level: slot.level,
            message: slot.message.clone(),
        })
        .collect()
}

/// RAII handle for one active warning. Drop clears the `kind` from the
/// banner. Obtained from [`set_warning`].
#[must_use = "dropping the guard immediately clears the warning; bind it for as long as the condition holds"]
pub struct WarningGuard {
    kind: &'static str,
    id: u64,
}

impl WarningGuard {
    /// Refresh the held warning's `level` + `message` in place — e.g. a
    /// disk-usage percentage ticking up, or a `warn` escalating to `crit`.
    /// A no-op once a later [`set_warning`] for the same kind has superseded
    /// this guard (its entry is owned by the newer id).
    pub fn update(&self, level: &'static str, message: impl Into<String>) {
        let mut reg = lock();
        if let Some(slot) = reg.get_mut(self.kind)
            && slot.owner == self.id
        {
            slot.level = level;
            slot.message = message.into();
        }
    }
}

impl Drop for WarningGuard {
    fn drop(&mut self) {
        let mut reg = lock();
        // Only remove if we still own the slot — a later `set_warning` for
        // the same kind takes over the entry (fresh owner id), so this
        // superseded guard's drop must not clobber the live one.
        let ours = reg.get(self.kind).is_some_and(|slot| slot.owner == self.id);
        if ours {
            reg.remove(self.kind);
        }
    }
}

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

    /// Find the snapshot entry for `kind`, if present. Tests filter by
    /// their own unique kinds instead of asserting global emptiness, so
    /// they stay correct under cargo's parallel test runner (the registry
    /// is process-global and shared across tests).
    fn find(kind: &str) -> Option<ServerWarning> {
        snapshot().into_iter().find(|w| w.kind == kind)
    }

    #[test]
    fn set_surfaces_in_snapshot_then_drop_clears() {
        {
            let _g = set_warning("t_setdrop", "warn", "boom");
            let w = find("t_setdrop").expect("warning present while guard held");
            assert_eq!(w.level, "warn");
            assert_eq!(w.message, "boom");
        }
        assert!(find("t_setdrop").is_none(), "dropping the guard clears it");
    }

    #[test]
    fn update_refreshes_level_and_message_in_place() {
        let g = set_warning("t_update", "warn", "85% full");
        g.update("crit", "96% full");
        let w = find("t_update").expect("present");
        assert_eq!(w.level, "crit");
        assert_eq!(w.message, "96% full");
        drop(g);
        assert!(find("t_update").is_none());
    }

    #[test]
    fn resetting_same_kind_survives_stale_guard_drop() {
        let a = set_warning("t_reset", "warn", "first");
        let b = set_warning("t_reset", "crit", "second");
        // `b` took over the entry; dropping the superseded `a` must not
        // remove `b`'s live warning.
        drop(a);
        let w = find("t_reset").expect("b's warning survives a's drop");
        assert_eq!(w.level, "crit");
        assert_eq!(w.message, "second");
        // A stale guard's update is also a no-op — prove by dropping b and
        // confirming the kind is gone (only the owner clears it).
        drop(b);
        assert!(find("t_reset").is_none());
    }

    #[test]
    fn stale_guard_update_is_noop() {
        let a = set_warning("t_staleupd", "warn", "old");
        let b = set_warning("t_staleupd", "warn", "new");
        a.update("crit", "should not apply");
        let w = find("t_staleupd").expect("present");
        assert_eq!(w.message, "new", "stale guard update ignored");
        drop(b);
        drop(a);
        assert!(find("t_staleupd").is_none());
    }

    #[test]
    fn distinct_kinds_coexist_and_snapshot_is_sorted() {
        let _z = set_warning("t_zzz", "warn", "z");
        let _a = set_warning("t_aaa", "warn", "a");
        let kinds: Vec<&str> = snapshot()
            .iter()
            .map(|w| w.kind)
            .filter(|k| k.starts_with("t_zzz") || k.starts_with("t_aaa"))
            .collect();
        assert_eq!(kinds, ["t_aaa", "t_zzz"], "kind-sorted order");
    }

    #[test]
    fn concurrent_unique_kinds_never_race_and_clear_out() {
        // Each thread owns a distinct kind and hammers set/update/drop in a
        // loop. The shared resource under test is the one mutex; a data
        // race would trip under the thread sanitizer / show a wrong final
        // state. After every guard has dropped, none of the kinds remain.
        const KINDS: [&str; 6] = ["cc0", "cc1", "cc2", "cc3", "cc4", "cc5"];
        std::thread::scope(|s| {
            for kind in KINDS {
                s.spawn(move || {
                    for i in 0..500 {
                        let g = set_warning(kind, "warn", format!("iter {i}"));
                        g.update("crit", format!("iter {i} upd"));
                        drop(g);
                    }
                });
            }
        });
        for kind in KINDS {
            assert!(find(kind).is_none(), "{kind} cleared after all drops");
        }
    }

    #[test]
    fn concurrent_same_kind_resolves_to_absent_after_all_drop() {
        // Many threads contend on ONE kind, each holding its own guard
        // briefly. The id-guard invariant means whichever set ran last owns
        // the entry, and once every guard has dropped the kind is absent —
        // no stale guard leaves a phantom warning behind.
        std::thread::scope(|s| {
            for t in 0..8 {
                s.spawn(move || {
                    for i in 0..300 {
                        let g = set_warning("cc_shared", "warn", format!("{t}:{i}"));
                        drop(g);
                    }
                });
            }
        });
        assert!(find("cc_shared").is_none());
    }
}