perf(desktop): remove X11 staging copies and preserve receipt progress

crates/rds-desktop/src/capture/x11.rs

@@ -4,7 +4,7 @@
 //!
 //! MIT-SHM removes the dominant capture cost: the pixmap is written
 //! into a shared segment the client mmaps, so a 1080p frame costs one
-//! completion event plus a memcpy instead of ~8 MiB serialized through
+//! reply and a synchronous borrowed view instead of ~8 MiB serialized through
 //! the X socket. `CreateSegment` (SHM ≥ 1.2 fd passing) is preferred —
 //! the server allocates and owns nothing persists on the client side.
 //! `ext-image-copy-capture` and the DRM/KMS tap remain scheduled behind
@@ -15,6 +15,7 @@
 //! module (workspace convention: FFI paths allow `unsafe_code` locally).
 #![allow(unsafe_code)]
 
+use std::borrow::Cow;
 use std::os::fd::OwnedFd;
 
 use bytes::Bytes;
@@ -27,7 +28,7 @@ use x11rb::protocol::shm::{self, ConnectionExt as _, Seg};
 use x11rb::protocol::xproto::{ConnectionExt, GetImageReply, ImageFormat};
 use x11rb::rust_connection::RustConnection;
 
-use crate::{Capturer, DesktopError, RawFrame};
+use crate::{BgraFrame, Capturer, DesktopError, RawFrame};
 
 /// MIT-SHM segment the server fills in place — the reply is a
 /// completion event, not a serialized pixmap.
@@ -170,11 +171,29 @@ impl X11Capturer {
     }
 }
 
-impl Capturer for X11Capturer {
-    fn capture(&mut self) -> Result<RawFrame, DesktopError> {
+impl X11Capturer {
+    /// The synchronous callback is the lifetime boundary: the next capture
+    /// cannot overwrite the mapping until conversion/encode returns. Owned
+    /// capture consumers retain the separate snapshot-producing trait method.
+    pub(crate) fn capture_with<T>(
+        &mut self,
+        process: impl FnOnce(BgraFrame<'_>) -> T,
+    ) -> Result<T, DesktopError> {
+        let (width, height) = (u32::from(self.width), u32::from(self.height));
+        let pixels = self.capture_pixels()?;
+        Ok(process(BgraFrame {
+            width,
+            height,
+            stride: width * 4,
+            data: pixels.as_ref(),
+        }))
+    }
+
+    fn capture_pixels(&mut self) -> Result<Cow<'_, [u8]>, DesktopError> {
         if let Some(shm) = &self.shm {
-            // `reply()` waits for the ShmCompletion event the server
-            // sends after filling the segment.
+            // The GetImage reply arrives after the server has finished
+            // writing the segment. No new request can reuse it while the
+            // callback holds the capturer's exclusive borrow.
             let filled = self
                 .conn
                 .shm_get_image(
@@ -192,12 +211,11 @@ impl Capturer for X11Capturer {
                 .and_then(|c| c.reply().map_err(|e| DesktopError::Capture(e.to_string())));
             match filled {
                 Ok(_) => {
-                    return Ok(RawFrame {
-                        width: u32::from(self.width),
-                        height: u32::from(self.height),
-                        stride: u32::from(self.width) * 4,
-                        data: Bytes::copy_from_slice(&shm.map[..shm.len]),
-                    });
+                    let shm = self
+                        .shm
+                        .as_ref()
+                        .expect("successful capture retains its segment");
+                    return Ok(Cow::Borrowed(&shm.map[..shm.len]));
                 }
                 Err(e) => {
                     tracing::warn!("MIT-SHM capture failed ({e}); falling back to GetImage");
@@ -230,11 +248,19 @@ impl Capturer for X11Capturer {
                 self.height
             )));
         }
+        Ok(Cow::Owned(reply.data))
+    }
+}
+
+impl Capturer for X11Capturer {
+    fn capture(&mut self) -> Result<RawFrame, DesktopError> {
+        let (width, height) = (u32::from(self.width), u32::from(self.height));
+        let pixels = self.capture_pixels()?;
         Ok(RawFrame {
-            width: u32::from(self.width),
-            height: u32::from(self.height),
-            stride: u32::from(self.width) * 4,
-            data: Bytes::from(reply.data),
+            width,
+            height,
+            stride: width * 4,
+            data: Bytes::from(pixels.into_owned()),
         })
     }
 
@@ -347,6 +373,55 @@ mod tests {
             &[0x56, 0x34, 0x12],
             "captured pixel did not match the native paint"
         );
+        // The producer borrows the completed mapping directly; the owned
+        // snapshot above must survive another capture and a different paint.
+        let mapped = cap.shm.as_ref().unwrap().map.as_ptr();
+        cap.conn
+            .change_gc(
+                gc,
+                &x11rb::protocol::xproto::ChangeGCAux::new().foreground(0xAB_CD_EF),
+            )
+            .unwrap()
+            .check()
+            .unwrap();
+        cap.conn
+            .poly_fill_rectangle(
+                cap.root,
+                gc,
+                &[Rectangle {
+                    x: 0,
+                    y: 0,
+                    width: 100,
+                    height: 100,
+                }],
+            )
+            .unwrap()
+            .check()
+            .unwrap();
+        cap.capture_with(|borrowed| {
+            assert_eq!(
+                borrowed.data.as_ptr(),
+                mapped,
+                "borrowed capture copied pixels"
+            );
+            assert_eq!(&borrowed.data[pixel..pixel + 3], &[0xEF, 0xCD, 0xAB]);
+        })
+        .unwrap();
+        assert_eq!(&frame.data[pixel..pixel + 3], &[0x56, 0x34, 0x12]);
+
+        // A server without usable SHM still provides the same borrowed input
+        // contract from its owned GetImage reply; no extra staging copy.
+        let shm = cap.shm.take().unwrap();
+        cap.conn.shm_detach(shm.seg).unwrap().check().unwrap();
+        drop(shm);
+        cap.capture_with(|borrowed| {
+            assert_eq!(&borrowed.data[pixel..pixel + 3], &[0xEF, 0xCD, 0xAB]);
+        })
+        .unwrap();
+        assert_eq!(
+            &cap.capture().unwrap().data[pixel..pixel + 3],
+            &[0xEF, 0xCD, 0xAB]
+        );
         cap.conn.free_gc(gc).unwrap().check().unwrap();
     }
 }

crates/rds-desktop/src/codec/openh264.rs

@@ -13,7 +13,7 @@ use openh264::formats::{BGRA8Source, RGB8Source, RGBSource, YUVBuffer, YUVSource
 use openh264::{Error as OhError, OpenH264API};
 use rds_core::Codec;
 
-use crate::{Decoder, DesktopError, EncodedFrame, Encoder, RawFrame};
+use crate::{BgraFrame, Decoder, DesktopError, EncodedFrame, Encoder, RawFrame};
 
 // A large recovery picture can leave timestamp rate control in debt for
 // seconds at its low-rate floor. Preserve occasional delta progress while
@@ -84,7 +84,7 @@ impl H264Encoder {
 
 impl Encoder for H264Encoder {
     fn encode(&mut self, frame: &RawFrame) -> Result<EncodedFrame, DesktopError> {
-        self.encode_timed(frame, self.started.elapsed().as_millis() as u64)
+        self.encode_bgra(frame.into())
     }
     fn request_idr(&mut self) {
         self.want_idr = true;
@@ -99,10 +99,26 @@ impl Encoder for H264Encoder {
 }
 
 impl H264Encoder {
+    #[cfg(test)]
     fn encode_timed(
         &mut self,
         frame: &RawFrame,
         now_ms: u64,
+    ) -> Result<EncodedFrame, DesktopError> {
+        self.encode_bgra_timed(frame.into(), now_ms)
+    }
+
+    pub(crate) fn encode_bgra(
+        &mut self,
+        frame: BgraFrame<'_>,
+    ) -> Result<EncodedFrame, DesktopError> {
+        self.encode_bgra_timed(frame, self.started.elapsed().as_millis() as u64)
+    }
+
+    fn encode_bgra_timed(
+        &mut self,
+        frame: BgraFrame<'_>,
+        now_ms: u64,
     ) -> Result<EncodedFrame, DesktopError> {
         let dimensions = (frame.width, frame.height);
         let target = self
@@ -340,14 +356,14 @@ impl Decoder for H264Decoder {
 pub fn bgra_to_i420(frame: &RawFrame) -> YUVBuffer {
     let stride = frame.stride as usize;
     if stride.is_multiple_of(4) && frame.data.len() >= stride * frame.height as usize {
-        return YUVBuffer::from_bgra8_source(StridedBgra(frame));
+        return YUVBuffer::from_bgra8_source(StridedBgra(frame.into()));
     }
-    bgra_to_i420_scalar(frame)
+    bgra_to_i420_scalar(frame.into())
 }
 
 /// `RawFrame` as an `openh264` BGRA source, carrying its real stride so
 /// padded capture buffers need no intermediate copy.
-struct StridedBgra<'a>(&'a RawFrame);
+struct StridedBgra<'a>(BgraFrame<'a>);
 
 impl RGBSource for StridedBgra<'_> {
     fn dimensions(&self) -> (usize, usize) {
@@ -367,7 +383,7 @@ impl RGB8Source for StridedBgra<'_> {
     }
 
     fn rgb8_data(&self) -> &[u8] {
-        &self.0.data
+        self.0.data
     }
 
     fn pixel_stride(&self) -> usize {
@@ -383,7 +399,7 @@ impl BGRA8Source for StridedBgra<'_> {}
 
 /// Scalar fallback for strides that are not whole pixels — also the
 /// test reference for the SIMD path's output.
-fn bgra_to_i420_scalar(frame: &RawFrame) -> YUVBuffer {
+fn bgra_to_i420_scalar(frame: BgraFrame<'_>) -> YUVBuffer {
     let (w, h) = (frame.width as usize, frame.height as usize);
     let stride = frame.stride as usize;
     let mut yuv = vec![0u8; w * h * 3 / 2];
@@ -424,6 +440,34 @@ mod tests {
         }
     }
 
+    #[test]
+    fn borrowed_input_preserves_owned_bitstream_and_recovery() {
+        let mut owned = H264Encoder::new(1_000_000, 30.0).unwrap();
+        let mut borrowed = H264Encoder::new(1_000_000, 30.0).unwrap();
+        let mut decoder = H264Decoder::new().unwrap();
+        for sequence in 0..8 {
+            let mut raw = frame();
+            raw.data = Bytes::from(vec![sequence as u8 * 20; raw.data.len()]);
+            if sequence == 4 {
+                owned.request_idr();
+                borrowed.request_idr();
+            }
+            let a = owned.encode_timed(&raw, sequence * 34).unwrap();
+            let b = borrowed
+                .encode_bgra_timed((&raw).into(), sequence * 34)
+                .unwrap();
+            assert_eq!(
+                a.data, b.data,
+                "borrowed frame changed bitstream at {sequence}"
+            );
+            assert_eq!(a.keyframe, b.keyframe);
+            if sequence == 0 || sequence == 4 {
+                assert!(a.keyframe, "explicit recovery did not emit IDR");
+            }
+            assert!(decoder.decode(&b).unwrap().is_some());
+        }
+    }
+
     /// The keyframe flag must describe the emitted bitstream: IDR on
     /// frame 0 (fresh encoder) and after `request_idr`, deltas between.
     #[test]
@@ -633,7 +677,7 @@ mod tests {
             data: Bytes::from(data),
         };
         let fast = bgra_to_i420(&raw);
-        let slow = bgra_to_i420_scalar(&raw);
+        let slow = bgra_to_i420_scalar((&raw).into());
         assert_eq!(fast.dimensions(), slow.dimensions());
         for (a, b) in fast.y().iter().zip(slow.y()) {
             assert!((i32::from(*a) - i32::from(*b)).abs() <= 2, "Y diverges");

crates/rds-desktop/src/lib.rs

@@ -58,6 +58,29 @@ pub struct RawFrame {
     pub data: Bytes,
 }
 
+/// Borrowed BGRA input for synchronous capture/conversion/encode work. A
+/// native mapping can supply pixels without first allocating an owned frame.
+#[cfg(any(feature = "x11", test))]
+#[derive(Clone, Copy)]
+pub(crate) struct BgraFrame<'a> {
+    pub width: u32,
+    pub height: u32,
+    pub stride: u32,
+    pub data: &'a [u8],
+}
+
+#[cfg(any(feature = "x11", test))]
+impl<'a> From<&'a RawFrame> for BgraFrame<'a> {
+    fn from(frame: &'a RawFrame) -> Self {
+        Self {
+            width: frame.width,
+            height: frame.height,
+            stride: frame.stride,
+            data: &frame.data,
+        }
+    }
+}
+
 /// Local software receive limit: up to 8K pixels, no side above 8192.
 /// Check before allocating BGRA output; native codec internal memory is separate.
 pub(crate) fn frame_bytes(width: usize, height: usize) -> Option<usize> {

crates/rds-desktop/src/scaling.rs

@@ -1,12 +1,9 @@
 //! Owned low-latency BGRA downscaling for a configured serving profile. The
 //! display capabilities keep original coordinates; only video dimensions change.
-use crate::{DesktopError, RawFrame};
+use crate::{BgraFrame, DesktopError, RawFrame};
 use bytes::Bytes;
 
-pub(crate) fn downscale(frame: RawFrame, height: Option<u32>) -> Result<RawFrame, DesktopError> {
-    let Some(height) = height.filter(|height| *height < frame.height) else {
-        return Ok(frame);
-    };
+pub(crate) fn downscale(frame: BgraFrame<'_>, height: u32) -> Result<RawFrame, DesktopError> {
     let height = height & !1;
     let width =
         ((u64::from(frame.width) * u64::from(height) / u64::from(frame.height)) as u32) & !1;
@@ -57,7 +54,7 @@ mod tests {
             stride: 72,
             data: Bytes::from(bytes),
         };
-        let resized = downscale(frame, Some(4)).unwrap();
+        let resized = downscale((&frame).into(), 4).unwrap();
         assert_eq!((resized.width, resized.height, resized.stride), (8, 4, 32));
         assert_eq!(&resized.data[..], &vec![12; 128]);
     }