perf: MIT-SHM capture + codec hot-path allocations

CHANGELOG.md

@@ -2,6 +2,16 @@
 
 ## [Unreleased]
 
+- Desktop capture/encode throughput:
+  - X11 capture uses MIT-SHM (`CreateSegment` fd-passing, server ≥1.2)
+    when available — the pixmap lands in a shared segment instead of an
+    ~8 MiB serialized `GetImage` reply per 1080p frame; plain `GetImage`
+    remains the fallback for remote/older servers.
+  - The encoder recycles its I420 input buffer across frames (~3 MiB
+    per frame at 1080p — ~190 MB/s of alloc churn at 60 fps removed).
+  - Decode output goes I420→RGBA in one SIMD pass (`write_rgba8`,
+    AVX2 on x86-64) plus an in-place R↔B swap — replacing an RGB8
+    scratch buffer plus a scalar expand.
 - Desktop encode path:
   - OpenH264 now runs `ScreenContentRealTime` — the correct usage type
     for desktop content (text/sharp edges, not camera footage); adaptive

crates/rds-desktop/Cargo.toml

@@ -8,10 +8,11 @@ repository.workspace = true
 
 [features]
 default = []
-x11 = ["dep:x11rb", "dep:openh264"]
+x11 = ["dep:x11rb", "dep:openh264", "dep:memmap2"]
 
 [dependencies]
 bytes.workspace = true
+memmap2 = { version = "0.9", optional = true }
 openh264 = { version = "0.9", optional = true }
 rds-core.workspace = true
 rds-net.workspace = true

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

@@ -1,25 +1,51 @@
-//! X11 capture via `x11rb`: `GetImage` polling for frames.
+//! X11 capture via `x11rb`: MIT-SHM shared-memory `GetImage` where the
+//! server supports it, plain `GetImage` polling as the portable
+//! fallback.
 //!
-//! This is the portable baseline backend. The MIT-SHM zero-copy path,
-//! `ext-image-copy-capture` and the DRM/KMS tap are scheduled behind the
-//! same `Capturer` trait; GetImage polling is what every X server already
-//! supports. Input injection lives in `crate::input::x11`.
+//! 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
+//! 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
+//! the same `Capturer` trait. Input injection lives in
+//! `crate::input::x11`.
+//!
+//! `mmap` on the server-provided segment is the one unsafe call in this
+//! module (workspace convention: FFI paths allow `unsafe_code` locally).
+#![allow(unsafe_code)]
+
+use std::os::fd::OwnedFd;
 
 use bytes::Bytes;
+use memmap2::{MmapMut, MmapOptions};
 use rds_core::{DesktopCaps, DisplayInfo};
-use x11rb::connection::Connection as _;
+use x11rb::connection::{Connection as _, RequestConnection as _};
+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};
 
+/// MIT-SHM segment the server fills in place — the reply is a
+/// completion event, not a serialized pixmap.
+struct ShmPath {
+    seg: Seg,
+    map: MmapMut,
+    len: usize,
+    /// The mapping stays valid without it, but keeping the fd makes
+    /// ownership/lifetime explicit.
+    _fd: OwnedFd,
+}
+
 /// Polls one X11 screen (root window) into `RawFrame`s.
 pub struct X11Capturer {
     conn: RustConnection,
     root: x11rb::protocol::xproto::Window,
     width: u16,
     height: u16,
     screen: usize,
+    shm: Option<ShmPath>,
 }
 
 impl X11Capturer {
@@ -35,18 +61,84 @@ impl X11Capturer {
             setup.roots[idx].height_in_pixels,
         );
         let _ = default_screen;
+        let shm = Self::try_shm(&conn, width, height);
         Ok(Self {
             conn,
             root,
             width,
             height,
             screen: idx,
+            shm,
+        })
+    }
+
+    /// MIT-SHM ≥ 1.2 (fd passing): the server creates the segment and
+    /// hands back an fd we mmap. `None` on older servers, remote
+    /// displays, or any setup failure — the caller falls back to plain
+    /// `GetImage`.
+    fn try_shm(conn: &RustConnection, width: u16, height: u16) -> Option<ShmPath> {
+        conn.extension_information(shm::X11_EXTENSION_NAME).ok()??;
+        let ver = conn.shm_query_version().ok()?.reply().ok()?;
+        if (ver.major_version, ver.minor_version) < (1, 2) {
+            return None;
+        }
+        let len = usize::from(width) * usize::from(height) * 4;
+        let seg = conn.generate_id().ok()?;
+        let reply = conn
+            .shm_create_segment(seg, len as u32, false)
+            .ok()?
+            .reply()
+            .ok()?;
+        let fd: OwnedFd = reply.shm_fd;
+        // SAFETY: `fd` is a fresh server-created segment of exactly
+        // `len` bytes; nothing else mutates its size.
+        let map = unsafe { MmapOptions::new().len(len).map_mut(&fd) }.ok()?;
+        Some(ShmPath {
+            seg,
+            map,
+            len,
+            _fd: fd,
         })
     }
 }
 
 impl Capturer for X11Capturer {
     fn capture(&mut self) -> Result<RawFrame, DesktopError> {
+        if let Some(shm) = &self.shm {
+            // `reply()` waits for the ShmCompletion event the server
+            // sends after filling the segment.
+            let filled = self
+                .conn
+                .shm_get_image(
+                    self.root,
+                    0,
+                    0,
+                    self.width,
+                    self.height,
+                    !0,
+                    ImageFormat::Z_PIXMAP.into(),
+                    shm.seg,
+                    0,
+                )
+                .map_err(|e| DesktopError::Capture(e.to_string()))
+                .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]),
+                    });
+                }
+                Err(e) => {
+                    tracing::warn!("MIT-SHM capture failed ({e}); falling back to GetImage");
+                }
+            }
+        }
+        // Reached only when the shm attempt failed at runtime — stop
+        // paying for attempts that can no longer succeed.
+        self.shm = None;
         let reply: GetImageReply = self
             .conn
             .get_image(
@@ -79,6 +171,14 @@ impl Capturer for X11Capturer {
     }
 }
 
+impl Drop for X11Capturer {
+    fn drop(&mut self) {
+        if let Some(shm) = self.shm.take() {
+            let _ = self.conn.shm_detach(shm.seg);
+        }
+    }
+}
+
 /// Displays visible over X11.
 pub fn capabilities() -> Result<DesktopCaps, DesktopError> {
     match X11Capturer::new(0) {
@@ -89,3 +189,30 @@ pub fn capabilities() -> Result<DesktopCaps, DesktopError> {
         Err(e) => Err(e),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Real capture needs an X server — skipped silently when `$DISPLAY`
+    /// is unset (CI has none). On a local server (`:N`) MIT-SHM must be
+    /// present; remote `host:N` displays legitimately lack it.
+    #[test]
+    fn capture_roundtrip() {
+        let Some(display) = std::env::var_os("DISPLAY") else {
+            return;
+        };
+        let local = display.to_string_lossy().starts_with(':');
+        let Ok(mut cap) = X11Capturer::new(0) else {
+            return;
+        };
+        if local {
+            assert!(cap.shm.is_some(), "local X server without MIT-SHM 1.2?");
+        }
+        for _ in 0..3 {
+            let frame = cap.capture().unwrap();
+            assert_eq!(frame.data.len(), (frame.width * frame.height * 4) as usize);
+            assert_eq!(frame.stride, frame.width * 4);
+        }
+    }
+}

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

@@ -27,6 +27,10 @@ pub struct H264Encoder {
     pending_bitrate: Option<u32>,
     fps: f32,
     want_idr: bool,
+    /// Recycled I420 input buffer — a 1080p frame is ~3 MiB, so a fresh
+    /// allocation per frame at 60 fps is ~190 MB/s of pure alloc churn.
+    /// Rebuilt only when frame dimensions change.
+    yuv_buf: Option<YUVBuffer>,
 }
 
 impl H264Encoder {
@@ -38,6 +42,7 @@ impl H264Encoder {
             pending_bitrate: None,
             fps,
             want_idr: true,
+            yuv_buf: None,
         })
     }
 
@@ -84,11 +89,24 @@ impl Encoder for H264Encoder {
             self.inner.force_intra_frame();
             self.want_idr = false;
         }
-        let yuv = bgra_to_i420(frame);
+        let (w, h) = (frame.width as usize, frame.height as usize);
+        let stride = frame.stride as usize;
+        let yuv = if stride.is_multiple_of(4) && frame.data.len() >= stride * h {
+            match self.yuv_buf.take() {
+                Some(mut buf) if buf.dimensions() == (w, h) => {
+                    buf.read_bgra8(StridedBgra(frame));
+                    buf
+                }
+                _ => YUVBuffer::from_bgra8_source(StridedBgra(frame)),
+            }
+        } else {
+            bgra_to_i420_scalar(frame)
+        };
         let stream = self
             .inner
             .encode(&yuv)
             .map_err(|e| DesktopError::Encode(e.to_string()))?;
+        self.yuv_buf = Some(yuv);
         // The flag the writer's collapse trusts must be the truth on
         // the wire, not a schedule assumption: OpenH264 decides when
         // forced and periodic IDRs actually land, so read the NALs.
@@ -171,20 +189,13 @@ impl Decoder for H264Decoder {
             return Ok(None);
         };
         let (w, h) = yuv.dimensions();
-        let mut rgb = vec![0u8; w * h * 3];
-        yuv.write_rgb8(&mut rgb);
-        // Expand RGB to BGRA for a uniform RawFrame layout.
+        // I420→RGBA in one SIMD pass (AVX2 on x86-64), then swap R↔B
+        // in place for the RawFrame BGRA contract — the swap vectorizes
+        // trivially and avoids a separate 3-byte-per-pixel scratch.
         let mut bgra = vec![0u8; w * h * 4];
-        for (dst, src) in bgra
-            .as_chunks_mut::<4>()
-            .0
-            .iter_mut()
-            .zip(rgb.as_chunks::<3>().0)
-        {
-            dst[0] = src[2];
-            dst[1] = src[1];
-            dst[2] = src[0];
-            dst[3] = 0xFF;
+        yuv.write_rgba8(&mut bgra);
+        for px in bgra.as_chunks_mut::<4>().0.iter_mut() {
+            px.swap(0, 2);
         }
         Ok(Some(RawFrame {
             width: w as u32,