jay_screenshot: add support for multiple planes

src/cli/screenshot.rs

@@ -11,13 +11,14 @@ use {
         },
         wire::{
             jay_compositor::TakeScreenshot,
-            jay_screenshot::{Dmabuf, Error},
+            jay_screenshot::{Dmabuf, Error, Format, Plane},
         },
     },
     chrono::Local,
     jay_algorithms::qoi::xrgb8888_encode_qoi,
     png::{BitDepth, ColorType, Encoder, SrgbRenderingIntent},
-    std::rc::Rc,
+    std::{cell::RefCell, mem, rc::Rc},
+    uapi::OwnedFd,
 };
 
 pub fn main(global: GlobalArgs, args: ScreenshotArgs) {
@@ -48,16 +49,60 @@ async fn run(screenshot: Rc<Screenshot>) {
         res.push(Err(err.msg.to_owned()));
     });
     Dmabuf::handle(tc, sid, result.clone(), |res, buf| {
-        res.push(Ok(buf));
+        let mut planes = PlaneVec::new();
+        planes.push(DmaBufPlane {
+            offset: buf.offset,
+            stride: buf.stride,
+            fd: buf.fd,
+        });
+        let dmabuf = DmaBuf {
+            id: DmaBufIds::default().next(),
+            width: buf.width as _,
+            height: buf.height as _,
+            format: XRGB8888,
+            modifier: buf.modifier_lo as u64 | ((buf.modifier_hi as u64) << 32),
+            planes,
+        };
+        res.push(Ok(ScreenshotWithDevice {
+            dev: buf.drm_dev,
+            buf: dmabuf,
+        }));
+    });
+    let planes = Rc::new(RefCell::new(PlaneVec::new()));
+    Plane::handle(tc, sid, planes.clone(), |planes, buf| {
+        planes.borrow_mut().push(DmaBufPlane {
+            offset: buf.offset,
+            stride: buf.stride,
+            fd: buf.fd,
+        });
     });
-    let buf = match result.pop().await {
+    Format::handle(
+        tc,
+        sid,
+        (planes.clone(), result.clone()),
+        |(planes, res), buf| {
+            let dmabuf = DmaBuf {
+                id: DmaBufIds::default().next(),
+                width: buf.width as _,
+                height: buf.height as _,
+                format: XRGB8888,
+                modifier: buf.modifier_lo as u64 | ((buf.modifier_hi as u64) << 32),
+                planes: mem::take(&mut *planes.borrow_mut()),
+            };
+            res.push(Ok(ScreenshotWithDevice {
+                dev: buf.drm_dev,
+                buf: dmabuf,
+            }));
+        },
+    );
+    let shot = match result.pop().await {
         Ok(b) => b,
         Err(e) => {
             fatal!("Could not take a screenshot: {}", e);
         }
     };
     let format = screenshot.args.format;
-    let data = buf_to_bytes(&DmaBufIds::default(), &buf, format);
+    let data = buf_to_bytes(&shot.dev, &shot.buf, format);
     let filename = match &screenshot.args.filename {
         Some(f) => f.clone(),
         _ => {
@@ -74,8 +119,13 @@ async fn run(screenshot: Rc<Screenshot>) {
     }
 }
 
-pub fn buf_to_bytes(dma_buf_ids: &DmaBufIds, buf: &Dmabuf, format: ScreenshotFormat) -> Vec<u8> {
-    let drm = match Drm::reopen(buf.drm_dev.raw(), false) {
+pub struct ScreenshotWithDevice {
+    pub dev: Rc<OwnedFd>,
+    pub buf: DmaBuf,
+}
+
+pub fn buf_to_bytes(dev: &Rc<OwnedFd>, buf: &DmaBuf, format: ScreenshotFormat) -> Vec<u8> {
+    let drm = match Drm::reopen(dev.raw(), false) {
         Ok(drm) => drm,
         Err(e) => {
             fatal!("Could not open the drm device: {}", ErrorFmt(e));
@@ -87,21 +137,7 @@ pub fn buf_to_bytes(dma_buf_ids: &DmaBufIds, buf: &Dmabuf, format: ScreenshotFor
             fatal!("Could not create a gbm device: {}", ErrorFmt(e));
         }
     };
-    let mut planes = PlaneVec::new();
-    planes.push(DmaBufPlane {
-        offset: buf.offset,
-        stride: buf.stride,
-        fd: buf.fd.clone(),
-    });
-    let dmabuf = DmaBuf {
-        id: dma_buf_ids.next(),
-        width: buf.width as _,
-        height: buf.height as _,
-        format: XRGB8888,
-        modifier: (buf.modifier_hi as u64) << 32 | (buf.modifier_lo as u64),
-        planes,
-    };
-    let bo = match gbm.import_dmabuf(&dmabuf, GBM_BO_USE_LINEAR | GBM_BO_USE_RENDERING) {
+    let bo = match gbm.import_dmabuf(&buf, GBM_BO_USE_LINEAR | GBM_BO_USE_RENDERING) {
         Ok(bo) => Rc::new(bo),
         Err(e) => {
             fatal!("Could not import screenshot dmabuf: {}", ErrorFmt(e));
@@ -115,7 +151,12 @@ pub fn buf_to_bytes(dma_buf_ids: &DmaBufIds, buf: &Dmabuf, format: ScreenshotFor
     };
     let data = unsafe { bo_map.data() };
     if format == ScreenshotFormat::Qoi {
-        return xrgb8888_encode_qoi(data, buf.width, buf.height, bo_map.stride() as u32);
+        return xrgb8888_encode_qoi(
+            data,
+            buf.width as _,
+            buf.height as _,
+            bo_map.stride() as u32,
+        );
     }
 
     let mut out = vec![];
@@ -128,7 +169,7 @@ pub fn buf_to_bytes(dma_buf_ids: &DmaBufIds, buf: &Dmabuf, format: ScreenshotFor
                 image_data.extend_from_slice(&[pixel[2], pixel[1], pixel[0], 255])
             }
         }
-        let mut encoder = Encoder::new(&mut out, buf.width, buf.height);
+        let mut encoder = Encoder::new(&mut out, buf.width as _, buf.height as _);
         encoder.set_color(ColorType::Rgba);
         encoder.set_depth(BitDepth::Eight);
         encoder.set_srgb(SrgbRenderingIntent::Perceptual);

src/ifs/jay_compositor.rs

@@ -13,7 +13,7 @@ use {
             jay_randr::JayRandr,
             jay_render_ctx::JayRenderCtx,
             jay_screencast::JayScreencast,
-            jay_screenshot::JayScreenshot,
+            jay_screenshot::{JayScreenshot, PLANES_SINCE},
             jay_seat_events::JaySeatEvents,
             jay_select_toplevel::{JaySelectToplevel, JayToplevelSelector},
             jay_select_workspace::{JaySelectWorkspace, JayWorkspaceSelector},
@@ -69,7 +69,7 @@ impl Global for JayCompositorGlobal {
     }
 
     fn version(&self) -> u32 {
-        5
+        6
     }
 
     fn required_caps(&self) -> ClientCaps {
@@ -111,29 +111,41 @@ impl JayCompositor {
             id,
             client: self.client.clone(),
             tracker: Default::default(),
+            version: self.version,
+            bo: Cell::new(None),
         });
         track!(self.client, ss);
         self.client.add_client_obj(&ss)?;
         match take_screenshot(&self.client.state, include_cursor) {
             Ok(s) => {
                 let dmabuf = s.bo.dmabuf();
-                let plane = &dmabuf.planes[0];
-                ss.send_dmabuf(
-                    &s.drm,
-                    &plane.fd,
-                    dmabuf.width,
-                    dmabuf.height,
-                    plane.offset,
-                    plane.stride,
-                    dmabuf.modifier,
-                );
+                if ss.version < PLANES_SINCE {
+                    let plane = &dmabuf.planes[0];
+                    ss.send_dmabuf(
+                        &s.drm,
+                        &plane.fd,
+                        dmabuf.width,
+                        dmabuf.height,
+                        plane.offset,
+                        plane.stride,
+                        dmabuf.modifier,
+                    );
+                } else {
+                    for plane in &dmabuf.planes {
+                        ss.send_plane(&plane);
+                    }
+                    ss.send_format(&s.drm, dmabuf);
+                }
+                ss.bo.set(Some(s.bo));
             }
             Err(e) => {
                 let msg = ErrorFmt(e).to_string();
                 ss.send_error(&msg);
             }
         }
-        self.client.remove_obj(ss.deref())?;
+        if ss.version < PLANES_SINCE {
+            self.client.remove_obj(ss.deref())?;
+        }
         Ok(())
     }
 }

src/ifs/jay_screenshot.rs

@@ -1,18 +1,26 @@
 use {
     crate::{
-        client::Client,
+        client::{Client, ClientError},
         leaks::Tracker,
         object::{Object, Version},
+        video::{
+            dmabuf::{DmaBuf, DmaBufPlane},
+            gbm::GbmBo,
+        },
         wire::{jay_screenshot::*, JayScreenshotId},
     },
-    std::{convert::Infallible, rc::Rc},
+    std::{cell::Cell, rc::Rc},
     uapi::OwnedFd,
 };
 
+pub const PLANES_SINCE: Version = Version(6);
+
 pub struct JayScreenshot {
     pub id: JayScreenshotId,
     pub client: Rc<Client>,
     pub tracker: Tracker<Self>,
+    pub version: Version,
+    pub bo: Cell<Option<GbmBo>>,
 }
 
 impl JayScreenshot {
@@ -45,10 +53,36 @@ impl JayScreenshot {
             msg,
         });
     }
+
+    pub fn send_plane(&self, plane: &DmaBufPlane) {
+        self.client.event(Plane {
+            self_id: self.id,
+            fd: plane.fd.clone(),
+            offset: plane.offset,
+            stride: plane.stride,
+        });
+    }
+
+    pub fn send_format(&self, drm_dev: &Rc<OwnedFd>, dmabuf: &DmaBuf) {
+        self.client.event(Format {
+            self_id: self.id,
+            drm_dev: drm_dev.clone(),
+            format: dmabuf.format.drm,
+            width: dmabuf.width,
+            height: dmabuf.height,
+            modifier_lo: dmabuf.modifier as u32,
+            modifier_hi: (dmabuf.modifier >> 32) as u32,
+        });
+    }
 }
 
 impl JayScreenshotRequestHandler for JayScreenshot {
-    type Error = Infallible;
+    type Error = ClientError;
+
+    fn destroy(&self, _req: Destroy, _slf: &Rc<Self>) -> Result<(), Self::Error> {
+        self.client.remove_obj(self)?;
+        Ok(())
+    }
 }
 
 object_base! {

src/it/test_client.rs

@@ -25,7 +25,7 @@ use {
 
 pub struct TestClient {
     pub run: Rc<TestRun>,
-    pub server: Rc<Client>,
+    pub _server: Rc<Client>,
     pub tran: Rc<TestTransport>,
     pub registry: Rc<TestRegistry>,
     pub jc: Rc<TestJayCompositor>,
@@ -92,12 +92,8 @@ impl TestClient {
     }
 
     pub async fn take_screenshot(&self, include_cursor: bool) -> Result<Vec<u8>, TestError> {
-        let dmabuf = self.jc.take_screenshot(include_cursor).await?;
-        let qoi = buf_to_bytes(
-            &self.server.state.dma_buf_ids,
-            &dmabuf,
-            ScreenshotFormat::Qoi,
-        );
+        let (shot, _ts) = self.jc.take_screenshot(include_cursor).await?;
+        let qoi = buf_to_bytes(&shot.dev, &shot.buf, ScreenshotFormat::Qoi);
         Ok(qoi)
     }
 

src/it/test_ifs/test_jay_compositor.rs

@@ -1,5 +1,6 @@
 use {
     crate::{
+        cli::screenshot::ScreenshotWithDevice,
         client::ClientId,
         it::{
             test_error::{TestError, TestResult},
@@ -11,7 +12,6 @@ use {
         utils::{buffd::MsgParser, cell_ext::CellExt},
         wire::{
             jay_compositor::{self, *},
-            jay_screenshot::Dmabuf,
             JayCompositorId,
         },
     },
@@ -49,10 +49,15 @@ impl TestJayCompositor {
         Ok(())
     }
 
-    pub async fn take_screenshot(&self, include_cursor: bool) -> Result<Dmabuf, TestError> {
+    pub async fn take_screenshot(
+        &self,
+        include_cursor: bool,
+    ) -> Result<(ScreenshotWithDevice, Rc<TestJayScreenshot>), TestError> {
         let js = Rc::new(TestJayScreenshot {
+            tran: self.tran.clone(),
             id: self.tran.id(),
             result: Cell::new(None),
+            planes: Default::default(),
         });
         self.tran.send(TakeScreenshot2 {
             self_id: self.id,
@@ -62,7 +67,7 @@ impl TestJayCompositor {
         self.tran.add_obj(js.clone())?;
         self.tran.sync().await;
         match js.result.take() {
-            Some(Ok(res)) => Ok(res),
+            Some(Ok(res)) => Ok((res, js)),
             Some(Err(res)) => bail!("Compositor could not take a screenshot: {}", res),
             None => bail!("Compositor did not send a screenshot"),
         }

src/it/test_ifs/test_registry.rs

@@ -165,7 +165,7 @@ impl TestRegistry {
         get_jay_compositor,
         jay_compositor,
         jay_compositor,
-        1,
+        6,
         TestJayCompositor
     );
     create_singleton!(get_compositor, compositor, wl_compositor, 6, TestCompositor);