✨ Add LRU image cache eviction to headless exporter

backend/scripts/_env

@@ -81,6 +81,7 @@ export PENPOT_INTERNAL_URI=http://localhost:3450
 # defaulting to ../frontend/resources/public/js (render-wasm.js/.wasm).
 # export PENPOT_WASM_HEADLESS=true
 # export PENPOT_WASM_DIR=../frontend/resources/public/js
+# export PENPOT_WASM_IMAGE_CACHE_MB=256
 
 export JAVA_OPTS="\
        -Djava.util.logging.manager=org.apache.logging.log4j.jul.LogManager \

exporter/src/app/config.cljs

@@ -47,7 +47,10 @@
    ;; via the in-process Skia/WASM pipeline. `:wasm-dir` points at the built
    ;; render-wasm artifact (render-wasm.js/.wasm).
    [:wasm-headless {:optional true} :boolean]
-   [:wasm-dir {:optional true} :string]])
+   [:wasm-dir {:optional true} :string]
+   ;; Byte budget (in MB) for the WASM image cache; least-recently-used
+   ;; images are evicted between requests once the store exceeds it.
+   [:wasm-image-cache-mb {:optional true} ::sm/int]])
 
 (def ^:private decode-config
   (sm/decoder schema:config sm/string-transformer))

exporter/src/app/renderer/wasm.cljs

@@ -430,6 +430,14 @@
                               ;; thenable so `p/mcat` doesn't throw "expected thenable".
                               (p/do (on-object (assoc object :path path)))))
                           objects))))))
+       (p/fmap (fn [result]
+                 ;; Trim the image store AFTER the request (never mid-render, so
+                 ;; an image can't disappear under a running export). Images the
+                 ;; next request needs again are simply re-provisioned.
+                 (let [evicted (wasm/evict-images! (cf/get :wasm-image-cache-mb 256))]
+                   (when (pos? evicted)
+                     (l/info :hint "wasm render: evicted cached images" :count evicted)))
+                 result))
        (p/merr (fn [cause]
                  (l/error :hint "wasm render: failed" :cause cause)
                  ;; A panic/abort can leave the shared module's buffer allocated

exporter/src/app/wasm.cljs

@@ -182,6 +182,14 @@
     (.set heap img-u8 (+ ptr 36))
     (h/call module "_store_image")))
 
+(defn evict-images!
+  "Evicts least-recently-used images until the module's image store retains at
+  most `max-mb` megabytes. Called between requests, so an image can never
+  disappear under a running render; evicted images are re-provisioned by any
+  later request that needs them. Returns the number of evicted images."
+  [max-mb]
+  (h/call wasm/internal-module "_evict_images_to_budget" max-mb))
+
 (defn provision-fonts!
   "Resolves and uploads every font needed by `shape-id`. `resolve-font` is an
   injected fn of the family map -> promise of TTF bytes (or nil to skip). This

render-wasm/src/main.rs

@@ -741,6 +741,19 @@ pub extern "C" fn is_image_cached(
     Ok(result)
 }
 
+/// Evicts least-recently-used images until the store retains at most
+/// `max_mb` megabytes of image data. Called by the headless exporter between
+/// requests — never mid-render, so an image can't disappear under a running
+/// export; evicted images are re-provisioned by later requests that need
+/// them. Returns the number of evicted images.
+#[no_mangle]
+#[wasm_error]
+pub extern "C" fn evict_images_to_budget(max_mb: u32) -> Result<u32> {
+    let max_bytes = (max_mb as usize) * 1024 * 1024;
+    let evicted = get_resources().images.evict_to_budget(max_bytes);
+    Ok(evicted as u32)
+}
+
 #[no_mangle]
 #[wasm_error]
 pub extern "C" fn set_shape_svg_raw_content() -> Result<()> {

render-wasm/src/render/images.rs

@@ -6,6 +6,7 @@ use crate::error::Result;
 use crate::get_gpu_state;
 use skia_safe::gpu::{surfaces, Budgeted, DirectContext};
 use skia_safe::{self as skia, Codec, ISize, Size};
+use std::cell::Cell;
 use std::collections::HashMap;
 
 pub type Image = skia::Image;
@@ -66,8 +67,19 @@ enum StoredImage {
     },
 }
 
+struct StoredEntry {
+    image: StoredImage,
+    /// Approximate retained cost: encoded byte length (raw/svg) or the
+    /// decoded RGBA size for images registered from a GL texture.
+    bytes: usize,
+    /// LRU tick; `Cell` so read paths can touch it without `&mut self`.
+    last_used: Cell<u64>,
+}
+
 pub struct ImageStore {
-    images: HashMap<(Uuid, bool), StoredImage>,
+    images: HashMap<(Uuid, bool), StoredEntry>,
+    total_bytes: usize,
+    tick: Cell<u64>,
     /// gpu-only
     context: Option<Box<DirectContext>>,
 }
@@ -212,6 +224,8 @@ impl ImageStore {
         let context = &gpu_state.context;
         Self {
             images: HashMap::with_capacity(2048),
+            total_bytes: 0,
+            tick: Cell::new(0),
             context: Some(Box::new(context.clone())),
         }
     }
@@ -222,10 +236,65 @@ impl ImageStore {
     pub fn new_without_gpu() -> Self {
         Self {
             images: HashMap::with_capacity(16),
+            total_bytes: 0,
+            tick: Cell::new(0),
             context: None,
         }
     }
 
+    /// Bumps the LRU clock and returns the new tick.
+    fn next_tick(&self) -> u64 {
+        let t = self.tick.get() + 1;
+        self.tick.set(t);
+        t
+    }
+
+    fn insert_entry(&mut self, key: (Uuid, bool), image: StoredImage, bytes: usize) {
+        let last_used = Cell::new(self.next_tick());
+        self.total_bytes += bytes;
+        self.images.insert(
+            key,
+            StoredEntry {
+                image,
+                bytes,
+                last_used,
+            },
+        );
+    }
+
+    /// Evicts least-recently-used images until the store retains at most
+    /// `max_bytes`. Meant to be called by the headless exporter *between*
+    /// requests, so an image can never disappear under a running render;
+    /// evicted images are simply re-provisioned by a later request that
+    /// needs them (`is_image_cached` reports them as missing). Returns the
+    /// number of evicted images.
+    pub fn evict_to_budget(&mut self, max_bytes: usize) -> usize {
+        if self.total_bytes <= max_bytes {
+            return 0;
+        }
+
+        // Order the keys once instead of rescanning the whole store for the
+        // minimum on every eviction.
+        let mut keys: Vec<_> = self
+            .images
+            .iter()
+            .map(|(key, entry)| (entry.last_used.get(), *key))
+            .collect();
+        keys.sort_unstable_by_key(|(last_used, _)| *last_used);
+
+        let mut evicted = 0;
+        for (_, key) in keys {
+            if self.total_bytes <= max_bytes {
+                break;
+            }
+            if let Some(entry) = self.images.remove(&key) {
+                self.total_bytes -= entry.bytes;
+                evicted += 1;
+            }
+        }
+        evicted
+    }
+
     pub fn add(
         &mut self,
         id: Uuid,
@@ -239,40 +308,43 @@ impl ImageStore {
         }
 
         let raw_data = image_data.to_vec();
+        let bytes = raw_data.len();
 
         match self.context.as_mut() {
             Some(context) => {
                 if let Some(gpu_image) = decode_image(context, &raw_data) {
-                    self.images.insert(key, StoredImage::Gpu(gpu_image));
+                    self.insert_entry(key, StoredImage::Gpu(gpu_image), bytes);
                 } else if let Some((dom, size)) = parse_svg(&raw_data) {
-                    self.images.insert(
+                    self.insert_entry(
                         key,
                         StoredImage::Svg {
                             dom,
                             size,
                             raster: None,
                         },
+                        bytes,
                     );
                 } else {
                     // The lazy re-decode in `get_internal` only retries raster codecs,
                     // so SVGs that fail to parse here stay raw.
-                    self.images.insert(key, StoredImage::Raw(raw_data));
+                    self.insert_entry(key, StoredImage::Raw(raw_data), bytes);
                 }
             }
             // GPU-free: keep the encoded bytes; decoded on the CPU at draw time.
             // SVGs still get parsed up front since that needs no GPU context.
             None => {
                 if let Some((dom, size)) = parse_svg(&raw_data) {
-                    self.images.insert(
+                    self.insert_entry(
                         key,
                         StoredImage::Svg {
                             dom,
                             size,
                             raster: None,
                         },
+                        bytes,
                     );
                 } else {
-                    self.images.insert(key, StoredImage::Raw(raw_data));
+                    self.insert_entry(key, StoredImage::Raw(raw_data), bytes);
                 }
             }
         }
@@ -303,7 +375,8 @@ impl ImageStore {
             ));
         };
         let image = create_image_from_gl_texture(context, texture_id, width, height)?;
-        self.images.insert(key, StoredImage::Gpu(image));
+        let bytes = (width as usize) * (height as usize) * 4;
+        self.insert_entry(key, StoredImage::Gpu(image), bytes);
 
         Ok(())
     }
@@ -337,7 +410,9 @@ impl ImageStore {
     }
 
     fn decode_raw_cpu_image(&self, id: &Uuid, is_thumbnail: bool) -> Option<Image> {
-        match self.images.get(&(*id, is_thumbnail))? {
+        let entry = self.images.get(&(*id, is_thumbnail))?;
+        entry.last_used.set(self.next_tick());
+        match &entry.image {
             StoredImage::Raw(raw_data) => {
                 let data = unsafe { skia::Data::new_bytes(raw_data) };
                 Image::from_encoded(&data)
@@ -361,7 +436,8 @@ impl ImageStore {
             .images
             .get(&(*id, false))
             .or_else(|| self.images.get(&(*id, true)))?;
-        match entry {
+        entry.last_used.set(self.next_tick());
+        match &entry.image {
             StoredImage::Svg { dom, size, .. } => Some((dom, *size)),
             _ => None,
         }
@@ -369,16 +445,20 @@ impl ImageStore {
 
     fn get_internal(&mut self, id: &Uuid, is_thumbnail: bool) -> Option<&Image> {
         let key = (*id, is_thumbnail);
+        let tick = self.tick.get() + 1;
+        self.tick.set(tick);
         // Use entry API to mutate the HashMap in-place if needed
         if let Some(entry) = self.images.get_mut(&key) {
-            match entry {
+            entry.last_used.set(tick);
+            let image = &mut entry.image;
+            match image {
                 StoredImage::Gpu(ref img) => Some(img),
                 StoredImage::Raw(raw_data) => {
                     let context = self.context.as_mut()?;
                     let gpu_image = decode_image(context, raw_data)?;
-                    *entry = StoredImage::Gpu(gpu_image);
+                    *image = StoredImage::Gpu(gpu_image);
 
-                    if let StoredImage::Gpu(ref img) = entry {
+                    if let StoredImage::Gpu(ref img) = image {
                         Some(img)
                     } else {
                         None