Watch
0
0
Fork
You've already forked hyperhive
0
hyperhive/swarm-controller/src/agent_limits.rs
atlas 738dd1c95b swarm-controller: split main.rs into jobq, routes and startup modules
no behaviour change: a pure move of items and tests into new modules, plus
visibility edits (pub(crate)) and path updates in sibling modules that named
crate-root items. same routes, same OpenAPI, 561 tests before and after.

why: #5470, new features stop landing in one 6.4k-line file by default.

files outside swarm-controller/src: docs/agent-lifecycle/agent-roster.md,
swarm-ui JobsPage.tsx and AgentsPage.tsx, nix/host-modules/swarm-controller.nix
and swarm-grafana.nix hold comments that named main.rs for items that moved.
2026-10-18 13:44:43 +02:00

842 lines
30 KiB
Rust

//! `/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 systemd ignores a value
//! it can't parse, leaving the agent uncapped, so they are checked here against
//! systemd's own syntax (`Caps`) 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::routes::{AppState, error_problem, wanted_error_status};
use crate::schedules::{live, queue, read_or_default, refused, write_revision};
/// 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.
///
/// Each value is stored as given, and only in a form systemd applies: a value
/// systemd ignores ("Invalid …, ignoring") would leave the agent uncapped,
/// since a swarm cap replaces the hive's answer.
#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize, ToSchema)]
#[serde(deny_unknown_fields)]
pub struct Caps {
/// systemd `CPUQuota=`: digits, optionally `.` and one or two digits, then
/// `%`; above zero and at most `21474836.47%`. `"100%"` is one core.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schema(example = "400%")]
cpu_quota: Option<String>,
/// systemd `MemoryMax=`: `infinity`; a percentage of physical memory
/// written as for `cpu_quota`, above zero and at most `100%`; or a size:
/// digits, optionally `.` and digits, then nothing, `B` or one of
/// `K`/`M`/`G`/`T` (1024-based, upper case), at least one byte and below
/// 2^64 - 1 bytes.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[schema(example = "8G")]
memory_max: Option<String>,
}
impl Caps {
/// The value to store, or why it is refused.
pub(crate) fn into_stored(self) -> Result<AgentCaps, String> {
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<AgentCaps> for Caps {
fn from(v: AgentCaps) -> Self {
Self {
cpu_quota: v.cpu_quota,
memory_max: v.memory_max,
}
}
}
// The accepted forms are a subset of what systemd parses (checked against
// `systemd-analyze verify`, systemd 260) and stricter than hive-c0re's own
// `validate_cpu_quota` / `validate_memory_max`, so the hive never refuses a
// stored value.
/// systemd keeps `CPUQuota=` in hundredths of a percent in a C `int`.
const CPU_QUOTA_MAX_HUNDREDTHS: u64 = 2_147_483_647;
fn cpu_quota(value: &str) -> Result<String, String> {
match hundredths(value) {
Some(h) if (1..=CPU_QUOTA_MAX_HUNDREDTHS).contains(&h) => Ok(value.to_owned()),
_ => Err(format!(
"invalid cpu_quota {value:?}: expected a percentage above zero with at most two \
decimals, such as \"200%\" (100% = one full core), up to \"21474836.47%\""
)),
}
}
fn memory_max(value: &str) -> Result<String, String> {
let applies = value == "infinity"
|| if value.ends_with('%') {
hundredths(value).is_some_and(|h| (1..=10_000).contains(&h))
} else {
size_bytes(value).is_some_and(|b| b >= 1 && b < u128::from(u64::MAX))
};
if applies {
Ok(value.to_owned())
} else {
Err(format!(
"invalid memory_max {value:?}: expected a size above zero such as \"8G\" \
(suffix B, K, M, G or T, upper case), a percentage up to \"100%\" with at \
most two decimals, or \"infinity\""
))
}
}
/// `value` as systemd reads a percentage, in hundredths of a percent: digits,
/// optionally `.` and one or two digits, then `%`. `None` for anything else,
/// and for a whole part with a leading zero other than `0` itself: systemd
/// ignores `08%` and may read `010%` as octal.
fn hundredths(value: &str) -> Option<u64> {
let number = value.strip_suffix('%')?;
let (int, frac) = match number.split_once('.') {
Some((int, frac)) if (1..=2).contains(&frac.len()) && digits(frac) => (int, frac),
Some(_) => return None,
None => (number, ""),
};
if !digits(int) || (int.len() > 1 && int.starts_with('0')) {
return None;
}
let frac = match frac.len() {
0 => 0,
1 => frac.parse::<u64>().ok()? * 10,
_ => frac.parse::<u64>().ok()?,
};
int.parse::<u64>().ok()?.checked_mul(100)?.checked_add(frac)
}
/// `value` as one term of systemd's `parse_size` reads it, in bytes: digits,
/// optionally `.` and digits, then nothing, `B`, or `K`/`M`/`G`/`T` (1024-based,
/// case-sensitive). `P` and `E`, which systemd also reads, are left out
/// because hive-c0re's validator refuses them. `None` for anything else, and for what
/// `parse_size` refuses as out of range: a whole part or fraction digits past
/// `u64`, or a whole part that, plus one when there is a fraction, is past
/// `u64` once scaled.
fn size_bytes(value: &str) -> Option<u128> {
let (number, exponent) = match value.char_indices().last()? {
(i, 'B') => (&value[..i], 0),
(i, c) => match "KMGT".find(c) {
Some(k) => (&value[..i], k + 1),
None => (value, 0),
},
};
let scale = 1024u128.pow(u32::try_from(exponent).ok()?);
let (int, frac) = number.split_once('.').unwrap_or((number, ""));
if !digits(int) || (number.contains('.') && !digits(frac)) {
return None;
}
let whole = u128::from(int.parse::<u64>().ok()?);
let frac_value = if frac.is_empty() {
0
} else {
u128::from(frac.parse::<u64>().ok()?)
};
if whole + u128::from(frac_value > 0) > u128::from(u64::MAX) / scale {
return None;
}
let frac_bytes = frac_value * scale / 10u128.checked_pow(u32::try_from(frac.len()).ok()?)?;
Some(whole * scale + frac_bytes)
}
fn digits(s: &str) -> bool {
!s.is_empty() && s.bytes().all(|b| b.is_ascii_digit())
}
fn agent_name(name: &str) -> Result<String, (StatusCode, String)> {
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_or_create(
client: &async_nats::Client,
) -> Result<async_nats::jetstream::kv::Store, ProblemDetails> {
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<AppState>,
Path(name): Path<String>,
) -> Result<Json<Caps>, ProblemDetails> {
let agent = agent_name(&name).map_err(refused)?;
let client = queue(&state).map_err(refused)?;
let Some(store) = agent_limits::open(&client)
.await
.map_err(|e| queue_problem(&anyhow::Error::new(e)))?
else {
return Ok(Json(Caps::default()));
};
let (current, _, _) = read_or_default::<AgentCaps>(&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<AppState>,
Path(name): Path<String>,
Json(req): Json<Caps>,
) -> Result<Json<Caps>, 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"),
));
}
store(&client, &agent, &value).await?;
Ok(Json(value.into()))
}
/// What [`store`] replaced, or `None` when it wrote nothing because the caps
/// were already stored.
pub(crate) type Replaced = Option<crate::schedules::Replaced<AgentCaps>>;
/// Replace `agent`'s stored caps with `value`.
pub(crate) async fn store(
client: &async_nats::Client,
agent: &str,
value: &AgentCaps,
) -> Result<Replaced, ProblemDetails> {
let store = open_or_create(client).await?;
for _ in 0..MAX_ATTEMPTS {
let (current, revision, undecodable) = read_or_default::<AgentCaps>(&store, agent)
.await
.map_err(|e| queue_problem(&e))?;
if current == *value && !undecodable {
tracing::info!(%agent, ?value, "set agent limits");
return Ok(None);
}
// A deleted key reads as present at its marker's revision, so ask
// whether it holds a value before keeping one to put back.
let held = store
.get(agent)
.await
.map_err(|e| queue_problem(&anyhow::Error::new(e)))?
.is_some();
if let Some(written) = write_revision(&store, agent, value, revision)
.await
.map_err(|e| queue_problem(&e))?
{
tracing::info!(%agent, ?value, "set agent limits");
return Ok(Some(crate::schedules::Replaced::new(
held.then_some(current),
written,
)));
}
}
Err(queue_problem(&anyhow::anyhow!(
"gave up writing the limits of {agent} after {MAX_ATTEMPTS} conflicting writes"
)))
}
/// Undo what [`store`] did, unless another write has landed since. Best
/// effort: a failure is logged, since the caller is already on its way to an
/// error of its own.
pub(crate) async fn restore(client: &async_nats::Client, agent: &str, replaced: Replaced) {
let Some(replaced) = replaced else { return };
let undone = match open_or_create(client).await {
Ok(store) => replaced.undo(&store, agent).await,
Err(e) => Err(anyhow::anyhow!("{e:?}")),
};
match undone {
Ok(true) => {}
Ok(false) => {
tracing::info!(%agent, "left the agent's limits: another write landed since");
}
Err(e) => {
tracing::warn!(%agent, error = %format!("{e:#}"), "could not put back the agent's limits");
}
}
}
#[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<AppState>,
Path(name): Path<String>,
) -> Result<StatusCode, ProblemDetails> {
let agent = agent_name(&name).map_err(refused)?;
let client = queue(&state).map_err(refused)?;
let store = open_or_create(&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)]
pub(crate) 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;
/// Of the `bad` values, systemd parses `""` (no quota) and `+5%`, and the
/// last adds a directive on a second line. It ignores every other one
/// ("Invalid CPU quota …, ignoring", checked with `systemd-analyze
/// verify`), which would leave the agent uncapped.
#[test]
fn cpu_quota_takes_a_percentage_above_zero() {
for ok in [
"400%",
"100%",
"0.5%",
"0.01%",
"1500%",
"400.12%",
"21474836.47%",
] {
assert_eq!(cpu_quota(ok).as_deref(), Ok(ok));
}
for bad in [
"0%",
"0.0%",
"0.00%",
"-50%",
"400",
"4 00%",
"%",
"lots",
"",
"1e3%",
"+5%",
"400.123%",
"0.001%",
"99999999999999999999%",
"21474836.48%",
"05%",
"08%",
"085480%",
"010%",
"00.5%",
"1.%",
".5%",
"400%\nExecStartPre=/bin/sh -c id #%",
] {
assert!(cpu_quota(bad).is_err(), "{bad:?} was accepted");
}
}
/// As above for `MemoryMax=`. Of the `bad` values, systemd also parses `3P`
/// and `2E`, which hive-c0re's validator refuses, and `8 G` and `1.G`.
#[test]
fn memory_max_takes_a_size_a_share_or_infinity_that_systemd_applies() {
for ok in [
"8G",
"512M",
"4096K",
"1048576",
"1048576B",
"1.5B",
"0.5K",
"1.5T",
"1.23456789G",
"16383T",
"16777215T",
"18446744073709551614",
"18446744073709551614.5",
"1.00000000000000000001G",
"50%",
"12.34%",
"0.01%",
"100%",
"100.00%",
"infinity",
] {
assert_eq!(memory_max(ok).as_deref(), Ok(ok), "{ok:?}");
}
for bad in [
"0",
"0G",
"0.0G",
"0.5",
"0%",
"-8G",
"150%",
"100.01%",
"12.345%",
"08%",
"085480%",
"010%",
"00.5%",
"8X",
"G",
"8 G",
"",
"eight",
"8GG",
"8gb",
"8g",
"8GB",
"4096k",
"1048576b",
"3P",
"2E",
".5G",
"1.G",
"Infinity",
"17000000T",
"16777216T",
"16777215.5T",
"99999999999999999999T",
"18446744073709551615",
"18446744073709551616",
"1.123456789012345678901234567890G",
"8G\nDevicePolicy=auto",
] {
assert!(memory_max(bad).is_err(), "{bad:?} was accepted");
}
}
/// One hive, `pr1ma`, on `client`.
pub(crate) fn state(client: &async_nats::Client) -> crate::routes::AppState {
crate::routes::AppState {
hives: Arc::new(vec![crate::routes::hives::HiveEntry {
name: "pr1ma".to_owned(),
domain: "pr1ma.example".to_owned(),
dashboard: "pr1ma.hives.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,
agents_md: None,
reported_inference: None,
agent_login: 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()),
reviews: Arc::default(),
try_outs: crate::review_classifier::TryOuts::default(),
runs: crate::review_classifier::runs::Runs::Live(None),
builds: crate::builder::runs::Builds::Live(None),
alerts: None,
}
}
/// A state on a real server whose `pr1ma` declares `atlas`.
pub(crate) async fn declared(client: &async_nats::Client) -> crate::routes::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::routes::AppState,
agent: &str,
value: Caps,
) -> Result<axum::Json<Caps>, problem_details::ProblemDetails> {
put_agent_limits(
State(state.clone()),
Path(agent.to_owned()),
axum::Json(value),
)
.await
}
async fn get(state: &crate::routes::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.
pub(crate) async fn published(client: &async_nats::Client, agent: &str) -> Option<String> {
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("8G")))
.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 read before any cap was set answers `{}` and leaves the bucket
/// uncreated.
#[tokio::test]
async fn a_read_before_any_cap_answers_empty_and_creates_no_bucket() {
let nats = TestNats::start();
let client = nats.client().await;
let state = state(&client);
assert_eq!(get(&state, "atlas").await, Caps::default());
assert!(
swarm_queue_client::agent_limits::open(&client)
.await
.expect("the queue answers")
.is_none()
);
}
/// 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::startup::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 value carrying a directive on a second line is a 400 through the real
/// router, and the caps already set stay.
#[tokio::test]
async fn a_cap_with_a_second_line_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":"400%\nExecStartPre=/bin/sh -c id #%"}"#,
r#"{"memory_max":"8G\nDevicePolicy=auto"}"#,
] {
assert_eq!(put_raw(&app, body).await, StatusCode::BAD_REQUEST, "{body}");
assert_eq!(
published(&client, "atlas").await.as_deref(),
Some(CAPPED),
"{body}"
);
}
}
/// 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}"
);
}
}
/// A stored value this build cannot decode, such as one a newer build
/// wrote, reads as no caps, and a `PUT`, even of no caps, replaces it.
#[tokio::test]
async fn an_undecodable_value_reads_as_no_caps_until_a_put() {
let nats = TestNats::start();
let client = nats.client().await;
let state = declared(&client).await;
let store = swarm_queue_client::agent_limits::open_or_create(&client)
.await
.expect("bucket");
store
.put(
"atlas",
br#"{"cpu_quota":{"later":"400%"}}"#.to_vec().into(),
)
.await
.expect("written");
assert_eq!(get(&state, "atlas").await, Caps::default());
let answer = put(&state, "atlas", Caps::default())
.await
.expect("caps are set");
assert_eq!(answer.0, Caps::default());
assert_eq!(published(&client, "atlas").await.as_deref(), Some("{}"));
}
/// `restore` puts back the caps a `store` replaced, or takes the key out
/// again when there was none.
#[tokio::test]
async fn restore_undoes_a_store() {
let nats = TestNats::start();
let client = nats.client().await;
let caps = |json: &str| {
serde_json::from_str::<Caps>(json)
.expect("caps")
.into_stored()
.expect("valid")
};
let first = super::store(&client, "atlas", &caps(CAPPED))
.await
.expect("stored");
super::restore(&client, "atlas", first).await;
assert_eq!(published(&client, "atlas").await, None);
super::store(&client, "atlas", &caps(CAPPED))
.await
.expect("stored");
let second = super::store(&client, "atlas", &caps(r#"{"cpu_quota":"100%"}"#))
.await
.expect("replaced");
super::restore(&client, "atlas", second).await;
assert_eq!(published(&client, "atlas").await.as_deref(), Some(CAPPED));
}
/// A write that landed after the one being undone is newer, so `restore`
/// leaves it, whether the undone write replaced a value or made the key.
#[tokio::test]
async fn restore_leaves_a_write_that_landed_since() {
let nats = TestNats::start();
let client = nats.client().await;
let caps = |json: &str| {
serde_json::from_str::<Caps>(json)
.expect("caps")
.into_stored()
.expect("valid")
};
let other = caps(r#"{"cpu_quota":"100%"}"#);
let made = super::store(&client, "atlas", &caps(CAPPED))
.await
.expect("stored");
super::store(&client, "atlas", &other).await.expect("a PUT");
super::restore(&client, "atlas", made).await;
assert_eq!(
published(&client, "atlas").await.as_deref(),
Some(r#"{"cpu_quota":"100%"}"#)
);
let replaced = super::store(&client, "atlas", &caps(r#"{"cpu_quota":"200%"}"#))
.await
.expect("replaced");
super::store(&client, "atlas", &caps(CAPPED))
.await
.expect("a PUT");
super::restore(&client, "atlas", replaced).await;
assert_eq!(published(&client, "atlas").await.as_deref(), Some(CAPPED));
}
}