Share concurrent cold mask builds (#1653)

Concurrent requests for the same cold combined mask each rasterized and feathered the whole
image. A per-key OnceLock slot now lets one caller build while the others wait for the same
tiles; different keys build independently outside the cache mutex, and the result is admitted
through the byte-budgeted cache with the generation taken at lookup, so a purge mid-build still
wins. The vector mask renders sequentially inside the slot so a Rayon worker can't wait on itself.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

crates/compose/src/masks.rs

@@ -7,7 +7,7 @@
 
 use std::collections::HashMap;
 use std::hash::{Hash, Hasher};
-use std::sync::{Arc, Mutex, OnceLock};
+use std::sync::{Arc, Mutex, OnceLock, Weak};
 
 use photocraft_color::{ColorMode, PixelFormat, SampleType};
 use photocraft_doc::Layer;
@@ -28,6 +28,8 @@ struct CacheEntry {
 #[derive(Default)]
 struct Cache {
     map: HashMap<u64, CacheEntry>,
+    /// Builds in progress, shared by every caller that asks for the same cold mask.
+    in_flight: HashMap<u64, Weak<OnceLock<Surface>>>,
     tick: u64,
     bytes: usize,
     generation: u64,
@@ -83,6 +85,52 @@ const CAPACITY: usize = 64;
 /// fit; larger individual results remain usable but are not retained in the cache.
 pub const CACHE_BUDGET: usize = 512 << 20;
 
+enum CachedMask {
+    Hit(Surface),
+    /// A build slot, and the cache generation it may fill.
+    Pending(Arc<OnceLock<Surface>>, u64),
+}
+
+fn lookup(c: &Mutex<Cache>, k: u64) -> CachedMask {
+    let mut c = c.lock().unwrap_or_else(|e| e.into_inner());
+    if let Some(surface) = c.get(k) {
+        return CachedMask::Hit(surface);
+    }
+    let generation = c.generation;
+    // The callers own each slot, so completed builds do not retain an extra
+    // surface after their last caller returns. Sweep dead weak slots on misses.
+    c.in_flight.retain(|_, slot| slot.strong_count() > 0);
+    if let Some(slot) = c.in_flight.get(&k).and_then(Weak::upgrade) {
+        return CachedMask::Pending(slot, generation);
+    }
+    let slot = Arc::new(OnceLock::new());
+    c.in_flight.insert(k, Arc::downgrade(&slot));
+    CachedMask::Pending(slot, generation)
+}
+
+/// The slot's surface, built once by whichever caller gets there first. `build` returns the
+/// mask and the pixel-mask input its key was made from (kept alive by the cache entry).
+fn finish_mask(c: &Mutex<Cache>, k: u64, slot: &Arc<OnceLock<Surface>>, generation: u64, build: impl FnOnce() -> (Surface, Option<Surface>)) -> Surface {
+    slot.get_or_init(|| {
+        // Never hold the global cache mutex while rasterizing or feathering.
+        let (surface, pixel_input) = build();
+        let bytes = entry_bytes(&surface, pixel_input.as_ref());
+        let mut c = c.lock().unwrap_or_else(|e| e.into_inner());
+        if c.in_flight.get(&k).and_then(Weak::upgrade).is_some_and(|active| Arc::ptr_eq(&active, slot)) {
+            c.insert(k, surface.clone(), pixel_input, bytes, CACHE_BUDGET, generation);
+        }
+        surface
+    })
+    .clone()
+}
+
+fn cached_mask(c: &Mutex<Cache>, k: u64, build: impl FnOnce() -> (Surface, Option<Surface>)) -> Surface {
+    match lookup(c, k) {
+        CachedMask::Hit(surface) => surface,
+        CachedMask::Pending(slot, generation) => finish_mask(c, k, &slot, generation, build),
+    }
+}
+
 fn surface_bytes(surface: &Surface) -> usize {
     surface.tiles().fold(surface.default_bytes().len(), |bytes, (_, tile)| bytes.saturating_add(tile.bytes().len()))
 }
@@ -174,23 +222,20 @@ fn gaussian(v: &mut [f32], w: usize, h: usize, sigma: f32) {
 /// `None` when the layer has no enabled vector mask and no feathered pixel mask (use the pixel
 /// mask directly). Values outside the computed area are the surface's default.
 pub fn combined_mask(layer: &Layer, canvas: Rect) -> Option<Surface> {
+    if !layer.vector_mask.as_ref().is_some_and(|v| v.enabled) && !has_feather(layer) {
+        return None;
+    }
+    Some(cached_mask(cache(), key(layer, canvas), || (build_mask(layer, canvas), layer.mask.as_ref().map(|m| m.surface.clone()))))
+}
+
+fn build_mask(layer: &Layer, canvas: Rect) -> Surface {
     let vm = layer.vector_mask.as_ref().filter(|v| v.enabled);
     let pixel = layer.mask.as_ref().filter(|m| m.enabled);
     let (sv, sp) = (vm.map_or(0.0, |v| feather_sigma(v.feather)), pixel.map_or(0.0, |m| feather_sigma(m.feather)));
-    if vm.is_none() && sp <= 0.0 {
-        return None;
-    }
-    let k = key(layer, canvas);
-    let generation = {
-        let mut c = cache().lock().unwrap_or_else(|e| e.into_inner());
-        if let Some(surface) = c.get(k) {
-            return Some(surface);
-        }
-        c.generation
-    };
+    let compiled = vm.map(photocraft_vector::CompiledVectorMask::new);
     // Far outside the path the vector mask is constant.
     let far = Rect::from_xywh(canvas.x0 - 1_000_000, canvas.y0 - 1_000_000, 1, 1);
-    let v_out = vm.map_or(1.0, |vm| photocraft_vector::vector_mask_values(vm, far)[0]);
+    let v_out = compiled.as_ref().map_or(1.0, |vm| vm.render_sequential(far)[0]);
     let p_def = pixel.map_or(1.0, |m| {
         let d = m.surface.default_pixel().first().copied().unwrap_or(1.0);
         1.0 - m.density * (1.0 - d)
@@ -222,8 +267,10 @@ pub fn combined_mask(layer: &Layer, canvas: Rect) -> Option<Surface> {
     let mut s = Surface::with_default(FORMAT, &[p_def * v_out]);
     if !area.is_empty() {
         let (w, h) = (area.width() as usize, area.height() as usize);
-        let mut v = match vm {
-            Some(vm) => photocraft_vector::vector_mask_values(vm, area),
+        // A slot initializer must not enter Rayon: another tile requesting the
+        // same mask could be stolen onto its stack and wait for its own build.
+        let mut v = match &compiled {
+            Some(vm) => vm.render_sequential(area),
             None => vec![1.0; w * h],
         };
         gaussian(&mut v, w, h, sv);
@@ -237,10 +284,7 @@ pub fn combined_mask(layer: &Layer, canvas: Rect) -> Option<Surface> {
         }
         s.write_region(area, &v);
     }
-    let pixel_input = layer.mask.as_ref().map(|m| m.surface.clone());
-    let bytes = entry_bytes(&s, pixel_input.as_ref());
-    cache().lock().unwrap_or_else(|e| e.into_inner()).insert(k, s.clone(), pixel_input, bytes, CACHE_BUDGET, generation);
-    Some(s)
+    s
 }
 
 #[cfg(test)]
@@ -249,6 +293,164 @@ mod tests {
     use photocraft_doc::LayerMask;
     use photocraft_doc::vector::{Knot, Path, Subpath, VectorMask};
 
+    fn pending(c: &Mutex<Cache>, k: u64) -> Arc<OnceLock<Surface>> {
+        match lookup(c, k) {
+            CachedMask::Pending(slot, _) => slot,
+            CachedMask::Hit(_) => panic!("expected a cold mask"),
+        }
+    }
+
+    /// A built mask (and no pixel-mask input), as `finish_mask`'s build returns it.
+    fn test_surface() -> (Surface, Option<Surface>) {
+        let mut surface = Surface::new(FORMAT);
+        surface.fill_rect(Rect::new(0, 0, 2, 2), &[0.5]);
+        (surface, None)
+    }
+
+    #[cfg(not(target_arch = "wasm32"))]
+    #[test]
+    fn cold_callers_share_one_build_without_retaining_a_second_cache() {
+        use std::sync::Barrier;
+        use std::sync::atomic::{AtomicUsize, Ordering};
+
+        let c = Mutex::new(Cache::default());
+        // Acquire all callers' slots before allowing any builder to finish.
+        let slots: Vec<_> = (0..8).map(|_| pending(&c, 7)).collect();
+        assert!(slots.iter().all(|slot| Arc::ptr_eq(slot, &slots[0])));
+        let barrier = Barrier::new(slots.len());
+        let builds = AtomicUsize::new(0);
+        let outputs = std::thread::scope(|scope| {
+            let handles: Vec<_> = slots
+                .iter()
+                .map(|slot| {
+                    let (c, barrier, builds) = (&c, &barrier, &builds);
+                    scope.spawn(move || {
+                        barrier.wait();
+                        finish_mask(c, 7, slot, 0, || {
+                            builds.fetch_add(1, Ordering::SeqCst);
+                            test_surface()
+                        })
+                    })
+                })
+                .collect();
+            handles.into_iter().map(|handle| handle.join().unwrap()).collect::<Vec<_>>()
+        });
+        assert_eq!(builds.load(Ordering::SeqCst), 1);
+        assert!(outputs.iter().all(|surface| Arc::ptr_eq(surface.tiles().next().unwrap().1, outputs[0].tiles().next().unwrap().1)));
+
+        // Evicting the retained result cannot make existing callers build again.
+        c.lock().unwrap().map.clear();
+        let again = finish_mask(&c, 7, &slots[0], 0, || panic!("rebuilt an active slot"));
+        assert!(Arc::ptr_eq(again.tiles().next().unwrap().1, outputs[0].tiles().next().unwrap().1));
+        drop(slots);
+        let next = pending(&c, 8);
+        assert_eq!(c.lock().unwrap().in_flight.len(), 1);
+        drop(next);
+    }
+
+    #[cfg(not(target_arch = "wasm32"))]
+    #[test]
+    fn distinct_keys_build_while_another_key_is_paused() {
+        use std::sync::mpsc;
+        use std::time::Duration;
+
+        let c = Mutex::new(Cache::default());
+        let (started_tx, started_rx) = mpsc::channel();
+        let (release_tx, release_rx) = mpsc::channel();
+        let (second_tx, second_rx) = mpsc::channel();
+        std::thread::scope(|scope| {
+            let c = &c;
+            let first = scope.spawn(move || {
+                cached_mask(c, 1, || {
+                    started_tx.send(()).unwrap();
+                    release_rx.recv_timeout(Duration::from_secs(10)).unwrap();
+                    test_surface()
+                })
+            });
+            started_rx.recv_timeout(Duration::from_secs(5)).unwrap();
+            let second = scope.spawn(move || {
+                let surface = cached_mask(c, 2, test_surface);
+                second_tx.send(surface.sample_channel(0, 0, 0)).unwrap();
+            });
+            assert_eq!(second_rx.recv_timeout(Duration::from_secs(5)).unwrap(), 0.5);
+            release_tx.send(()).unwrap();
+            first.join().unwrap();
+            second.join().unwrap();
+        });
+    }
+
+    #[test]
+    fn obsolete_in_flight_build_does_not_repopulate_the_cache() {
+        let c = Mutex::new(Cache::default());
+        let old = pending(&c, 1);
+        c.lock().unwrap().in_flight.clear();
+        let current = pending(&c, 1);
+        assert!(!Arc::ptr_eq(&old, &current));
+        let old_surface = finish_mask(&c, 1, &old, 0, test_surface);
+        assert!(c.lock().unwrap().map.is_empty());
+        let current_surface = finish_mask(&c, 1, &current, 0, test_surface);
+        assert!(!Arc::ptr_eq(old_surface.tiles().next().unwrap().1, current_surface.tiles().next().unwrap().1));
+        assert_eq!(c.lock().unwrap().map.len(), 1);
+    }
+
+    #[test]
+    fn a_purge_during_a_shared_build_keeps_the_cache_empty() {
+        let c = Mutex::new(Cache::default());
+        let CachedMask::Pending(slot, generation) = lookup(&c, 1) else { panic!("expected a cold mask") };
+        c.lock().unwrap().clear();
+        finish_mask(&c, 1, &slot, generation, test_surface);
+        assert!(c.lock().unwrap().map.is_empty());
+        assert_eq!(c.lock().unwrap().bytes, 0);
+    }
+
+    #[cfg(not(target_arch = "wasm32"))]
+    #[test]
+    fn parallel_feathered_pixel_mask_returns_the_same_tiles() {
+        let canvas = Rect::new(0, 0, 600, 400);
+        let mut layer = Layer::raster("cold shared mask", PixelFormat::RGBA8);
+        let mut mask = photocraft_doc::LayerMask::reveal_all();
+        mask.surface.fill_rect(canvas, &[0.5]);
+        mask.feather = 8.0;
+        layer.mask = Some(mask);
+        let barrier = std::sync::Barrier::new(8);
+        let outputs = std::thread::scope(|scope| {
+            let handles: Vec<_> = (0..8)
+                .map(|_| {
+                    let (layer, barrier) = (&layer, &barrier);
+                    scope.spawn(move || {
+                        barrier.wait();
+                        combined_mask(layer, canvas).unwrap()
+                    })
+                })
+                .collect();
+            handles.into_iter().map(|handle| handle.join().unwrap()).collect::<Vec<_>>()
+        });
+        assert!(outputs.iter().all(|surface| Arc::ptr_eq(surface.tiles().next().unwrap().1, outputs[0].tiles().next().unwrap().1)));
+    }
+
+    #[cfg(not(target_arch = "wasm32"))]
+    #[test]
+    fn vector_mask_build_inside_a_small_rayon_pool_completes() {
+        use rayon::prelude::*;
+        use std::sync::mpsc;
+        use std::time::Duration;
+
+        let (tx, rx) = mpsc::channel();
+        let handle = std::thread::spawn(move || {
+            let canvas = Rect::new(0, 0, 512, 512);
+            let mut layer = Layer::raster("rayon vector mask", PixelFormat::RGBA8);
+            let mut mask = VectorMask::new(photocraft_vector::shapes::ellipse(8.0, 8.0, 496.0, 496.0));
+            mask.feather = 8.0;
+            layer.vector_mask = Some(mask);
+            let pool = rayon::ThreadPoolBuilder::new().num_threads(2).build().unwrap();
+            let outputs: Vec<_> = pool.install(|| (0..8).into_par_iter().map(|_| combined_mask(&layer, canvas).unwrap()).collect());
+            assert!(outputs.iter().all(|surface| Arc::ptr_eq(surface.tiles().next().unwrap().1, outputs[0].tiles().next().unwrap().1)));
+            tx.send(()).unwrap();
+        });
+        rx.recv_timeout(Duration::from_secs(10)).expect("mask initializer deadlocked inside Rayon");
+        handle.join().unwrap();
+    }
+
     fn one_tile(sample: SampleType) -> Surface {
         let mut surface = Surface::with_default(PixelFormat { mode: ColorMode::Grayscale, sample, alpha: false }, &[1.0]);
         surface.fill_rect(Rect::new(0, 0, 1, 1), &[0.0]);

crates/vector/src/lib.rs

@@ -107,12 +107,33 @@ impl CompiledVectorMask {
             MaskGeometry::Constant(value) => vec![*value; rect.width() as usize * rect.height() as usize],
             MaskGeometry::Path(rasterizer) => rasterizer.render(rect),
         };
+        self.apply_density(&mut v);
+        v
+    }
+
+    /// Effective mask values without entering a thread pool. Use this while
+    /// initializing a shared result that other Rayon workers may wait for.
+    pub fn render_sequential(&self, rect: Rect) -> Vec<f32> {
+        let mut v = match &self.geometry {
+            MaskGeometry::Constant(value) => vec![*value; rect.width() as usize * rect.height() as usize],
+            MaskGeometry::Path(rasterizer) => {
+                let mut v = vec![0.0; rect.width() as usize * rect.height() as usize];
+                if !v.is_empty() {
+                    rasterizer.render_into(rect, &mut v);
+                }
+                v
+            }
+        };
+        self.apply_density(&mut v);
+        v
+    }
+
+    fn apply_density(&self, v: &mut [f32]) {
         if let Some(d) = self.density {
-            for x in &mut v {
+            for x in v {
                 *x = 1.0 - d * (1.0 - *x);
             }
         }
-        v
     }
 }
 

crates/vector/src/tests.rs

@@ -337,6 +337,37 @@ fn coverage_surface_and_vector_mask_values() {
     assert!(vector_mask_values(&empty, Rect::new(0, 0, 4, 4)).iter().all(|x| *x == 0.0));
 }
 
+#[test]
+fn sequential_vector_masks_match_parallel_rendering() {
+    let rect = Rect::new(-3, -7, 510, 250);
+    let mut path = shapes::ellipse(10.25, 8.75, 470.0, 220.0);
+    path.subpaths.push(shapes::rect(200.5, 90.25, 45.0, 50.0).subpaths.remove(0).with_op(PathOp::Subtract));
+    for inverted in [false, true] {
+        for density in [0.0, 0.25, 1.0] {
+            path.inverted = inverted;
+            let mut mask = VectorMask::new(path.clone());
+            mask.density = density;
+            let compiled = CompiledVectorMask::new(&mask);
+            assert_eq!(compiled.render_sequential(rect), compiled.render(rect));
+        }
+    }
+    let compiled = CompiledVectorMask::new(&VectorMask::new(path.clone()));
+    for empty in [Rect::EMPTY, Rect::new(1, 1, 1, 10), Rect::new(1, 1, 10, 1)] {
+        assert_eq!(compiled.render_sequential(empty), compiled.render(empty));
+    }
+    let mut empty = VectorMask::new(Path::default());
+    for inverted in [false, true] {
+        empty.path.inverted = inverted;
+        let compiled = CompiledVectorMask::new(&empty);
+        assert_eq!(compiled.render_sequential(rect), compiled.render(rect));
+    }
+    let mut disabled = VectorMask::new(path);
+    disabled.enabled = false;
+    let compiled = CompiledVectorMask::new(&disabled);
+    assert_eq!(compiled.render_sequential(rect), compiled.render(rect));
+    assert_eq!(compiled.render_sequential(Rect::EMPTY), compiled.render(Rect::EMPTY));
+}
+
 fn iou(a: &[f32], b: &[f32]) -> f64 {
     let (mut i, mut u) = (0.0, 0.0);
     for (x, y) in a.iter().zip(b) {