From c88d2b3a7f95a91dcbbc6528d30d93c2e3ec66fe Mon Sep 17 00:00:00 2001
From: Mara Johnson <mjohnson7cd@gmail.com>
Date: Sun, 11 Oct 2026 00:36:38 -0400
Subject: [PATCH] fix(render): prevent render hang and system lag (#2081)

---
 crates/gpu/src/context.rs | 17 ++++++++++++++++-
 crates/gpu/src/tests.rs   | 18 ++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/crates/gpu/src/context.rs b/crates/gpu/src/context.rs
index 788757bf..11f28c70 100644
--- a/crates/gpu/src/context.rs
+++ b/crates/gpu/src/context.rs
@@ -43,6 +43,9 @@ const UPLOAD_BUDGET: usize = 1536 << 20;
 /// Working textures kept for reuse once their last image is dropped (bytes).
 const POOL_BUDGET: usize = 768 << 20;
 
+/// Readback staging buffers kept for reuse across sizes (bytes).
+const STAGING_BUDGET: u64 = 256 << 20;
+
 /// A GPU image: premultiplied RGBA f32 in a texture (cheap to clone; immutable once written).
 /// Working images come from the context's texture pool and go back to it when the last clone
 /// is dropped.
@@ -774,6 +777,12 @@ impl GpuContext {
         self.unavailable.push((entry, "dropped by a test".into()));
     }
 
+    /// Bytes of readback staging buffers kept for reuse (tests).
+    #[cfg(test)]
+    pub(crate) fn staging_bytes(&self) -> u64 {
+        self.staging.lock().map(|s| s.iter().map(|(k, v)| k * v.len() as u64).sum()).unwrap_or(0)
+    }
+
     pub fn can_readback(&self) -> bool {
         self.readbacks.check().is_ok() && (self.can_wait() || self.deferred.is_some())
     }
@@ -1153,9 +1162,15 @@ impl<'g> Enc<'g> {
         }
         out.copy_from_slice(&bytes);
         // Keep a few staging buffers per size (frames read back at the same size every time).
+        // Layers whose bounds change every frame read back at a new size each time: drop the
+        // other sizes once over budget, or mapped host memory grows for the whole render (#2081).
         if let Ok(mut s) = self.g.staging.lock() {
+            let held: u64 = s.iter().map(|(k, v)| k.saturating_mul(v.len() as u64)).sum();
+            if held.saturating_add(size) > STAGING_BUDGET {
+                s.retain(|k, _| *k == size);
+            }
             let v = s.entry(size).or_default();
-            if v.len() < 2 && size <= 256 << 20 {
+            if v.len() < 2 && size <= STAGING_BUDGET {
                 v.push(buf);
             }
         }
diff --git a/crates/gpu/src/tests.rs b/crates/gpu/src/tests.rs
index 150f29ad..784a393d 100644
--- a/crates/gpu/src/tests.rs
+++ b/crates/gpu/src/tests.rs
@@ -839,3 +839,21 @@ fn effects_whose_kernels_failed_render_on_the_cpu() {
         check(&format!("missing kernels {depth:?}"), Some(d), 0.0);
     }
 }
+
+/// A layer whose bounds change every frame reads back at a new size each time: the staging
+/// buffers kept for reuse stay within budget instead of growing for the whole render (#2081).
+#[test]
+fn readback_staging_stays_bounded_across_sizes() {
+    let Some(g) = gpu() else { return };
+    if !g.ctx.can_wait() {
+        return;
+    }
+    let mut e = crate::context::Enc::new(&g.ctx);
+    // 1024 × (256 + i) RGBA f32: ~4 MiB each, ~480 MiB over all sizes.
+    for i in 0..120 {
+        let img = g.ctx.image(1024, 256 + i);
+        assert!(e.download(&img).is_some(), "readback {i}");
+    }
+    let held = g.ctx.staging_bytes();
+    assert!(held <= (256 << 20) + 2 * 1024 * 16 * 376, "staging buffers hold {held} bytes");
+}