From 4ad9579e9744307c021224ca7e6218ea9862eb8a Mon Sep 17 00:00:00 2001 From: iris Date: Thu, 15 Oct 2026 11:53:19 +0200 Subject: [PATCH] Guard the creation roll-back by revision A store answers the revision it wrote and what it replaced, and restoring compares against that revision, so a write that landed after the creation's own is kept instead of overwritten. A write that changed nothing rolls back nothing, and a key that held no value is removed rather than set to an empty one. A test fails the second store after the first and checks the earlier caps come back. --- swarm-controller/src/agent_limits.rs | 93 +++++++++++++++++++++------- swarm-controller/src/inference.rs | 82 +++++++++++++++++------- swarm-controller/src/main.rs | 73 +++++++++++++++++++--- swarm-controller/src/schedules.rs | 48 ++++++++++++-- 4 files changed, 237 insertions(+), 59 deletions(-) diff --git a/swarm-controller/src/agent_limits.rs b/swarm-controller/src/agent_limits.rs index 06563bff..9d7005a9 100644 --- a/swarm-controller/src/agent_limits.rs +++ b/swarm-controller/src/agent_limits.rs @@ -19,7 +19,7 @@ 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}; +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 @@ -257,8 +257,9 @@ pub(crate) async fn put_agent_limits( Ok(Json(value.into())) } -/// What [`store`] replaced: the caps that were stored, or `None` for no key. -pub(crate) struct Replaced(Option); +/// What [`store`] replaced, or `None` when it wrote nothing because the caps +/// were already stored. +pub(crate) type Replaced = Option>; /// Replace `agent`'s stored caps with `value`. pub(crate) async fn store( @@ -271,14 +272,26 @@ pub(crate) async fn store( let (current, revision, undecodable) = read_or_default::(&store, agent) .await .map_err(|e| queue_problem(&e))?; - let replaced = Replaced(revision.map(|_| current.clone())); - if (current == *value && !undecodable) - || write(&store, agent, value, revision) - .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(replaced); + return Ok(Some(crate::schedules::Replaced::new( + held.then_some(current), + written, + ))); } } Err(queue_problem(&anyhow::anyhow!( @@ -286,21 +299,23 @@ pub(crate) async fn store( ))) } -/// Put back what [`store`] replaced. Best effort: a failure is logged, since -/// the caller is already on its way to an error of its own. +/// 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 put_back = match replaced.0 { - Some(previous) => store(client, agent, &previous).await.map(|_| ()), - None => match open_or_create(client).await { - Ok(store) => store - .delete(agent) - .await - .map_err(|e| queue_problem(&anyhow::Error::new(e))), - Err(e) => Err(e), - }, + 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:?}")), }; - if let Err(e) = put_back { - tracing::warn!(%agent, error = ?e, "could not put back the agent's limits"); + 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"); + } } } @@ -786,4 +801,38 @@ pub(crate) mod tests { 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::(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)); + } } diff --git a/swarm-controller/src/inference.rs b/swarm-controller/src/inference.rs index 99f21fff..73cab8b1 100644 --- a/swarm-controller/src/inference.rs +++ b/swarm-controller/src/inference.rs @@ -26,7 +26,7 @@ use serde::{Deserialize, Serialize}; use swarm_queue_client::agent_inference::{self, AgentInference, Effective, Effort, Source}; use utoipa::ToSchema; -use crate::schedules::{live, queue, read_or_default, refused, write}; +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 @@ -317,9 +317,9 @@ pub(crate) async fn put_agent_inference( Ok(Json(value.into())) } -/// What [`store`] replaced: the settings that were stored, or `None` for no -/// key. -pub(crate) struct Replaced(Option); +/// What [`store`] replaced, or `None` when it wrote nothing because the +/// settings were already stored. +pub(crate) type Replaced = Option>; /// Replace `agent`'s stored inference settings with `value`. pub(crate) async fn store( @@ -332,14 +332,26 @@ pub(crate) async fn store( let (current, revision, undecodable) = read_or_default::(&store, agent) .await .map_err(|e| queue_problem(&e))?; - let replaced = Replaced(revision.map(|_| current.clone())); - if (current == *value && !undecodable) - || write(&store, agent, value, revision) - .await - .map_err(|e| queue_problem(&e))? + if current == *value && !undecodable { + tracing::info!(%agent, ?value, "set inference settings"); + 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 inference settings"); - return Ok(replaced); + return Ok(Some(crate::schedules::Replaced::new( + held.then_some(current), + written, + ))); } } Err(queue_problem(&anyhow::anyhow!( @@ -348,21 +360,23 @@ pub(crate) async fn store( ))) } -/// Put back what [`store`] replaced. Best effort: a failure is logged, since -/// the caller is already on its way to an error of its own. +/// 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 put_back = match replaced.0 { - Some(previous) => store(client, agent, &previous).await.map(|_| ()), - None => match open(client).await { - Ok(store) => store - .delete(agent) - .await - .map_err(|e| queue_problem(&anyhow::Error::new(e))), - Err(e) => Err(e), - }, + let Some(replaced) = replaced else { return }; + let undone = match open(client).await { + Ok(store) => replaced.undo(&store, agent).await, + Err(e) => Err(anyhow::anyhow!("{e:?}")), }; - if let Err(e) = put_back { - tracing::warn!(%agent, error = ?e, "could not put back the agent's inference settings"); + match undone { + Ok(true) => {} + Ok(false) => { + tracing::info!(%agent, "left the agent's inference settings: another write landed since"); + } + Err(e) => { + tracing::warn!(%agent, error = %format!("{e:#}"), "could not put back the agent's inference settings"); + } } } @@ -638,4 +652,26 @@ mod tests { let entry = bucket.get("atlas").await.expect("read").expect("a value"); assert_eq!(&entry[..], br#"{"model":"opus"}"#); } + + /// A write that landed after the one being undone is newer, so `restore` + /// leaves it. + #[tokio::test] + async fn restore_leaves_a_write_that_landed_since() { + let nats = crate::test_nats::TestNats::start(); + let client = nats.client().await; + let bucket = agent_inference::open_or_create(&client) + .await + .expect("bucket"); + let value = |json: &str| settings(json).unwrap().into_stored().unwrap(); + + let made = store(&client, "atlas", &value(r#"{"model":"opus"}"#)) + .await + .expect("stored"); + store(&client, "atlas", &value(r#"{"model":"haiku"}"#)) + .await + .expect("a PUT"); + restore(&client, "atlas", made).await; + let entry = bucket.get("atlas").await.expect("read").expect("a value"); + assert_eq!(&entry[..], br#"{"model":"haiku"}"#); + } } diff --git a/swarm-controller/src/main.rs b/swarm-controller/src/main.rs index 9c3502a9..52952535 100644 --- a/swarm-controller/src/main.rs +++ b/swarm-controller/src/main.rs @@ -1946,8 +1946,8 @@ async fn create_agent( /// creation is refused after the write. #[derive(Default)] struct StoredSettings { - limits: Option, - inference: Option, + limits: agent_limits::Replaced, + inference: inference::Replaced, } impl StoredSettings { @@ -1961,11 +1961,11 @@ impl StoredSettings { ) -> Result { let mut stored = Self::default(); if let Some(limits) = limits { - stored.limits = Some(agent_limits::store(client, agent, limits).await?); + stored.limits = agent_limits::store(client, agent, limits).await?; } if let Some(inference) = inference { match inference::store(client, agent, inference).await { - Ok(replaced) => stored.inference = Some(replaced), + Ok(replaced) => stored.inference = replaced, Err(e) => { stored.roll_back(client, agent).await; return Err(e); @@ -1976,12 +1976,8 @@ impl StoredSettings { } async fn roll_back(self, client: &async_nats::Client, agent: &str) { - if let Some(replaced) = self.limits { - agent_limits::restore(client, agent, replaced).await; - } - if let Some(replaced) = self.inference { - inference::restore(client, agent, replaced).await; - } + agent_limits::restore(client, agent, self.limits).await; + inference::restore(client, agent, self.inference).await; } } @@ -3913,6 +3909,63 @@ mod tests { ); } + /// When the inference store fails after the caps are stored, the caps go + /// back to what they were: the earlier value, or no key. Without the + /// roll-back this leaves a cap behind for an agent that was never created. + /// The inference bucket is made with a one-byte value limit, so any write + /// to it fails. + #[tokio::test] + async fn a_failed_second_store_puts_the_first_back() { + let nats = crate::test_nats::TestNats::start(); + let client = nats.client().await; + async_nats::jetstream::new(client.clone()) + .create_key_value(async_nats::jetstream::kv::Config { + bucket: swarm_queue_client::agent_inference::BUCKET.to_owned(), + history: 1, + max_value_size: 1, + ..Default::default() + }) + .await + .expect("a bucket that refuses any value"); + let caps = |json: &str| { + serde_json::from_str::(json) + .expect("caps") + .into_stored() + .expect("valid") + }; + let inference = + serde_json::from_str::(r#"{"model":"opus"}"#) + .expect("settings") + .into_stored() + .expect("valid"); + crate::agent_limits::store(&client, "atlas", &caps(r#"{"cpu_quota":"100%"}"#)) + .await + .expect("an earlier cap"); + + for agent in ["atlas", "iris"] { + let stored = super::StoredSettings::store( + &client, + agent, + Some(&caps(r#"{"cpu_quota":"400%"}"#)), + Some(&inference), + ) + .await; + assert!(stored.is_err(), "{agent}: the inference write must fail"); + } + assert_eq!( + crate::agent_limits::tests::published(&client, "atlas") + .await + .as_deref(), + Some(r#"{"cpu_quota":"100%"}"#), + "the earlier cap comes back" + ); + assert_eq!( + crate::agent_limits::tests::published(&client, "iris").await, + None, + "a cap stored for a name with none is removed" + ); + } + /// A valid cap beside an invalid model stores neither and queues nothing, /// on a server that could have stored the cap. #[tokio::test] diff --git a/swarm-controller/src/schedules.rs b/swarm-controller/src/schedules.rs index 4c6c0084..fa2c836a 100644 --- a/swarm-controller/src/schedules.rs +++ b/swarm-controller/src/schedules.rs @@ -539,21 +539,61 @@ pub(crate) async fn write( value: &T, revision: Option, ) -> anyhow::Result { + Ok(write_revision(store, key, value, revision).await?.is_some()) +} + +/// As [`write`], answering the revision written, or `None` when another write +/// landed first. +pub(crate) async fn write_revision( + store: &Store, + key: &str, + value: &T, + revision: Option, +) -> anyhow::Result> { let encoded = serde_json::to_vec(value)?; match revision { Some(revision) => match store.update(key, encoded.into(), revision).await { - Ok(_) => Ok(true), - Err(e) if matches!(e.kind(), UpdateErrorKind::WrongLastRevision) => Ok(false), + Ok(written) => Ok(Some(written)), + Err(e) if matches!(e.kind(), UpdateErrorKind::WrongLastRevision) => Ok(None), Err(e) => Err(e).with_context(|| format!("writing {key}")), }, None => match store.create(key, encoded.into()).await { - Ok(_) => Ok(true), - Err(e) if matches!(e.kind(), CreateErrorKind::AlreadyExists) => Ok(false), + Ok(written) => Ok(Some(written)), + Err(e) if matches!(e.kind(), CreateErrorKind::AlreadyExists) => Ok(None), Err(e) => Err(e).with_context(|| format!("writing {key}")), }, } } +/// A write to a key, with what it replaced, so [`Replaced::undo`] can reverse +/// it. +pub(crate) struct Replaced { + /// The value the key held, or `None` when it held none. + previous: Option, + /// The revision the write made. + written: u64, +} + +impl Replaced { + pub(crate) fn new(previous: Option, written: u64) -> Self { + Self { previous, written } + } + + /// Put `previous` back, or remove the key if it held none, against the + /// revision the write made. `Ok(false)` when another write has landed + /// since: that one is newer than what is being undone and stays. + pub(crate) async fn undo(self, store: &Store, key: &str) -> anyhow::Result { + match self.previous { + Some(previous) => write(store, key, &previous, Some(self.written)).await, + None => match store.delete_expect_revision(key, Some(self.written)).await { + Ok(()) => Ok(true), + Err(e) if matches!(e.kind(), UpdateErrorKind::WrongLastRevision) => Ok(false), + Err(e) => Err(e).with_context(|| format!("removing {key}")), + }, + } + } +} + /// Apply `change` to the definitions and write them back, re-reading and /// re-applying when another write landed in between. async fn commit(