Fix VRAM leak: use weak references in ToplevelHandleState and stop capture sessions

Cargo.toml

@@ -148,4 +148,4 @@ cosmic-protocols = { git = "https://github.com/pop-os//cosmic-protocols", branch
 cosmic-client-toolkit = { git = "https://github.com/pop-os//cosmic-protocols", branch = "main" }
 
 [patch.crates-io]
-smithay = { git = "https://github.com/smithay/smithay.git", rev = "3d3f9e3" }
+smithay = { git = "https://github.com/smithay/smithay.git", rev = "599857c" }

src/shell/element/mod.rs

@@ -47,7 +47,7 @@ use std::{
 
 pub mod surface;
 use self::stack::MoveResult;
-pub use self::surface::CosmicSurface;
+pub use self::surface::{CosmicSurface, WeakCosmicSurface};
 pub mod stack;
 pub use self::stack::CosmicStack;
 pub mod window;

src/shell/element/surface.rs

@@ -22,7 +22,7 @@ use smithay::{
         },
     },
     desktop::{
-        PopupManager, Window, WindowSurface, WindowSurfaceType, space::SpaceElement,
+        PopupManager, WeakWindow, Window, WindowSurface, WindowSurfaceType, space::SpaceElement,
         utils::OutputPresentationFeedback,
     },
     input::{
@@ -67,6 +67,25 @@ use crate::{
 #[derive(Debug, Clone, PartialEq, Hash, Eq)]
 pub struct CosmicSurface(pub Window);
 
+/// A weak reference to a [`CosmicSurface`] that does not keep it alive.
+#[derive(Debug, Clone)]
+pub struct WeakCosmicSurface(WeakWindow);
+
+impl WeakCosmicSurface {
+    /// Attempt to upgrade to a strong [`CosmicSurface`] reference.
+    /// Returns `None` if the window has already been dropped.
+    pub fn upgrade(&self) -> Option<CosmicSurface> {
+        self.0.upgrade().map(CosmicSurface)
+    }
+}
+
+impl CosmicSurface {
+    /// Create a weak reference to this surface.
+    pub fn downgrade(&self) -> WeakCosmicSurface {
+        WeakCosmicSurface(self.0.downgrade())
+    }
+}
+
 impl From<ToplevelSurface> for CosmicSurface {
     fn from(s: ToplevelSurface) -> Self {
         CosmicSurface(Window::new_wayland_window(s))

src/shell/mod.rs

@@ -89,7 +89,7 @@ pub mod layout;
 mod seats;
 mod workspace;
 pub mod zoom;
-pub use self::element::{CosmicMapped, CosmicMappedRenderElement, CosmicSurface};
+pub use self::element::{CosmicMapped, CosmicMappedRenderElement, CosmicSurface, WeakCosmicSurface};
 pub use self::seats::*;
 pub use self::workspace::*;
 use self::zoom::{OutputZoomState, ZoomState};

src/wayland/handlers/image_capture_source.rs

@@ -45,7 +45,7 @@ impl ToplevelCaptureSourceHandler for State {
         toplevel: &ForeignToplevelHandle,
     ) {
         let data = match window_from_ext(self, toplevel) {
-            Some(toplevel) => ImageCaptureSourceKind::Toplevel(toplevel.clone()),
+            Some(toplevel) => ImageCaptureSourceKind::Toplevel(toplevel.downgrade()),
             None => ImageCaptureSourceKind::Destroyed,
         };
         source.user_data().insert_if_missing(|| data);

src/wayland/handlers/image_copy_capture/mod.rs

@@ -41,7 +41,9 @@ mod render;
 mod user_data;
 pub use self::render::*;
 use self::user_data::*;
-pub use self::user_data::{FrameHolder, ImageCopySessions, SessionData, SessionHolder};
+pub use self::user_data::{
+    FrameHolder, ImageCopySessions, SessionData, SessionHolder, stop_all_capture_sessions,
+};
 
 impl ImageCopyCaptureHandler for State {
     fn image_copy_capture_state(&mut self) -> &mut ImageCopyCaptureState {
@@ -59,9 +61,9 @@ impl ImageCopyCaptureHandler for State {
                 let output = shell.workspaces.space_for_handle(handle)?.output();
                 constraints_for_output(output, &mut self.backend)
             }
-            ImageCaptureSourceKind::Toplevel(window) => {
-                constraints_for_toplevel(window, &mut self.backend)
-            }
+            ImageCaptureSourceKind::Toplevel(weak) => weak
+                .upgrade()
+                .and_then(|window| constraints_for_toplevel(&window, &mut self.backend)),
             _ => None,
         }
     }
@@ -126,7 +128,11 @@ impl ImageCopyCaptureHandler for State {
                 });
                 workspace.add_session(session);
             }
-            ImageCaptureSourceKind::Toplevel(mut toplevel) => {
+            ImageCaptureSourceKind::Toplevel(weak) => {
+                let Some(mut toplevel) = weak.upgrade() else {
+                    session.stop();
+                    return;
+                };
                 let size = toplevel.geometry().size.to_physical(1);
                 session.user_data().insert_if_missing_threadsafe(|| {
                     Mutex::new(SessionUserData::new(OutputDamageTracker::new(
@@ -234,7 +240,10 @@ impl ImageCopyCaptureHandler for State {
 
                 workspace.add_cursor_session(session);
             }
-            ImageCaptureSourceKind::Toplevel(mut toplevel) => {
+            ImageCaptureSourceKind::Toplevel(weak) => {
+                let Some(mut toplevel) = weak.upgrade() else {
+                    return;
+                };
                 let shell = self.common.shell.read();
                 if let Some(element) = shell.element_for_surface(&toplevel) {
                     if element.has_active_window(&toplevel) {
@@ -280,7 +289,10 @@ impl ImageCopyCaptureHandler for State {
             ImageCaptureSourceKind::Workspace(handle) => {
                 render_workspace_to_buffer(self, session, frame, handle)
             }
-            ImageCaptureSourceKind::Toplevel(toplevel) => {
+            ImageCaptureSourceKind::Toplevel(weak) => {
+                let Some(toplevel) = weak.upgrade() else {
+                    return;
+                };
                 render_window_to_buffer(self, session, frame, &toplevel)
             }
             ImageCaptureSourceKind::Destroyed => unreachable!(),
@@ -328,7 +340,11 @@ impl ImageCopyCaptureHandler for State {
                     workspace.remove_session(&session)
                 }
             }
-            ImageCaptureSourceKind::Toplevel(mut toplevel) => toplevel.remove_session(&session),
+            ImageCaptureSourceKind::Toplevel(weak) => {
+                if let Some(mut toplevel) = weak.upgrade() {
+                    toplevel.remove_session(&session);
+                }
+            }
             ImageCaptureSourceKind::Destroyed => unreachable!(),
         }
     }
@@ -357,8 +373,10 @@ impl ImageCopyCaptureHandler for State {
                     workspace.remove_cursor_session(&session)
                 }
             }
-            ImageCaptureSourceKind::Toplevel(mut toplevel) => {
-                toplevel.remove_cursor_session(&session)
+            ImageCaptureSourceKind::Toplevel(weak) => {
+                if let Some(mut toplevel) = weak.upgrade() {
+                    toplevel.remove_cursor_session(&session);
+                }
             }
             ImageCaptureSourceKind::Destroyed => unreachable!(),
         }

src/wayland/handlers/image_copy_capture/user_data.rs

@@ -10,6 +10,8 @@ use smithay::{
     },
 };
 
+use smithay::utils::user_data::UserDataMap;
+
 use crate::shell::{CosmicSurface, Workspace};
 
 type ImageCopySessionsData = RefCell<ImageCopySessions>;
@@ -33,6 +35,25 @@ pub struct ImageCopySessions {
     cursor_sessions: Vec<CursorSession>,
 }
 
+/// Drop all capture sessions stored in the given `UserDataMap`.
+///
+/// When a toplevel is destroyed, its owned `Session` objects must be dropped
+/// so that `Session::drop()` fires — this fails active frames and sends
+/// `stopped` to the client, releasing GPU buffers.
+///
+/// Smithay doesn't automatically stop sessions when their capture source's
+/// underlying toplevel dies. A cleaner long-term fix would be in smithay's
+/// `ImageCopyCaptureState::cleanup()` — e.g. stopping sessions whose
+/// `source().alive()` is false — so all compositors benefit without manual
+/// session management. For now, we handle it on the cosmic-comp side.
+pub fn stop_all_capture_sessions(user_data: &UserDataMap) {
+    if let Some(data) = user_data.get::<ImageCopySessionsData>() {
+        let mut data = data.borrow_mut();
+        data.sessions.clear();
+        data.cursor_sessions.clear();
+    }
+}
+
 pub trait SessionHolder {
     fn add_session(&mut self, session: Session);
     fn remove_session(&mut self, session: &SessionRef);

src/wayland/handlers/toplevel_info.rs

@@ -3,7 +3,7 @@
 use smithay::utils::{Rectangle, user_data::UserDataMap};
 
 use crate::{
-    shell::CosmicSurface,
+    shell::{CosmicSurface, element::surface::WeakCosmicSurface},
     state::State,
     utils::prelude::Global,
     wayland::protocols::toplevel_info::{
@@ -62,6 +62,16 @@ impl Window for CosmicSurface {
     fn user_data(&self) -> &UserDataMap {
         CosmicSurface::user_data(self)
     }
+
+    fn downgrade(&self) -> WeakCosmicSurface {
+        CosmicSurface::downgrade(self)
+    }
+
+    fn upgrade(weak: &WeakCosmicSurface) -> Option<Self> {
+        weak.upgrade()
+    }
+
+    type Weak = WeakCosmicSurface;
 }
 
 delegate_toplevel_info!(State, CosmicSurface);

src/wayland/protocols/image_capture_source.rs

@@ -2,7 +2,7 @@ use super::{
     workspace::{WorkspaceHandle, WorkspaceHandler},
 };
 use crate::{
-    shell::CosmicSurface,
+    shell::WeakCosmicSurface,
 };
 use cosmic_protocols::image_capture_source::v1::server::{
     zcosmic_workspace_image_capture_source_manager_v1::{
@@ -30,11 +30,17 @@ pub struct WorkspaceImageCaptureSourceManagerGlobalData {
     filter: Box<dyn for<'a> Fn(&'a Client) -> bool + Send + Sync>,
 }
 
-#[derive(Debug, Clone, PartialEq)]
+/// What a capture source targets.
+///
+/// `Toplevel` uses `WeakCosmicSurface` (not a strong `CosmicSurface`) so
+/// that capture sources created by cosmic-workspaces for window thumbnails
+/// don't prevent the underlying window from being dropped. `Output` already
+/// uses `WeakOutput` for the same reason.
+#[derive(Debug, Clone)]
 pub enum ImageCaptureSourceKind {
     Output(WeakOutput),
     Workspace(WorkspaceHandle),
-    Toplevel(CosmicSurface),
+    Toplevel(WeakCosmicSurface),
     Destroyed,
 }