Use image cache to stop leaking images (#29452)

This PR fixes several possible memory leaks due to loading images in
markdown files and the image viewer, using the new image cache APIs

TODO: 
- [x] Ensure this didn't break rendering in any of the affected
components.

Release Notes:

- Fixed several image related memory leaks

crates/gpui/examples/image_gallery.rs

@@ -1,9 +1,9 @@
 use futures::FutureExt;
 use gpui::{
     App, AppContext, Application, Asset as _, AssetLogger, Bounds, ClickEvent, Context, ElementId,
-    Entity, HashMapImageCache, ImageAssetLoader, ImageCache, ImageCacheProvider, KeyBinding, Menu,
+    Entity, ImageAssetLoader, ImageCache, ImageCacheProvider, KeyBinding, Menu, MenuItem,
-    MenuItem, SharedString, TitlebarOptions, Window, WindowBounds, WindowOptions, actions, div,
+    RetainAllImageCache, SharedString, TitlebarOptions, Window, WindowBounds, WindowOptions,
-    hash, image_cache, img, prelude::*, px, rgb, size,
+    actions, div, hash, image_cache, img, prelude::*, px, rgb, size,
 };
 use reqwest_client::ReqwestClient;
 use std::{collections::HashMap, sync::Arc};
@@ -14,7 +14,7 @@ struct ImageGallery {
     image_key: String,
     items_count: usize,
     total_count: usize,
-    image_cache: Entity<HashMapImageCache>,
+    image_cache: Entity<RetainAllImageCache>,
 }
 
 impl ImageGallery {
@@ -44,8 +44,8 @@ impl Render for ImageGallery {
             .text_color(gpui::white())
             .child("Manually managed image cache:")
             .child(
-                image_cache(self.image_cache.clone()).child(
                 div()
+                    .image_cache(self.image_cache.clone())
                     .id("main")
                     .font_family(".SystemUIFont")
                     .text_color(gpui::black())
@@ -95,7 +95,7 @@ impl Render for ImageGallery {
                                     .map(|ix| img(format!("{}-{}", image_url, ix)).size_20()),
                             ),
                     ),
-            ))
+            )
             .child(
                 "Automatically managed image cache:"
             )
@@ -282,7 +282,7 @@ fn main() {
                 image_key: "".into(),
                 items_count: IMAGES_IN_GALLERY,
                 total_count: 0,
-                image_cache: HashMapImageCache::new(ctx),
+                image_cache: RetainAllImageCache::new(ctx),
             })
         })
         .unwrap();

crates/gpui/src/elements/div.rs

@@ -40,6 +40,8 @@ use std::{
 use taffy::style::Overflow;
 use util::ResultExt;
 
+use super::ImageCacheProvider;
+
 const DRAG_THRESHOLD: f64 = 2.;
 const TOOLTIP_SHOW_DELAY: Duration = Duration::from_millis(500);
 const HOVERABLE_TOOLTIP_HIDE_DELAY: Duration = Duration::from_millis(500);
@@ -1134,6 +1136,7 @@ pub fn div() -> Div {
         interactivity,
         children: SmallVec::default(),
         prepaint_listener: None,
+        image_cache: None,
     }
 }
 
@@ -1142,6 +1145,7 @@ pub struct Div {
     interactivity: Interactivity,
     children: SmallVec<[AnyElement; 2]>,
     prepaint_listener: Option<Box<dyn Fn(Vec<Bounds<Pixels>>, &mut Window, &mut App) + 'static>>,
+    image_cache: Option<Box<dyn ImageCacheProvider>>,
 }
 
 impl Div {
@@ -1154,6 +1158,12 @@ impl Div {
         self.prepaint_listener = Some(Box::new(listener));
         self
     }
+
+    /// Add an image cache at the location of this div in the element tree.
+    pub fn image_cache(mut self, cache: impl ImageCacheProvider) -> Self {
+        self.image_cache = Some(Box::new(cache));
+        self
+    }
 }
 
 /// A frame state for a `Div` element, which contains layout IDs for its children.
@@ -1199,7 +1209,12 @@ impl Element for Div {
         cx: &mut App,
     ) -> (LayoutId, Self::RequestLayoutState) {
         let mut child_layout_ids = SmallVec::new();
-        let layout_id =
+        let image_cache = self
+            .image_cache
+            .as_mut()
+            .map(|provider| provider.provide(window, cx));
+
+        let layout_id = window.with_image_cache(image_cache, |window| {
             self.interactivity
                 .request_layout(global_id, window, cx, |style, window, cx| {
                     window.with_text_style(style.text_style().cloned(), |window| {
@@ -1210,7 +1225,9 @@ impl Element for Div {
                             .collect::<SmallVec<_>>();
                         window.request_layout(style, child_layout_ids.iter().copied(), cx)
                     })
-                });
+                })
+        });
+
         (layout_id, DivFrameState { child_layout_ids })
     }
 
@@ -1291,18 +1308,25 @@ impl Element for Div {
         window: &mut Window,
         cx: &mut App,
     ) {
-        self.interactivity.paint(
+        let image_cache = self
-            global_id,
+            .image_cache
-            bounds,
+            .as_mut()
-            hitbox.as_ref(),
+            .map(|provider| provider.provide(window, cx));
-            window,
+
-            cx,
+        window.with_image_cache(image_cache, |window| {
-            |_style, window, cx| {
+            self.interactivity.paint(
-                for child in &mut self.children {
+                global_id,
-                    child.paint(window, cx);
+                bounds,
-                }
+                hitbox.as_ref(),
-            },
+                window,
-        );
+                cx,
+                |_style, window, cx| {
+                    for child in &mut self.children {
+                        child.paint(window, cx);
+                    }
+                },
+            )
+        });
     }
 }
 

crates/gpui/src/elements/image_cache.rs

@@ -109,7 +109,7 @@ impl Element for ImageCacheElement {
         cx: &mut App,
     ) -> (LayoutId, Self::RequestLayoutState) {
         let image_cache = self.image_cache_provider.provide(window, cx);
-        window.with_image_cache(image_cache, |window| {
+        window.with_image_cache(Some(image_cache), |window| {
             let child_layout_ids = self
                 .children
                 .iter_mut()
@@ -145,7 +145,7 @@ impl Element for ImageCacheElement {
         cx: &mut App,
     ) {
         let image_cache = self.image_cache_provider.provide(window, cx);
-        window.with_image_cache(image_cache, |window| {
+        window.with_image_cache(Some(image_cache), |window| {
             for child in &mut self.children {
                 child.paint(window, cx);
             }
@@ -217,9 +217,9 @@ impl<T: ImageCache> ImageCacheProvider for Entity<T> {
 }
 
 /// An implementation of ImageCache, that uses an LRU caching strategy to unload images when the cache is full
-pub struct HashMapImageCache(HashMap<u64, ImageCacheItem>);
+pub struct RetainAllImageCache(HashMap<u64, ImageCacheItem>);
 
-impl fmt::Debug for HashMapImageCache {
+impl fmt::Debug for RetainAllImageCache {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("HashMapImageCache")
             .field("num_images", &self.0.len())
@@ -227,11 +227,11 @@ impl fmt::Debug for HashMapImageCache {
     }
 }
 
-impl HashMapImageCache {
+impl RetainAllImageCache {
     /// Create a new image cache.
     #[inline]
     pub fn new(cx: &mut App) -> Entity<Self> {
-        let e = cx.new(|_cx| HashMapImageCache(HashMap::new()));
+        let e = cx.new(|_cx| RetainAllImageCache(HashMap::new()));
         cx.observe_release(&e, |image_cache, cx| {
             for (_, mut item) in std::mem::replace(&mut image_cache.0, HashMap::new()) {
                 if let Some(Ok(image)) = item.get() {
@@ -307,13 +307,39 @@ impl HashMapImageCache {
     }
 }
 
-impl ImageCache for HashMapImageCache {
+impl ImageCache for RetainAllImageCache {
     fn load(
         &mut self,
         resource: &Resource,
         window: &mut Window,
         cx: &mut App,
     ) -> Option<Result<Arc<RenderImage>, ImageCacheError>> {
-        HashMapImageCache::load(self, resource, window, cx)
+        RetainAllImageCache::load(self, resource, window, cx)
+    }
+}
+
+/// Constructs a retain-all image cache that uses the element state associated with the given ID.
+pub fn retain_all(id: impl Into<ElementId>) -> RetainAllImageCacheProvider {
+    RetainAllImageCacheProvider { id: id.into() }
+}
+
+/// A provider struct for creating a retain-all image cache inline
+pub struct RetainAllImageCacheProvider {
+    id: ElementId,
+}
+
+impl ImageCacheProvider for RetainAllImageCacheProvider {
+    fn provide(&mut self, window: &mut Window, cx: &mut App) -> AnyImageCache {
+        window
+            .with_global_id(self.id.clone(), |global_id, window| {
+                window.with_element_state::<Entity<RetainAllImageCache>, _>(
+                    global_id,
+                    |cache, _window| {
+                        let mut cache = cache.unwrap_or_else(|| RetainAllImageCache::new(cx));
+                        (cache.clone(), cache)
+                    },
+                )
+            })
+            .into()
     }
 }

crates/gpui/src/elements/img.rs

@@ -522,6 +522,37 @@ impl ImageSource {
             ImageSource::Image(data) => window.use_asset::<AssetLogger<ImageDecoder>>(data, cx),
         }
     }
+
+    pub(crate) fn get_data(
+        &self,
+        cache: Option<AnyImageCache>,
+        window: &mut Window,
+        cx: &mut App,
+    ) -> Option<Result<Arc<RenderImage>, ImageCacheError>> {
+        match self {
+            ImageSource::Resource(resource) => {
+                if let Some(cache) = cache {
+                    cache.load(resource, window, cx)
+                } else {
+                    window.get_asset::<ImgResourceLoader>(resource, cx)
+                }
+            }
+            ImageSource::Custom(loading_fn) => loading_fn(window, cx),
+            ImageSource::Render(data) => Some(Ok(data.to_owned())),
+            ImageSource::Image(data) => window.get_asset::<AssetLogger<ImageDecoder>>(data, cx),
+        }
+    }
+
+    /// Remove this image source from the asset system
+    pub fn remove_asset(&self, cx: &mut App) {
+        match self {
+            ImageSource::Resource(resource) => {
+                cx.remove_asset::<ImgResourceLoader>(resource);
+            }
+            ImageSource::Custom(_) | ImageSource::Render(_) => {}
+            ImageSource::Image(data) => cx.remove_asset::<AssetLogger<ImageDecoder>>(data),
+        }
+    }
 }
 
 #[derive(Clone)]

crates/gpui/src/platform.rs

@@ -1537,6 +1537,22 @@ impl Image {
             .and_then(|result| result.ok())
     }
 
+    /// Use the GPUI `get_asset` API to make this image renderable
+    pub fn get_render_image(
+        self: Arc<Self>,
+        window: &mut Window,
+        cx: &mut App,
+    ) -> Option<Arc<RenderImage>> {
+        ImageSource::Image(self)
+            .get_data(None, window, cx)
+            .and_then(|result| result.ok())
+    }
+
+    /// Use the GPUI `remove_asset` API to drop this image, if possible.
+    pub fn remove_asset(self: Arc<Self>, cx: &mut App) {
+        ImageSource::Image(self).remove_asset(cx);
+    }
+
     /// Convert the clipboard image to an `ImageData` object.
     pub fn to_image_data(&self, svg_renderer: SvgRenderer) -> Result<Arc<RenderImage>> {
         fn frames_for_image(

crates/gpui/src/window.rs

@@ -2117,6 +2117,16 @@ impl Window {
             None
         })
     }
+
+    /// Asynchronously load an asset, if the asset hasn't finished loading or doesn't exist this will return None.
+    /// Your view will not be re-drawn once the asset has finished loading.
+    ///
+    /// Note that the multiple calls to this method will only result in one `Asset::load` call at a
+    /// time.
+    pub fn get_asset<A: Asset>(&mut self, source: &A::Source, cx: &mut App) -> Option<A::Output> {
+        let (task, _) = cx.fetch_asset::<A>(source);
+        task.clone().now_or_never()
+    }
     /// Obtain the current element offset. This method should only be called during the
     /// prepaint phase of element drawing.
     pub fn element_offset(&self) -> Point<Pixels> {
@@ -2876,14 +2886,18 @@ impl Window {
     }
 
     /// Executes the provided function with the specified image cache.
-    pub(crate) fn with_image_cache<F, R>(&mut self, image_cache: AnyImageCache, f: F) -> R
+    pub fn with_image_cache<F, R>(&mut self, image_cache: Option<AnyImageCache>, f: F) -> R
     where
         F: FnOnce(&mut Self) -> R,
     {
-        self.image_cache_stack.push(image_cache);
+        if let Some(image_cache) = image_cache {
-        let result = f(self);
+            self.image_cache_stack.push(image_cache);
-        self.image_cache_stack.pop();
+            let result = f(self);
-        result
+            self.image_cache_stack.pop();
+            result
+        } else {
+            f(self)
+        }
     }
 
     /// Sets an input handler, such as [`ElementInputHandler`][element_input_handler], which interfaces with the

crates/image_viewer/src/image_viewer.rs

@@ -35,9 +35,19 @@ impl ImageView {
     pub fn new(
         image_item: Entity<ImageItem>,
         project: Entity<Project>,
+        window: &mut Window,
         cx: &mut Context<Self>,
     ) -> Self {
         cx.subscribe(&image_item, Self::on_image_event).detach();
+        cx.on_release_in(window, |this, window, cx| {
+            let image_data = this.image_item.read(cx).image.clone();
+            if let Some(image) = image_data.clone().get_render_image(window, cx) {
+                cx.drop_image(image, None);
+            }
+            image_data.remove_asset(cx);
+        })
+        .detach();
+
         Self {
             image_item,
             project,
@@ -234,7 +244,9 @@ impl SerializableItem for ImageView {
                 .update(cx, |project, cx| project.open_image(project_path, cx))?
                 .await?;
 
-            cx.update(|_, cx| Ok(cx.new(|cx| ImageView::new(image_item, project, cx))))?
+            cx.update(
+                |window, cx| Ok(cx.new(|cx| ImageView::new(image_item, project, window, cx))),
+            )?
         })
     }
 
@@ -365,13 +377,13 @@ impl ProjectItem for ImageView {
         project: Entity<Project>,
         _: &Pane,
         item: Entity<Self::Item>,
-        _: &mut Window,
+        window: &mut Window,
         cx: &mut Context<Self>,
     ) -> Self
     where
         Self: Sized,
     {
-        Self::new(item, project, cx)
+        Self::new(item, project, window, cx)
     }
 }
 

crates/markdown_preview/src/markdown_preview_view.rs

@@ -7,8 +7,8 @@ use editor::scroll::Autoscroll;
 use editor::{Editor, EditorEvent};
 use gpui::{
     App, ClickEvent, Context, Entity, EventEmitter, FocusHandle, Focusable, InteractiveElement,
-    IntoElement, ListState, ParentElement, Render, Styled, Subscription, Task, WeakEntity, Window,
+    IntoElement, ListState, ParentElement, Render, RetainAllImageCache, Styled, Subscription, Task,
-    list,
+    WeakEntity, Window, list,
 };
 use language::LanguageRegistry;
 use settings::Settings;
@@ -30,6 +30,7 @@ const REPARSE_DEBOUNCE: Duration = Duration::from_millis(200);
 
 pub struct MarkdownPreviewView {
     workspace: WeakEntity<Workspace>,
+    image_cache: Entity<RetainAllImageCache>,
     active_editor: Option<EditorState>,
     focus_handle: FocusHandle,
     contents: Option<ParsedMarkdown>,
@@ -264,6 +265,7 @@ impl MarkdownPreviewView {
                 tab_content_text,
                 language_registry,
                 parsing_markdown_task: None,
+                image_cache: RetainAllImageCache::new(cx),
             };
 
             this.set_editor(active_editor, window, cx);
@@ -506,7 +508,9 @@ impl Render for MarkdownPreviewView {
     fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
         let buffer_size = ThemeSettings::get_global(cx).buffer_font_size(cx);
         let buffer_line_height = ThemeSettings::get_global(cx).buffer_line_height;
+
         v_flex()
+            .image_cache(self.image_cache.clone())
             .id("MarkdownPreview")
             .key_context("MarkdownPreview")
             .track_focus(&self.focus_handle(cx))

crates/markdown_preview/src/markdown_renderer.rs

@@ -139,7 +139,6 @@ pub fn render_parsed_markdown(
             .map(|block| render_markdown_block(block, &mut cx)),
     )
 }
-
 pub fn render_markdown_block(block: &ParsedMarkdownElement, cx: &mut RenderContext) -> AnyElement {
     use ParsedMarkdownElement::*;
     match block {

crates/repl/src/notebook/cell.rs

@@ -3,7 +3,9 @@ use std::sync::Arc;
 
 use editor::{Editor, EditorMode, MultiBuffer};
 use futures::future::Shared;
-use gpui::{App, Entity, Hsla, Task, TextStyleRefinement, prelude::*};
+use gpui::{
+    App, Entity, Hsla, RetainAllImageCache, Task, TextStyleRefinement, image_cache, prelude::*,
+};
 use language::{Buffer, Language, LanguageRegistry};
 use markdown_preview::{markdown_parser::parse_markdown, markdown_renderer::render_markdown_block};
 use nbformat::v4::{CellId, CellMetadata, CellType};
@@ -148,6 +150,7 @@ impl Cell {
 
                     MarkdownCell {
                         markdown_parsing_task,
+                        image_cache: RetainAllImageCache::new(cx),
                         languages: languages.clone(),
                         id: id.clone(),
                         metadata: metadata.clone(),
@@ -329,6 +332,7 @@ pub trait RunnableCell: RenderableCell {
 pub struct MarkdownCell {
     id: CellId,
     metadata: CellMetadata,
+    image_cache: Entity<RetainAllImageCache>,
     source: String,
     parsed_markdown: Option<markdown_preview::markdown_elements::ParsedMarkdown>,
     markdown_parsing_task: Task<()>,
@@ -403,12 +407,12 @@ impl Render for MarkdownCell {
                     .child(self.gutter(window, cx))
                     .child(
                         v_flex()
+                            .image_cache(self.image_cache.clone())
                             .size_full()
                             .flex_1()
                             .p_3()
                             .font_ui(cx)
                             .text_size(TextSize::Default.rems(cx))
-                            //
                             .children(parsed.children.iter().map(|child| {
                                 div().relative().child(div().relative().child(
                                     render_markdown_block(child, &mut markdown_render_context),

crates/repl/src/outputs/markdown.rs

@@ -1,5 +1,7 @@
 use anyhow::Result;
-use gpui::{App, ClipboardItem, Context, Entity, Task, Window, div, prelude::*};
+use gpui::{
+    App, ClipboardItem, Context, Entity, RetainAllImageCache, Task, Window, div, prelude::*,
+};
 use language::Buffer;
 use markdown_preview::{
     markdown_elements::ParsedMarkdown, markdown_parser::parse_markdown,
@@ -11,6 +13,7 @@ use crate::outputs::OutputContent;
 
 pub struct MarkdownView {
     raw_text: String,
+    image_cache: Entity<RetainAllImageCache>,
     contents: Option<ParsedMarkdown>,
     parsing_markdown_task: Option<Task<Result<()>>>,
 }
@@ -33,6 +36,7 @@ impl MarkdownView {
 
         Self {
             raw_text: text.clone(),
+            image_cache: RetainAllImageCache::new(cx),
             contents: None,
             parsing_markdown_task: Some(task),
         }
@@ -74,6 +78,7 @@ impl Render for MarkdownView {
             markdown_preview::markdown_renderer::RenderContext::new(None, window, cx);
 
         v_flex()
+            .image_cache(self.image_cache.clone())
             .gap_3()
             .py_4()
             .children(parsed.children.iter().map(|child| {

crates/zed/src/zed.rs

@@ -26,8 +26,8 @@ use git_ui::project_diff::ProjectDiffToolbar;
 use gpui::{
     Action, App, AppContext as _, AsyncWindowContext, Context, DismissEvent, Element, Entity,
     Focusable, KeyBinding, ParentElement, PathPromptOptions, PromptLevel, ReadGlobal, SharedString,
-    Styled, Task, TitlebarOptions, UpdateGlobal, Window, WindowKind, WindowOptions, actions, point,
+    Styled, Task, TitlebarOptions, UpdateGlobal, Window, WindowKind, WindowOptions, actions,
-    px,
+    image_cache, point, px, retain_all,
 };
 use image_viewer::ImageInfo;
 use migrate::{MigrationBanner, MigrationEvent, MigrationNotification, MigrationType};
@@ -1336,7 +1336,7 @@ fn show_markdown_app_notification<F>(
                 let primary_button_on_click = primary_button_on_click.clone();
                 cx.new(move |cx| {
                     MessageNotification::new_from_builder(cx, move |window, cx| {
-                        gpui::div()
+                        image_cache(retain_all("notification-cache"))
                             .text_xs()
                             .child(markdown_preview::markdown_renderer::render_parsed_markdown(
                                 &parsed_markdown.clone(),