x11: Implement image transfer using the MIT-SHM extension (#46)

* Implement shared memory transfer for X11

* Change to x11rb

* Fix leak + review from psychon

* Add waits to prevent illegal writes to SHM

* @ids1024 code review

Cargo.toml

@@ -16,7 +16,7 @@ rust-version = "1.60.0"
 default = ["x11", "wayland", "wayland-dlopen"]
 wayland = ["wayland-backend", "wayland-client", "nix", "fastrand"]
 wayland-dlopen = ["wayland-sys/dlopen"]
-x11 = ["bytemuck", "x11rb", "x11-dl"]
+x11 = ["bytemuck", "nix", "x11rb", "x11-dl"]
 
 [dependencies]
 log = "0.4.17"
@@ -30,7 +30,7 @@ wayland-backend = { version = "0.1.0", features = ["client_system"], optional =
 wayland-client = { version = "0.30.0", optional = true }
 wayland-sys = "0.30.0"
 x11-dl = { version  = "2.19.1", optional = true }
-x11rb = { version = "0.11.0", features = ["allow-unsafe-code", "dl-libxcb"], optional = true }
+x11rb = { version = "0.11.0", features = ["allow-unsafe-code", "dl-libxcb", "shm"], optional = true }
 
 [target.'cfg(all(unix, not(any(target_vendor = "apple", target_os = "android", target_os = "redox", target_os = "linux", target_os = "freebsd"))))'.dependencies]
 fastrand = { version = "1.8.0", optional = true }

src/x11.rs

@@ -1,18 +1,24 @@
 //! Implementation of software buffering for X11.
 //!
 //! This module converts the input buffer into an XImage and then sends it over the wire to be
-//! drawn. A more effective implementation would use shared memory instead of the wire. In
-//! addition, we may also want to blit to a pixmap instead of a window.
+//! drawn by the X server. The SHM extension is used if available.
+
+#![allow(clippy::uninlined_format_args)]
 
 use crate::SoftBufferError;
+use nix::libc::{shmat, shmctl, shmdt, shmget, IPC_PRIVATE, IPC_RMID};
 use raw_window_handle::{XcbDisplayHandle, XcbWindowHandle, XlibDisplayHandle, XlibWindowHandle};
-use std::fmt;
+use std::ptr::{null_mut, NonNull};
+use std::{fmt, io};
 
 use x11_dl::xlib::Display;
 use x11_dl::xlib_xcb::Xlib_xcb;
 
-use x11rb::connection::Connection;
-use x11rb::protocol::xproto::{self, ConnectionExt as _, Gcontext, Window};
+use x11rb::connection::{Connection, RequestConnection, SequenceNumber};
+use x11rb::cookie::Cookie;
+use x11rb::errors::{ConnectionError, ReplyError, ReplyOrIdError};
+use x11rb::protocol::shm::{self, ConnectionExt as _};
+use x11rb::protocol::xproto::{self, ConnectionExt as _};
 use x11rb::xcb_ffi::XCBConnection;
 
 /// The handle to an X11 drawing context.
@@ -21,13 +27,35 @@ pub struct X11Impl {
     connection: XCBConnection,
 
     /// The window to draw to.
-    window: Window,
+    window: xproto::Window,
 
     /// The graphics context to use when drawing.
-    gc: Gcontext,
+    gc: xproto::Gcontext,
 
     /// The depth (bits per pixel) of the drawing context.
     depth: u8,
+
+    /// Information about SHM, if it is available.
+    shm: Option<ShmInfo>,
+}
+
+struct ShmInfo {
+    /// The shared memory segment, paired with its ID.
+    seg: Option<(ShmSegment, shm::Seg)>,
+
+    /// A cookie indicating that the shared memory segment is ready to be used.
+    ///
+    /// We can't soundly read from or write to the SHM segment until the X server is done processing the
+    /// `shm::PutImage` request. However, the X server handles requests in order, which means that, if
+    /// we send a very small request after the `shm::PutImage` request, then the X server will have to
+    /// process that request before it can process the `shm::PutImage` request. Therefore, we can use
+    /// the reply to that small request to determine when the `shm::PutImage` request is done.
+    ///
+    /// In this case, we use `GetInputFocus` since it is a very small request.
+    ///
+    /// We store the sequence number instead of the `Cookie` since we cannot hold a self-referential
+    /// reference to the `connection` field.
+    done_processing: Option<SequenceNumber>,
 }
 
 impl X11Impl {
@@ -97,10 +125,13 @@ impl X11Impl {
 
         let window = window_handle.window;
 
-        // Start getting the depth of the window.
+        // Run in parallel: start getting the window depth and the SHM extension.
         let geometry_token = connection
             .get_geometry(window)
             .swbuf_err("Failed to send geometry request")?;
+        connection
+            .prefetch_extension_information(shm::X11_EXTENSION_NAME)
+            .swbuf_err("Failed to send SHM query request")?;
 
         // Create a new graphics context to draw to.
         let gc = connection
@@ -121,38 +152,292 @@ impl X11Impl {
             .reply()
             .swbuf_err("Failed to get geometry reply")?;
 
+        // See if SHM is available.
+        let shm_info = {
+            let present = is_shm_available(&connection);
+
+            if present {
+                // SHM is available.
+                Some(ShmInfo {
+                    seg: None,
+                    done_processing: None,
+                })
+            } else {
+                None
+            }
+        };
+
         Ok(Self {
             connection,
             window,
             gc,
             depth: geometry_reply.depth,
+            shm: shm_info,
         })
     }
 
     pub(crate) unsafe fn set_buffer(&mut self, buffer: &[u32], width: u16, height: u16) {
         // Draw the image to the buffer.
-        let result = self.connection.put_image(
-            xproto::ImageFormat::Z_PIXMAP,
-            self.window,
-            self.gc,
-            width,
-            height,
-            0,
-            0,
-            0,
-            self.depth,
-            bytemuck::cast_slice(buffer),
-        );
+        let result = unsafe { self.set_buffer_shm(buffer, width, height) }.and_then(|had_shm| {
+            if had_shm {
+                Ok(())
+            } else {
+                log::debug!("Falling back to non-SHM method");
+                self.set_buffer_fallback(buffer, width, height)
+            }
+        });
 
-        match result {
-            Err(e) => log::error!("Failed to draw image to window: {}", e),
-            Ok(token) => token.ignore_error(),
+        if let Err(e) = result {
+            log::error!("Failed to draw image to window: {}", e);
+        }
+    }
+
+    /// Put the given buffer into the window using the SHM method.
+    ///
+    /// Returns `false` if SHM is not available.
+    ///
+    /// # Safety
+    ///
+    /// The buffer's length must be `width * height`.
+    unsafe fn set_buffer_shm(
+        &mut self,
+        buffer: &[u32],
+        width: u16,
+        height: u16,
+    ) -> Result<bool, PushBufferError> {
+        let shm_info = match self.shm {
+            Some(ref mut info) => info,
+            None => return Ok(false),
+        };
+
+        // If the X server is still processing the last image, wait for it to finish.
+        shm_info.finish_wait(&self.connection)?;
+
+        // Get the SHM segment to use.
+        let necessary_size = (width as usize) * (height as usize) * 4;
+        let (segment, segment_id) = shm_info.segment(&self.connection, necessary_size)?;
+
+        // Copy the buffer into the segment.
+        // SAFETY: The buffer is properly sized and we've ensured that the X server isn't reading from it.
+        unsafe {
+            segment.copy(buffer);
+        }
+
+        // Put the image into the window.
+        self.connection
+            .shm_put_image(
+                self.window,
+                self.gc,
+                width,
+                height,
+                0,
+                0,
+                width,
+                height,
+                0,
+                0,
+                self.depth,
+                xproto::ImageFormat::Z_PIXMAP.into(),
+                false,
+                segment_id,
+                0,
+            )?
+            .ignore_error();
+
+        // Send a short request to act as a notification for when the X server is done processing the image.
+        shm_info.begin_wait(&self.connection)?;
+
+        Ok(true)
+    }
+
+    /// Put the given buffer into the window using the fallback wire transfer method.
+    fn set_buffer_fallback(
+        &mut self,
+        buffer: &[u32],
+        width: u16,
+        height: u16,
+    ) -> Result<(), PushBufferError> {
+        self.connection
+            .put_image(
+                xproto::ImageFormat::Z_PIXMAP,
+                self.window,
+                self.gc,
+                width,
+                height,
+                0,
+                0,
+                0,
+                self.depth,
+                bytemuck::cast_slice(buffer),
+            )?
+            .ignore_error();
+
+        Ok(())
+    }
+}
+
+impl ShmInfo {
+    /// Allocate a new `ShmSegment` of the given size.
+    fn segment(
+        &mut self,
+        conn: &impl Connection,
+        size: usize,
+    ) -> Result<(&mut ShmSegment, shm::Seg), PushBufferError> {
+        // Round the size up to the next power of two to prevent frequent reallocations.
+        let size = size.next_power_of_two();
+
+        // Get the size of the segment currently in use.
+        let needs_realloc = match self.seg {
+            Some((ref seg, _)) => seg.size() < size,
+            None => true,
+        };
+
+        // Reallocate if necessary.
+        if needs_realloc {
+            let new_seg = ShmSegment::new(size)?;
+            self.associate(conn, new_seg)?;
+        }
+
+        // Get the segment and ID.
+        Ok(self
+            .seg
+            .as_mut()
+            .map(|(ref mut seg, id)| (seg, *id))
+            .unwrap())
+    }
+
+    /// Associate an SHM segment with the server.
+    fn associate(
+        &mut self,
+        conn: &impl Connection,
+        seg: ShmSegment,
+    ) -> Result<(), PushBufferError> {
+        // Register the guard.
+        let new_id = conn.generate_id()?;
+        conn.shm_attach(new_id, seg.id(), true)?.ignore_error();
+
+        // Take out the old one and detach it.
+        if let Some((old_seg, old_id)) = self.seg.replace((seg, new_id)) {
+            conn.shm_detach(old_id)?.ignore_error();
+
+            // Drop the old segment.
+            drop(old_seg);
+        }
+
+        Ok(())
+    }
+
+    /// Begin waiting for the SHM processing to finish.
+    fn begin_wait(&mut self, c: &impl Connection) -> Result<(), PushBufferError> {
+        let cookie = c.get_input_focus()?.sequence_number();
+        let old_cookie = self.done_processing.replace(cookie);
+        debug_assert!(old_cookie.is_none());
+        Ok(())
+    }
+
+    /// Wait for the SHM processing to finish.
+    fn finish_wait(&mut self, c: &impl Connection) -> Result<(), PushBufferError> {
+        if let Some(done_processing) = self.done_processing.take() {
+            // Cast to a cookie and wait on it.
+            let cookie = Cookie::<_, xproto::GetInputFocusReply>::new(c, done_processing);
+            cookie.reply()?;
+        }
+
+        Ok(())
+    }
+}
+
+struct ShmSegment {
+    id: i32,
+    ptr: NonNull<i8>,
+    size: usize,
+}
+
+impl ShmSegment {
+    /// Create a new `ShmSegment` with the given size.
+    fn new(size: usize) -> io::Result<Self> {
+        unsafe {
+            // Create the shared memory segment.
+            let id = shmget(IPC_PRIVATE, size, 0o600);
+            if id == -1 {
+                return Err(io::Error::last_os_error());
+            }
+
+            // Map the SHM to our memory space.
+            let ptr = {
+                let ptr = shmat(id, null_mut(), 0);
+                match NonNull::new(ptr as *mut i8) {
+                    Some(ptr) => ptr,
+                    None => {
+                        shmctl(id, IPC_RMID, null_mut());
+                        return Err(io::Error::last_os_error());
+                    }
+                }
+            };
+
+            Ok(Self { id, ptr, size })
+        }
+    }
+
+    /// Copy data into this shared memory segment.
+    ///
+    /// # Safety
+    ///
+    /// This function assumes that the size of `self`'s buffer is larger than or equal to `data.len()`.
+    /// In addition, no other processes should be reading from or writing to this memory.
+    unsafe fn copy<T: bytemuck::NoUninit>(&mut self, data: &[T]) {
+        debug_assert!(data.len() * std::mem::size_of::<T>() <= self.size,);
+        let incoming_data = bytemuck::cast_slice::<_, u8>(data);
+
+        unsafe {
+            std::ptr::copy_nonoverlapping(
+                incoming_data.as_ptr(),
+                self.ptr.as_ptr() as *mut u8,
+                incoming_data.len(),
+            )
+        }
+    }
+
+    /// Get the size of this shared memory segment.
+    fn size(&self) -> usize {
+        self.size
+    }
+
+    /// Get the shared memory ID.
+    fn id(&self) -> u32 {
+        self.id as _
+    }
+}
+
+impl Drop for ShmSegment {
+    fn drop(&mut self) {
+        unsafe {
+            // Detach the shared memory segment.
+            shmdt(self.ptr.as_ptr() as _);
+
+            // Delete the shared memory segment.
+            shmctl(self.id, IPC_RMID, null_mut());
         }
     }
 }
 
 impl Drop for X11Impl {
     fn drop(&mut self) {
+        // If we used SHM, make sure it's detached from the server.
+        if let Some(mut shm) = self.shm.take() {
+            // If we were in the middle of processing a buffer, wait for it to finish.
+            shm.finish_wait(&self.connection).ok();
+
+            if let Some((segment, seg_id)) = shm.seg.take() {
+                if let Ok(token) = self.connection.shm_detach(seg_id) {
+                    token.ignore_error();
+                }
+
+                // Drop the segment.
+                drop(segment);
+            }
+        }
+
         // Close the graphics context that we created.
         if let Ok(token) = self.connection.free_gc(self.gc) {
             token.ignore_error();
@@ -160,12 +445,93 @@ impl Drop for X11Impl {
     }
 }
 
+/// Test to see if SHM is available.
+fn is_shm_available(c: &impl Connection) -> bool {
+    // Create a small SHM segment.
+    let seg = match ShmSegment::new(0x1000) {
+        Ok(seg) => seg,
+        Err(_) => return false,
+    };
+
+    // Attach and detach it.
+    let seg_id = match c.generate_id() {
+        Ok(id) => id,
+        Err(_) => return false,
+    };
+
+    let (attach, detach) = {
+        let attach = c.shm_attach(seg_id, seg.id(), false);
+        let detach = c.shm_detach(seg_id);
+
+        match (attach, detach) {
+            (Ok(attach), Ok(detach)) => (attach, detach),
+            _ => return false,
+        }
+    };
+
+    // Check the replies.
+    matches!((attach.check(), detach.check()), (Ok(()), Ok(())))
+}
+
+/// An error that can occur when pushing a buffer to the window.
+#[derive(Debug)]
+enum PushBufferError {
+    /// We encountered an X11 error.
+    X11(ReplyError),
+
+    /// We exhausted the XID space.
+    XidExhausted,
+
+    /// A system error occurred while creating the shared memory segment.
+    System(io::Error),
+}
+
+impl fmt::Display for PushBufferError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::X11(e) => write!(f, "X11 error: {}", e),
+            Self::XidExhausted => write!(f, "XID space exhausted"),
+            Self::System(e) => write!(f, "System error: {}", e),
+        }
+    }
+}
+
+impl std::error::Error for PushBufferError {}
+
+impl From<ConnectionError> for PushBufferError {
+    fn from(e: ConnectionError) -> Self {
+        Self::X11(ReplyError::ConnectionError(e))
+    }
+}
+
+impl From<ReplyError> for PushBufferError {
+    fn from(e: ReplyError) -> Self {
+        Self::X11(e)
+    }
+}
+
+impl From<ReplyOrIdError> for PushBufferError {
+    fn from(e: ReplyOrIdError) -> Self {
+        match e {
+            ReplyOrIdError::ConnectionError(e) => Self::X11(ReplyError::ConnectionError(e)),
+            ReplyOrIdError::X11Error(e) => Self::X11(ReplyError::X11Error(e)),
+            ReplyOrIdError::IdsExhausted => Self::XidExhausted,
+        }
+    }
+}
+
+impl From<io::Error> for PushBufferError {
+    fn from(e: io::Error) -> Self {
+        Self::System(e)
+    }
+}
+
 /// Convenient wrapper to cast errors into SoftBufferError.
 trait ResultExt<T, E> {
     fn swbuf_err(self, msg: impl Into<String>) -> Result<T, SoftBufferError>;
 }
 
-impl<T, E: fmt::Debug + fmt::Display + 'static> ResultExt<T, E> for Result<T, E> {
+impl<T, E: std::error::Error + 'static> ResultExt<T, E> for Result<T, E> {
     fn swbuf_err(self, msg: impl Into<String>) -> Result<T, SoftBufferError> {
         self.map_err(|e| {
             SoftBufferError::PlatformError(Some(msg.into()), Some(Box::new(LibraryError(e))))