hyperhive/swarm-queue-client/src/agent_limits.rs

//! The per-agent resource-cap KV bucket: one key per agent, holding the CPU
//! and memory caps the swarm sets for it.
//!
//! **swarm-controller is the only writer**, and creates the bucket on its
//! first write. Every field is optional: an absent field leaves the agent on
//! its hive's own answer, so an agent with no key runs exactly as it would
//! with no swarm.
//!
//! **The hive running the agent reads the key**, not the agent: the caps go
//! into the container's systemd drop-in, which the host writes. A hive reads
//! with one `DIRECT.GET` per agent
//! ([`crate::agent_limits::direct_get_subject`]), and its grant covers every
//! key of this bucket, read-only. Scoping it to the hive's own agents would
//! need a roster in the auth responder.
//!
//! Keyed by the agent alone, like [`crate::agent_inference`]: an agent can
//! move between hives, and its caps go with it.

/// The KV bucket caps are published into, keyed by agent name.
pub const BUCKET: &str = "agent-limits";

/// The stream backing [`BUCKET`]. `create_key_value` prefixes `KV_`.
pub const STREAM: &str = "KV_agent-limits";

/// The subject `agent`'s entry is stored under.
#[must_use]
pub fn key_subject(agent: &str) -> String {
    format!("$KV.{BUCKET}.{agent}")
}

/// The request subject that reads `agent`'s current entry.
#[must_use]
pub fn direct_get_subject(agent: &str) -> String {
    format!("$JS.API.DIRECT.GET.{STREAM}.{}", key_subject(agent))
}

/// One agent's caps: the value under `BUCKET`/`<agent>`, as systemd values
/// carried verbatim. `None` is "not capped by the swarm", never "uncapped".
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct AgentCaps {
    /// systemd `CPUQuota=`, e.g. `"400%"`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cpu_quota: Option<String>,
    /// systemd `MemoryMax=`, e.g. `"8G"`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub memory_max: Option<String>,
}

impl AgentCaps {
    /// True when no field is set, which a hive reads the same as no key.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.cpu_quota.is_none() && self.memory_max.is_none()
    }
}

/// Open the caps bucket for writing, creating it if nothing has yet.
///
/// **Controller-side only.** `history: 1`: a hive wants the current value.
#[cfg(feature = "kv")]
pub async fn open_or_create(
    client: &async_nats::Client,
) -> Result<async_nats::jetstream::kv::Store, crate::Error> {
    let js = async_nats::jetstream::new(client.clone());
    crate::kv::open_or_create(
        &js,
        BUCKET,
        "agent-limits",
        "Resource caps the swarm controller sets for each agent".to_owned(),
    )
    .await
}

/// Open the caps bucket for reading, or `None` when nothing has created it
/// yet, which reads the same as no caps.
///
/// **Controller-side**, for reads that must not create the bucket.
///
/// # Errors
///
/// [`crate::Error::OpenBucket`] for every failure other than the bucket not
/// existing.
#[cfg(feature = "kv")]
pub async fn open(
    client: &async_nats::Client,
) -> Result<Option<async_nats::jetstream::kv::Store>, crate::Error> {
    let js = async_nats::jetstream::new(client.clone());
    match js.get_key_value(BUCKET).await {
        Ok(store) => Ok(Some(store)),
        Err(e) if crate::wanted::is_bucket_not_found(&e) => Ok(None),
        Err(source) => Err(crate::Error::OpenBucket {
            bucket: BUCKET.to_owned(),
            source,
        }),
    }
}

/// What reading one agent's key found.
#[cfg(feature = "agent-limits")]
#[derive(Debug)]
pub enum Read {
    /// The bucket answered; an absent or deleted key has every field unset.
    Current(AgentCaps),
    /// Nothing answers for the bucket's stream: no cap has been written yet,
    /// or `JetStream` is not serving it right now. Not an unset value, since
    /// the second case says nothing about the agent's caps.
    NoBucket,
}

/// Read `agent`'s current caps with one direct get.
///
/// # Errors
///
/// A failed request or an undecodable value. A request the queue refused for
/// permissions reaches the client as a timeout, not as a refusal.
#[cfg(feature = "agent-limits")]
pub async fn read(
    client: &async_nats::Client,
    agent: &str,
) -> Result<Read, Box<dyn std::error::Error + Send + Sync>> {
    use async_nats::jetstream::stream::DirectGetErrorKind;

    let js = async_nats::jetstream::new(client.clone());
    let stream = js.get_stream_no_info(STREAM).await?;
    let message = match stream.direct_get_last_for_subject(key_subject(agent)).await {
        Ok(message) => message,
        Err(e) if matches!(e.kind(), DirectGetErrorKind::NotFound) => {
            return Ok(Read::Current(AgentCaps::default()));
        }
        Err(e) if is_no_responders(&e) => return Ok(Read::NoBucket),
        Err(e) => return Err(e.into()),
    };
    // A KV delete or purge is stored as an empty marker carrying this header.
    let removed = message
        .headers
        .get("KV-Operation")
        .is_some_and(|op| matches!(op.as_str(), "DEL" | "PURGE"));
    if removed {
        return Ok(Read::Current(AgentCaps::default()));
    }
    Ok(Read::Current(serde_json::from_slice(&message.payload)?))
}

#[cfg(feature = "agent-limits")]
fn is_no_responders(e: &async_nats::jetstream::stream::DirectGetError) -> bool {
    matches!(
        e.kind(),
        async_nats::jetstream::stream::DirectGetErrorKind::ErrorResponse(status, _)
            if status == async_nats::StatusCode::NO_RESPONDERS
    )
}

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

    /// The names the hive's grant in `swarm-nats-auth` is built from.
    #[test]
    fn names_match_the_queue_grant() {
        assert_eq!(STREAM, format!("KV_{BUCKET}"));
        assert_eq!(key_subject("iris"), "$KV.agent-limits.iris");
        assert_eq!(
            direct_get_subject("iris"),
            "$JS.API.DIRECT.GET.KV_agent-limits.$KV.agent-limits.iris"
        );
    }

    /// An unset field is left out on the wire, and `{}` is every field unset.
    #[test]
    fn unset_fields_are_absent() {
        let value = AgentCaps {
            memory_max: Some("8G".to_owned()),
            ..AgentCaps::default()
        };
        assert_eq!(
            serde_json::to_string(&value).unwrap(),
            r#"{"memory_max":"8G"}"#
        );
        let empty: AgentCaps = serde_json::from_str("{}").unwrap();
        assert!(empty.is_empty());
    }

    /// Fields a newer controller adds must not cost an older hive its caps.
    #[test]
    fn decoding_tolerates_unknown_fields() {
        let decoded: AgentCaps =
            serde_json::from_str(r#"{"cpu_quota":"400%","io_weight":"10"}"#).unwrap();
        assert_eq!(decoded.cpu_quota.as_deref(), Some("400%"));
    }
}