offscreen: Release the render target on close and close dropped windows

Cargo.toml

@@ -146,8 +146,9 @@ opt-level = 0
 debug     = true
 
 # `gpui-pre-wgpu` 0.3.8 with its headless renderer behind an `offscreen` feature
-# instead of `test-support` / `bench-support`, and a constructor that takes the
-# wgpu instance (iOS needs Metal). Applications that enable `offscreen` must
-# patch it the same way.
+# instead of `test-support` / `bench-support`, a constructor that takes the wgpu
+# instance (iOS needs Metal), and methods that release its render target and
+# clear its sprite atlas between captures. Applications that enable `offscreen`
+# must patch it the same way.
 [patch.crates-io]
 gpui-pre-wgpu = { path = "vendor/gpui-pre-wgpu" }

README.md

@@ -106,10 +106,17 @@ let image = window.render(size(px(390.), content_height))?; // image::RgbaImage
 window.close(cx);
 ```
 
+Dropping an `OffscreenWindow` without calling `close` (for example when the task
+holding it is cancelled) still closes the window, on the next turn of the main
+thread's executor. The GPU device and pipelines are kept for later captures;
+once the last offscreen window closes, the render target sized to the last
+frame and the glyph and image tiles the offscreen windows uploaded are released.
+
 It needs `gpui-pre-wgpu`'s headless renderer, which upstream only builds for
 tests and benchmarks and which picks a Vulkan/GL device that iOS lacks. This
 repository vendors `gpui-pre-wgpu` 0.3.8 with that renderer behind an
-`offscreen` feature and a constructor that takes the wgpu instance
+`offscreen` feature, a constructor that takes the wgpu instance, and methods
+that release its render target and clear its sprite atlas
 (`vendor/gpui-pre-wgpu`). An application enabling `offscreen` must patch it the
 same way:
 

src/offscreen.rs

@@ -9,17 +9,24 @@
 //!
 //! The renderer (one GPU device plus its pipelines) is created on first use and
 //! kept for the rest of the process on the main thread, so later captures skip
-//! device creation and shader compilation.
-
-use std::{cell::RefCell, rc::Rc, sync::Arc};
+//! device creation and shader compilation. What a capture leaves on the GPU (the
+//! render target, sized to the last frame, and the glyph and image tiles of the
+//! atlas the offscreen windows share) is released once the last offscreen window
+//! closes.
+
+use std::{
+    cell::{Cell, RefCell},
+    rc::Rc,
+    sync::Arc,
+};
 
 use anyhow::{Context as _, Result};
 use gpui::{
-    point, px, App, Bounds, Capslock, DevicePixels, DispatchEventResult, Entity, GpuSpecs,
-    Modifiers, Pixels, PlatformAtlas, PlatformDisplay, PlatformInput, PlatformInputHandler,
-    PlatformWindow, Point, PromptButton, PromptLevel, Render, RequestFrameOptions, Scene, Size,
-    Window, WindowAppearance, WindowBackgroundAppearance, WindowBounds, WindowControlArea,
-    WindowHandle, WindowOptions, WindowParams, WindowVisibility,
+    point, px, AnyWindowHandle, App, AppContext as _, AsyncApp, Bounds, Capslock, DevicePixels,
+    DispatchEventResult, Entity, GpuSpecs, Modifiers, Pixels, PlatformAtlas, PlatformDisplay,
+    PlatformInput, PlatformInputHandler, PlatformWindow, Point, PromptButton, PromptLevel, Render,
+    RequestFrameOptions, Scene, Size, Window, WindowAppearance, WindowBackgroundAppearance,
+    WindowBounds, WindowControlArea, WindowHandle, WindowOptions, WindowParams, WindowVisibility,
 };
 use gpui_wgpu::WgpuHeadlessRenderer;
 pub use image::RgbaImage;
@@ -31,14 +38,21 @@ thread_local! {
     static PENDING: RefCell<Option<Rc<RefCell<State>>>> = const { RefCell::new(None) };
     static RENDERER: RefCell<Option<Rc<RefCell<WgpuHeadlessRenderer>>>> =
         const { RefCell::new(None) };
+    /// Offscreen platform windows GPUI has not dropped yet.
+    static OPEN_WINDOWS: Cell<usize> = const { Cell::new(0) };
 }
 
 /// A GPUI window rendered into images instead of onto the screen.
 ///
-/// Dropping the handle leaves the window open; call [`Self::close`].
+/// Close it with [`Self::close`]. Dropping the handle without closing it (for
+/// example when the task capturing it is cancelled) closes the window on the
+/// next turn of the main thread's executor instead, since `Drop` has no `App`.
 pub struct OffscreenWindow<V> {
     handle: WindowHandle<V>,
     state: Rc<RefCell<State>>,
+    /// Closes the window when the handle is dropped; `None` once [`Self::close`]
+    /// has closed it.
+    close_on_drop: Option<AsyncApp>,
 }
 
 impl<V: Render + 'static> OffscreenWindow<V> {
@@ -81,6 +95,7 @@ impl<V: Render + 'static> OffscreenWindow<V> {
         Ok(Self {
             handle: handle?,
             state,
+            close_on_drop: Some(cx.to_async()),
         })
     }
 
@@ -150,17 +165,39 @@ impl<V: Render + 'static> OffscreenWindow<V> {
     }
 
     /// Closes the window, dropping its view tree.
-    pub fn close(self, cx: &mut App) {
-        let _ = self
+    pub fn close(mut self, cx: &mut App) {
+        // Within the window's own update it can't be reached; the drop below
+        // then closes it from a fresh task instead.
+        if self
             .handle
-            .update(cx, |_, window, _| window.remove_window());
+            .update(cx, |_, window, _| window.remove_window())
+            .is_ok()
+        {
+            self.close_on_drop = None;
+        }
+    }
+}
+
+impl<V> Drop for OffscreenWindow<V> {
+    fn drop(&mut self) {
+        let Some(cx) = self.close_on_drop.take() else {
+            return;
+        };
+        // The handle can be dropped while the `App` is borrowed (inside an update,
+        // or by a task being cancelled), so close the window from a fresh task.
+        let handle: AnyWindowHandle = *self.handle;
+        cx.spawn(async move |cx| {
+            let _ = cx.update_window(handle, |_, window, _| window.remove_window());
+        })
+        .detach();
     }
 }
 
 /// Answers the platform's `open_window` with an offscreen window when
 /// [`OffscreenWindow::open`] asked for one.
 pub(crate) fn take_pending(_params: &WindowParams) -> Option<Box<dyn PlatformWindow>> {
     let state = PENDING.with(|pending| pending.borrow_mut().take())?;
+    OPEN_WINDOWS.with(|open| open.set(open.get() + 1));
     Some(Box::new(OffscreenPlatformWindow(state)))
 }
 
@@ -216,6 +253,30 @@ struct State {
 
 struct OffscreenPlatformWindow(Rc<RefCell<State>>);
 
+/// GPUI drops the platform window with the window itself, after its view tree.
+/// Once no offscreen window is left, nothing holds tiles in the shared atlas, so
+/// it and the render target are released; the device and pipelines stay.
+impl Drop for OffscreenPlatformWindow {
+    fn drop(&mut self) {
+        let open = OPEN_WINDOWS.with(|open| {
+            open.set(open.get().saturating_sub(1));
+            open.get()
+        });
+        if open > 0 {
+            return;
+        }
+        let Some(renderer) = RENDERER.with(|slot| slot.borrow().clone()) else {
+            return;
+        };
+        let Ok(mut renderer) = renderer.try_borrow_mut() else {
+            log::warn!("offscreen renderer busy; keeping its textures");
+            return;
+        };
+        renderer.release_render_target();
+        renderer.clear_sprite_atlas();
+    }
+}
+
 // No native window backs it, so there is no handle to report; `NotSupported`
 // is `raw_window_handle`'s variant for exactly this.
 impl HasWindowHandle for OffscreenPlatformWindow {

vendor/gpui-pre-wgpu/src/wgpu_renderer.rs

@@ -2376,6 +2376,24 @@ impl WgpuHeadlessRenderer {
         self.core.max_texture_size
     }
 
+    /// Drops the textures sized to the last render (the render target and the path
+    /// intermediates), keeping the device, pipelines and atlas. The next render
+    /// recreates them at its own size.
+    #[cfg(feature = "offscreen")]
+    pub fn release_render_target(&mut self) {
+        self.render_target = None;
+        self.core.resources.invalidate_intermediate_textures();
+    }
+
+    /// Drops every glyph and image tile in [`Self::sprite_atlas`], with the bind
+    /// groups that keep their textures alive. Only call it while no window draws
+    /// through the atlas: tiles such a window already holds would dangle.
+    #[cfg(feature = "offscreen")]
+    pub fn clear_sprite_atlas(&mut self) {
+        self.core.resources.atlas_texture_bind_groups.clear();
+        self.core.atlas.clear();
+    }
+
     fn from_headless_context(context: WgpuContext, target_format: wgpu::TextureFormat) -> Self {
         // The device may be shared with earlier renderers; observe only errors raised from
         // here on, including any while this renderer builds its pipelines.
@@ -2798,6 +2816,47 @@ mod tests {
         Ok(())
     }
 
+    #[cfg(all(target_os = "linux", feature = "offscreen"))]
+    #[test]
+    fn headless_renderer_releases_frame_textures() -> anyhow::Result<()> {
+        let mut renderer = WgpuHeadlessRenderer::new()?;
+        let key = gpui::AtlasKey::Image(gpui::RenderImageParams {
+            image_id: gpui::ImageId(1),
+            frame_index: 0,
+        });
+        let pixel = device_size(1, 1);
+        let mut build = || Ok(Some((pixel, std::borrow::Cow::Owned(vec![0, 0, 0, 255]))));
+        renderer
+            .sprite_atlas()
+            .get_or_insert_with(key.clone(), &mut build)?
+            .expect("tile should be created");
+        let mut scene = Scene::default();
+        scene.insert_primitive(solid_quad(2.0, 2.0, 4.0, 3.0, gpui::red()));
+        scene.finish();
+        renderer.render_to_image(&scene, device_size(13, 7))?;
+        assert!(renderer.render_target.is_some());
+        assert!(renderer.core.resources.path_intermediate_texture.is_some());
+
+        renderer.release_render_target();
+        assert!(renderer.render_target.is_none());
+        assert!(renderer.core.resources.path_intermediate_texture.is_none());
+        assert!(renderer.core.resources.path_msaa_texture.is_none());
+
+        renderer.clear_sprite_atlas();
+        assert!(renderer.core.resources.atlas_texture_bind_groups.is_empty());
+        let mut rebuilt = false;
+        renderer.sprite_atlas().get_or_insert_with(key, &mut || {
+            rebuilt = true;
+            build()
+        })?;
+        assert!(rebuilt, "the cleared atlas should rasterize the tile again");
+
+        // Both are recreated on demand.
+        let image = renderer.render_to_image(&scene, device_size(13, 7))?;
+        assert_pixel(&image, 3, 3, RED);
+        Ok(())
+    }
+
     #[cfg(target_os = "linux")]
     #[test]
     fn headless_renderer_rejects_invalid_sizes() -> anyhow::Result<()> {