hyperhive/hive-agent/src/swarm_agent_icon.rs

//! Publishing this agent's icon into the swarm's `agent-icons` bucket.
//!
//! The bytes are the file the harness's own `GET /icon` serves
//! ([`ICON_PATH`]), unmodified. The key is the agent name, so the swarm can
//! show the icon while the agent is stopped or after it has moved hives.
//!
//! **Once per harness start.** The icon comes from the agent's NixOS config,
//! and a config change is applied by stopping the container, switching it and
//! starting it again, so a new icon always arrives with a new harness process.
//! An agent with no icon deletes its key, so an icon removed from the config
//! also disappears from the swarm.
//!
//! The agent's own queue credential is granted exactly this agent's key.
//!
//! The write is a `JetStream` publish straight to the key's subject, the same
//! message `kv::Store::put` sends, without opening the bucket first. That
//! keeps the grant to the one subject. swarm-controller creates the bucket,
//! so a write can arrive before it exists; that write and any other failed
//! one are retried with backoff until one is acked.

use std::path::Path;
use std::time::Duration;

use async_nats::jetstream::context::PublishErrorKind;
use backon::{ExponentialBuilder, Retryable};

/// Where the agent's NixOS config puts the icon
/// (`services.hyperhive.agent.icon`), and the file `GET /icon` serves. Absent
/// when none is configured.
const ICON_PATH: &str = "/etc/hyperhive/icon.svg";

/// First retry delay after a failed write, doubled per failure up to
/// [`MAX_RETRY`], with no cap on the number of tries.
const FIRST_RETRY: Duration = Duration::from_secs(5);

/// Longest wait between retries.
const MAX_RETRY: Duration = Duration::from_mins(5);

/// The KV header and value that mark an entry deleted. `kv::Store::get`
/// answers `None` for an entry carrying them.
const KV_OPERATION: &str = "KV-Operation";
const KV_OPERATION_DELETE: &str = "DEL";

/// What to write under this agent's key.
#[derive(Debug, PartialEq, Eq)]
enum Icon {
    /// The SVG's bytes.
    Set(Vec<u8>),
    /// No icon is configured, so the key is deleted.
    Unset,
}

/// Read the icon at `path`. A missing file is [`Icon::Unset`]; any other read
/// error is returned, and the key is left as it is.
fn read_icon(path: &Path) -> std::io::Result<Icon> {
    match std::fs::read(path) {
        Ok(bytes) => Ok(Icon::Set(bytes)),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Icon::Unset),
        Err(e) => Err(e),
    }
}

/// Start the publish task, if this agent has a label.
pub fn spawn() {
    let agent = crate::identity::label();
    if agent.is_empty() {
        tracing::warn!("this agent has no label; not publishing its icon upward");
        return;
    }
    tokio::spawn(run(agent));
}

async fn run(agent: String) {
    let client = crate::swarm_queue::client().await;
    let subject = swarm_queue_client::agent_icon::subject(&agent);
    let icon = match read_icon(Path::new(ICON_PATH)) {
        Ok(icon) => icon,
        Err(e) => {
            tracing::warn!(path = ICON_PATH, error = %e, "reading the agent icon failed; not publishing it");
            return;
        }
    };
    let js = async_nats::jetstream::new(client.clone());
    let written = (|| write(&client, &js, &subject, &icon))
        .retry(
            ExponentialBuilder::default()
                .with_min_delay(FIRST_RETRY)
                .with_max_delay(MAX_RETRY)
                .without_max_times(),
        )
        .when(|failure| matches!(failure, Failure::Retry(_)))
        .notify(|failure, delay| {
            if let Failure::Retry(e) = failure {
                tracing::warn!(error = e, retry_in = ?delay, "publishing the agent icon failed");
            }
        })
        .await;
    match written {
        Ok(()) => tracing::info!(
            subject,
            set = matches!(icon, Icon::Set(_)),
            "published this agent's icon"
        ),
        Err(Failure::TooLarge(e)) => {
            tracing::warn!(error = %e, "the agent icon is over the queue's payload limit; not publishing it");
        }
        // Only a `TooLarge` ends the retries.
        Err(Failure::Retry(e)) => {
            tracing::error!(error = e, "gave up publishing the agent icon");
        }
    }
}

/// Why a write did not land.
enum Failure {
    /// The payload can never fit; retrying cannot help.
    TooLarge(async_nats::jetstream::context::PublishError),
    /// Anything else, rendered for the log.
    Retry(String),
}

/// One acked write of `icon` to `subject`.
async fn write(
    client: &async_nats::Client,
    js: &async_nats::jetstream::Context,
    subject: &str,
    icon: &Icon,
) -> Result<(), Failure> {
    // An unconnected client does not fail a publish, it buffers it, and the
    // payload limit it checks against is the library default until the
    // server has announced its own.
    swarm_queue_client::ensure_connected(client)
        .map_err(|e| Failure::Retry(swarm_queue_client::chain(&e)))?;
    let sent = match icon {
        Icon::Set(bytes) => js.publish(subject.to_owned(), bytes.clone().into()).await,
        Icon::Unset => {
            let mut headers = async_nats::HeaderMap::new();
            headers.insert(KV_OPERATION, KV_OPERATION_DELETE);
            js.publish_with_headers(subject.to_owned(), headers, Vec::new().into())
                .await
        }
    };
    let ack = sent.map_err(|e| match e.kind() {
        PublishErrorKind::MaxPayloadExceeded => Failure::TooLarge(e),
        _ => Failure::Retry(e.to_string()),
    })?;
    // `StreamNotFound` here is the bucket not existing yet.
    ack.await.map_err(|e| Failure::Retry(e.to_string()))?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::{Icon, read_icon};

    #[test]
    fn a_configured_icon_is_published_byte_for_byte() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("icon.svg");
        std::fs::write(&path, b"<svg/>\n").expect("write");
        assert_eq!(
            read_icon(&path).expect("read"),
            Icon::Set(b"<svg/>\n".to_vec())
        );
    }

    /// No file is the unconfigured state, which clears the key rather than
    /// leaving a removed icon in the swarm.
    #[test]
    fn no_icon_file_clears_the_key() {
        let dir = tempfile::tempdir().expect("tempdir");
        assert_eq!(
            read_icon(&dir.path().join("icon.svg")).expect("read"),
            Icon::Unset
        );
    }

    /// Any other read failure is an error, so a transient one does not delete
    /// an icon the agent still has.
    #[test]
    fn an_unreadable_icon_is_an_error_not_an_unset_icon() {
        let dir = tempfile::tempdir().expect("tempdir");
        assert!(read_icon(dir.path()).is_err());
    }
}