diff --git a/nix/module-eval/core-toggle.nix b/nix/module-eval/core-toggle.nix index f99c1a11..48326b2f 100644 --- a/nix/module-eval/core-toggle.nix +++ b/nix/module-eval/core-toggle.nix @@ -286,6 +286,22 @@ let && !(lib.hasInfix "mint swarmctl" s) && lib.hasInfix "mint swarm-authelia-bridge" s; } + { + # swarm-controller checks no caller itself: this location's authelia + # subrequest is the operator gate for every route under `/api/`. The + # webhook location is the control, the one path to the same socket that + # skips the gate, so the first clause reads a string that can lack it. + name = "the swarm UI's API location asks authelia before reaching the controller"; + ok = + let + l = allLocal.services.nginx.virtualHosts."t.local".locations; + socket = allLocal.services.hyperhive.deploy.swarm-controller.socketPath; + in + lib.hasInfix "auth_request /__hive_authelia;" l."/api/".extraConfig + && lib.hasInfix socket l."/api/".proxyPass + && lib.hasInfix socket l."/webhook/forge/".proxyPass + && !(lib.hasInfix "auth_request" l."/webhook/forge/".extraConfig); + } { name = "a hive with matrix off serves no matrix discovery endpoint"; ok = diff --git a/swarm-controller/README.md b/swarm-controller/README.md index 036ff283..4c7e12d4 100644 --- a/swarm-controller/README.md +++ b/swarm-controller/README.md @@ -52,6 +52,11 @@ The swarm's control plane. The swarm UI and `swarmctl` are its clients. agent applies them itself. `GET` also answers the model and effort the agent last reported in its turn-state header, kept in memory → [`docs/swarm/inference.md`](../docs/swarm/inference.md). +- **Resource caps** — each agent's CPU and memory cap, one key per agent in + the `agent-limits` KV bucket, which this daemon alone writes + (`GET`/`PUT`/`DELETE /api/agents/{name}/limits`); the hive running the + agent reads it, and each cap replaces the hive's own answer for that field → + [`docs/scheduler/coordinator.md`](../docs/scheduler/coordinator.md#container-resource-limits). - **Relays** — each agent's terminal and turn-state header as SSE, agent icons, the cross-repo issue report, and the UI's quick links. diff --git a/swarm-controller/src/agent_limits.rs b/swarm-controller/src/agent_limits.rs new file mode 100644 index 00000000..478a92db --- /dev/null +++ b/swarm-controller/src/agent_limits.rs @@ -0,0 +1,527 @@ +//! `/api/agents/{name}/limits` — the CPU and memory caps the swarm sets for +//! one agent. +//! +//! Each agent's caps are one value under its own key of `agent-limits` +//! (`swarm_queue_client::agent_limits`), which this daemon alone writes. The +//! hive running the agent reads the key; on the hive, a cap replaces the +//! hive's own answer for that field, and a field left unset stays with the +//! hive. +//! +//! The hive hands these values to systemd verbatim, and a value systemd +//! rejects stops the agent's container from starting, so they are validated +//! here as well as on the hive. + +use axum::Json; +use axum::extract::{Path, State}; +use axum::http::StatusCode; +use problem_details::ProblemDetails; +use serde::{Deserialize, Serialize}; +use swarm_queue_client::agent_limits::{self, AgentCaps}; +use utoipa::ToSchema; + +use crate::schedules::{live, queue, read, refused, write}; +use crate::{AppState, error_problem, wanted_error_status}; + +/// How many times a losing writer re-reads and re-applies, as in +/// `wanted::MAX_ATTEMPTS`. +const MAX_ATTEMPTS: usize = 5; + +/// One agent's caps. A field left out is not capped by the swarm. +#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize, ToSchema)] +#[serde(deny_unknown_fields)] +pub struct Caps { + /// systemd `CPUQuota=`: a percentage above zero, `"100%"` being one core. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[schema(example = "400%")] + cpu_quota: Option, + /// systemd `MemoryMax=`: a size above zero with an optional `K`/`M`/`G`/`T` + /// suffix, a percentage of physical memory up to `100%`, or `infinity`. + /// Stored in the one spelling systemd accepts (`8gb` is stored as `8G`). + #[serde(default, skip_serializing_if = "Option::is_none")] + #[schema(example = "8G")] + memory_max: Option, +} + +impl Caps { + /// The value to store, or why it is refused. + fn into_stored(self) -> Result { + Ok(AgentCaps { + cpu_quota: self.cpu_quota.as_deref().map(cpu_quota).transpose()?, + memory_max: self.memory_max.as_deref().map(memory_max).transpose()?, + }) + } +} + +impl From for Caps { + fn from(v: AgentCaps) -> Self { + Self { + cpu_quota: v.cpu_quota, + memory_max: v.memory_max, + } + } +} + +// The accepted shapes are hive-c0re's `resource_limits::validate_cpu_quota` / +// `validate_memory_max`, the operator's own input path for the same drop-in, +// with zero and an over-100% memory share refused as well. + +fn cpu_quota(value: &str) -> Result { + match value.strip_suffix('%') { + Some(n) if is_positive_number(n) => Ok(value.to_owned()), + _ => Err(format!( + "invalid cpu_quota {value:?}: expected a percentage above zero such as \"200%\" \ + (100% = one full core)" + )), + } +} + +fn memory_max(value: &str) -> Result { + if value == "infinity" { + return Ok(value.to_owned()); + } + let normalized = match value.strip_suffix('%') { + Some(n) => (is_positive_number(n) && n.parse::().is_ok_and(|p| p <= 100.0)) + .then(|| value.to_owned()), + None => byte_size(value), + }; + normalized.ok_or_else(|| { + format!( + "invalid memory_max {value:?}: expected a size above zero such as \"8G\", a \ + percentage up to \"100%\", or \"infinity\"" + ) + }) +} + +/// A byte count with an optional `K`/`M`/`G`/`T` multiplier, in any case and +/// with an optional trailing `B`, in the form systemd takes: `8gb` → `8G`, a +/// bare `1048576B` kept as it is. +fn byte_size(value: &str) -> Option { + const UNITS: [char; 8] = ['K', 'M', 'G', 'T', 'k', 'm', 'g', 't']; + let (body, b) = match value.strip_suffix(['B', 'b']) { + Some(body) => (body, true), + None => (value, false), + }; + let (mantissa, unit) = match body.strip_suffix(UNITS) { + Some(mantissa) => (mantissa, body[mantissa.len()..].to_ascii_uppercase()), + None if b => (body, "B".to_owned()), + None => (body, String::new()), + }; + is_positive_number(mantissa).then(|| format!("{mantissa}{unit}")) +} + +/// Digits with at most one `.` between digits, and not zero. No sign or +/// exponent, which systemd does not take either. +fn is_positive_number(value: &str) -> bool { + let (int, frac) = value.split_once('.').unwrap_or((value, "1")); + let digits = |s: &str| !s.is_empty() && s.bytes().all(|b| b.is_ascii_digit()); + digits(int) && digits(frac) && value.bytes().any(|b| (b'1'..=b'9').contains(&b)) +} + +fn agent_name(name: &str) -> Result { + hive_types::Ident::parse(name) + .map(hive_types::Ident::into_string) + .map_err(|reason| (StatusCode::BAD_REQUEST, reason.to_owned())) +} + +fn queue_problem(e: &anyhow::Error) -> ProblemDetails { + tracing::warn!(error = %format!("{e:#}"), "agent limits: queue failure"); + error_problem(wanted_error_status(e), &format!("{e:#}")) +} + +async fn open( + client: &async_nats::Client, +) -> Result { + agent_limits::open_or_create(client) + .await + .map_err(|e| queue_problem(&anyhow::Error::new(e))) +} + +#[utoipa::path( + get, + path = "/api/agents/{name}/limits", + params(("name" = String, Path, description = "agent name")), + responses( + (status = 200, description = "the caps the swarm sets for the agent; `{}` when it \ + sets none", body = Caps), + (status = 400, description = "not an identifier (problem+json)", body = String), + (status = 503, description = "no swarm queue is configured on this host, or it is \ + not connected (problem+json)", body = String), + (status = 500, description = "the caps could not be read (problem+json)", + body = String), + ), + tag = "agents" +)] +pub(crate) async fn get_agent_limits( + State(state): State, + Path(name): Path, +) -> Result, ProblemDetails> { + let agent = agent_name(&name).map_err(refused)?; + let client = queue(&state).map_err(refused)?; + let store = open(&client).await?; + let (current, _) = read::(&store, &agent) + .await + .map_err(|e| queue_problem(&e))?; + Ok(Json(current.into())) +} + +/// Set the CPU and memory caps the swarm declares for one agent. +/// +/// The value replaces the stored one whole. Whichever hive runs the agent +/// applies it; a running agent picks the caps up only at its next spawn or +/// rebuild. +#[utoipa::path( + put, + path = "/api/agents/{name}/limits", + params(("name" = String, Path, description = "agent name")), + request_body = Caps, + responses( + (status = 200, description = "the caps as stored", body = Caps), + (status = 400, description = "not an identifier, not declared on any hive, or a cap \ + is not a value systemd takes or is zero (problem+json)", body = String), + (status = 422, description = "an unknown field"), + (status = 503, description = "no swarm queue is configured on this host, or it is \ + not connected (problem+json)", body = String), + (status = 500, description = "the caps could not be written (problem+json)", + body = String), + ), + tag = "agents" +)] +pub(crate) async fn put_agent_limits( + State(state): State, + Path(name): Path, + Json(req): Json, +) -> Result, ProblemDetails> { + let agent = agent_name(&name).map_err(refused)?; + let value = req + .into_stored() + .map_err(|why| error_problem(StatusCode::BAD_REQUEST, &why))?; + let client = queue(&state).map_err(refused)?; + if !live(&state).await?.contains(&agent) { + return Err(error_problem( + StatusCode::BAD_REQUEST, + &format!("no hive's wanted state declares an agent {agent:?} that is not destroyed"), + )); + } + let store = open(&client).await?; + for _ in 0..MAX_ATTEMPTS { + let (current, revision) = read::(&store, &agent) + .await + .map_err(|e| queue_problem(&e))?; + if current == value + || write(&store, &agent, &value, revision) + .await + .map_err(|e| queue_problem(&e))? + { + tracing::info!(%agent, ?value, "set agent limits"); + return Ok(Json(value.into())); + } + } + Err(queue_problem(&anyhow::anyhow!( + "gave up writing the limits of {agent} after {MAX_ATTEMPTS} conflicting writes" + ))) +} + +#[utoipa::path( + delete, + path = "/api/agents/{name}/limits", + params(("name" = String, Path, description = "agent name")), + responses( + (status = 204, description = "cleared; the hive's own answer applies again"), + (status = 400, description = "not an identifier (problem+json)", body = String), + (status = 503, description = "no swarm queue is configured on this host, or it is \ + not connected (problem+json)", body = String), + (status = 500, description = "the caps could not be cleared (problem+json)", + body = String), + ), + tag = "agents" +)] +pub(crate) async fn delete_agent_limits( + State(state): State, + Path(name): Path, +) -> Result { + let agent = agent_name(&name).map_err(refused)?; + let client = queue(&state).map_err(refused)?; + let store = open(&client).await?; + store + .delete(&agent) + .await + .map_err(|e| queue_problem(&anyhow::Error::new(e)))?; + tracing::info!(%agent, "cleared agent limits"); + Ok(StatusCode::NO_CONTENT) +} + +#[cfg(test)] +mod tests { + use super::{ + Caps, cpu_quota, delete_agent_limits, get_agent_limits, memory_max, put_agent_limits, + }; + use crate::test_nats::TestNats; + use crate::wanted::WantedWriter; + use axum::extract::{Path, State}; + use axum::http::StatusCode; + use std::sync::Arc; + use swarm_queue_client::wanted::AgentState; + + #[test] + fn cpu_quota_takes_a_percentage_above_zero() { + for ok in ["400%", "100%", "0.5%", "1500%"] { + assert_eq!(cpu_quota(ok).as_deref(), Ok(ok)); + } + for bad in [ + "0%", "0.0%", "-50%", "400", "4 00%", "%", "lots", "", "1e3%", "+5%", + ] { + assert!(cpu_quota(bad).is_err(), "{bad:?} was accepted"); + } + } + + #[test] + fn memory_max_takes_a_size_a_share_or_infinity_and_stores_the_systemd_spelling() { + for (given, stored) in [ + ("8G", "8G"), + ("8gb", "8G"), + ("512m", "512M"), + ("1048576", "1048576"), + ("1048576B", "1048576B"), + ("1.5T", "1.5T"), + ("50%", "50%"), + ("100%", "100%"), + ("infinity", "infinity"), + ] { + assert_eq!(memory_max(given).as_deref(), Ok(stored), "{given:?}"); + } + for bad in [ + "0", "0G", "0.0G", "0%", "-8G", "150%", "8X", "G", "8 G", "", "eight", "8GG", + ] { + assert!(memory_max(bad).is_err(), "{bad:?} was accepted"); + } + } + + /// One hive, `pr1ma`, on `client`. + fn state(client: &async_nats::Client) -> crate::AppState { + crate::AppState { + hives: Arc::new(vec![crate::HiveEntry { + name: "pr1ma".to_owned(), + domain: "pr1ma.example".to_owned(), + }]), + links: Arc::new(Vec::new()), + status: Some(Arc::new(crate::status::StatusReader::new( + client.clone(), + std::time::Duration::from_mins(1), + ))), + wanted: Some(Arc::new(WantedWriter::new(client.clone()))), + agent_status: None, + agent_icons: None, + reported_inference: None, + jobq: Arc::new(std::sync::Mutex::new(hive_jobq::scheduler::Scheduler::new( + hive_jobq::Graph::new(), + hive_jobq::resources::ResourceTable::new(), + ))), + webhook_secret: None, + config_prs: Arc::default(), + swarm_name: None, + auth: None, + forge: Arc::default(), + create_gate: Arc::default(), + accounts: crate::linked_accounts::Accounts::Fake(Arc::default()), + try_outs: crate::review_classifier::TryOuts::default(), + } + } + + /// A state on a real server whose `pr1ma` declares `atlas`. + async fn declared(client: &async_nats::Client) -> crate::AppState { + WantedWriter::new(client.clone()) + .set("pr1ma", "atlas", AgentState::Up) + .await + .expect("declare atlas"); + state(client) + } + + fn caps(cpu: Option<&str>, mem: Option<&str>) -> Caps { + Caps { + cpu_quota: cpu.map(str::to_owned), + memory_max: mem.map(str::to_owned), + } + } + + async fn put( + state: &crate::AppState, + agent: &str, + value: Caps, + ) -> Result, problem_details::ProblemDetails> { + put_agent_limits( + State(state.clone()), + Path(agent.to_owned()), + axum::Json(value), + ) + .await + } + + async fn get(state: &crate::AppState, agent: &str) -> Caps { + get_agent_limits(State(state.clone()), Path(agent.to_owned())) + .await + .expect("caps are read") + .0 + } + + /// The raw value a hive reads for `agent`, or `None` for no key or a + /// deleted one. + async fn published(client: &async_nats::Client, agent: &str) -> Option { + let store = swarm_queue_client::agent_limits::open_or_create(client) + .await + .expect("bucket"); + store + .get(agent) + .await + .expect("read") + .map(|v| String::from_utf8(v.to_vec()).expect("utf-8")) + } + + /// Through the real handlers and a real `nats-server`: a cap lands under + /// the agent's own key as the hive reads it, `GET` answers it, and + /// `DELETE` takes it out again. + /// + /// The published bytes are pinned literally because hive-c0re's + /// `workers::swarm_limits` tests decode this same string; a change to + /// either side's reading of it fails one of the two. + #[tokio::test] + async fn the_route_writes_the_caps_a_hive_reads_and_delete_clears_them() { + let nats = TestNats::start(); + let client = nats.client().await; + let state = declared(&client).await; + + let answer = put(&state, "atlas", caps(Some("400%"), Some("8gb"))) + .await + .expect("caps are set"); + assert_eq!(answer.0, caps(Some("400%"), Some("8G"))); + assert_eq!( + published(&client, "atlas").await.as_deref(), + Some(r#"{"cpu_quota":"400%","memory_max":"8G"}"#) + ); + assert_eq!(get(&state, "atlas").await, caps(Some("400%"), Some("8G"))); + + assert_eq!( + delete_agent_limits(State(state.clone()), Path("atlas".to_owned())).await, + Ok(StatusCode::NO_CONTENT) + ); + assert_eq!(published(&client, "atlas").await, None); + assert_eq!(get(&state, "atlas").await, Caps::default()); + } + + /// A name no hive declares gets no key, so a typo cannot cap an agent + /// that does not exist yet. + #[tokio::test] + async fn an_agent_no_hive_declares_is_refused_and_nothing_is_written() { + let nats = TestNats::start(); + let client = nats.client().await; + let state = declared(&client).await; + + let problem = put(&state, "iris", caps(Some("400%"), None)) + .await + .expect_err("no hive declares iris"); + assert_eq!(problem.status, Some(StatusCode::BAD_REQUEST), "{problem:?}"); + assert_eq!(published(&client, "iris").await, None); + } + + /// A refused cap is answered before the queue is asked: the client here + /// was never connected, and would answer 503. + #[tokio::test] + async fn a_refused_cap_is_a_400_before_the_queue() { + let client = async_nats::ConnectOptions::new() + .retry_on_initial_connect() + .connect("127.0.0.1:1") + .await + .expect("returns without a connection"); + let state = state(&client); + + let problem = put(&state, "atlas", caps(Some("0%"), None)) + .await + .expect_err("a zero cap is refused"); + assert_eq!(problem.status, Some(StatusCode::BAD_REQUEST), "{problem:?}"); + + // The control: a valid cap on the same state does reach the queue. + let problem = put(&state, "atlas", caps(Some("400%"), None)) + .await + .expect_err("the queue is not connected"); + assert_eq!( + problem.status, + Some(StatusCode::SERVICE_UNAVAILABLE), + "{problem:?}" + ); + } + + const CAPPED: &str = r#"{"cpu_quota":"400%","memory_max":"8G"}"#; + + /// `atlas` capped, behind the real router. + async fn capped_router(client: &async_nats::Client) -> axum::Router { + let state = declared(client).await; + let _ = put(&state, "atlas", caps(Some("400%"), Some("8G"))) + .await + .expect("cap atlas"); + crate::build_app(state) + } + + async fn put_raw(app: &axum::Router, body: &'static str) -> StatusCode { + use tower::ServiceExt as _; + app.clone() + .oneshot( + axum::http::Request::put("/api/agents/atlas/limits") + .header(axum::http::header::CONTENT_TYPE, "application/json") + .body(axum::body::Body::from(body)) + .expect("request"), + ) + .await + .expect("router answers") + .status() + } + + /// A body that is not JSON is refused before the handler runs, and the + /// caps already set stay. The well-formed `{}` after it is the control: + /// the same router and request do reach the bucket. + #[tokio::test] + async fn a_malformed_body_is_a_400_and_the_caps_stay() { + let nats = TestNats::start(); + let client = nats.client().await; + let app = capped_router(&client).await; + + for body in [r#"{"cpu_quota":"#, "cpu_quota=400%", ""] { + assert_eq!( + put_raw(&app, body).await, + StatusCode::BAD_REQUEST, + "{body:?}" + ); + assert_eq!( + published(&client, "atlas").await.as_deref(), + Some(CAPPED), + "{body:?}" + ); + } + + assert_eq!(put_raw(&app, "{}").await, StatusCode::OK); + assert_eq!(published(&client, "atlas").await.as_deref(), Some("{}")); + } + + /// A field the body does not define is a 422 and writes nothing: a + /// misspelt `memory_max` must not read as "no memory cap". + #[tokio::test] + async fn an_unknown_field_is_a_422_and_the_caps_stay() { + let nats = TestNats::start(); + let client = nats.client().await; + let app = capped_router(&client).await; + + for body in [ + r#"{"cpu_quota":"100%","memory":"2G"}"#, + r#"{"limits":{"cpu_quota":"100%"}}"#, + ] { + assert_eq!( + put_raw(&app, body).await, + StatusCode::UNPROCESSABLE_ENTITY, + "{body}" + ); + assert_eq!( + published(&client, "atlas").await.as_deref(), + Some(CAPPED), + "{body}" + ); + } + } +} diff --git a/swarm-controller/src/main.rs b/swarm-controller/src/main.rs index 63c3f03e..cde67d31 100644 --- a/swarm-controller/src/main.rs +++ b/swarm-controller/src/main.rs @@ -44,6 +44,7 @@ use utoipa_axum::{router::OpenApiRouter, routes}; mod agent_icon; mod agent_identity; mod agent_input; +mod agent_limits; mod agent_renewal; mod agent_state_stream; mod agent_status; @@ -3300,6 +3301,11 @@ fn build_app(state: AppState) -> axum::Router { inference::put_agent_inference, inference::delete_agent_inference )) + .routes(routes!( + agent_limits::get_agent_limits, + agent_limits::put_agent_limits, + agent_limits::delete_agent_limits + )) .routes(routes!(get_config_prs)) .routes(routes!(create_agent)) .routes(routes!(mint_agent_identity)) diff --git a/swarm-queue-client/Cargo.toml b/swarm-queue-client/Cargo.toml index 1123f9a3..5496599d 100644 --- a/swarm-queue-client/Cargo.toml +++ b/swarm-queue-client/Cargo.toml @@ -47,6 +47,9 @@ agent-inference = ["async-nats/jetstream"] # `ToSchema` on the review classifier's shared types, for swarm-controller, # whose HTTP API serves the same shapes the queue carries. schema = ["dep:utoipa"] +# `agent_limits::read`, for hive-c0re, which reads the key of each agent it +# runs the same way. Explicit for the same reason as `notices`. +agent-limits = ["async-nats/jetstream"] [dependencies] # Bare (no `kv`/`jetstream`) unless a consumer opts into the `kv` feature diff --git a/swarm-queue-client/src/agent_limits.rs b/swarm-queue-client/src/agent_limits.rs new file mode 100644 index 00000000..91b857af --- /dev/null +++ b/swarm-queue-client/src/agent_limits.rs @@ -0,0 +1,166 @@ +//! 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`/``, 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, + /// systemd `MemoryMax=`, e.g. `"8G"`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub memory_max: Option, +} + +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 { + 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 +} + +/// 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> { + 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%")); + } +} diff --git a/swarm-queue-client/src/lib.rs b/swarm-queue-client/src/lib.rs index fa5b991b..6cd0da3a 100644 --- a/swarm-queue-client/src/lib.rs +++ b/swarm-queue-client/src/lib.rs @@ -216,6 +216,11 @@ pub mod agent_schedules; /// unconditional; writing needs `kv` and reading needs `agent-inference`. pub mod agent_inference; +/// The bucket swarm-controller sets each agent's CPU and memory caps in — one +/// key per agent, read by the hive that runs it. The names are unconditional; +/// writing needs `kv` and reading needs `agent-limits`. +pub mod agent_limits; + #[cfg(feature = "kv")] mod kv;