Fix deadlock when misusing the Caches (#2318)

### What

I introduced this bug in my "refactor" PR
#2303 (oops!), specifically commit
16a9aed

Learnings: returning locks is a bad idea, as the user will hold them
longer than you or they intend, and then you will have a double-lock.

An alternative here is to use more fine-grained locking, one per cache,
and then use a lot more `Arc`s, i.e. returning `Arc<Cache>` from
`Caches::entry`.

### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)

PR Build Summary: https://build.rerun.io/pr/2318

Docs preview: https://rerun.io/preview/637c88f/docs

crates/re_data_ui/src/image.rs

@@ -26,7 +26,9 @@ impl EntityDataUi for Tensor {
     ) {
         re_tracing::profile_function!();
 
-        let decoded = ctx.cache.entry::<TensorDecodeCache>().entry(self.clone());
+        let decoded = ctx
+            .cache
+            .entry(|c: &mut TensorDecodeCache| c.entry(self.clone()));
         match decoded {
             Ok(decoded) => {
                 let annotations = crate::annotations(ctx, query, entity_path);
@@ -58,7 +60,7 @@ fn tensor_ui(
 ) {
     // See if we can convert the tensor to a GPU texture.
     // Even if not, we will show info about the tensor.
-    let tensor_stats = *ctx.cache.entry::<TensorStatsCache>().entry(tensor);
+    let tensor_stats = ctx.cache.entry(|c: &mut TensorStatsCache| c.entry(tensor));
     let debug_name = entity_path.to_string();
     let texture_result = gpu_bridge::tensor_to_gpu(
         ctx.render_ctx,

crates/re_space_view_spatial/src/scene/parts/images.rs

@@ -61,7 +61,7 @@ fn to_textured_rect(
     let Some([height, width, _]) = tensor.image_height_width_channels() else { return None; };
 
     let debug_name = ent_path.to_string();
-    let tensor_stats = *ctx.cache.entry::<TensorStatsCache>().entry(tensor);
+    let tensor_stats = ctx.cache.entry(|c: &mut TensorStatsCache| c.entry(tensor));
 
     match gpu_bridge::tensor_to_gpu(
         ctx.render_ctx,
@@ -206,7 +206,7 @@ impl ImagesPart {
                 return Ok(());
             }
 
-            let tensor = match ctx.cache.entry::<TensorDecodeCache>().entry(tensor) {
+            let tensor = match ctx.cache.entry(|c: &mut TensorDecodeCache| c.entry(tensor)) {
                 Ok(tensor) => tensor,
                 Err(err) => {
                     re_log::warn_once!(
@@ -374,18 +374,18 @@ impl ImagesPart {
         let radius_scale = *properties.backproject_radius_scale.get();
         let point_radius_from_world_depth = radius_scale * pixel_width_from_depth;
 
-        let max_data_value =
-            if let Some((_min, max)) = ctx.cache.entry::<TensorStatsCache>().entry(tensor).range {
-                max as f32
-            } else {
-                // This could only happen for Jpegs, and we should never get here.
-                // TODO(emilk): refactor the code so that we can always calculate a range for the tensor
-                re_log::warn_once!("Couldn't calculate range for a depth tensor!?");
-                match tensor.data {
-                    TensorData::U16(_) => u16::MAX as f32,
-                    _ => 10.0,
-                }
-            };
+        let tensor_stats = ctx.cache.entry(|c: &mut TensorStatsCache| c.entry(tensor));
+        let max_data_value = if let Some((_min, max)) = tensor_stats.range {
+            max as f32
+        } else {
+            // This could only happen for Jpegs, and we should never get here.
+            // TODO(emilk): refactor the code so that we can always calculate a range for the tensor
+            re_log::warn_once!("Couldn't calculate range for a depth tensor!?");
+            match tensor.data {
+                TensorData::U16(_) => u16::MAX as f32,
+                _ => 10.0,
+            }
+        };
 
         Ok(DepthCloud {
             world_from_obj: world_from_obj.into(),

crates/re_space_view_spatial/src/scene/parts/meshes.rs

@@ -42,11 +42,10 @@ impl MeshPart {
 
             let outline_mask_ids = ent_context.highlight.index_outline_mask(instance_key);
 
-            if let Some(mesh) =
-                ctx.cache
-                    .entry::<MeshCache>()
-                    .entry(&ent_path.to_string(), &mesh, ctx.render_ctx)
-            {
+            let mesh = ctx
+                .cache
+                .entry(|c: &mut MeshCache| c.entry(&ent_path.to_string(), &mesh, ctx.render_ctx));
+            if let Some(mesh) = mesh {
                 instances.extend(mesh.mesh_instances.iter().map(move |mesh_instance| {
                     MeshInstance {
                         gpu_mesh: mesh_instance.gpu_mesh.clone(),

crates/re_space_view_spatial/src/ui.rs

@@ -816,20 +816,21 @@ pub fn picking(
 
                                 let tensor_name = instance_path.to_string();
 
-                                match ctx.cache
-                                        .entry::<TensorDecodeCache>()
-                                        .entry(tensor) {
-                                    Ok(decoded_tensor) =>
+                                let decoded_tensor = ctx.cache.entry(|c: &mut TensorDecodeCache| c.entry(tensor));
+                                match decoded_tensor {
+                                    Ok(decoded_tensor) => {
+                                        let tensor_stats = ctx.cache.entry(|c: &mut TensorStatsCache| c.entry(&decoded_tensor));
                                         show_zoomed_image_region(
                                             ctx.render_ctx,
                                             ui,
                                             &decoded_tensor,
-                                            ctx.cache.entry::<TensorStatsCache>().entry(&decoded_tensor),
+                                            &tensor_stats,
                                             &scene.annotation_map.find(&instance_path.entity_path),
                                             decoded_tensor.meter,
                                             &tensor_name,
                                             [coords[0] as _, coords[1] as _],
-                                        ),
+                                        );
+                                    }
                                     Err(err) => re_log::warn_once!(
                                             "Encountered problem decoding tensor at path {tensor_name}: {err}"
                                         ),

crates/re_viewer_context/src/caches.rs

@@ -24,19 +24,18 @@ impl Caches {
         }
     }
 
-    /// Retrieves a cache for reading and writing.
+    /// Accesses a cache for reading and writing.
     ///
     /// Adds the cache lazily if it wasn't already there.
-    pub fn entry<C: Cache + Default>(&self) -> parking_lot::MappedMutexGuard<'_, C> {
-        parking_lot::MutexGuard::map(self.0.lock(), |map| {
-            map.entry(TypeId::of::<C>())
-                .or_insert(Box::<C>::default())
-                .as_any_mut()
-                .downcast_mut::<C>()
-                .expect(
-                    "Downcast failed, this indicates a bug in how `Caches` adds new cache types.",
-                )
-        })
+    pub fn entry<C: Cache + Default, R>(&self, f: impl FnOnce(&mut C) -> R) -> R {
+        f(self
+            .0
+            .lock()
+            .entry(TypeId::of::<C>())
+            .or_insert(Box::<C>::default())
+            .as_any_mut()
+            .downcast_mut::<C>()
+            .expect("Downcast failed, this indicates a bug in how `Caches` adds new cache types."))
     }
 }
 

crates/re_viewer_context/src/tensor/tensor_stats_cache.rs

@@ -7,8 +7,9 @@ use crate::Cache;
 pub struct TensorStatsCache(nohash_hasher::IntMap<TensorId, TensorStats>);
 
 impl TensorStatsCache {
-    pub fn entry(&mut self, tensor: &Tensor) -> &TensorStats {
-        self.0
+    pub fn entry(&mut self, tensor: &Tensor) -> TensorStats {
+        *self
+            .0
             .entry(tensor.tensor_id)
             .or_insert_with(|| TensorStats::new(tensor))
     }

crates/re_viewport/src/view_tensor/scene.rs

@@ -33,7 +33,7 @@ impl SceneTensor {
         tensor: Tensor,
     ) {
         if !tensor.is_shaped_like_an_image() {
-            match ctx.cache.entry::<TensorDecodeCache>().entry(tensor) {
+            match ctx.cache.entry(|c: &mut TensorDecodeCache| c.entry(tensor)) {
                 Ok(tensor) => {
                     let instance_path = InstancePath::instance(ent_path.clone(), InstanceKey(0));
                     self.tensors.insert(instance_path, tensor);

crates/re_viewport/src/view_tensor/ui.rs

@@ -67,15 +67,11 @@ impl ViewTensorState {
             return;
         };
 
+        let tensor_stats = ctx.cache.entry(|c: &mut TensorStatsCache| c.entry(tensor));
         ctx.re_ui
             .selection_grid(ui, "tensor_selection_ui")
             .show(ui, |ui| {
-                tensor_summary_ui_grid_contents(
-                    ctx.re_ui,
-                    ui,
-                    tensor,
-                    ctx.cache.entry::<TensorStatsCache>().entry(tensor),
-                );
+                tensor_summary_ui_grid_contents(ctx.re_ui, ui, tensor, &tensor_stats);
                 self.texture_settings.ui(ctx.re_ui, ui);
                 self.color_mapping.ui(ctx.render_ctx, ctx.re_ui, ui);
             });
@@ -178,7 +174,7 @@ fn paint_tensor_slice(
 ) -> anyhow::Result<(egui::Response, egui::Painter, egui::Rect)> {
     re_tracing::profile_function!();
 
-    let tensor_stats = *ctx.cache.entry::<TensorStatsCache>().entry(tensor);
+    let tensor_stats = ctx.cache.entry(|c: &mut TensorStatsCache| c.entry(tensor));
     let colormapped_texture = super::tensor_slice_to_gpu::colormapped_texture(
         ctx.render_ctx,
         tensor,