hyperhive/hive-c0re/src/workers/swarm_limits.rs

//! Reading the caps the swarm sets for this hive's agents.
//!
//! swarm-controller keeps one key per agent in the `agent-limits` bucket
//! ([`swarm_wire::agent_limits`]). Every [`READ_INTERVAL`] this reads
//! the key of each agent this hive has and stores what it finds in
//! `swarm-resource-limits.json`, where the drop-in reads it.
//! A running agent gets a changed cap at its first spawn or rebuild after
//! that read.
//!
//! A failed read, or a bucket that does not answer, leaves the stored caps as
//! they are: dropping them would read exactly like the swarm clearing them.

use std::collections::BTreeMap;
use std::time::Duration;

use anyhow::{Context, Result};
use swarm_queue_client::agent_limits::{self, Read};
use swarm_wire::agent_limits::AgentCaps;

use crate::resource_limits::AgentLimits;
use crate::stats::sweep_health::{self, SweepHealth};

/// How often the caps are read.
pub const READ_INTERVAL: Duration = Duration::from_secs(30);

/// At [`READ_INTERVAL`] this is ~90 seconds of failing reads.
const FAILURES_BEFORE_BANNER: u32 = 3;

/// Read the caps now and every [`READ_INTERVAL`] until shutdown.
pub async fn follow(client: async_nats::Client, mut shutdown: tokio::sync::watch::Receiver<bool>) {
    let mut health = SweepHealth::new("swarm_limits_read", "warn", FAILURES_BEFORE_BANNER);
    loop {
        match sync(&client).await {
            Ok(()) => health.record_ok(),
            Err(e) => {
                tracing::warn!(error = ?e, "swarm limits: reading the caps failed");
                let err = format!("{e:#}");
                health.record_err(|ctx| {
                    let age = ctx.since_last_ok.map_or_else(
                        || "no success this session".to_owned(),
                        |d| format!("last ok {} ago", sweep_health::fmt_age(d)),
                    );
                    format!(
                        "reading the swarm's resource caps is failing ({} consecutive, {age}) \
                         — agents keep the caps last read: {err}",
                        ctx.consecutive
                    )
                });
            }
        }
        tokio::select! {
            () = tokio::time::sleep(READ_INTERVAL) => {}
            _ = shutdown.changed() => {
                tracing::info!("swarm limits: shutdown signal received");
                return;
            }
        }
    }
}

/// Read every agent's key and store the result, unless a read fails or a
/// value is refused.
async fn sync(client: &async_nats::Client) -> Result<()> {
    // An unconnected client does not fail a JetStream request, it hangs on it.
    swarm_queue_client::ensure_connected(client)?;
    let agents = crate::lifecycle::agents_for_meta_listing()
        .await
        .context("enumerating this hive's agents")?;
    let mut declared = BTreeMap::new();
    for spec in agents {
        let read = agent_limits::read(client, &spec.name)
            .await
            .map_err(|e| anyhow::anyhow!(e))
            .with_context(|| format!("reading the swarm's caps for {}", spec.name))?;
        match read {
            Read::Current(caps) => {
                declared.insert(spec.name, caps);
            }
            Read::NoBucket => {
                tracing::debug!("swarm limits: the agent-limits bucket does not answer");
                return Ok(());
            }
        }
    }
    let caps = validated(&declared)?;
    if caps != crate::resource_limits::read_swarm().context("reading the stored swarm caps")? {
        crate::resource_limits::write_swarm(&caps).context("storing the swarm caps")?;
        tracing::info!(?caps, "swarm limits: stored the swarm's resource caps");
    }
    Ok(())
}

/// The caps to store for `declared`, by agent. An agent with no field set
/// gets no entry.
///
/// Every value goes through the `resource_limits` validators, because it is
/// rendered verbatim into the container's systemd drop-in. One bad value
/// refuses the whole set.
fn validated(declared: &BTreeMap<String, AgentCaps>) -> Result<BTreeMap<String, AgentLimits>> {
    let mut out = BTreeMap::new();
    for (agent, caps) in declared {
        if caps.is_empty() {
            continue;
        }
        let limits = validated_caps(caps)
            .map_err(|why| anyhow::anyhow!("the swarm set an invalid cap for {agent}: {why}"))?;
        out.insert(agent.clone(), limits);
    }
    Ok(out)
}

fn validated_caps(caps: &AgentCaps) -> Result<AgentLimits, String> {
    if let Some(cpu) = &caps.cpu_quota {
        crate::resource_limits::validate_cpu_quota(cpu)?;
    }
    Ok(AgentLimits {
        cpu_quota: caps.cpu_quota.clone(),
        memory_max: caps
            .memory_max
            .as_deref()
            .map(crate::resource_limits::validate_memory_max)
            .transpose()?,
    })
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use super::validated;
    use swarm_wire::agent_limits::AgentCaps;

    /// `atlas`'s answer with `atlas_key` as the value read for it, over the
    /// hive-wide defaults of 200% / 4G.
    fn atlas_with(atlas_key: &str) -> (String, String) {
        let declared = BTreeMap::from([(
            "atlas".to_owned(),
            serde_json::from_str::<AgentCaps>(atlas_key).expect("decodes"),
        )]);
        let stored = validated(&declared).expect("valid caps");
        crate::resource_limits::effective_from(&stored, "atlas", "200%", "4G")
    }

    /// The other half of swarm-controller's
    /// `agent_limits::the_route_writes_the_caps_a_hive_reads_and_delete_clears_them`:
    /// the first value is the bytes that route stores. The cap replaces the
    /// hive default; a deleted key reads as `{}` and gives the default back.
    #[test]
    fn the_caps_the_route_stores_replace_the_hive_default_and_a_cleared_key_restores_it() {
        assert_eq!(
            atlas_with(r#"{"cpu_quota":"400%","memory_max":"8G"}"#),
            ("400%".to_owned(), "8G".to_owned())
        );
        assert_eq!(atlas_with("{}"), ("200%".to_owned(), "4G".to_owned()));
    }

    /// Per field: a swarm memory cap leaves CPU on the hive default.
    #[test]
    fn a_field_the_swarm_leaves_unset_stays_with_the_hive() {
        assert_eq!(
            atlas_with(r#"{"memory_max":"2G"}"#),
            ("200%".to_owned(), "2G".to_owned())
        );
    }

    #[test]
    fn an_agent_with_no_caps_gets_no_entry() {
        let declared = BTreeMap::from([("atlas".to_owned(), AgentCaps::default())]);
        assert!(validated(&declared).expect("valid").is_empty());
    }

    /// The values are rendered verbatim into a systemd drop-in, so a value
    /// the hive's own validators refuse, a newline smuggling in a directive
    /// among them, refuses the whole set, the valid cap beside it included.
    #[test]
    fn an_invalid_cap_is_refused_and_nothing_is_stored() {
        for bad in [
            r#"{"cpu_quota":"400%\nExecStartPre=/bin/sh -c id #%"}"#,
            r#"{"memory_max":"8G\nDevicePolicy=auto"}"#,
            r#"{"cpu_quota":"lots"}"#,
            r#"{"memory_max":"8X"}"#,
        ] {
            let declared = BTreeMap::from([
                ("atlas".to_owned(), AgentCaps::default()),
                (
                    "iris".to_owned(),
                    serde_json::from_str(r#"{"memory_max":"2G"}"#).expect("decodes"),
                ),
                (
                    "sock".to_owned(),
                    serde_json::from_str(bad).expect("decodes"),
                ),
            ]);
            let err = validated(&declared).expect_err(bad);
            assert!(err.to_string().contains("sock"), "{err:#}");
        }
    }
}