swarm-controller: KV-store + serve per-agent status (#3341 item 2)

swarm-controller/src/agent_status.rs

@@ -0,0 +1,304 @@
+//! Swarm-wide view of what each agent last said about itself.
+//!
+//! Sibling of [`crate::status`], not a mode of it — see
+//! `swarm_queue_client::agent_status`'s module doc for why the two buckets
+//! stay separate. Everything about *how* a snapshot is read (queue is the
+//! store, no cache, freshness derived at read time, rows come from the
+//! roster so an empty bucket cannot render as "every agent is fine") is
+//! identical to [`crate::status`]; this module differs only in whose
+//! roster it reads and how it recovers a hive name for a reporting agent.
+//!
+//! **The roster here is agent *names* only** ([`crate::auth::AuthBridge::list_agent_identities`]) —
+//! it does not say which hive an agent lives on. That is not a gap this
+//! module needs to close: a *reporting* agent's hive comes straight out of
+//! its bucket key (`swarm_queue_client::agent_status::split_key`), and a
+//! *never-reported* agent's hive is unknowable regardless of what any
+//! roster might claim, so [`AgentStatusRow::hive`] is `None` in exactly
+//! that one case rather than a second lookup this module would have to
+//! keep in sync with the bucket.
+
+use std::collections::BTreeMap;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+use anyhow::{Context, Result};
+use futures_util::TryStreamExt as _;
+use serde::{Deserialize, Serialize};
+use utoipa::ToSchema;
+
+use crate::status::Freshness;
+
+/// One row of the aggregate.
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, ToSchema)]
+pub struct AgentStatusRow {
+    pub name: String,
+    /// The hive this agent reported from, taken from its bucket key.
+    /// `None` when the agent has never reported — see the module doc for
+    /// why that is not answered from anywhere else.
+    pub hive: Option<String>,
+    pub freshness: Freshness,
+    /// When the snapshot was stored, unix seconds — the bucket's own
+    /// stamp, same rule as [`crate::status::HiveStatus::last_seen_unix`].
+    pub last_seen_unix: Option<i64>,
+    pub age_seconds: Option<u64>,
+    /// The agent's status text + running flag, unopened. `None` exactly
+    /// when `freshness` is `NeverReported`, or when the stored bytes
+    /// failed to decode as [`swarm_queue_client::agent_status::AgentStatus`]
+    /// (logged as a warning naming the agent).
+    pub snapshot: Option<serde_json::Value>,
+}
+
+/// A snapshot as retained by the bucket, with the hive it was published
+/// under already split out of the key.
+#[derive(Clone, Debug)]
+struct Offered {
+    hive: String,
+    received_at: SystemTime,
+    payload: Option<serde_json::Value>,
+}
+
+/// Reads the per-agent status aggregate out of the KV bucket. Same shape
+/// as [`crate::status::StatusReader`] — see that type for why the bucket
+/// is resolved lazily and cached, and why a resolution failure is not.
+pub struct AgentStatusReader {
+    client: async_nats::Client,
+    store: tokio::sync::OnceCell<async_nats::jetstream::kv::Store>,
+    stale_after: Duration,
+}
+
+impl AgentStatusReader {
+    #[must_use]
+    pub fn new(client: async_nats::Client, stale_after: Duration) -> Self {
+        Self {
+            client,
+            store: tokio::sync::OnceCell::new(),
+            stale_after,
+        }
+    }
+
+    async fn store(
+        &self,
+    ) -> std::result::Result<&async_nats::jetstream::kv::Store, swarm_queue_client::Error> {
+        self.store
+            .get_or_try_init(|| swarm_queue_client::agent_status::open_or_create(&self.client))
+            .await
+    }
+
+    /// The aggregate, rendered against `now`.
+    ///
+    /// Every roster agent produces a row whether or not it has ever
+    /// reported; a reporting agent outside the roster produces one too —
+    /// same two-sided shape as [`crate::status::StatusReader::view`].
+    pub async fn view(&self, roster: &[String], now: SystemTime) -> Result<Vec<AgentStatusRow>> {
+        swarm_queue_client::ensure_connected(&self.client)?;
+
+        let store = self.store().await?;
+
+        let mut keys = store
+            .keys()
+            .await
+            .context("listing agent-status bucket keys")?;
+        let mut entries: BTreeMap<String, Offered> = BTreeMap::new();
+        while let Some(key) = keys
+            .try_next()
+            .await
+            .context("reading the agent-status bucket's key list")?
+        {
+            let Some((hive, agent)) = swarm_queue_client::agent_status::split_key(&key) else {
+                tracing::warn!(%key, "agent-status bucket key doesn't split into hive/agent, skipping");
+                continue;
+            };
+            let Some(entry) = store
+                .entry(&key)
+                .await
+                .with_context(|| format!("reading agent-status entry {key}"))?
+            else {
+                continue;
+            };
+            let payload = match serde_json::from_slice::<
+                swarm_queue_client::agent_status::AgentStatus,
+            >(&entry.value)
+            {
+                Ok(status) => serde_json::to_value(status).ok(),
+                Err(e) => {
+                    tracing::warn!(
+                        %agent, %hive, error = %e,
+                        "agent-status snapshot didn't decode; reporting the timestamp without it"
+                    );
+                    None
+                }
+            };
+            entries.insert(
+                agent.to_owned(),
+                Offered {
+                    hive: hive.to_owned(),
+                    received_at: to_system_time(entry.created.unix_timestamp()),
+                    payload,
+                },
+            );
+        }
+
+        Ok(render(roster, &entries, now, self.stale_after))
+    }
+}
+
+fn to_system_time(unix_seconds: i64) -> SystemTime {
+    u64::try_from(unix_seconds).map_or(UNIX_EPOCH, |secs| UNIX_EPOCH + Duration::from_secs(secs))
+}
+
+/// Turn a roster plus whatever the bucket held into the rendered rows. Pure
+/// for the same testing reason as [`crate::status::render`].
+fn render(
+    roster: &[String],
+    entries: &BTreeMap<String, Offered>,
+    now: SystemTime,
+    stale_after: Duration,
+) -> Vec<AgentStatusRow> {
+    let mut rows: Vec<AgentStatusRow> = roster
+        .iter()
+        .map(|name| match entries.get(name) {
+            Some(offered) => row(name.clone(), Freshness::Fresh, offered, now, stale_after),
+            None => AgentStatusRow {
+                name: name.clone(),
+                hive: None,
+                freshness: Freshness::NeverReported,
+                last_seen_unix: None,
+                age_seconds: None,
+                snapshot: None,
+            },
+        })
+        .collect();
+
+    rows.extend(
+        entries
+            .iter()
+            .filter(|(name, _)| !roster.iter().any(|r| r == *name))
+            .map(|(name, offered)| {
+                row(name.clone(), Freshness::Unknown, offered, now, stale_after)
+            }),
+    );
+    rows
+}
+
+/// `default_freshness` is `Unknown` for an agent outside the roster and
+/// `Fresh` for one inside it — `row` then downgrades to `Stale` on age,
+/// same two-step [`crate::status::row`] uses.
+fn row(
+    name: String,
+    default_freshness: Freshness,
+    offered: &Offered,
+    now: SystemTime,
+    stale_after: Duration,
+) -> AgentStatusRow {
+    let age = now.duration_since(offered.received_at).ok();
+    let freshness = match (default_freshness, age) {
+        (Freshness::Unknown, _) => Freshness::Unknown,
+        (_, Some(age)) if age > stale_after => Freshness::Stale,
+        _ => Freshness::Fresh,
+    };
+    AgentStatusRow {
+        name,
+        hive: Some(offered.hive.clone()),
+        freshness,
+        last_seen_unix: offered
+            .received_at
+            .duration_since(UNIX_EPOCH)
+            .ok()
+            .and_then(|d| i64::try_from(d.as_secs()).ok()),
+        age_seconds: age.map(|age| age.as_secs()),
+        snapshot: offered.payload.clone(),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{Freshness, Offered, render};
+    use std::collections::BTreeMap;
+    use std::time::{Duration, SystemTime};
+
+    fn roster() -> Vec<String> {
+        vec!["iris".to_owned(), "atlas".to_owned()]
+    }
+
+    fn entries(rows: &[(&str, &str, SystemTime)]) -> BTreeMap<String, Offered> {
+        rows.iter()
+            .map(|(agent, hive, at)| {
+                (
+                    (*agent).to_owned(),
+                    Offered {
+                        hive: (*hive).to_owned(),
+                        received_at: *at,
+                        payload: Some(serde_json::json!({ "running": true })),
+                    },
+                )
+            })
+            .collect()
+    }
+
+    fn t0() -> SystemTime {
+        SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000)
+    }
+
+    #[test]
+    fn an_empty_bucket_renders_every_roster_agent_as_never_reported() {
+        let rows = render(&roster(), &BTreeMap::new(), t0(), Duration::from_mins(2));
+        assert_eq!(rows.len(), 2);
+        assert!(rows.iter().all(|r| r.freshness == Freshness::NeverReported));
+        assert!(rows.iter().all(|r| r.hive.is_none()));
+    }
+
+    #[test]
+    fn a_reporting_agent_gets_its_hive_from_the_bucket_key() {
+        let rows = render(
+            &roster(),
+            &entries(&[("iris", "pr1ma", t0())]),
+            t0(),
+            Duration::from_mins(2),
+        );
+        let iris = rows.iter().find(|r| r.name == "iris").expect("present");
+        assert_eq!(iris.hive.as_deref(), Some("pr1ma"));
+        assert_eq!(iris.freshness, Freshness::Fresh);
+    }
+
+    #[test]
+    fn a_stale_report_still_carries_its_hive() {
+        let rows = render(
+            &roster(),
+            &entries(&[("iris", "pr1ma", t0())]),
+            t0() + Duration::from_mins(10),
+            Duration::from_mins(2),
+        );
+        let iris = rows.iter().find(|r| r.name == "iris").expect("present");
+        assert_eq!(iris.freshness, Freshness::Stale);
+        assert_eq!(iris.hive.as_deref(), Some("pr1ma"), "stale is not absent");
+    }
+
+    #[test]
+    fn an_agent_outside_the_roster_is_surfaced_as_unknown() {
+        let rows = render(
+            &roster(),
+            &entries(&[("ghost", "pr1ma", t0())]),
+            t0(),
+            Duration::from_mins(2),
+        );
+        assert_eq!(rows.len(), 3, "two roster agents plus the stranger");
+        let ghost = rows.iter().find(|r| r.name == "ghost").expect("present");
+        assert_eq!(ghost.freshness, Freshness::Unknown);
+        assert_eq!(
+            ghost.hive.as_deref(),
+            Some("pr1ma"),
+            "still known from its key"
+        );
+    }
+
+    #[test]
+    fn one_agent_reporting_does_not_vouch_for_a_silent_one() {
+        let rows = render(
+            &roster(),
+            &entries(&[("iris", "pr1ma", t0())]),
+            t0(),
+            Duration::from_mins(2),
+        );
+        let atlas = rows.iter().find(|r| r.name == "atlas").expect("present");
+        assert_eq!(atlas.freshness, Freshness::NeverReported);
+    }
+}

swarm-controller/src/main.rs

@@ -40,6 +40,7 @@ use swarm_authelia_bridge_sock::BridgeResponse;
 use utoipa::{OpenApi, ToSchema};
 use utoipa_axum::{router::OpenApiRouter, routes};
 
+mod agent_status;
 mod auth;
 mod config_pr;
 mod forge;
@@ -436,6 +437,10 @@ struct AppState {
     /// up, so there is nowhere to publish a declaration to. Shares that
     /// reader's connection rather than opening a second one.
     wanted: Option<Arc<wanted::WantedWriter>>,
+    /// Per-agent status, sharing `status`'s connection — same "one queue
+    /// connection, several consumers" rationale as `wanted` above. `None`
+    /// in exactly the state `status` is.
+    agent_status: Option<Arc<agent_status::AgentStatusReader>>,
     /// The swarm-level job graph, wrapped in its
     /// [`hive_jobq::scheduler::Scheduler`] now that something drives it
     /// (`spawn_jobq_worker`) — the graph alone was enough for the
@@ -726,6 +731,22 @@ fn wanted_writer(status: Option<&Arc<status::StatusReader>>) -> Option<Arc<wante
     status.map(|s| Arc::new(wanted::WantedWriter::new(s.queue_client())))
 }
 
+/// The per-agent status reader, sharing the status reader's connection and
+/// staleness threshold — same rationale as [`wanted_writer`], and the same
+/// threshold on purpose: both buckets are fed by the same publisher cadence
+/// (`hive-c0re`'s `PUBLISH_INTERVAL`), so two separately-configured
+/// thresholds would just be two ways to get out of sync with one cadence.
+fn agent_status_reader(
+    status: Option<&Arc<status::StatusReader>>,
+) -> Option<Arc<agent_status::AgentStatusReader>> {
+    status.map(|s| {
+        Arc::new(agent_status::AgentStatusReader::new(
+            s.queue_client(),
+            status::StatusReader::stale_after_from_env(),
+        ))
+    })
+}
+
 /// Both handlers below take the same hive name and reject it the same way.
 ///
 /// Reports the status and the detail rather than a rendered
@@ -868,6 +889,49 @@ async fn get_hives_status(
     }
 }
 
+/// What each agent last said about itself, read from the swarm queue at
+/// request time.
+///
+/// Every agent in the roster gets a row whether or not it has ever
+/// reported — see the `agent_status` module for why, and for why its hive
+/// comes from its own bucket key rather than a second lookup.
+#[utoipa::path(
+    get,
+    path = "/api/agents/status",
+    responses(
+        (status = 200, description = "a row per agent, freshness derived now", body = Vec<agent_status::AgentStatusRow>),
+        (status = 503, description = "no swarm queue is configured here, no identity bridge is configured, or either store could not be read", body = String),
+    ),
+    tag = "agents"
+)]
+async fn get_agents_status(
+    State(state): State<AppState>,
+) -> Result<Json<Vec<agent_status::AgentStatusRow>>, StatusUnavailable> {
+    let Some(reader) = state.agent_status.as_ref() else {
+        return Err(StatusUnavailable(
+            "no swarm queue is configured on this host".to_owned(),
+        ));
+    };
+    let Some(bridge) = state.auth.as_ref() else {
+        return Err(StatusUnavailable(
+            "no identity bridge is configured on this host".to_owned(),
+        ));
+    };
+    let roster = bridge.list_agent_identities().await.map_err(|e| {
+        let detail = format!("{e:#}");
+        tracing::warn!(error = %detail, "reading the agent roster failed");
+        StatusUnavailable(detail)
+    })?;
+    match reader.view(&roster, std::time::SystemTime::now()).await {
+        Ok(rows) => Ok(Json(rows)),
+        Err(e) => {
+            let detail = format!("{e:#}");
+            tracing::warn!(error = %detail, "reading the agent-status bucket failed");
+            Err(StatusUnavailable(detail))
+        }
+    }
+}
+
 /// Body of `POST /api/agents` — the agent name to create, and the hive the
 /// creation is aimed at. The repo name inside `forge::CONFIG_ORG` is the
 /// same string as `name`: one repo per agent, named after it, same
@@ -1480,6 +1544,7 @@ async fn main() -> Result<()> {
         hives: Arc::new(load_hives()),
         links: Arc::new(load_links()),
         wanted: wanted_writer(status.as_ref()),
+        agent_status: agent_status_reader(status.as_ref()),
         status,
         jobq,
         webhook_secret,
@@ -1514,6 +1579,7 @@ fn build_app(state: AppState) -> axum::Router {
         .routes(routes!(get_config_prs))
         .routes(routes!(create_agent))
         .routes(routes!(get_agents))
+        .routes(routes!(get_agents_status))
         .routes(routes!(set_agent_state))
         .routes(routes!(get_hive_wanted))
         .routes(routes!(issue_report::get_repos))
@@ -1622,6 +1688,7 @@ mod tests {
             // No queue, for the same reason as `status`: these tests drive
             // agent creation, which publishes no declaration.
             wanted: None,
+            agent_status: None,
             jobq: std::sync::Arc::clone(&sched),
             webhook_secret: None,
             config_prs: None,

swarm-queue-client/src/agent_status.rs

@@ -49,6 +49,25 @@ pub fn key(hive: &str, agent: &str) -> String {
     format!("{hive}/{agent}")
 }
 
+/// The inverse of [`key`]: split a bucket key back into `(hive, agent)`.
+///
+/// `None` for a key with zero or more than one `/` — a hive publishing
+/// under [`key`] never produces one, so a malformed key means something
+/// else wrote this bucket. Splitting on the *first* `/` would be equally
+/// valid today (neither name can contain one), but this rejects rather
+/// than guesses, so a future name-charset change can't silently start
+/// misreading old keys.
+#[must_use]
+pub fn split_key(key: &str) -> Option<(&str, &str)> {
+    let mut parts = key.splitn(3, '/');
+    let hive = parts.next()?;
+    let agent = parts.next()?;
+    if parts.next().is_some() {
+        return None;
+    }
+    Some((hive, agent))
+}
+
 /// One agent's status snapshot, as published under [`key`].
 ///
 /// Mirrors the tuple `container_view::read_agent_status_live` returns on
@@ -106,13 +125,28 @@ pub async fn open_or_create(
 
 #[cfg(test)]
 mod tests {
-    use super::{AgentStatus, key};
+    use super::{AgentStatus, key, split_key};
 
     #[test]
     fn key_joins_hive_and_agent_with_a_slash() {
         assert_eq!(key("prod", "iris"), "prod/iris");
     }
 
+    #[test]
+    fn split_key_is_the_inverse_of_key() {
+        assert_eq!(split_key(&key("prod", "iris")), Some(("prod", "iris")));
+    }
+
+    #[test]
+    fn split_key_rejects_a_key_with_no_slash() {
+        assert_eq!(split_key("iris"), None);
+    }
+
+    #[test]
+    fn split_key_rejects_a_key_with_two_slashes() {
+        assert_eq!(split_key("prod/iris/extra"), None);
+    }
+
     #[test]
     fn agent_status_round_trips() {
         let status = AgentStatus {