Merge latest changes from pop-os/xdg-desktop-portal-cosmic

Author: meocoder31099

src/screencast_thread.rs

@@ -481,11 +481,14 @@ fn start_stream(
 
     let framerate = 0; // default not limit the frame rate.
 
-    let (width, height) =
-        match block_on(wayland_helper.capture_source_shm(capture_source.clone(), overlay_cursor)) {
-            Some(frame) => (frame.width, frame.height),
-            None => return Err(anyhow::anyhow!("failed to use shm capture to get size")),
-        };
+    let session = wayland_helper.capture_source_session(capture_source, overlay_cursor);
+
+    let Some((width, height)) = block_on(session.wait_for_formats(|formats| formats.buffer_size))
+    else {
+        return Err(anyhow::anyhow!(
+            "failed to get formats for image copy; session stopped"
+        ));
+    };
 
     let dmabuf_helper = wayland_helper.dmabuf();
 
@@ -513,8 +516,6 @@ fn start_stream(
         &mut initial_params,
     )?;
 
-    let session = wayland_helper.capture_source_session(capture_source, overlay_cursor);
-
     let data = StreamData {
         wayland_helper,
         dmabuf_helper,

src/screenshot.rs

@@ -14,6 +14,7 @@ use cosmic::iced_runtime::platform_specific::wayland::layer_surface::{
 use cosmic::iced_winit::commands::layer_surface::{destroy_layer_surface, get_layer_surface};
 use cosmic::widget::horizontal_space;
 use cosmic_client_toolkit::sctk::shell::wlr_layer::{Anchor, KeyboardInteractivity, Layer};
+use futures::stream::{FuturesUnordered, StreamExt};
 use image::RgbaImage;
 use rustix::fd::AsFd;
 use std::borrow::Cow;
@@ -140,20 +141,23 @@ impl Screenshot {
         outputs: &[Output],
     ) -> anyhow::Result<HashMap<String, Vec<(u32, u32, Bytes)>>> {
         let wayland_helper = self.wayland_helper.clone();
-
-        let mut map: HashMap<String, _> = HashMap::with_capacity(outputs.len());
-        for Output { output, name, .. } in outputs {
-            let frame = wayland_helper
-                .capture_output_toplevels_shm(&output, false)
-                .await
-                .into_iter()
-                .filter_map(|img| img.image_transformed().ok())
-                .map(|img| (img.width(), img.height(), img.into_vec().into()))
-                .collect();
-            map.insert(name.clone(), frame);
-        }
-
-        Ok(map)
+        Ok(outputs
+            .iter()
+            .map(move |Output { output, name, .. }| {
+                let wayland_helper = wayland_helper.clone();
+                async move {
+                    let frame = wayland_helper
+                        .capture_output_toplevels_shm(&output, false)
+                        .filter_map(|img| async move { img.image_transformed().ok() })
+                        .map(|img| (img.width(), img.height(), img.into_vec().into()))
+                        .collect()
+                        .await;
+                    (name.clone(), frame)
+                }
+            })
+            .collect::<FuturesUnordered<_>>()
+            .collect::<HashMap<String, _>>()
+            .await)
     }
 
     async fn interactive_output_images(

src/wayland/mod.rs

@@ -14,7 +14,10 @@ use cosmic_client_toolkit::{
     toplevel_info::{ToplevelInfo, ToplevelInfoState},
     workspace::WorkspaceState,
 };
-use futures::channel::oneshot;
+use futures::{
+    channel::oneshot,
+    stream::{FuturesOrdered, Stream, StreamExt},
+};
 use rustix::fd::{FromRawFd, RawFd};
 use std::{
     collections::HashMap,
@@ -158,6 +161,8 @@ impl AppData {
 #[derive(Default)]
 struct SessionState {
     formats: Option<Formats>,
+    stopped: bool,
+    wakers: Vec<std::task::Waker>,
 }
 
 struct SessionInner {
@@ -175,17 +180,32 @@ impl Session {
     }
 
     fn update<F: FnOnce(&mut SessionState)>(&self, f: F) {
-        f(&mut self.0.state.lock().unwrap());
+        let mut state = self.0.state.lock().unwrap();
+        f(&mut state);
+        for waker in std::mem::take(&mut state.wakers) {
+            waker.wake();
+        }
         self.0.condvar.notify_all();
     }
 
-    fn wait_for_formats<T, F: FnMut(&Formats) -> T>(&self, mut cb: F) -> T {
-        let data = self
-            .0
-            .condvar
-            .wait_while(self.0.state.lock().unwrap(), |data| data.formats.is_none())
-            .unwrap();
-        cb(data.formats.as_ref().unwrap())
+    /// Wait for the `Formats` to be sent from the compositor for the stream, and run
+    /// a callback with the state mutex locked.
+    ///
+    /// If formats has not been sent, this will wait until it is received. It returns
+    /// `None` if the server has sent `stopped`.
+    pub async fn wait_for_formats<T, F: FnMut(&Formats) -> T>(&self, mut cb: F) -> Option<T> {
+        std::future::poll_fn(|context| {
+            let mut state = self.0.state.lock().unwrap();
+            if state.stopped {
+                std::task::Poll::Ready(None)
+            } else if let Some(formats) = &state.formats {
+                std::task::Poll::Ready(Some(cb(formats)))
+            } else {
+                state.wakers.push(context.waker().clone());
+                std::task::Poll::Pending
+            }
+        })
+        .await
     }
 
     /// Capture to `wl_buffer`, blocking until capture either succeeds or fails
@@ -289,35 +309,32 @@ impl WaylandHelper {
         }
     }
 
-    pub async fn capture_output_toplevels_shm(
-        &self,
+    pub fn capture_output_toplevels_shm<'a>(
+        &'a self,
         output: &wl_output::WlOutput,
         overlay_cursor: bool,
-    ) -> Vec<ShmImage<OwnedFd>> {
+    ) -> impl Stream<Item = ShmImage<OwnedFd>> + 'a {
         // get the active workspace for this output
         // get the toplevels for that workspace
         // capture each toplevel
 
-        let Some(toplevels) = self
+        let toplevels = self
             .inner
             .output_toplevels
             .lock()
             .unwrap()
             .get(output)
             .cloned()
-        else {
-            return Vec::new();
-        };
+            .unwrap_or_default();
 
-        // TODO is `FuturesOrdered` more optimal?
-        let mut images = Vec::new();
-        for foreign_toplevel in toplevels.into_iter() {
-            let source = CaptureSource::Toplevel(foreign_toplevel.clone());
-            if let Some(image) = self.capture_source_shm(source, overlay_cursor).await {
-                images.push(image);
-            }
-        }
-        images
+        toplevels
+            .into_iter()
+            .map(|foreign_toplevel| {
+                let source = CaptureSource::Toplevel(foreign_toplevel.clone());
+                self.capture_source_shm(source, overlay_cursor)
+            })
+            .collect::<FuturesOrdered<_>>()
+            .filter_map(|x| async { x })
     }
 
     pub fn capture_source_session(&self, source: CaptureSource, overlay_cursor: bool) -> Session {
@@ -364,7 +381,9 @@ impl WaylandHelper {
         let session = self.capture_source_session(source, overlay_cursor);
 
         // TODO: Check that format has been advertised in `Formats`
-        let (width, height) = session.wait_for_formats(|formats| formats.buffer_size);
+        let (width, height) = session
+            .wait_for_formats(|formats| formats.buffer_size)
+            .await?;
 
         let fd = buffer::create_memfd(width, height);
         let buffer =
@@ -577,8 +596,13 @@ impl ScreencopyHandler for AppData {
         }
     }
 
-    fn stopped(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _session: &CaptureSession) {
-        // TODO
+    fn stopped(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, session: &CaptureSession) {
+        if let Some(session) = Session::for_session(session) {
+            session.update(|data| {
+                data.stopped = true;
+            });
+        }
+        // TODO signal users of session in some way?
     }
 
     fn ready(