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