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

@@ -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));
+    }
 }