perf(fleet): cache grouped views with revision-based invalidation

wakey-control-plane/src/api/control.rs

@@ -12,5 +12,6 @@ pub use enroll::{
     enroll, healthz, issue_enroll_token, list_enroll_tokens, revoke_agent, revoke_enroll_token,
     set_agent_nickname,
 };
+pub(crate) use fleet::FleetCache;
 pub use fleet::{list_fleet_devices, refresh_fleet_devices, wake_fleet_device};
 pub use stats::{StateStatsResponse, state_stats};

wakey-control-plane/src/api/control/fleet/mod.rs

@@ -27,6 +27,28 @@ use types::{
     RefreshFleetDevicesResponse, WakeFleetDeviceRequest, WakeFleetDeviceResponse,
 };
 
+pub(crate) type FleetCache =
+    crate::runtime::view_cache::ViewCache<(u64, BTreeSet<String>), Vec<FleetDevice>>;
+
+async fn fleet_view(state: &AppState) -> anyhow::Result<std::sync::Arc<Vec<FleetDevice>>> {
+    state
+        .fleet
+        .get_or_try_init(
+            || async {
+                let connected = state
+                    .sessions
+                    .read()
+                    .await
+                    .keys()
+                    .cloned()
+                    .collect::<BTreeSet<_>>();
+                (state.store.fleet_revision(), connected)
+            },
+            || build_fleet_view(state),
+        )
+        .await
+}
+
 const REFRESH_CONCURRENCY: usize = 8;
 
 #[tracing::instrument(name = "fleet_list", skip_all)]
@@ -289,6 +311,21 @@ async fn load_fleet_devices(
     state: &AppState,
     query: &ListFleetDevicesQuery,
 ) -> anyhow::Result<Vec<FleetDevice>> {
+    let cached = fleet_view(state).await?;
+    // Cache the complete view. Query-specific visibility, search and limits
+    // must never affect what another client (or a wake request) can retrieve.
+    let mut devices = cached.as_ref().clone();
+    tracing::Span::current().record("grouped", devices.len());
+    filter_fleet_devices(&mut devices, query);
+    tracing::Span::current().record("filtered", devices.len());
+    let limit = query.limit.unwrap_or(500).clamp(1, 1000);
+    devices.truncate(limit);
+    tracing::Span::current().record("returned", devices.len());
+    Ok(devices)
+}
+
+#[tracing::instrument(name = "fleet_rebuild", skip_all)]
+async fn build_fleet_view(state: &AppState) -> anyhow::Result<Vec<FleetDevice>> {
     let known_devices = state.store.list_known_devices().await?;
     let agent_devices = state.store.list_agent_device_rows().await?;
     let connected = {
@@ -322,13 +359,5 @@ async fn load_fleet_devices(
         agent_status,
         identifier_map,
     };
-    let mut devices = build_fleet_devices(known_devices, agent_devices, &context);
-    tracing::Span::current().record("grouped", devices.len());
-    filter_fleet_devices(&mut devices, query);
-    tracing::Span::current().record("filtered", devices.len());
-    let limit = query.limit.unwrap_or(500).clamp(1, 1000);
-    // this is going to be 10 million times harder if i decide to do this lazily
-    devices.truncate(limit);
-    tracing::Span::current().record("returned", devices.len());
-    Ok(devices)
+    Ok(build_fleet_devices(known_devices, agent_devices, &context))
 }

wakey-control-plane/src/api/control/fleet/tests.rs

@@ -12,6 +12,216 @@ use super::normalize_refresh_agent_ids;
 use super::types::ListFleetDevicesQuery;
 use crate::state::{AgentDeviceRow, AgentDeviceWithChildren, DeviceIdentifier, KnownDevice};
 
+#[tokio::test]
+async fn fleet_cache_tracks_store_edits_and_connections_without_caching_filters() {
+    use crate::runtime::{AgentSession, AppState};
+    use crate::state::{DeviceIdentifierInput, KnownDeviceInput, Store};
+    use std::sync::Arc;
+    use std::time::Duration;
+    use tokio::sync::{Mutex, RwLock, mpsc};
+
+    let dir = std::env::temp_dir().join(format!("wakey-fleet-cache-{}", uuid::Uuid::new_v4()));
+    let state = AppState {
+        store: Arc::new(
+            Store::load_or_init(&dir.join("state.sqlite3"), vec![], Duration::from_secs(60))
+                .await
+                .unwrap(),
+        ),
+        sessions: Arc::new(RwLock::new(HashMap::new())),
+        pending: Arc::new(Mutex::new(HashMap::new())),
+        public_url: "http://localhost".into(),
+        command_timeout: Duration::from_secs(30),
+        enroll_token_ttl: Duration::from_secs(60),
+        terminals: crate::runtime::terminals::TerminalRegistry::new(),
+        alerts: crate::api::AlertsHub::new(),
+        fleet: Arc::new(super::FleetCache::default()),
+    };
+    let store = &state.store;
+    let token = store
+        .issue_enroll_token(Duration::from_secs(60))
+        .await
+        .unwrap();
+    let agent = store.enroll(&token.enroll_token).await.unwrap();
+    let observed = wakey_core::Device::from_parts(
+        vec![wakey_core::NeighborEntry {
+            ip: "192.168.1.10".parse().unwrap(),
+            mac: Some("aa:bb:cc:dd:ee:01".parse().unwrap()),
+            dev: Some("br-lan".into()),
+            state: wakey_core::NeighborState::Reachable,
+        }],
+        vec![],
+    );
+    store
+        .replace_agent_device_snapshot(&agent.agent_id, &[observed])
+        .await
+        .unwrap();
+    let initial = super::fleet_view(&state).await.unwrap();
+    assert_eq!(initial.len(), 1);
+    assert!(Arc::ptr_eq(
+        &initial,
+        &super::fleet_view(&state).await.unwrap()
+    ));
+    assert!(initial[0].recommended_route.is_none());
+
+    let (tx, _rx) = mpsc::unbounded_channel();
+    state.sessions.write().await.insert(
+        agent.agent_id.clone(),
+        AgentSession {
+            connection_id: "test".into(),
+            tx,
+            capabilities: vec![],
+            capability_options: Default::default(),
+        },
+    );
+    let connected = super::fleet_view(&state).await.unwrap();
+    assert!(connected[0].recommended_route.is_some());
+    store
+        .set_agent_nickname(&agent.agent_id, Some("router"))
+        .await
+        .unwrap();
+    assert_eq!(
+        super::fleet_view(&state).await.unwrap()[0].agents[0]
+            .nickname
+            .as_deref(),
+        Some("router")
+    );
+
+    let saved = store
+        .create_known_device(KnownDeviceInput {
+            display_name: "my pc".into(),
+            pinned: true,
+            notes: None,
+            identifiers: vec![DeviceIdentifierInput {
+                kind: "mac".into(),
+                value: "aa:bb:cc:dd:ee:01".into(),
+            }],
+        })
+        .await
+        .unwrap();
+    assert_eq!(
+        super::fleet_view(&state).await.unwrap()[0].display_name,
+        "my pc"
+    );
+    store
+        .rename_known_device(&saved.device_id, "renamed pc")
+        .await
+        .unwrap();
+    let before_error = super::fleet_view(&state).await.unwrap();
+    assert!(
+        store
+            .rename_known_device(&saved.device_id, " ")
+            .await
+            .is_err()
+    );
+    assert!(Arc::ptr_eq(
+        &before_error,
+        &super::fleet_view(&state).await.unwrap()
+    ));
+    assert_eq!(
+        super::fleet_view(&state).await.unwrap()[0].display_name,
+        "renamed pc"
+    );
+    let query = ListFleetDevicesQuery {
+        query: Some("does-not-exist".into()),
+        visibility: Some("all".into()),
+        ..Default::default()
+    };
+    assert!(
+        super::load_fleet_devices(&state, &query)
+            .await
+            .unwrap()
+            .is_empty()
+    );
+    let all = ListFleetDevicesQuery {
+        visibility: Some("all".into()),
+        ..Default::default()
+    };
+    assert_eq!(
+        super::load_fleet_devices(&state, &all).await.unwrap().len(),
+        1
+    );
+
+    store
+        .detach_device_identifier(&saved.device_id, &saved.identifiers[0].identifier_key)
+        .await
+        .unwrap();
+    assert_eq!(super::fleet_view(&state).await.unwrap().len(), 2);
+    store
+        .attach_device_identifier(
+            &saved.device_id,
+            DeviceIdentifierInput {
+                kind: "mac".into(),
+                value: "aa:bb:cc:dd:ee:01".into(),
+            },
+        )
+        .await
+        .unwrap();
+    assert_eq!(super::fleet_view(&state).await.unwrap().len(), 1);
+    let other = store
+        .create_known_device(KnownDeviceInput {
+            display_name: "other".into(),
+            pinned: false,
+            notes: None,
+            identifiers: vec![],
+        })
+        .await
+        .unwrap();
+    assert_eq!(super::fleet_view(&state).await.unwrap().len(), 2);
+    let limited = ListFleetDevicesQuery {
+        visibility: Some("all".into()),
+        limit: Some(1),
+        ..Default::default()
+    };
+    assert_eq!(
+        super::load_fleet_devices(&state, &limited)
+            .await
+            .unwrap()
+            .len(),
+        1
+    );
+    assert_eq!(
+        super::load_fleet_devices(&state, &all).await.unwrap().len(),
+        2
+    );
+    store
+        .merge_known_devices(&saved.device_id, &other.device_id)
+        .await
+        .unwrap();
+    assert_eq!(super::fleet_view(&state).await.unwrap().len(), 1);
+    store.forget_known_device(&saved.device_id).await.unwrap();
+    assert!(
+        super::fleet_view(&state).await.unwrap()[0]
+            .known_device
+            .is_none()
+    );
+
+    state.sessions.write().await.clear();
+    assert!(
+        super::fleet_view(&state).await.unwrap()[0]
+            .recommended_route
+            .is_none()
+    );
+    let before_reload = super::fleet_view(&state).await.unwrap();
+    store.reload_from_disk().await.unwrap();
+    assert!(!Arc::ptr_eq(
+        &before_reload,
+        &super::fleet_view(&state).await.unwrap()
+    ));
+    store.revoke_agent(&agent.agent_id).await.unwrap();
+    assert!(
+        super::fleet_view(&state).await.unwrap()[0].agents[0]
+            .nickname
+            .is_none()
+    );
+    store
+        .replace_agent_device_snapshot(&agent.agent_id, &[])
+        .await
+        .unwrap();
+    assert!(super::fleet_view(&state).await.unwrap().is_empty());
+    drop(state);
+    std::fs::remove_dir_all(dir).unwrap();
+}
+
 #[test]
 fn refresh_agent_selection_keeps_all_agents_beyond_previous_cap() {
     let ids = (0..140)

wakey-control-plane/src/api/mod.rs

@@ -6,6 +6,7 @@ mod audit;
 mod commands;
 mod control;
 mod terminals;
+pub(crate) use control::FleetCache;
 
 pub use alerts::{AlertsHub, active_alerts, alert_history, alerts_stream, run_alert_evaluator};
 pub use audit::list_audit_events;

wakey-control-plane/src/runtime/mod.rs

@@ -26,6 +26,7 @@ use crate::ws;
 mod admin;
 mod process;
 pub mod terminals;
+pub(crate) mod view_cache;
 pub use admin::revoke_agent;
 pub use admin::{
     issue_enroll_token, list_enroll_tokens, migrate_sqlite_state, revoke_enroll_token, state_stats,
@@ -44,6 +45,7 @@ pub struct AppState {
     pub enroll_token_ttl: Duration,
     pub terminals: terminals::TerminalRegistry,
     pub alerts: api::AlertsHub,
+    pub(crate) fleet: Arc<api::FleetCache>,
 }
 
 #[derive(Clone)]
@@ -190,6 +192,7 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
         enroll_token_ttl: daemon.enroll_token_ttl,
         terminals: terminals::TerminalRegistry::new(),
         alerts: api::AlertsHub::new(),
+        fleet: Arc::new(api::FleetCache::default()),
     };
 
     tokio::spawn(api::run_alert_evaluator(app_state.clone()));

wakey-control-plane/src/runtime/view_cache.rs

@@ -0,0 +1,157 @@
+//! Lazy derived views keyed by their inputs, with one rebuild at a time.
+
+use std::future::Future;
+use std::sync::Arc;
+
+use tokio::sync::Mutex;
+
+pub(crate) struct ViewCache<K, V> {
+    entry: Mutex<Option<(K, Arc<V>)>>,
+}
+
+impl<K, V> Default for ViewCache<K, V> {
+    fn default() -> Self {
+        Self {
+            entry: Mutex::new(None),
+        }
+    }
+}
+
+impl<K: Eq, V> ViewCache<K, V> {
+    /// Reuse a matching view or rebuild it. `key` must reflect every input to
+    /// `build`; rechecking it prevents an in-flight old build being published.
+    /// Errors are not cached, and cancellation releases the rebuild lock.
+    pub(crate) async fn get_or_try_init<E, KF, BF>(
+        &self,
+        key: impl Fn() -> KF,
+        build: impl Fn() -> BF,
+    ) -> Result<Arc<V>, E>
+    where
+        KF: Future<Output = K>,
+        BF: Future<Output = Result<V, E>>,
+    {
+        // Keep publication and rebuilding under one lock. Waiting readers then
+        // reuse the completed view rather than launching duplicate DB work.
+        let mut entry = self.entry.lock().await;
+        loop {
+            let before = key().await;
+            if let Some((stored_key, value)) = &*entry
+                && *stored_key == before
+            {
+                return Ok(Arc::clone(value));
+            }
+            let value = build().await?;
+            if before == key().await {
+                let value = Arc::new(value);
+                *entry = Some((before, Arc::clone(&value)));
+                return Ok(value);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::atomic::{AtomicU64, Ordering};
+    use tokio::sync::Notify;
+
+    #[tokio::test]
+    async fn concurrent_readers_share_a_build_and_changed_keys_rebuild() {
+        let cache = ViewCache::default();
+        let revision = AtomicU64::new(0);
+        let builds = AtomicU64::new(0);
+        let key = || async { revision.load(Ordering::SeqCst) };
+        let build = || async {
+            builds.fetch_add(1, Ordering::SeqCst);
+            tokio::task::yield_now().await;
+            Ok::<_, ()>(42)
+        };
+        let (a, b) = tokio::join!(
+            cache.get_or_try_init(key, build),
+            cache.get_or_try_init(key, build)
+        );
+        assert!(Arc::ptr_eq(&a.unwrap(), &b.unwrap()));
+        assert_eq!(builds.load(Ordering::SeqCst), 1);
+        revision.fetch_add(1, Ordering::SeqCst);
+        cache.get_or_try_init(key, build).await.unwrap();
+        assert_eq!(builds.load(Ordering::SeqCst), 2);
+    }
+
+    #[tokio::test]
+    async fn invalidation_during_a_build_discards_the_old_result() {
+        let cache = ViewCache::default();
+        let revision = AtomicU64::new(0);
+        let builds = AtomicU64::new(0);
+        let started = Notify::new();
+        let release = Notify::new();
+        let reader = cache.get_or_try_init(
+            || async { revision.load(Ordering::SeqCst) },
+            || async {
+                let captured = revision.load(Ordering::SeqCst);
+                if builds.fetch_add(1, Ordering::SeqCst) == 0 {
+                    started.notify_one();
+                    release.notified().await;
+                }
+                Ok::<_, ()>(captured)
+            },
+        );
+        let writer = async {
+            started.notified().await;
+            revision.store(1, Ordering::SeqCst);
+            release.notify_one();
+        };
+        let (result, ()) = tokio::join!(reader, writer);
+        assert_eq!(*result.unwrap(), 1);
+        assert_eq!(builds.load(Ordering::SeqCst), 2);
+    }
+
+    #[tokio::test]
+    async fn failed_builds_are_retried() {
+        let cache = ViewCache::default();
+        assert!(
+            cache
+                .get_or_try_init(|| async { 0 }, || async { Err::<u8, _>("db failed") })
+                .await
+                .is_err()
+        );
+        assert_eq!(
+            *cache
+                .get_or_try_init(|| async { 0 }, || async { Ok::<_, &str>(7) })
+                .await
+                .unwrap(),
+            7
+        );
+    }
+
+    #[tokio::test]
+    async fn cancelling_a_build_does_not_lock_out_readers() {
+        let cache = Arc::new(ViewCache::default());
+        let started = Arc::new(Notify::new());
+        let task = tokio::spawn({
+            let cache = Arc::clone(&cache);
+            let started = Arc::clone(&started);
+            async move {
+                cache
+                    .get_or_try_init(
+                        || async { 0 },
+                        || async {
+                            started.notify_one();
+                            std::future::pending::<Result<u8, ()>>().await
+                        },
+                    )
+                    .await
+            }
+        });
+        started.notified().await;
+        task.abort();
+        assert!(task.await.unwrap_err().is_cancelled());
+        assert_eq!(
+            *cache
+                .get_or_try_init(|| async { 0 }, || async { Ok::<_, ()>(7) })
+                .await
+                .unwrap(),
+            7
+        );
+    }
+}

wakey-control-plane/src/state/store.rs

@@ -16,8 +16,20 @@ use crate::state::types::{
 pub struct Store {
     db_path: PathBuf,
     pool: SqlitePool,
+    fleet_revision: std::sync::atomic::AtomicU64,
 }
 impl Store {
+    pub(crate) fn fleet_revision(&self) -> u64 {
+        self.fleet_revision
+            .load(std::sync::atomic::Ordering::Acquire)
+    }
+
+    /// Call after a successful fleet-relevant write, never before its commit.
+    fn touch_fleet(&self) {
+        self.fleet_revision
+            .fetch_add(1, std::sync::atomic::Ordering::Release);
+    }
+
     /// Begins a new transaction with an IMMEDIATE write lock.
     ///
     /// Use this for all transactions that perform writes to avoid SQLite upgrade deadlocks.

wakey-control-plane/src/state/store/agent_devices.rs

@@ -32,6 +32,9 @@ impl Store {
         tx.commit()
             .await
             .context("failed committing device retention cleanup")?;
+        if removed > 0 {
+            self.touch_fleet();
+        }
         Ok(removed)
     }
 
@@ -251,6 +254,7 @@ impl Store {
         tx.commit()
             .await
             .context("failed committing device snapshot")?;
+        self.touch_fleet();
 
         tracing::Span::current().record("accepted", incoming_keys.len());
         Ok(incoming_keys.len())
@@ -644,10 +648,12 @@ mod tests {
                 .await
                 .unwrap();
         }
+        let before_cleanup = store.fleet_revision();
         assert_eq!(
             store.prune_devices_at(Duration::ZERO, 1_000).await.unwrap(),
             0
         );
+        assert_eq!(store.fleet_revision(), before_cleanup);
         assert_eq!(
             store
                 .prune_devices_at(Duration::from_secs(900), 1_000)
@@ -655,6 +661,8 @@ mod tests {
                 .unwrap(),
             1
         );
+        let after_cleanup = store.fleet_revision();
+        assert_ne!(after_cleanup, before_cleanup);
         assert_eq!(
             store
                 .prune_devices_at(Duration::from_secs(900), 1_000)
@@ -662,6 +670,7 @@ mod tests {
                 .unwrap(),
             0
         );
+        assert_eq!(store.fleet_revision(), after_cleanup);
 
         for table in [
             "agent_devices",

wakey-control-plane/src/state/store/db.rs

@@ -25,7 +25,11 @@ impl Store {
             .await
             .with_context(|| format!("failed to migrate state db {}", db_path.display()))?;
 
-        let store = Self { db_path, pool };
+        let store = Self {
+            db_path,
+            pool,
+            fleet_revision: Default::default(),
+        };
 
         store.ensure_schema_version().await?;
         store
@@ -87,7 +91,8 @@ impl Store {
             .fetch_one(&self.pool)
             .await
             .context("failed validating SQLite state connection")?;
-        info!(path = %self.db_path.display(), "reload requested; SQLite backend does not require in-memory reload");
+        self.touch_fleet();
+        info!(path = %self.db_path.display(), "reload requested; invalidated derived fleet views");
         Ok(())
     }
 

wakey-control-plane/src/state/store/devices.rs

@@ -37,6 +37,7 @@ impl Store {
         tx.commit()
             .await
             .context("failed committing known device transaction")?;
+        self.touch_fleet();
         self.get_known_device(&device_id)
             .await?
             .ok_or_else(|| anyhow::anyhow!("created known device disappeared"))
@@ -84,6 +85,7 @@ impl Store {
         if result.rows_affected() == 0 {
             return Ok(None);
         }
+        self.touch_fleet();
         self.get_known_device(device_id).await
     }
 
@@ -93,6 +95,9 @@ impl Store {
             .execute(&self.pool)
             .await
             .context("failed deleting known device")?;
+        if result.rows_affected() > 0 {
+            self.touch_fleet();
+        }
         Ok(result.rows_affected() > 0)
     }
 
@@ -159,6 +164,7 @@ impl Store {
         tx.commit()
             .await
             .context("failed committing known device merge transaction")?;
+        self.touch_fleet();
         self.get_known_device(target_device_id).await
     }
 
@@ -195,6 +201,7 @@ impl Store {
         tx.commit()
             .await
             .context("failed committing device identifier transaction")?;
+        self.touch_fleet();
         self.get_known_device(device_id).await
     }
 
@@ -237,6 +244,7 @@ impl Store {
         tx.commit()
             .await
             .context("failed committing device identifier detach transaction")?;
+        self.touch_fleet();
         self.get_known_device(device_id).await
     }
 

wakey-control-plane/src/state/store/enrollment.rs

@@ -54,6 +54,7 @@ impl Store {
         tx.commit()
             .await
             .context("failed committing state db after enroll")?;
+        self.touch_fleet();
         info!(agent_id = %agent_id, "issued persistent agent credentials");
 
         Ok(IssuedAgent {
@@ -113,6 +114,9 @@ impl Store {
         tx.commit()
             .await
             .context("failed committing db after agent revoke")?;
+        if result.rows_affected() > 0 {
+            self.touch_fleet();
+        }
         Ok(result.rows_affected() > 0)
     }
 
@@ -146,6 +150,7 @@ impl Store {
                 .context("failed clearing agent nickname")?;
         }
 
+        self.touch_fleet();
         Ok(true)
     }