Caches per recording (#7513)

Moves `Caches` from `AppState` to the `StoreHub`, and make them
per-recording instead of viewer-wide in the process.

This is:
* a pre-requisite to fix the repeated static logging issues (#7404)
* a necessary step towards proper secondary caching
* a fix for a very old issue: cache memory is not reclaimed when closing
a recording

I'm not a fan of the implementation: it's brittle, confusing, and I
don't like it. Still, it's a net improvement over the status quo, and
fixes a very annoying "leak" for end users.

Before:


https://github.com/user-attachments/assets/a54c1923-a939-4700-a485-eeb2a2dd4966



After:


https://github.com/user-attachments/assets/dd5adb65-d328-474c-bc31-e65a59332f03



---

* Part of #7404

crates/viewer/re_viewer/src/app.rs

@@ -1232,7 +1232,6 @@ impl App {
                 );
             }
             store_hub.purge_fraction_of_ram(fraction_to_purge);
-            self.state.cache.purge_memory();
 
             let mem_use_after = MemoryUse::capture();
 
@@ -1600,7 +1599,7 @@ impl eframe::App for App {
 
             // We haven't called `begin_frame` at this point, so pretend we did and add one to the active frame index.
             let renderer_active_frame_idx = render_ctx.active_frame_idx().wrapping_add(1);
-            self.state.cache.begin_frame(renderer_active_frame_idx);
+            store_hub.begin_frame(renderer_active_frame_idx);
         }
 
         self.show_text_logs_as_notifications();

crates/viewer/re_viewer/src/app_state.rs

@@ -8,10 +8,9 @@ use re_smart_channel::ReceiveSet;
 use re_types::blueprint::components::PanelState;
 use re_ui::ContextExt as _;
 use re_viewer_context::{
-    blueprint_timeline, AppOptions, ApplicationSelectionState, Caches, CommandSender,
-    ComponentUiRegistry, PlayState, RecordingConfig, SpaceViewClassExt as _,
-    SpaceViewClassRegistry, StoreContext, StoreHub, SystemCommandSender as _, ViewStates,
-    ViewerContext,
+    blueprint_timeline, AppOptions, ApplicationSelectionState, CommandSender, ComponentUiRegistry,
+    PlayState, RecordingConfig, SpaceViewClassExt as _, SpaceViewClassRegistry, StoreContext,
+    StoreHub, SystemCommandSender as _, ViewStates, ViewerContext,
 };
 use re_viewport::Viewport;
 use re_viewport_blueprint::ui::add_space_view_or_container_modal_ui;
@@ -28,10 +27,6 @@ pub struct AppState {
     /// Global options for the whole viewer.
     pub(crate) app_options: AppOptions,
 
-    /// Things that need caching.
-    #[serde(skip)]
-    pub(crate) cache: Caches,
-
     /// Configuration for the current recording (found in [`EntityDb`]).
     recording_configs: HashMap<StoreId, RecordingConfig>,
     blueprint_cfg: RecordingConfig,
@@ -73,7 +68,6 @@ impl Default for AppState {
     fn default() -> Self {
         Self {
             app_options: Default::default(),
-            cache: Default::default(),
             recording_configs: Default::default(),
             blueprint_cfg: Default::default(),
             selection_panel: Default::default(),
@@ -149,7 +143,6 @@ impl AppState {
 
         let Self {
             app_options,
-            cache,
             recording_configs,
             blueprint_cfg,
             selection_panel,
@@ -254,7 +247,7 @@ impl AppState {
         let egui_ctx = ui.ctx().clone();
         let ctx = ViewerContext {
             app_options,
-            cache,
+            cache: store_context.caches,
             space_view_class_registry,
             reflection,
             component_ui_registry,
@@ -321,7 +314,7 @@ impl AppState {
         // but it's just a bunch of refs so not really that big of a deal in practice.
         let ctx = ViewerContext {
             app_options,
-            cache,
+            cache: store_context.caches,
             space_view_class_registry,
             reflection,
             component_ui_registry,

crates/viewer/re_viewer_context/src/store_context.rs

@@ -1,7 +1,7 @@
 use re_entity_db::{EntityDb, StoreBundle};
 use re_log_types::{ApplicationId, StoreId};
 
-use crate::StoreHub;
+use crate::{Caches, StoreHub};
 
 /// The current Blueprint and Recording being displayed by the viewer
 pub struct StoreContext<'a> {
@@ -24,6 +24,9 @@ pub struct StoreContext<'a> {
     /// This is the same bundle as is in [`Self::hub`], but extracted for ease-of-access.
     pub bundle: &'a StoreBundle,
 
+    /// Things that need caching.
+    pub caches: &'a Caches,
+
     /// The store hub, which keeps track of all the default and active blueprints, among other things.
     pub hub: &'a StoreHub,
 }

crates/viewer/re_viewer_context/src/store_hub.rs

@@ -1,4 +1,4 @@
-use ahash::{HashMap, HashMapExt};
+use ahash::{HashMap, HashMapExt, HashSet};
 
 use anyhow::Context as _;
 use itertools::Itertools as _;
@@ -8,7 +8,7 @@ use re_entity_db::{EntityDb, StoreBundle};
 use re_log_types::{ApplicationId, StoreId, StoreKind};
 use re_query::CachesStats;
 
-use crate::StoreContext;
+use crate::{Caches, StoreContext};
 
 /// Interface for accessing all blueprints and recordings
 ///
@@ -42,6 +42,9 @@ pub struct StoreHub {
     active_blueprint_by_app_id: HashMap<ApplicationId, StoreId>,
     store_bundle: StoreBundle,
 
+    /// Things that need caching.
+    caches_per_recording: HashMap<StoreId, Caches>,
+
     /// The [`ChunkStoreGeneration`] from when the [`EntityDb`] was last saved
     blueprint_last_save: HashMap<StoreId, ChunkStoreGeneration>,
 
@@ -137,6 +140,7 @@ impl StoreHub {
             active_blueprint_by_app_id: Default::default(),
             store_bundle,
 
+            caches_per_recording: Default::default(),
             blueprint_last_save: Default::default(),
             blueprint_last_gc: Default::default(),
         }
@@ -158,6 +162,8 @@ impl StoreHub {
     pub fn read_context(&mut self) -> Option<StoreContext<'_>> {
         static EMPTY_ENTITY_DB: once_cell::sync::Lazy<EntityDb> =
             once_cell::sync::Lazy::new(|| EntityDb::new(re_log_types::StoreId::empty_recording()));
+        static EMPTY_CACHES: once_cell::sync::Lazy<Caches> =
+            once_cell::sync::Lazy::new(Default::default);
 
         // If we have an app-id, then use it to look up the blueprint.
         let app_id = self.active_application_id.clone()?;
@@ -211,12 +217,15 @@ impl StoreHub {
             self.active_rec_id = None;
         }
 
+        let caches = self.active_caches();
+
         Some(StoreContext {
             app_id,
             blueprint: active_blueprint,
             default_blueprint,
             recording: recording.unwrap_or(&EMPTY_ENTITY_DB),
             bundle: &self.store_bundle,
+            caches: caches.unwrap_or(&EMPTY_CACHES),
             hub: self,
         })
     }
@@ -238,6 +247,7 @@ impl StoreHub {
     }
 
     pub fn remove(&mut self, store_id: &StoreId) {
+        _ = self.caches_per_recording.remove(store_id);
         let removed_store = self.store_bundle.remove(store_id);
 
         let Some(removed_store) = removed_store else {
@@ -296,8 +306,18 @@ impl StoreHub {
     /// Remove all open recordings and applications, and go to the welcome page.
     pub fn clear_recordings(&mut self) {
         // Keep only the welcome screen:
-        self.store_bundle
-            .retain(|db| db.app_id() == Some(&Self::welcome_screen_app_id()));
+        let mut store_ids_retained = HashSet::default();
+        self.store_bundle.retain(|db| {
+            if db.app_id() == Some(&Self::welcome_screen_app_id()) {
+                store_ids_retained.insert(db.store_id().clone());
+                true
+            } else {
+                false
+            }
+        });
+        self.caches_per_recording
+            .retain(|store_id, _| store_ids_retained.contains(store_id));
+
         self.active_rec_id = None;
         self.active_application_id = Some(Self::welcome_screen_app_id());
     }
@@ -342,7 +362,17 @@ impl StoreHub {
             re_log::warn!("Failed to save blueprints: {err}");
         }
 
-        self.store_bundle.retain(|db| db.app_id() != Some(app_id));
+        let mut store_ids_removed = HashSet::default();
+        self.store_bundle.retain(|db| {
+            if db.app_id() == Some(app_id) {
+                store_ids_removed.insert(db.store_id().clone());
+                false
+            } else {
+                true
+            }
+        });
+        self.caches_per_recording
+            .retain(|store_id, _| !store_ids_removed.contains(store_id));
 
         if self.active_application_id.as_ref() == Some(app_id) {
             self.active_application_id = None;
@@ -375,6 +405,24 @@ impl StoreHub {
             .and_then(|id| self.store_bundle.get(id))
     }
 
+    /// Directly access the [`Caches`] for the active recording.
+    ///
+    /// This returns `None` only if there is no active recording: the cache itself is always
+    /// present if there's an active recording.
+    #[inline]
+    pub fn active_caches(&self) -> Option<&Caches> {
+        self.active_rec_id.as_ref().and_then(|store_id| {
+            let caches = self.caches_per_recording.get(store_id);
+
+            debug_assert!(
+                caches.is_some(),
+                "active recordings should always have associated caches",
+            );
+
+            caches
+        })
+    }
+
     /// Change the active/visible recording id.
     ///
     /// This will also change the application-id to match the newly active recording.
@@ -392,7 +440,10 @@ impl StoreHub {
             self.set_active_app(app_id);
         }
 
-        self.active_rec_id = Some(recording_id);
+        self.active_rec_id = Some(recording_id.clone());
+
+        // Make sure the active recording has associated caches, always.
+        _ = self.caches_per_recording.entry(recording_id).or_default();
     }
 
     /// Activate a recording by its [`StoreId`].
@@ -552,6 +603,10 @@ impl StoreHub {
             return;
         };
 
+        if let Some(caches) = self.caches_per_recording.get_mut(&store_id) {
+            caches.purge_memory();
+        }
+
         let store_bundle = &mut self.store_bundle;
 
         let Some(entity_db) = store_bundle.get_mut(&store_id) else {
@@ -637,6 +692,15 @@ impl StoreHub {
         }
     }
 
+    /// See `re_viewer_context::Cache::begin_frame`.
+    pub fn begin_frame(&mut self, renderer_active_frame_idx: u64) {
+        if let Some(store_id) = self.active_recording_id().cloned() {
+            if let Some(caches) = self.caches_per_recording.get_mut(&store_id) {
+                caches.begin_frame(renderer_active_frame_idx);
+            }
+        }
+    }
+
     /// Persist any in-use blueprints to durable storage.
     pub fn save_app_blueprints(&mut self) -> anyhow::Result<()> {
         let Some(saver) = &self.persistence.saver else {

crates/viewer/re_viewer_context/src/test_context.rs

@@ -66,6 +66,7 @@ impl TestContext {
                     default_blueprint: None,
                     recording: &self.recording_store,
                     bundle: &Default::default(),
+                    caches: &Default::default(),
                     hub: &Default::default(),
                 };
 

crates/viewer/re_viewport_blueprint/src/space_view.rs

@@ -770,6 +770,7 @@ mod tests {
             default_blueprint: None,
             recording: &test_ctx.recording_store,
             bundle: &Default::default(),
+            caches: &Default::default(),
             hub: &re_viewer_context::StoreHub::test_hub(),
         };
 

crates/viewer/re_viewport_blueprint/src/space_view_contents.rs

@@ -616,6 +616,7 @@ mod tests {
             default_blueprint: None,
             recording: &recording,
             bundle: &Default::default(),
+            caches: &Default::default(),
             hub: &StoreHub::test_hub(),
         };