diff --git a/crates/viewer/re_viewer/src/app.rs b/crates/viewer/re_viewer/src/app.rs index e270663963ec..ac4c5e6e5044 100644 --- a/crates/viewer/re_viewer/src/app.rs +++ b/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(); diff --git a/crates/viewer/re_viewer/src/app_state.rs b/crates/viewer/re_viewer/src/app_state.rs index 0e95e539b6f1..4c39a8ae9951 100644 --- a/crates/viewer/re_viewer/src/app_state.rs +++ b/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, 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, diff --git a/crates/viewer/re_viewer_context/src/store_context.rs b/crates/viewer/re_viewer_context/src/store_context.rs index b401952315eb..35d851dc5023 100644 --- a/crates/viewer/re_viewer_context/src/store_context.rs +++ b/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, } diff --git a/crates/viewer/re_viewer_context/src/store_hub.rs b/crates/viewer/re_viewer_context/src/store_hub.rs index 848c5805dce4..25ed6d38ac20 100644 --- a/crates/viewer/re_viewer_context/src/store_hub.rs +++ b/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, store_bundle: StoreBundle, + /// Things that need caching. + caches_per_recording: HashMap, + /// The [`ChunkStoreGeneration`] from when the [`EntityDb`] was last saved blueprint_last_save: HashMap, @@ -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> { static EMPTY_ENTITY_DB: once_cell::sync::Lazy = once_cell::sync::Lazy::new(|| EntityDb::new(re_log_types::StoreId::empty_recording())); + static EMPTY_CACHES: once_cell::sync::Lazy = + 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 { diff --git a/crates/viewer/re_viewer_context/src/test_context.rs b/crates/viewer/re_viewer_context/src/test_context.rs index 5c982a97104d..608cc8ccdeb2 100644 --- a/crates/viewer/re_viewer_context/src/test_context.rs +++ b/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(), }; diff --git a/crates/viewer/re_viewport_blueprint/src/space_view.rs b/crates/viewer/re_viewport_blueprint/src/space_view.rs index 1c7ce8c4986c..b9fae7df1520 100644 --- a/crates/viewer/re_viewport_blueprint/src/space_view.rs +++ b/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(), }; diff --git a/crates/viewer/re_viewport_blueprint/src/space_view_contents.rs b/crates/viewer/re_viewport_blueprint/src/space_view_contents.rs index 13d5817c3997..4d4fb30c4e8e 100644 --- a/crates/viewer/re_viewport_blueprint/src/space_view_contents.rs +++ b/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(), };