Skip to content

Commit

Permalink
Caches per recording (#7513)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
teh-cmc authored Sep 26, 2024
1 parent 14ff694 commit 0034e6d
Showing 7 changed files with 83 additions and 21 deletions.
3 changes: 1 addition & 2 deletions crates/viewer/re_viewer/src/app.rs
Original file line number Diff line number Diff line change
@@ -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();
17 changes: 5 additions & 12 deletions crates/viewer/re_viewer/src/app_state.rs
Original file line number Diff line number Diff line change
@@ -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,
5 changes: 4 additions & 1 deletion crates/viewer/re_viewer_context/src/store_context.rs
Original file line number Diff line number Diff line change
@@ -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,
}
76 changes: 70 additions & 6 deletions crates/viewer/re_viewer_context/src/store_hub.rs
Original file line number Diff line number Diff line change
@@ -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 {
1 change: 1 addition & 0 deletions crates/viewer/re_viewer_context/src/test_context.rs
Original file line number Diff line number Diff line change
@@ -66,6 +66,7 @@ impl TestContext {
default_blueprint: None,
recording: &self.recording_store,
bundle: &Default::default(),
caches: &Default::default(),
hub: &Default::default(),
};

1 change: 1 addition & 0 deletions crates/viewer/re_viewport_blueprint/src/space_view.rs
Original file line number Diff line number Diff line change
@@ -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(),
};

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

0 comments on commit 0034e6d

Please sign in to comment.