hyperhive/swarm-controller/src/wanted.rs

//! Writes the agent set this swarm declares for each hive.
//!
//! The mirror of [`crate::status`]: that module reads what hives report,
//! this one writes what they are told, and both address the same queue.
//! Both hold a NATS client rather than a bucket handle, so a controller
//! that starts before a bucket exists picks it up without a restart.
//!
//! Where the mirror stops is the handle itself: `status` resolves one on
//! first use and caches it, while there is one wanted-state bucket **per
//! hive**, so no single handle serves them and `store` resolves per call.
//!
//! The bucket is the record. Nothing here keeps a second copy of the
//! declaration to reconcile against, because the current value can be read
//! back from the queue whenever it is needed.
//!
//! [`WantedWriter::set`] writes unconditionally, for every caller alike — no
//! per-caller `Intent`, no terminal-state refusal. A caller that needs to
//! decide *whether* to write at all (the agent-creation node's own
//! idempotency check) reads [`WantedWriter::view`] first and makes that call
//! itself.

use anyhow::{Context, Result};
use async_nats::jetstream::kv::{CreateErrorKind, UpdateErrorKind};
use swarm_queue_client::wanted::{AgentState, AgentWanted, HiveWanted};

/// The three outcomes of one write attempt, which the two KV verbs report
/// through separate error types.
enum Wrote {
    Ok,
    /// Another writer won the race; re-read and re-apply.
    LostRace,
    Failed(anyhow::Error),
}

/// How many times a losing writer re-reads and re-applies before giving up.
///
/// A conflict means another writer changed a *different* agent between this
/// one's read and its write, so a retry re-reads and re-applies onto the
/// winner. Bounded because an unbounded loop against a hot key is a spin,
/// and a caller that gets an error can ask again with fresh intent.
const MAX_ATTEMPTS: usize = 5;

/// Apply one agent's declared state to a hive's current declaration.
///
/// `current` is the hive's whole agent map — an undecodable value is an
/// **error**, not treated as absent, so a bad read never discards every
/// other agent's declaration under a fresh-start fallback.
fn apply(current: Option<&[u8]>, agent: &str, state: AgentState) -> Result<(HiveWanted, Vec<u8>)> {
    let mut declaration = match current {
        Some(raw) => serde_json::from_slice::<HiveWanted>(raw)
            .context("the hive's current declaration is not decodable")?,
        None => HiveWanted::default(),
    };
    declaration
        .agents
        .insert(agent.to_owned(), AgentWanted { state });
    let encoded = serde_json::to_vec(&declaration).context("encoding the new declaration")?;
    Ok((declaration, encoded))
}

/// Writes the wanted-state bucket, and reads it back.
pub struct WantedWriter {
    client: async_nats::Client,
}

impl WantedWriter {
    #[must_use]
    pub fn new(client: async_nats::Client) -> Self {
        Self { client }
    }

    /// The declaration currently published for `hive`, or `None`.
    pub async fn view(&self, hive: &str) -> Result<Option<HiveWanted>> {
        // An unconnected client does not fail a JetStream request, it hangs
        // on it — see `swarm_queue_client::ensure_connected`.
        swarm_queue_client::ensure_connected(&self.client)?;
        // One bucket per hive, resolved per call, created on first use — see
        // `swarm_queue_client::wanted` for why (only the controller creates
        // these; a hive opens its own read-only).
        let store = swarm_queue_client::wanted::open_or_create(&self.client, hive).await?;
        let Some(entry) = store
            .entry(hive)
            .await
            .with_context(|| format!("reading the declaration for {hive}"))?
        else {
            return Ok(None);
        };
        serde_json::from_slice(&entry.value)
            .map(Some)
            .with_context(|| format!("the declaration for {hive} is not decodable"))
    }

    /// Declare `agent` on `hive` to be in `state`, unconditionally (see the
    /// module doc), and return the whole declaration as published.
    ///
    /// Read-modify-write against the entry's revision, not a plain `put`:
    /// the value is the hive's whole agent map, so a blind write would drop
    /// a concurrent change to a different agent.
    pub async fn set(&self, hive: &str, agent: &str, state: AgentState) -> Result<HiveWanted> {
        swarm_queue_client::ensure_connected(&self.client)?;
        let store = swarm_queue_client::wanted::open_or_create(&self.client, hive).await?;

        for _ in 0..MAX_ATTEMPTS {
            let entry = store
                .entry(hive)
                .await
                .with_context(|| format!("reading the declaration for {hive}"))?;
            let revision = entry.as_ref().map(|e| e.revision);
            let (declaration, encoded) =
                apply(entry.as_ref().map(|e| e.value.as_ref()), agent, state)?;

            // A redeclare of the identical state re-encodes the declaration
            // unchanged (a caller that already sees this state and asks
            // again is not doing anything) — skip the network round-trip for
            // it entirely rather than writing identical bytes back.
            // `HiveWanted::agents` is a `BTreeMap`, so this byte comparison
            // is deterministic, not vulnerable to a map re-ordering itself
            // between reads.
            if entry.as_ref().map(|e| e.value.as_ref()) == Some(encoded.as_slice()) {
                tracing::debug!(hive, agent, "declaration already current, no write needed");
                return Ok(declaration);
            }

            // `update` and `create` have separate error types, and only one
            // variant of each means "someone else got there first". Every
            // other failure returns immediately: retrying a disconnect or a
            // permission error would spin the loop and then report a
            // conflict, blaming a concurrent writer that never existed.
            let written = match revision {
                Some(revision) => match store.update(hive, encoded.into(), revision).await {
                    Ok(_) => Wrote::Ok,
                    Err(e) if matches!(e.kind(), UpdateErrorKind::WrongLastRevision) => {
                        Wrote::LostRace
                    }
                    Err(e) => Wrote::Failed(anyhow::Error::new(e)),
                },
                None => match store.create(hive, encoded.into()).await {
                    Ok(_) => Wrote::Ok,
                    Err(e) if matches!(e.kind(), CreateErrorKind::AlreadyExists) => Wrote::LostRace,
                    Err(e) => Wrote::Failed(anyhow::Error::new(e)),
                },
            };
            match written {
                Wrote::Ok => {
                    tracing::info!(hive, agent, ?state, "declared agent state");
                    return Ok(declaration);
                }
                // Re-read and re-apply onto the winner's value, not over it.
                Wrote::LostRace => {
                    tracing::debug!(hive, agent, "declaration write lost a race, retrying");
                }
                Wrote::Failed(e) => {
                    return Err(e).with_context(|| format!("declaring {agent} on {hive}"));
                }
            }
        }
        anyhow::bail!("gave up declaring {agent} on {hive} after {MAX_ATTEMPTS} conflicting writes")
    }
}

#[cfg(test)]
mod tests {
    use super::apply;
    use swarm_queue_client::wanted::AgentState;

    #[test]
    fn declaring_one_agent_preserves_every_other() {
        let current = br#"{"agents":{"iris":{"state":"up"},"argus":{"state":"offline"}}}"#;
        let (declaration, _) = apply(Some(current), "atlas", AgentState::Up).unwrap();
        assert_eq!(declaration.agents.len(), 3);
        assert_eq!(declaration.agents["iris"].state, AgentState::Up);
        assert_eq!(declaration.agents["argus"].state, AgentState::Offline);
        assert_eq!(declaration.agents["atlas"].state, AgentState::Up);
    }

    #[test]
    fn redeclaring_an_agent_replaces_only_its_own_state() {
        let current = br#"{"agents":{"iris":{"state":"up"},"atlas":{"state":"up"}}}"#;
        let (declaration, _) = apply(Some(current), "atlas", AgentState::Offline).unwrap();
        assert_eq!(declaration.agents.len(), 2);
        assert_eq!(declaration.agents["iris"].state, AgentState::Up);
        assert_eq!(declaration.agents["atlas"].state, AgentState::Offline);
    }

    #[test]
    fn a_hive_with_no_declaration_yet_gets_a_one_agent_one() {
        let (declaration, _) = apply(None, "atlas", AgentState::Up).unwrap();
        assert_eq!(declaration.agents.len(), 1);
        assert_eq!(declaration.agents["atlas"].state, AgentState::Up);
    }

    // The failure this function exists to prevent: a fresh-start fallback
    // here would publish a one-agent document over a hive's whole set.
    #[test]
    fn an_undecodable_declaration_is_an_error_not_a_fresh_start() {
        let err = apply(Some(b"{not json"), "atlas", AgentState::Up).unwrap_err();
        assert!(
            err.to_string().contains("not decodable"),
            "unexpected error: {err}"
        );
    }

    #[test]
    fn an_unknown_state_in_the_current_value_is_also_an_error() {
        let current = br#"{"agents":{"iris":{"state":"sideways"}}}"#;
        assert!(apply(Some(current), "atlas", AgentState::Up).is_err());
    }

    #[test]
    fn the_encoded_form_round_trips() {
        let (_, encoded) = apply(None, "atlas", AgentState::Offline).unwrap();
        let (again, _) = apply(Some(&encoded), "iris", AgentState::Up).unwrap();
        assert_eq!(again.agents["atlas"].state, AgentState::Offline);
        assert_eq!(again.agents["iris"].state, AgentState::Up);
    }

    /// No terminal-state refusal any more (mara: "do not refuse to recreate
    /// an agent") — `apply` writes whatever it is asked to, including moving
    /// a `Destroyed` agent to any other state. Whether that write should
    /// happen at all is now the caller's own decision (see the module doc).
    #[test]
    fn a_destroyed_agent_can_be_declared_into_any_other_state() {
        let current = br#"{"agents":{"atlas":{"state":"destroyed"}}}"#;
        for state in [AgentState::Up, AgentState::Offline, AgentState::Paused] {
            let (declaration, _) = apply(Some(current), "atlas", state).unwrap();
            assert_eq!(declaration.agents["atlas"].state, state);
        }
    }

    /// The write is per-agent — a destroyed agent on the same hive as a live
    /// one must not block declaring the live one, and declaring one must not
    /// disturb the other.
    #[test]
    fn a_destroyed_agent_does_not_block_declaring_a_different_one() {
        let current = br#"{"agents":{"atlas":{"state":"destroyed"},"iris":{"state":"up"}}}"#;
        let (declaration, _) = apply(Some(current), "iris", AgentState::Offline).unwrap();
        assert_eq!(declaration.agents["atlas"].state, AgentState::Destroyed);
        assert_eq!(declaration.agents["iris"].state, AgentState::Offline);
    }
}