vinzenz/hyperhive
vinzenz/hyperhive
Watch
0
0
Fork
You've already forked hyperhive
 0
Code Activity

Guard the creation roll-back by revision

Browse source 
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.
This commit is contained in:
iris 2026-10-15 11:53:19 +02:00 • committed by mara
commit 4ad9579e97
4 changed files with 237 additions and 59 deletions
Download patch file Download diff file Expand all files Collapse all files

93
swarm-controller/src/agent_limits.rs
View file

@ -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<AgentCaps>);
/// 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(
@ -271,14 +272,26 @@ pub(crate) async fn store(
let (current, revision, undecodable) = read_or_default::<AgentCaps>(&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::<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));
}
}

82
swarm-controller/src/inference.rs
View file

@ -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<AgentInference>);
/// What [`store`] replaced, or `None` when it wrote nothing because the
/// settings were already stored.
pub(crate) type Replaced = Option<crate::schedules::Replaced<AgentInference>>;
/// 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::<AgentInference>(&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"}"#);
}
}

73
swarm-controller/src/main.rs
View file

@ -1946,8 +1946,8 @@ async fn create_agent(
/// creation is refused after the write.
#[derive(Default)]
struct StoredSettings {
limits: Option<agent_limits::Replaced>,
inference: Option<inference::Replaced>,
limits: agent_limits::Replaced,
inference: inference::Replaced,
}
impl StoredSettings {
@ -1961,11 +1961,11 @@ impl StoredSettings {
) -> Result<Self, problem_details::ProblemDetails> {
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::<crate::agent_limits::Caps>(json)
.expect("caps")
.into_stored()
.expect("valid")
};
let inference =
serde_json::from_str::<crate::inference::InferenceSettings>(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]

48
swarm-controller/src/schedules.rs
View file

@ -539,21 +539,61 @@ pub(crate) async fn write<T: Serialize>(
value: &T,
revision: Option<u64>,
) -> anyhow::Result<bool> {
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<T: Serialize>(
store: &Store,
key: &str,
value: &T,
revision: Option<u64>,
) -> anyhow::Result<Option<u64>> {
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<T> {
/// The value the key held, or `None` when it held none.
previous: Option<T>,
/// The revision the write made.
written: u64,
}
impl<T: Serialize> Replaced<T> {
pub(crate) fn new(previous: Option<T>, 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<bool> {
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<T>(