hyperhive/swarm-controller/src/pause_hold.rs

//! Holds an agent the hive paused, so the swarm's declaration stops undoing it.
//!
//! An agent whose harness paused its own turn loop says so in its turn-state
//! header (`paused_by_hive`, see [`crate::agent_login::Reported`]). The hive
//! resumes any paused agent that the swarm declares `Up`, so this declares it
//! `Paused` instead, which the hive's pause marker already agrees with. Only an
//! operator declaring the agent `Up` again resumes it.
//!
//! Each pause is held once. A header re-sent after the operator's resume and
//! before the hive cleared the marker still reports it, and must not pause the
//! agent a second time. The agent is only declared `Paused` while it is
//! declared `Up`: an operator's `Offline` or `Destroyed` stays what it is.

use std::sync::Arc;
use std::time::Duration;

use anyhow::Result;

use crate::agent_login::Reported;
use crate::wanted::WantedWriter;

/// How often the reported pauses are looked at.
const POLL: Duration = Duration::from_secs(5);

/// Hold the pauses `login` learns of, for as long as this daemon runs.
pub(crate) fn spawn(login: Arc<Reported>, wanted: Arc<WantedWriter>, hives: Vec<String>) {
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(POLL).await;
            hold_pending(&login, &wanted, &hives).await;
        }
    });
}

/// Hold every reported pause not held yet. A pause that could not be held is
/// tried again on the next call.
pub(crate) async fn hold_pending(login: &Reported, wanted: &WantedWriter, hives: &[String]) {
    for agent in login.unheld_pauses() {
        match hold(wanted, hives, &agent).await {
            Ok(()) => login.mark_pause_held(&agent),
            Err(e) => tracing::warn!(
                %agent,
                error = %format!("{e:#}"),
                "holding a hive-paused agent failed; will retry"
            ),
        }
    }
}

/// Declare `agent` `Paused` on whichever hive declares it `Up`.
async fn hold(wanted: &WantedWriter, hives: &[String], agent: &str) -> Result<()> {
    for hive in hives {
        if wanted.pause_if_up(hive, agent).await? {
            tracing::info!(%hive, %agent, "the hive paused this agent; declared it paused");
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::hold_pending;
    use crate::agent_login::Reported;
    use crate::test_nats::TestNats;
    use crate::wanted::WantedWriter;
    use swarm_wire::wanted::AgentState;

    fn header(reason: Option<&str>) -> Vec<u8> {
        let mut header = serde_json::json!({ "agent_state": "paused", "needs_login": false });
        if let Some(reason) = reason {
            header["paused_by_hive"] = reason.into();
        }
        serde_json::to_vec(&header).unwrap()
    }

    async fn declared(wanted: &WantedWriter, agent: &str) -> Option<AgentState> {
        let declaration = wanted.view("pr1ma").await.unwrap()?;
        declaration.agents.get(agent).map(|a| a.state)
    }

    /// The whole fix: a hive-paused agent declared `Up` is declared `Paused`,
    /// the operator's `Up` after that sticks, and the next pause is held too.
    #[tokio::test]
    async fn a_hive_pause_replaces_up_and_an_operator_up_sticks() {
        let nats = TestNats::start();
        let wanted = WantedWriter::new(nats.client().await);
        let hives = vec!["pr1ma".to_owned()];
        let login = Reported::default();
        let subject = swarm_queue_client::agent_state_subject("atlas");
        wanted.set("pr1ma", "atlas", AgentState::Up).await.unwrap();
        wanted.set("pr1ma", "iris", AgentState::Up).await.unwrap();

        login.record(&subject, &header(Some("skipped get_loose_ends")));
        hold_pending(&login, &wanted, &hives).await;
        assert_eq!(declared(&wanted, "atlas").await, Some(AgentState::Paused));
        assert_eq!(declared(&wanted, "iris").await, Some(AgentState::Up));

        // The operator resumes; the hive has not cleared its marker yet and
        // re-sends the header.
        wanted.set("pr1ma", "atlas", AgentState::Up).await.unwrap();
        login.record(&subject, &header(Some("skipped get_loose_ends")));
        hold_pending(&login, &wanted, &hives).await;
        assert_eq!(declared(&wanted, "atlas").await, Some(AgentState::Up));

        // The hive clears the marker, then pauses the agent again.
        login.record(&subject, &header(None));
        login.record(&subject, &header(Some("skipped get_loose_ends")));
        hold_pending(&login, &wanted, &hives).await;
        assert_eq!(declared(&wanted, "atlas").await, Some(AgentState::Paused));
    }

    /// An operator's `Offline` is not overruled by the agent's pause.
    #[tokio::test]
    async fn only_an_agent_declared_up_is_paused() {
        let nats = TestNats::start();
        let wanted = WantedWriter::new(nats.client().await);
        let hives = vec!["pr1ma".to_owned()];
        let login = Reported::default();
        wanted
            .set("pr1ma", "atlas", AgentState::Offline)
            .await
            .unwrap();

        login.record(
            &swarm_queue_client::agent_state_subject("atlas"),
            &header(Some("skipped get_loose_ends")),
        );
        hold_pending(&login, &wanted, &hives).await;
        assert_eq!(declared(&wanted, "atlas").await, Some(AgentState::Offline));
        assert!(login.unheld_pauses().is_empty(), "nothing left to hold");
    }

    /// An agent no hive declares has nothing to override: its pause is held at
    /// once and nothing is written to any declaration.
    #[tokio::test]
    async fn an_undeclared_agent_is_held_at_once_and_nothing_is_declared() {
        let nats = TestNats::start();
        let wanted = WantedWriter::new(nats.client().await);
        let hives = vec!["pr1ma".to_owned()];
        let login = Reported::default();
        login.record(
            &swarm_queue_client::agent_state_subject("atlas"),
            &header(Some("skipped get_loose_ends")),
        );
        assert_eq!(login.unheld_pauses(), ["atlas"], "control: pending first");

        hold_pending(&login, &wanted, &hives).await;
        assert!(login.unheld_pauses().is_empty(), "held at once");
        assert_eq!(declared(&wanted, "atlas").await, None);
        assert!(wanted.view("pr1ma").await.unwrap().is_none());
    }
}