protocols/screencopy: Make frame/session send stopped/fail on drop

Previously, `Frame` was stored in KMS frame udata, but in some cases the
udata was dropped without a capture happening, and `Frame` did not
implement `Drop`, so `fail` was never sent.

Instead, rename `DropableFrame` to `Frame` and `Frame` to `FrameRef`, so
we can have a single instance of `Frame`, that will send `fail` on drop.
This guarantees either `.success` or `.fail` are send, as long as its
not leaked.

This seems to fix https://github.com/pop-os/cosmic-comp/issues/1305.
xdg-desktop-portal-cosmic prints an error, buy retries (as it should for
an `Unknown` error; though maybe there should be a retry limit) and the
session continues working.

(Not sure if it should be sending `failed`, or queing it with the next
frame so it can send `success` to the client, but this works and is
desirable as a failsafe anyway.)

`Session` and `CursorSession` are similiarly updated.

`.fail()`, `.success()`, and `.stop()` now consume
`Frame`/`Session`/`CursorSession`. So to stop a session, it is now
necessary to call `.remove_session()`, but then simply dropping with
send `.stop()`.

Factoring out some `Request::Capture` handling into a `capture_frame`
function seems to clean up error handling and such a bit.

src/wayland/handlers/screencopy/mod.rs

@@ -27,8 +27,9 @@ use crate::{
     wayland::protocols::{
         image_capture_source::ImageCaptureSourceData,
         screencopy::{
-            delegate_screencopy, BufferConstraints, CursorSession, DmabufConstraints, Frame,
-            ScreencopyHandler, ScreencopyState, Session,
+            delegate_screencopy, BufferConstraints, CursorSession, CursorSessionRef,
+            DmabufConstraints, Frame, FrameRef, ScreencopyHandler, ScreencopyState, Session,
+            SessionRef,
         },
     },
 };
@@ -165,7 +166,6 @@ impl ScreencopyHandler for State {
         match session.source() {
             ImageCaptureSourceData::Output(weak) => {
                 let Some(mut output) = weak.upgrade() else {
-                    session.stop();
                     return;
                 };
 
@@ -196,7 +196,6 @@ impl ScreencopyHandler for State {
             ImageCaptureSourceData::Workspace(handle) => {
                 let mut shell = self.common.shell.write().unwrap();
                 let Some(workspace) = shell.workspaces.space_for_handle_mut(&handle) else {
-                    session.stop();
                     return;
                 };
 
@@ -252,11 +251,10 @@ impl ScreencopyHandler for State {
         }
     }
 
-    fn frame(&mut self, session: Session, frame: Frame) {
+    fn frame(&mut self, session: SessionRef, frame: Frame) {
         match session.source() {
             ImageCaptureSourceData::Output(weak) => {
                 let Some(mut output) = weak.upgrade() else {
-                    session.stop(); // will fail the frame as well
                     return;
                 };
 
@@ -273,7 +271,7 @@ impl ScreencopyHandler for State {
         }
     }
 
-    fn cursor_frame(&mut self, session: CursorSession, frame: Frame) {
+    fn cursor_frame(&mut self, session: CursorSessionRef, frame: Frame) {
         if !session.has_cursor() {
             frame.success(Transform::Normal, Vec::new(), self.common.clock.now());
             return;
@@ -290,18 +288,18 @@ impl ScreencopyHandler for State {
         render_cursor_to_buffer(self, &session, frame, &seat);
     }
 
-    fn frame_aborted(&mut self, frame: Frame) {
+    fn frame_aborted(&mut self, frame: FrameRef) {
         let shell = self.common.shell.read().unwrap();
         for mut output in shell.outputs().cloned() {
-            output.remove_frame(&frame)
+            output.remove_frame(&frame);
         }
     }
 
-    fn session_destroyed(&mut self, session: Session) {
+    fn session_destroyed(&mut self, session: SessionRef) {
         match session.source() {
             ImageCaptureSourceData::Output(weak) => {
                 if let Some(mut output) = weak.upgrade() {
-                    output.remove_session(session);
+                    output.remove_session(&session);
                 }
             }
             ImageCaptureSourceData::Workspace(handle) => {
@@ -313,19 +311,19 @@ impl ScreencopyHandler for State {
                     .workspaces
                     .space_for_handle_mut(&handle)
                 {
-                    workspace.remove_session(session)
+                    workspace.remove_session(&session)
                 }
             }
-            ImageCaptureSourceData::Toplevel(mut toplevel) => toplevel.remove_session(session),
+            ImageCaptureSourceData::Toplevel(mut toplevel) => toplevel.remove_session(&session),
             ImageCaptureSourceData::Destroyed => unreachable!(),
         }
     }
 
-    fn cursor_session_destroyed(&mut self, session: CursorSession) {
+    fn cursor_session_destroyed(&mut self, session: CursorSessionRef) {
         match session.source() {
             ImageCaptureSourceData::Output(weak) => {
                 if let Some(mut output) = weak.upgrade() {
-                    output.remove_cursor_session(session);
+                    output.remove_cursor_session(&session);
                 }
             }
             ImageCaptureSourceData::Workspace(handle) => {
@@ -337,11 +335,11 @@ impl ScreencopyHandler for State {
                     .workspaces
                     .space_for_handle_mut(&handle)
                 {
-                    workspace.remove_cursor_session(session)
+                    workspace.remove_cursor_session(&session)
                 }
             }
             ImageCaptureSourceData::Toplevel(mut toplevel) => {
-                toplevel.remove_cursor_session(session)
+                toplevel.remove_cursor_session(&session)
             }
             ImageCaptureSourceData::Destroyed => unreachable!(),
         }

src/wayland/handlers/screencopy/render.rs

@@ -46,13 +46,13 @@ use crate::{
             constraints_for_output, constraints_for_toplevel, SessionData, SessionUserData,
         },
         protocols::{
-            screencopy::{BufferConstraints, CursorSession, FailureReason, Frame, Session},
+            screencopy::{BufferConstraints, CursorSessionRef, FailureReason, Frame, SessionRef},
             workspace::WorkspaceHandle,
         },
     },
 };
 
-use super::super::data_device::get_dnd_icon;
+use super::{super::data_device::get_dnd_icon, user_data::SessionHolder};
 
 pub fn submit_buffer<R>(
     frame: Frame,
@@ -205,17 +205,16 @@ where
 
 pub fn render_workspace_to_buffer(
     state: &mut State,
-    session: Session,
+    session: SessionRef,
     frame: Frame,
     handle: WorkspaceHandle,
 ) {
     let shell = state.common.shell.read().unwrap();
     let Some(workspace) = shell.workspaces.space_for_handle(&handle) else {
-        session.stop();
         return;
     };
 
-    let output = workspace.output().clone();
+    let mut output = workspace.output().clone();
     let idx = shell.workspaces.idx_for_handle(&output, &handle).unwrap();
     std::mem::drop(shell);
 
@@ -227,7 +226,7 @@ pub fn render_workspace_to_buffer(
     let buffer_size = buffer_dimensions(&buffer).unwrap();
     if mode != Some(buffer_size) {
         let Some(constraints) = constraints_for_output(&output, &mut state.backend) else {
-            session.stop();
+            output.remove_session(&session);
             return;
         };
         session.update_constraints(constraints);
@@ -439,12 +438,12 @@ smithay::render_elements! {
 
 pub fn render_window_to_buffer(
     state: &mut State,
-    session: Session,
+    session: SessionRef,
     frame: Frame,
     toplevel: &CosmicSurface,
 ) {
     if !toplevel.alive() {
-        session.stop();
+        toplevel.clone().remove_session(&session);
         return;
     }
 
@@ -453,7 +452,7 @@ pub fn render_window_to_buffer(
     let buffer_size = buffer_dimensions(&buffer).unwrap();
     if buffer_size != geometry.size.to_buffer(1, Transform::Normal) {
         let Some(constraints) = constraints_for_toplevel(toplevel, &mut state.backend) else {
-            session.stop();
+            toplevel.clone().remove_session(&session);
             return;
         };
         session.update_constraints(constraints);
@@ -663,7 +662,7 @@ pub fn render_window_to_buffer(
 
 pub fn render_cursor_to_buffer(
     state: &mut State,
-    session: &CursorSession,
+    session: &CursorSessionRef,
     frame: Frame,
     seat: &Seat<State>,
 ) {

src/wayland/handlers/screencopy/user_data.rs

@@ -1,9 +1,4 @@
-use std::{
-    cell::RefCell,
-    collections::HashMap,
-    ops::{Deref, DerefMut},
-    sync::Mutex,
-};
+use std::{cell::RefCell, collections::HashMap, sync::Mutex};
 
 use smithay::{
     backend::renderer::{damage::OutputDamageTracker, utils::CommitCounter},
@@ -13,11 +8,13 @@ use smithay::{
 
 use crate::{
     shell::{CosmicSurface, Workspace},
-    wayland::protocols::screencopy::{CursorSession, FailureReason, Frame, Session},
+    wayland::protocols::screencopy::{
+        CursorSession, CursorSessionRef, Frame, FrameRef, Session, SessionRef,
+    },
 };
 
 type ScreencopySessionsData = RefCell<ScreencopySessions>;
-type PendingScreencopyBuffers = Mutex<Vec<(Session, DropableFrame)>>;
+type PendingScreencopyBuffers = Mutex<Vec<(SessionRef, Frame)>>;
 
 pub type SessionData = Mutex<SessionUserData>;
 
@@ -58,24 +55,24 @@ impl SessionUserData {
 
 #[derive(Debug, Default)]
 pub struct ScreencopySessions {
-    sessions: Vec<DropableSession>,
-    cursor_sessions: Vec<DropableCursorSession>,
+    sessions: Vec<Session>,
+    cursor_sessions: Vec<CursorSession>,
 }
 
 pub trait SessionHolder {
     fn add_session(&mut self, session: Session);
-    fn remove_session(&mut self, session: Session);
-    fn sessions(&self) -> Vec<Session>;
+    fn remove_session(&mut self, session: &SessionRef);
+    fn sessions(&self) -> Vec<SessionRef>;
 
     fn add_cursor_session(&mut self, session: CursorSession);
-    fn remove_cursor_session(&mut self, session: CursorSession);
-    fn cursor_sessions(&self) -> Vec<CursorSession>;
+    fn remove_cursor_session(&mut self, session: &CursorSessionRef);
+    fn cursor_sessions(&self) -> Vec<CursorSessionRef>;
 }
 
 pub trait FrameHolder {
-    fn add_frame(&mut self, session: Session, frame: Frame);
-    fn remove_frame(&mut self, frame: &Frame);
-    fn take_pending_frames(&self) -> Vec<(Session, Frame)>;
+    fn add_frame(&mut self, session: SessionRef, frame: Frame);
+    fn remove_frame(&mut self, frame: &FrameRef);
+    fn take_pending_frames(&self) -> Vec<(SessionRef, Frame)>;
 }
 
 impl SessionHolder for Output {
@@ -87,19 +84,19 @@ impl SessionHolder for Output {
             .unwrap()
             .borrow_mut()
             .sessions
-            .push(DropableSession(Some(session)));
+            .push(session);
     }
 
-    fn remove_session(&mut self, session: Session) {
+    fn remove_session(&mut self, session: &SessionRef) {
         self.user_data()
             .get::<ScreencopySessionsData>()
             .unwrap()
             .borrow_mut()
             .sessions
-            .retain(|s| *s != session);
+            .retain(|s| s != session);
     }
 
-    fn sessions(&self) -> Vec<Session> {
+    fn sessions(&self) -> Vec<SessionRef> {
         self.user_data()
             .get::<ScreencopySessionsData>()
             .map_or(Vec::new(), |sessions| {
@@ -107,7 +104,7 @@ impl SessionHolder for Output {
                     .borrow()
                     .sessions
                     .iter()
-                    .flat_map(|s| s.0.clone())
+                    .map(|s| (*s).clone())
                     .collect()
             })
     }
@@ -120,19 +117,19 @@ impl SessionHolder for Output {
             .unwrap()
             .borrow_mut()
             .cursor_sessions
-            .push(DropableCursorSession(Some(session)));
+            .push(session);
     }
 
-    fn remove_cursor_session(&mut self, session: CursorSession) {
+    fn remove_cursor_session(&mut self, session: &CursorSessionRef) {
         self.user_data()
             .get::<ScreencopySessionsData>()
             .unwrap()
             .borrow_mut()
             .cursor_sessions
-            .retain(|s| *s != session);
+            .retain(|s| s != session);
     }
 
-    fn cursor_sessions(&self) -> Vec<CursorSession> {
+    fn cursor_sessions(&self) -> Vec<CursorSessionRef> {
         self.user_data()
             .get::<ScreencopySessionsData>()
             .map_or(Vec::new(), |sessions| {
@@ -140,14 +137,14 @@ impl SessionHolder for Output {
                     .borrow()
                     .cursor_sessions
                     .iter()
-                    .flat_map(|s| s.0.clone())
+                    .map(|s| (*s).clone())
                     .collect()
             })
     }
 }
 
 impl FrameHolder for Output {
-    fn add_frame(&mut self, session: Session, frame: Frame) {
+    fn add_frame(&mut self, session: SessionRef, frame: Frame) {
         self.user_data()
             .insert_if_missing_threadsafe(PendingScreencopyBuffers::default);
         self.user_data()
@@ -155,61 +152,49 @@ impl FrameHolder for Output {
             .unwrap()
             .lock()
             .unwrap()
-            .push((session, DropableFrame(Some(frame))));
+            .push((session, frame));
     }
-    fn remove_frame(&mut self, frame: &Frame) {
+    fn remove_frame(&mut self, frame: &FrameRef) {
         if let Some(pending) = self.user_data().get::<PendingScreencopyBuffers>() {
             pending.lock().unwrap().retain(|(_, f)| f != frame);
         }
     }
-    fn take_pending_frames(&self) -> Vec<(Session, Frame)> {
+    fn take_pending_frames(&self) -> Vec<(SessionRef, Frame)> {
         self.user_data()
             .get::<PendingScreencopyBuffers>()
-            .map(|pending| {
-                pending
-                    .lock()
-                    .unwrap()
-                    .split_off(0)
-                    .into_iter()
-                    .map(|(s, mut f)| (s, f.0.take().unwrap()))
-                    .collect()
-            })
+            .map(|pending| std::mem::take(&mut *pending.lock().unwrap()))
             .unwrap_or_default()
     }
 }
 
 impl SessionHolder for Workspace {
     fn add_session(&mut self, session: Session) {
-        self.screencopy
-            .sessions
-            .push(DropableSession(Some(session)));
+        self.screencopy.sessions.push(session);
     }
 
-    fn remove_session(&mut self, session: Session) {
-        self.screencopy.sessions.retain(|s| *s != session);
+    fn remove_session(&mut self, session: &SessionRef) {
+        self.screencopy.sessions.retain(|s| s != session);
     }
-    fn sessions(&self) -> Vec<Session> {
+    fn sessions(&self) -> Vec<SessionRef> {
         self.screencopy
             .sessions
             .iter()
-            .flat_map(|s| s.0.clone())
+            .map(|s| (*s).clone())
             .collect()
     }
 
     fn add_cursor_session(&mut self, session: CursorSession) {
-        self.screencopy
-            .cursor_sessions
-            .push(DropableCursorSession(Some(session)));
+        self.screencopy.cursor_sessions.push(session);
     }
 
-    fn remove_cursor_session(&mut self, session: CursorSession) {
-        self.screencopy.cursor_sessions.retain(|s| *s != session);
+    fn remove_cursor_session(&mut self, session: &CursorSessionRef) {
+        self.screencopy.cursor_sessions.retain(|s| s != session);
     }
-    fn cursor_sessions(&self) -> Vec<CursorSession> {
+    fn cursor_sessions(&self) -> Vec<CursorSessionRef> {
         self.screencopy
             .cursor_sessions
             .iter()
-            .flat_map(|s| s.0.clone())
+            .map(|s| (*s).clone())
             .collect()
     }
 }
@@ -223,18 +208,18 @@ impl SessionHolder for CosmicSurface {
             .unwrap()
             .borrow_mut()
             .sessions
-            .push(DropableSession(Some(session)));
+            .push(session);
     }
 
-    fn remove_session(&mut self, session: Session) {
+    fn remove_session(&mut self, session: &SessionRef) {
         self.user_data()
             .get::<ScreencopySessionsData>()
             .unwrap()
             .borrow_mut()
             .sessions
-            .retain(|s| *s != session);
+            .retain(|s| s != session);
     }
-    fn sessions(&self) -> Vec<Session> {
+    fn sessions(&self) -> Vec<SessionRef> {
         self.user_data()
             .get::<ScreencopySessionsData>()
             .map_or(Vec::new(), |sessions| {
@@ -242,7 +227,7 @@ impl SessionHolder for CosmicSurface {
                     .borrow()
                     .sessions
                     .iter()
-                    .flat_map(|s| s.0.clone())
+                    .map(|s| (*s).clone())
                     .collect()
             })
     }
@@ -255,19 +240,19 @@ impl SessionHolder for CosmicSurface {
             .unwrap()
             .borrow_mut()
             .cursor_sessions
-            .push(DropableCursorSession(Some(session)));
+            .push(session);
     }
 
-    fn remove_cursor_session(&mut self, session: CursorSession) {
+    fn remove_cursor_session(&mut self, session: &CursorSessionRef) {
         self.user_data()
             .get::<ScreencopySessionsData>()
             .unwrap()
             .borrow_mut()
             .cursor_sessions
-            .retain(|s| *s != session);
+            .retain(|s| s != session);
     }
 
-    fn cursor_sessions(&self) -> Vec<CursorSession> {
+    fn cursor_sessions(&self) -> Vec<CursorSessionRef> {
         self.user_data()
             .get::<ScreencopySessionsData>()
             .map_or(Vec::new(), |sessions| {
@@ -275,86 +260,8 @@ impl SessionHolder for CosmicSurface {
                     .borrow()
                     .cursor_sessions
                     .iter()
-                    .flat_map(|s| s.0.clone())
+                    .map(|s| (*s).clone())
                     .collect()
             })
     }
 }
-
-#[derive(Debug)]
-struct DropableSession(Option<Session>);
-impl Deref for DropableSession {
-    type Target = Session;
-    fn deref(&self) -> &Self::Target {
-        self.0.as_ref().unwrap()
-    }
-}
-impl DerefMut for DropableSession {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        self.0.as_mut().unwrap()
-    }
-}
-impl PartialEq<Session> for DropableSession {
-    fn eq(&self, other: &Session) -> bool {
-        self.0.as_ref().map(|s| s == other).unwrap_or(false)
-    }
-}
-impl Drop for DropableSession {
-    fn drop(&mut self) {
-        if let Some(s) = self.0.take() {
-            s.stop();
-        }
-    }
-}
-
-#[derive(Debug)]
-struct DropableCursorSession(Option<CursorSession>);
-impl Deref for DropableCursorSession {
-    type Target = CursorSession;
-    fn deref(&self) -> &Self::Target {
-        self.0.as_ref().unwrap()
-    }
-}
-impl DerefMut for DropableCursorSession {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        self.0.as_mut().unwrap()
-    }
-}
-impl PartialEq<CursorSession> for DropableCursorSession {
-    fn eq(&self, other: &CursorSession) -> bool {
-        self.0.as_ref().map(|s| s == other).unwrap_or(false)
-    }
-}
-impl Drop for DropableCursorSession {
-    fn drop(&mut self) {
-        if let Some(s) = self.0.take() {
-            s.stop();
-        }
-    }
-}
-
-#[derive(Debug)]
-pub struct DropableFrame(Option<Frame>);
-impl Deref for DropableFrame {
-    type Target = Frame;
-    fn deref(&self) -> &Self::Target {
-        self.0.as_ref().unwrap()
-    }
-}
-impl DerefMut for DropableFrame {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        self.0.as_mut().unwrap()
-    }
-}
-impl PartialEq<Frame> for DropableFrame {
-    fn eq(&self, other: &Frame) -> bool {
-        self.0.as_ref().map(|f| f == other).unwrap_or(false)
-    }
-}
-impl Drop for DropableFrame {
-    fn drop(&mut self) {
-        if let Some(f) = self.0.take() {
-            f.fail(FailureReason::Unknown);
-        }
-    }
-}