Optimize performance for hundreds of entities (#7115)

joelreymont · joelreymont · commit 129d04ee06c6 · 2025-11-07T20:23:40.000+02:00
This commit addresses the performance bottlenecks described in issue #7115 where the Rerun viewer becomes slow with several hundred entities. Key improvements: 1. Share visualizable_entities_per_view across the frame - Added visualizable_entities_per_view to ViewerContext to avoid redundant computation - Previously, determine_visualizable_entities() was called multiple times per frame for each view: * Once in app_state.rs * Once in update_overrides() * Once or more in create_entity_add_info() for UI updates - Now computed once per frame and shared via ViewerContext 2. Update create_entity_add_info to use cached results - Modified entity_add_info.rs to use pre-computed visualizable entities from ViewerContext instead of recomputing 3. Improve cache invalidation strategy - Changed cache keys from latest_row_id to ChunkStoreGeneration - latest_row_id changes on every data insertion, invalidating cache every frame for streaming data
diff --git a/crates/viewer/re_test_context/src/lib.rs b/crates/viewer/re_test_context/src/lib.rs
@@ -494,6 +494,9 @@ impl TestContext {
         let drag_and_drop_manager =
             re_viewer_context::DragAndDropManager::new(ItemCollection::default());
 
+        // Tests don't compute visualizable entities per view, so we provide an empty map
+        let visualizable_entities_per_view: HashMap<ViewId, re_viewer_context::PerVisualizer<re_viewer_context::VisualizableEntities>> = HashMap::default();
+
         let mut context_render_state = self.egui_render_state.lock();
         let render_state = context_render_state.get_or_insert_with(create_egui_renderstate);
         let mut egui_renderer = render_state.renderer.write();
@@ -531,6 +534,7 @@ impl TestContext {
             storage_context: &storage_context,
             maybe_visualizable_entities_per_visualizer: &maybe_visualizable_entities_per_visualizer,
             indicated_entities_per_visualizer: &indicated_entities_per_visualizer,
+            visualizable_entities_per_view: &visualizable_entities_per_view,
             query_results: &self.query_results,
             time_ctrl: &self.time_ctrl.read(),
             blueprint_time_ctrl: &Default::default(),
diff --git a/crates/viewer/re_viewer/src/app_state.rs b/crates/viewer/re_viewer/src/app_state.rs
@@ -1,12 +1,15 @@
-use std::{borrow::Cow, str::FromStr as _};
+use std::{
+    borrow::Cow,
+    str::FromStr as _,
+};
 
 use ahash::HashMap;
 use egui::{Ui, text_edit::TextEditState, text_selection::LabelSelectionState};
 
 use re_chunk::{Timeline, TimelineName};
 use re_chunk_store::LatestAtQuery;
 use re_entity_db::EntityDb;
-use re_log_types::{AbsoluteTimeRangeF, DataSourceMessage, StoreId, TableId};
+use re_log_types::{AbsoluteTimeRangeF, DataSourceMessage, EntityPath, StoreId, TableId};
 use re_redap_browser::RedapServers;
 use re_redap_client::ConnectionRegistryHandle;
 use re_smart_channel::ReceiveSet;
@@ -18,7 +21,7 @@ use re_viewer_context::{
     DragAndDropManager, FallbackProviderRegistry, GlobalContext, IndicatedEntities, Item,
     MaybeVisualizableEntities, PerVisualizer, SelectionChange, StorageContext, StoreContext,
     StoreHub, SystemCommand, SystemCommandSender as _, TableStore, TimeControl, TimeControlCommand,
-    ViewClassRegistry, ViewId, ViewStates, ViewerContext, blueprint_timeline,
+    ViewClassRegistry, ViewId, ViewStates, ViewerContext, VisualizableEntities, blueprint_timeline,
     open_url::{self, ViewerOpenUrl},
 };
 use re_viewport::ViewportUi;
@@ -72,6 +75,12 @@ pub struct AppState {
     #[serde(skip)]
     pub(crate) share_modal: crate::ui::ShareModal,
 
+    #[serde(skip)]
+    view_query_cache: HashMap<ViewId, ViewQueryCacheEntry>,
+
+    #[serde(skip)]
+    visualizable_entities_cache: HashMap<ViewId, VisualizableEntitiesCacheEntry>,
+
     /// Test-only: single-shot callback to run at the end of the frame. Used in integration tests
     /// to interact with the `ViewerContext`.
     #[cfg(feature = "testing")]
@@ -127,6 +136,8 @@ impl Default for AppState {
             view_states: Default::default(),
             selection_state: Default::default(),
             focused_item: Default::default(),
+            view_query_cache: Default::default(),
+            visualizable_entities_cache: Default::default(),
 
             #[cfg(feature = "testing")]
             test_hook: None,
@@ -142,6 +153,58 @@ pub(crate) struct WelcomeScreenState {
     pub opacity: f32,
 }
 
+struct ViewQueryCacheEntry {
+    key: ViewQueryCacheKey,
+    result: DataQueryResult,
+}
+
+#[derive(Clone, PartialEq, Eq)]
+struct ViewQueryCacheKey {
+    recording_generation: re_chunk_store::ChunkStoreGeneration,
+    blueprint_generation: re_chunk_store::ChunkStoreGeneration,
+    blueprint_query: LatestAtQuery,
+}
+
+impl ViewQueryCacheKey {
+    fn new(
+        recording_generation: re_chunk_store::ChunkStoreGeneration,
+        blueprint_generation: re_chunk_store::ChunkStoreGeneration,
+        blueprint_query: &LatestAtQuery,
+    ) -> Self {
+        Self {
+            recording_generation,
+            blueprint_generation,
+            blueprint_query: blueprint_query.clone(),
+        }
+    }
+}
+
+struct VisualizableEntitiesCacheEntry {
+    key: VisualizableEntitiesCacheKey,
+    result: PerVisualizer<VisualizableEntities>,
+}
+
+#[derive(Clone, PartialEq, Eq, Hash)]
+struct VisualizableEntitiesCacheKey {
+    recording_generation: re_chunk_store::ChunkStoreGeneration,
+    blueprint_generation: re_chunk_store::ChunkStoreGeneration,
+    space_origin: EntityPath,
+}
+
+impl VisualizableEntitiesCacheKey {
+    fn new(
+        recording_generation: re_chunk_store::ChunkStoreGeneration,
+        blueprint_generation: re_chunk_store::ChunkStoreGeneration,
+        space_origin: &EntityPath,
+    ) -> Self {
+        Self {
+            recording_generation,
+            blueprint_generation,
+            space_origin: space_origin.clone(),
+        }
+    }
+}
+
 impl AppState {
     pub fn set_examples_manifest_url(&mut self, egui_ctx: &egui::Context, url: String) {
         self.welcome_screen.set_examples_manifest_url(egui_ctx, url);
@@ -291,39 +354,97 @@ impl AppState {
                     .maybe_visualizable_entities_for_visualizer_systems(recording.store_id());
                 let indicated_entities_per_visualizer =
                     view_class_registry.indicated_entities_per_visualizer(recording.store_id());
+                let recording_generation = recording.generation();
+                let blueprint_generation = store_context.blueprint.generation();
+
+                self.view_query_cache
+                    .retain(|view_id, _| viewport_ui.blueprint.views.contains_key(view_id));
+                self.visualizable_entities_cache
+                    .retain(|view_id, _| viewport_ui.blueprint.views.contains_key(view_id));
+
+                // Determine visualizable entities for each view (cached)
+                let visualizable_entities_per_view: HashMap<ViewId, PerVisualizer<VisualizableEntities>> = {
+                    re_tracing::profile_scope!("visualizable_entities_per_view");
+                    viewport_ui
+                        .blueprint
+                        .views
+                        .values()
+                        .map(|view| {
+                            let cache_key = VisualizableEntitiesCacheKey::new(
+                                recording_generation.clone(),
+                                blueprint_generation.clone(),
+                                &view.space_origin,
+                            );
+
+                            let visualizable_entities = if let Some(entry) = self.visualizable_entities_cache.get(&view.id)
+                                && entry.key == cache_key
+                            {
+                                PerVisualizer(entry.result.0.clone())
+                            } else {
+                                // This is the expensive operation we're caching
+                                view.class(view_class_registry)
+                                    .determine_visualizable_entities(
+                                        &maybe_visualizable_entities_per_visualizer,
+                                        recording,
+                                        &view_class_registry
+                                            .new_visualizer_collection(view.class_identifier()),
+                                        &view.space_origin,
+                                    )
+                            };
+
+                            self.visualizable_entities_cache.insert(
+                                view.id,
+                                VisualizableEntitiesCacheEntry {
+                                    key: cache_key,
+                                    result: PerVisualizer(visualizable_entities.0.clone()),
+                                },
+                            );
+
+                            (view.id, visualizable_entities)
+                        })
+                        .collect()
+                };
 
                 // Execute the queries for every `View`
-                let query_results = {
+                let query_results: HashMap<_, _> = {
                     re_tracing::profile_scope!("query_results");
                     viewport_ui
                         .blueprint
                         .views
                         .values()
                         .map(|view| {
-                            // TODO(andreas): This needs to be done in a store subscriber that exists per view (instance, not class!).
-                            // Note that right now we determine *all* visualizable entities, not just the queried ones.
-                            // In a store subscriber set this is fine, but on a per-frame basis it's wasteful.
-                            let visualizable_entities = view
-                                .class(view_class_registry)
-                                .determine_visualizable_entities(
-                                    &maybe_visualizable_entities_per_visualizer,
-                                    recording,
-                                    &view_class_registry
-                                        .new_visualizer_collection(view.class_identifier()),
-                                    &view.space_origin,
-                                );
-
-                            (
-                                view.id,
+                            let visualizable_entities = &visualizable_entities_per_view[&view.id];
+
+                            let cache_key = ViewQueryCacheKey::new(
+                                recording_generation.clone(),
+                                blueprint_generation.clone(),
+                                &blueprint_query,
+                            );
+
+                            let result = if let Some(entry) = self.view_query_cache.get(&view.id)
+                                && entry.key == cache_key
+                            {
+                                entry.result.clone()
+                            } else {
                                 view.contents.execute_query(
                                     store_context,
                                     view_class_registry,
                                     &blueprint_query,
-                                    &visualizable_entities,
-                                ),
-                            )
+                                    visualizable_entities,
+                                )
+                            };
+
+                            self.view_query_cache.insert(
+                                view.id,
+                                ViewQueryCacheEntry {
+                                    key: cache_key,
+                                    result: result.clone(),
+                                },
+                            );
+
+                            (view.id, result)
                         })
-                        .collect::<_>()
+                        .collect()
                 };
 
                 let app_blueprint_ctx = AppBlueprintCtx {
@@ -361,6 +482,7 @@ impl AppState {
                     maybe_visualizable_entities_per_visualizer:
                         &maybe_visualizable_entities_per_visualizer,
                     indicated_entities_per_visualizer: &indicated_entities_per_visualizer,
+                    visualizable_entities_per_view: &visualizable_entities_per_view,
                     query_results: &query_results,
                     time_ctrl,
                     blueprint_time_ctrl: blueprint_time_control,
@@ -389,9 +511,16 @@ impl AppState {
                     time_ctrl.timeline(),
                     &maybe_visualizable_entities_per_visualizer,
                     &indicated_entities_per_visualizer,
+                    &visualizable_entities_per_view,
                     view_states,
                 );
 
+                for (view_id, result) in &query_results {
+                    if let Some(entry) = self.view_query_cache.get_mut(view_id) {
+                        entry.result = result.clone();
+                    }
+                }
+
                 // We need to recreate the context to appease the borrow checker. It is a bit annoying, but
                 // it's just a bunch of refs so not really that big of a deal in practice.
                 let ctx = ViewerContext {
@@ -417,6 +546,7 @@ impl AppState {
                     maybe_visualizable_entities_per_visualizer:
                         &maybe_visualizable_entities_per_visualizer,
                     indicated_entities_per_visualizer: &indicated_entities_per_visualizer,
+                    visualizable_entities_per_view: &visualizable_entities_per_view,
                     query_results: &query_results,
                     time_ctrl,
                     blueprint_time_ctrl: blueprint_time_control,
@@ -806,6 +936,7 @@ fn update_overrides(
     active_timeline: &Timeline,
     maybe_visualizable_entities_per_visualizer: &PerVisualizer<MaybeVisualizableEntities>,
     indicated_entities_per_visualizer: &PerVisualizer<IndicatedEntities>,
+    visualizable_entities_per_view: &HashMap<ViewId, PerVisualizer<VisualizableEntities>>,
     view_states: &mut ViewStates,
 ) -> HashMap<ViewId, DataQueryResult> {
     use rayon::iter::{IntoParallelIterator as _, ParallelIterator as _};
@@ -814,6 +945,7 @@ fn update_overrides(
         view: &'a re_viewport_blueprint::ViewBlueprint,
         view_state: &'a dyn re_viewer_context::ViewState,
         query_result: DataQueryResult,
+        visualizable_entities: &'a PerVisualizer<VisualizableEntities>,
     }
 
     for view in viewport_blueprint.views.values() {
@@ -828,10 +960,12 @@ fn update_overrides(
                 let view_state = view_states
                     .get(view.id)
                     .expect("View state should exist, we just called ensure_state_exists on it.");
+                let visualizable_entities = &visualizable_entities_per_view[&view.id];
                 OverridesUpdateTask {
                     view,
                     view_state,
                     query_result,
+                    visualizable_entities,
                 }
             })
         })
@@ -844,24 +978,15 @@ fn update_overrides(
                  view,
                  view_state,
                  mut query_result,
+                 visualizable_entities,
              }| {
-                // TODO(andreas): This needs to be done in a store subscriber that exists per view (instance, not class!).
-                // Note that right now we determine *all* visualizable entities, not just the queried ones.
-                // In a store subscriber set this is fine, but on a per-frame basis it's wasteful.
-                let visualizable_entities = view
-                    .class(view_class_registry)
-                    .determine_visualizable_entities(
-                        maybe_visualizable_entities_per_visualizer,
-                        store_context.recording,
-                        &view_class_registry.new_visualizer_collection(view.class_identifier()),
-                        &view.space_origin,
-                    );
-
+                // Now using the pre-computed visualizable_entities instead of recomputing them.
+                // This avoids the expensive determine_visualizable_entities call that was happening twice per frame.
                 let resolver = re_viewport_blueprint::DataQueryPropertyResolver::new(
                     view,
                     view_class_registry,
                     maybe_visualizable_entities_per_visualizer,
-                    &visualizable_entities,
+                    visualizable_entities,
                     indicated_entities_per_visualizer,
                 );
 
diff --git a/crates/viewer/re_viewer_context/src/viewer_context.rs b/crates/viewer/re_viewer_context/src/viewer_context.rs
@@ -47,6 +47,11 @@ pub struct ViewerContext<'a> {
     ///                 or are we ever interested in a (definitely-)non-visualizable but archetype-matching entity?
     pub indicated_entities_per_visualizer: &'a PerVisualizer<IndicatedEntities>,
 
+    /// For each view, the set of entities that can be visualized by that view.
+    ///
+    /// This is computed once per frame and cached to avoid redundant computation.
+    pub visualizable_entities_per_view: &'a HashMap<ViewId, PerVisualizer<crate::VisualizableEntities>>,
+
     /// All the query results for this frame.
     pub query_results: &'a HashMap<ViewId, DataQueryResult>,
 
diff --git a/crates/viewer/re_viewport_blueprint/src/entity_add_info.rs b/crates/viewer/re_viewport_blueprint/src/entity_add_info.rs
@@ -71,15 +71,16 @@ pub fn create_entity_add_info(
 ) -> IntMap<EntityPath, EntityAddInfo> {
     let mut meta_data: IntMap<EntityPath, EntityAddInfo> = IntMap::default();
 
-    // TODO(andreas): This should be state that is already available because it's part of the view's state.
-    let class = view.class(ctx.view_class_registry());
-    let visualizable_entities = class.determine_visualizable_entities(
-        ctx.maybe_visualizable_entities_per_visualizer,
-        ctx.recording(),
-        &ctx.view_class_registry()
-            .new_visualizer_collection(view.class_identifier()),
-        &view.space_origin,
-    );
+    // Use pre-computed visualizable entities from the context to avoid recomputation.
+    // This was previously computed in app_state.rs and is now shared across the frame.
+    let Some(visualizable_entities) = ctx.visualizable_entities_per_view.get(&view.id) else {
+        // This should not happen in normal operation - if it does, something is wrong with the frame setup
+        re_log::error_once!(
+            "visualizable_entities not found for view {:?} - this is a bug",
+            view.id
+        );
+        return meta_data; // Return empty metadata
+    };
 
     tree.visit_children_recursively(|entity_path| {
         let can_add: CanAddToView =
