Merge branch 'main' into antoine/light-arrows

Author: abey79

crates/store/re_chunk/src/chunk.rs

@@ -414,6 +414,32 @@ impl Chunk {
         self
     }
 
+    /// Clones the chunk into a new chunk where all descriptors are untagged.
+    ///
+    /// Only useful as a migration tool while the Rerun ecosystem slowly moves over
+    /// to always using tags for everything.
+    #[doc(hidden)]
+    #[inline]
+    pub fn clone_as_untagged(&self) -> Self {
+        let mut chunk = self.clone();
+
+        let per_component_name = &mut chunk.components;
+        for (component_name, per_desc) in per_component_name.iter_mut() {
+            if per_desc.len() != 1 {
+                // If there are more than one entry, then we're in the land of UB anyway (for now).
+                continue;
+            }
+
+            let untagged_descriptor = ComponentDescriptor::new(*component_name);
+            *per_desc = std::mem::take(per_desc)
+                .into_values()
+                .map(|list_array| (untagged_descriptor.clone(), list_array))
+                .collect();
+        }
+
+        chunk
+    }
+
     /// Clones the chunk into a new chunk where all [`RowId`]s are [`RowId::ZERO`].
     pub fn zeroed(self) -> Self {
         let row_ids = std::iter::repeat(RowId::ZERO)

crates/store/re_chunk_store/src/lib.rs

@@ -39,7 +39,9 @@ pub use self::stats::{ChunkStoreChunkStats, ChunkStoreStats};
 pub use self::store::{
     ChunkStore, ChunkStoreConfig, ChunkStoreGeneration, ChunkStoreHandle, ColumnMetadata,
 };
-pub use self::subscribers::{ChunkStoreSubscriber, ChunkStoreSubscriberHandle};
+pub use self::subscribers::{
+    ChunkStoreSubscriber, ChunkStoreSubscriberHandle, PerStoreChunkSubscriber,
+};
 
 pub(crate) use self::store::ColumnMetadataState;
 

crates/store/re_chunk_store/src/query.rs

@@ -592,7 +592,10 @@ impl ChunkStore {
             })
             .unwrap_or_default();
 
-        debug_assert!(chunks.iter().map(|chunk| chunk.id()).all_unique());
+        debug_assert!(
+            chunks.iter().map(|chunk| chunk.id()).all_unique(),
+            "{entity_path}:{component_name} @ {query:?}",
+        );
 
         chunks
     }

crates/store/re_chunk_store/src/subscribers.rs

@@ -1,4 +1,7 @@
+use ahash::HashMap;
+use itertools::Itertools as _;
 use parking_lot::RwLock;
+use re_log_types::StoreId;
 
 use crate::{ChunkStore, ChunkStoreEvent};
 
@@ -57,6 +60,20 @@ pub trait ChunkStoreSubscriber: std::any::Any + Send + Sync {
     fn on_events(&mut self, events: &[ChunkStoreEvent]);
 }
 
+/// A [`ChunkStoreSubscriber`] that is instantiated for each unique [`StoreId`].
+pub trait PerStoreChunkSubscriber: Send + Sync + Default {
+    /// Arbitrary name for the subscriber.
+    ///
+    /// Does not need to be unique.
+    fn name() -> String;
+
+    /// Get notified of changes happening in a [`ChunkStore`], see [`ChunkStoreSubscriber::on_events`].
+    ///
+    /// Unlike [`ChunkStoreSubscriber::on_events`], all items are guaranteed to have the same [`StoreId`]
+    /// which does not change per invocation.
+    fn on_events<'a>(&mut self, events: impl Iterator<Item = &'a ChunkStoreEvent>);
+}
+
 /// All registered [`ChunkStoreSubscriber`]s.
 static SUBSCRIBERS: once_cell::sync::Lazy<RwLock<Vec<SharedStoreSubscriber>>> =
     once_cell::sync::Lazy::new(|| RwLock::new(Vec::new()));
@@ -143,6 +160,79 @@ impl ChunkStore {
         })
     }
 
+    /// Registers a [`PerStoreChunkSubscriber`] type so it gets automatically notified when data gets added and/or
+    /// removed to/from a [`ChunkStore`].
+    pub fn register_per_store_subscriber<S: PerStoreChunkSubscriber + Default + 'static>(
+    ) -> ChunkStoreSubscriberHandle {
+        let mut subscribers = SUBSCRIBERS.write();
+        subscribers.push(RwLock::new(Box::new(
+            PerStoreStoreSubscriberWrapper::<S>::default(),
+        )));
+        ChunkStoreSubscriberHandle(subscribers.len() as u32 - 1)
+    }
+
+    /// Passes a reference to the downcasted per-store subscriber to the given `FnMut` callback.
+    ///
+    /// Returns `None` if the subscriber doesn't exist or downcasting failed.
+    pub fn with_per_store_subscriber<S: PerStoreChunkSubscriber + 'static, T, F: FnMut(&S) -> T>(
+        ChunkStoreSubscriberHandle(handle): ChunkStoreSubscriberHandle,
+        store_id: &StoreId,
+        mut f: F,
+    ) -> Option<T> {
+        let subscribers = SUBSCRIBERS.read();
+        subscribers.get(handle as usize).and_then(|subscriber| {
+            let subscriber = subscriber.read();
+            subscriber
+                .as_any()
+                .downcast_ref::<PerStoreStoreSubscriberWrapper<S>>()
+                .and_then(|wrapper| wrapper.get(store_id).map(&mut f))
+        })
+    }
+
+    /// Passes a reference to the downcasted per-store subscriber to the given `FnOnce` callback.
+    ///
+    /// Returns `None` if the subscriber doesn't exist or downcasting failed.
+    pub fn with_per_store_subscriber_once<
+        S: PerStoreChunkSubscriber + 'static,
+        T,
+        F: FnOnce(&S) -> T,
+    >(
+        ChunkStoreSubscriberHandle(handle): ChunkStoreSubscriberHandle,
+        store_id: &StoreId,
+        f: F,
+    ) -> Option<T> {
+        let subscribers = SUBSCRIBERS.read();
+        subscribers.get(handle as usize).and_then(|subscriber| {
+            let subscriber = subscriber.read();
+            subscriber
+                .as_any()
+                .downcast_ref::<PerStoreStoreSubscriberWrapper<S>>()
+                .and_then(|wrapper| wrapper.get(store_id).map(f))
+        })
+    }
+
+    /// Passes a mutable reference to the downcasted per-store subscriber to the given callback.
+    ///
+    /// Returns `None` if the subscriber doesn't exist or downcasting failed.
+    pub fn with_per_store_subscriber_mut<
+        S: PerStoreChunkSubscriber + 'static,
+        T,
+        F: FnMut(&mut S) -> T,
+    >(
+        ChunkStoreSubscriberHandle(handle): ChunkStoreSubscriberHandle,
+        store_id: &StoreId,
+        mut f: F,
+    ) -> Option<T> {
+        let subscribers = SUBSCRIBERS.read();
+        subscribers.get(handle as usize).and_then(|subscriber| {
+            let mut subscriber = subscriber.write();
+            subscriber
+                .as_any_mut()
+                .downcast_mut::<PerStoreStoreSubscriberWrapper<S>>()
+                .and_then(|wrapper| wrapper.get_mut(store_id).map(&mut f))
+        })
+    }
+
     /// Called by [`ChunkStore`]'s mutating methods to notify subscriber subscribers of upcoming events.
     pub(crate) fn on_events(events: &[ChunkStoreEvent]) {
         re_tracing::profile_function!();
@@ -154,6 +244,47 @@ impl ChunkStore {
     }
 }
 
+/// Utility that makes a [`PerStoreChunkSubscriber`] a [`ChunkStoreSubscriber`].
+#[derive(Default)]
+struct PerStoreStoreSubscriberWrapper<S: PerStoreChunkSubscriber> {
+    subscribers: HashMap<StoreId, Box<S>>,
+}
+
+impl<S: PerStoreChunkSubscriber + 'static> PerStoreStoreSubscriberWrapper<S> {
+    fn get(&self, store_id: &StoreId) -> Option<&S> {
+        self.subscribers.get(store_id).map(|s| s.as_ref())
+    }
+
+    fn get_mut(&mut self, store_id: &StoreId) -> Option<&mut S> {
+        self.subscribers.get_mut(store_id).map(|s| s.as_mut())
+    }
+}
+
+impl<S: PerStoreChunkSubscriber + 'static> ChunkStoreSubscriber
+    for PerStoreStoreSubscriberWrapper<S>
+{
+    fn name(&self) -> String {
+        S::name()
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+
+    fn on_events(&mut self, events: &[ChunkStoreEvent]) {
+        for (store_id, events) in &events.iter().chunk_by(|e| e.store_id.clone()) {
+            self.subscribers
+                .entry(store_id)
+                .or_default()
+                .on_events(events);
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

crates/store/re_chunk_store/src/writes.rs

@@ -37,21 +37,6 @@ impl ChunkStore {
             return Ok(Vec::new());
         }
 
-        // NOTE: It is very easy to end in a situation where one pulls an old blueprint from
-        // somewhere, and then modifies it at runtime, therefore ending with both tagged and
-        // untagged components in the data.
-        // I'd like to keep this debug assertion a tiny bit longer just in case, so for we just
-        // ignore blueprints.
-        #[cfg(debug_assertions)]
-        if self.id.kind == re_log_types::StoreKind::Recording {
-            for (component_name, per_desc) in chunk.components().iter() {
-                assert!(
-                    per_desc.len() <= 1,
-                    "[DEBUG ONLY] Insert Chunk with multiple values for component named `{component_name}`: this is currently UB",
-                );
-            }
-        }
-
         if !chunk.is_sorted() {
             return Err(ChunkStoreError::UnsortedChunk);
         }
@@ -64,7 +49,31 @@ impl ChunkStore {
 
         self.insert_id += 1;
 
-        let non_compacted_chunk = Arc::clone(chunk); // we'll need it to create the store event
+        let mut chunk = Arc::clone(chunk);
+
+        // We're in a transition period during which the Rerun ecosystem is slowly moving over to tagged data.
+        //
+        // During that time, it is common to end up in situations where the blueprint intermixes both tagged
+        // and untagged components, which invariably leads to undefined behavior.
+        // To prevent that, we just always hot-patch it to untagged, for now.
+        //
+        // Examples:
+        // * An SDK logs a blueprint (tagged), which is then updated by the viewer (which uses untagged log calls).
+        // * Somebody loads an old .rbl from somewhere and starts logging new blueprint data to it.
+        // * Etc.
+        if self.id.kind == re_log_types::StoreKind::Blueprint {
+            chunk = Arc::new(chunk.clone_as_untagged());
+        }
+
+        #[cfg(debug_assertions)]
+        for (component_name, per_desc) in chunk.components().iter() {
+            assert!(
+                per_desc.len() <= 1,
+                "[DEBUG ONLY] Insert Chunk with multiple values for component named `{component_name}`: this is currently UB",
+            );
+        }
+
+        let non_compacted_chunk = Arc::clone(&chunk); // we'll need it to create the store event
 
         let (chunk, diffs) = if chunk.is_static() {
             // Static data: make sure to keep the most recent chunk available for each component column.
@@ -140,7 +149,7 @@ impl ChunkStore {
                     .or_insert_with(|| chunk.id());
             }
 
-            self.static_chunks_stats += ChunkStoreChunkStats::from_chunk(chunk);
+            self.static_chunks_stats += ChunkStoreChunkStats::from_chunk(&chunk);
 
             let mut diffs = vec![ChunkStoreDiff::addition(
                 non_compacted_chunk, /* added */
@@ -194,23 +203,23 @@ impl ChunkStore {
                 }
             }
 
-            (Arc::clone(chunk), diffs)
+            (Arc::clone(&chunk), diffs)
         } else {
             // Temporal data: just index the chunk on every dimension of interest.
             re_tracing::profile_scope!("temporal");
 
             let (elected_chunk, chunk_or_compacted) = {
                 re_tracing::profile_scope!("election");
 
-                let elected_chunk = self.find_and_elect_compaction_candidate(chunk);
+                let elected_chunk = self.find_and_elect_compaction_candidate(&chunk);
 
                 let chunk_or_compacted = if let Some(elected_chunk) = &elected_chunk {
                     let chunk_rowid_min = chunk.row_id_range().map(|(min, _)| min);
                     let elected_rowid_min = elected_chunk.row_id_range().map(|(min, _)| min);
 
                     let mut compacted = if elected_rowid_min < chunk_rowid_min {
                         re_tracing::profile_scope!("concat");
-                        elected_chunk.concatenated(chunk)?
+                        elected_chunk.concatenated(&chunk)?
                     } else {
                         re_tracing::profile_scope!("concat");
                         chunk.concatenated(elected_chunk)?
@@ -233,7 +242,7 @@ impl ChunkStore {
 
                     Arc::new(compacted)
                 } else {
-                    Arc::clone(chunk)
+                    Arc::clone(&chunk)
                 };
 
                 (elected_chunk, chunk_or_compacted)

crates/viewer/re_time_panel/src/data_density_graph.rs

@@ -14,7 +14,7 @@ use re_chunk_store::RangeQuery;
 use re_log_types::{ComponentPath, ResolvedTimeRange, TimeInt, Timeline};
 use re_viewer_context::{Item, TimeControl, UiLayout, ViewerContext};
 
-use crate::recursive_chunks_per_timeline_subscriber::PathRecursiveChunksPerTimeline;
+use crate::recursive_chunks_per_timeline_subscriber::PathRecursiveChunksPerTimelineStoreSubscriber;
 use crate::TimePanelItem;
 
 use super::time_ranges_ui::TimeRangesUi;
@@ -472,27 +472,33 @@ pub fn build_density_graph<'a>(
                 total_num_events,
             )
         } else {
-            PathRecursiveChunksPerTimeline::access(&store.id(), |chunks_per_timeline| {
-                let Some(info) = chunks_per_timeline
-                    .path_recursive_chunks_for_entity_and_timeline(&item.entity_path, &timeline)
-                else {
-                    return Default::default();
-                };
-
-                (
-                    info.recursive_chunks_info
-                        .values()
-                        .map(|info| {
-                            (
-                                info.chunk.clone(),
-                                info.resolved_time_range,
-                                info.num_events,
-                            )
-                        })
-                        .collect(),
-                    info.total_num_events,
-                )
-            })
+            PathRecursiveChunksPerTimelineStoreSubscriber::access(
+                &store.id(),
+                |chunks_per_timeline| {
+                    let Some(info) = chunks_per_timeline
+                        .path_recursive_chunks_for_entity_and_timeline(
+                            &item.entity_path,
+                            &timeline,
+                        )
+                    else {
+                        return Default::default();
+                    };
+
+                    (
+                        info.recursive_chunks_info
+                            .values()
+                            .map(|info| {
+                                (
+                                    info.chunk.clone(),
+                                    info.resolved_time_range,
+                                    info.num_events,
+                                )
+                            })
+                            .collect(),
+                        info.total_num_events,
+                    )
+                },
+            )
             .unwrap_or_default()
         }
     };

crates/viewer/re_time_panel/src/lib.rs

@@ -33,7 +33,7 @@ use re_viewer_context::{
 };
 use re_viewport_blueprint::ViewportBlueprint;
 
-use recursive_chunks_per_timeline_subscriber::PathRecursiveChunksPerTimeline;
+use recursive_chunks_per_timeline_subscriber::PathRecursiveChunksPerTimelineStoreSubscriber;
 use time_axis::TimelineAxis;
 use time_control_ui::TimeControlUi;
 use time_ranges_ui::TimeRangesUi;
@@ -147,7 +147,7 @@ impl TimePanel {
     /// This is implicitly called by [`Self::default`], but may need to be explicitly called in,
     /// e.g., testing context.
     pub fn ensure_registered_subscribers() {
-        PathRecursiveChunksPerTimeline::ensure_registered();
+        PathRecursiveChunksPerTimelineStoreSubscriber::ensure_registered();
     }
 
     pub fn new_blueprint_panel() -> Self {