841 lines
30 KiB
Rust
841 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::schedules::{live, queue, read_or_default, refused, write_revision};
|
|
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.
|
|
///
|
|
/// 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::AppState {
|
|
crate::AppState {
|
|
hives: Arc::new(vec![crate::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,
|
|
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::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<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::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::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));
|
|
}
|
|
}
|