wayland: optimize receiving messages

src/client/error.rs

@@ -19,10 +19,6 @@ pub enum ClientError {
     InvalidMethod,
     #[error("Client tried to access non-existent object {0}")]
     InvalidObject(ObjectId),
-    #[error("The message size is < 8")]
-    MessageSizeTooSmall,
-    #[error("The size of the message is not a multiple of 4")]
-    UnalignedMessage,
     #[error("The requested client {0} does not exist")]
     ClientDoesNotExist(ClientId),
     #[error("Server tried to allocate more than 0x1_00_00_00 ids")]

src/client/tasks.rs

@@ -2,11 +2,9 @@ use {
     crate::{
         async_engine::Phase,
         client::{Client, ClientError},
-        object::ObjectId,
         utils::{
-            buffd::{BufFdIn, BufFdOut, MsgParser},
+            buffd::{BufFdOut, MsgParser, WlBufFdIn, WlMessage},
             errorfmt::ErrorFmt,
-            vec_ext::VecExt,
         },
     },
     futures_util::{FutureExt, select},
@@ -49,14 +47,14 @@ async fn receive(data: Rc<Client>) {
     });
     let display = data.display().unwrap();
     let recv = async {
-        let mut buf = BufFdIn::new(&data.socket, &data.state.ring);
+        let mut buf = WlBufFdIn::new(&data.socket, &data.state.ring);
-        let mut data_buf = Vec::<u32>::new();
         loop {
-            let mut hdr = [0u32, 0];
+            let WlMessage {
-            buf.read_full(&mut hdr[..]).await?;
+                obj_id,
-            let obj_id = ObjectId::from_raw(hdr[0]);
+                message,
-            let len = (hdr[1] >> 16) as usize;
+                body,
-            let request = hdr[1] & 0xffff;
+                fds,
+            } = buf.read_message().await?;
             let obj = match data.objects.get_obj(obj_id) {
                 Ok(obj) => obj,
                 _ => {
@@ -65,28 +63,12 @@ async fn receive(data: Rc<Client>) {
                     return Err(ClientError::InvalidObject(obj_id));
                 }
             };
-            // log::trace!("obj: {}, request: {}, len: {}", obj_id, request, len);
+            let parser = MsgParser::new(fds, body);
-            if len < 8 {
+            if let Err(e) = obj.handle_request(&data, message, parser) {
-                return Err(ClientError::MessageSizeTooSmall);
-            }
-            if len % 4 != 0 {
-                return Err(ClientError::UnalignedMessage);
-            }
-            let len = len / 4 - 2;
-            data_buf.clear();
-            data_buf.reserve(len);
-            let unused = data_buf.split_at_spare_mut_ext().1;
-            buf.read_full(&mut unused[..len]).await?;
-            unsafe {
-                data_buf.set_len(len);
-            }
-            // log::trace!("{:x?}", data_buf);
-            let parser = MsgParser::new(&mut buf, &data_buf[..]);
-            if let Err(e) = obj.handle_request(&data, request, parser) {
                 if let ClientError::InvalidMethod = e
                     && let Ok(obj) = data.objects.get_obj(obj_id)
                 {
-                    data.invalid_request(&*obj, request);
+                    data.invalid_request(&*obj, message);
                     return Err(e);
                 }
                 return Err(ClientError::RequestError(Box::new(e)));

src/it/test_transport.rs

@@ -13,7 +13,10 @@ use {
         utils::{
             asyncevent::AsyncEvent,
             bitfield::Bitfield,
-            buffd::{BufFdIn, BufFdOut, MsgFormatter, MsgParser, OutBuffer, OutBufferSwapchain},
+            buffd::{
+                BufFdError, BufFdIn, BufFdOut, MsgFormatter, MsgParser, OutBuffer,
+                OutBufferSwapchain, WlBufFdIn, WlMessage,
+            },
             copyhashmap::CopyHashMap,
             hash_map_ext::HashMapExt,
             stack::Stack,
@@ -36,7 +39,6 @@ pub struct TestTransport {
     pub run: Rc<TestRun>,
     pub socket: Rc<OwnedFd>,
     pub client_id: Cell<ClientId>,
-    pub bufs: Stack<Vec<u32>>,
     pub swapchain: Rc<RefCell<OutBufferSwapchain>>,
     pub flush_request: AsyncEvent,
     pub incoming: Cell<Option<SpawnedFuture<()>>>,
@@ -153,7 +155,7 @@ impl TestTransport {
                 "",
                 Incoming {
                     tc: self.clone(),
-                    buf: BufFdIn::new(&self.socket, &self.run.state.ring),
+                    buf: WlBufFdIn::new(&self.socket, &self.run.state.ring),
                 }
                 .run(),
             ),
@@ -246,7 +248,7 @@ impl Outgoing {
 
 struct Incoming {
     tc: Rc<TestTransport>,
-    buf: BufFdIn,
+    buf: WlBufFdIn,
 }
 
 impl Incoming {
@@ -267,30 +269,15 @@ impl Incoming {
     }
 
     async fn handle_msg(&mut self) -> Result<(), TestError> {
-        let mut hdr = [0u32, 0];
+        let WlMessage {
-        if let Err(e) = self.buf.read_full(&mut hdr[..]).await {
+            obj_id,
-            return Err(e.with_context("Could not read from wayland socket"));
+            message,
-        }
+            body,
-        let obj_id = ObjectId::from_raw(hdr[0]);
+            fds,
-        let len = (hdr[1] >> 16) as usize;
+        } = match self.buf.read_message().await {
-        let request = hdr[1] & 0xffff;
+            Ok(m) => m,
-        if len < 8 {
+            Err(e) => return Err(e.with_context("Could not read from wayland socket")),
-            bail!("Message size is < 8");
+        };
-        }
-        if len % 4 != 0 {
-            bail!("Message size is not a multiple of 4");
-        }
-        let len = len / 4 - 2;
-        let mut data_buf = self.tc.bufs.pop().unwrap_or_default();
-        data_buf.clear();
-        data_buf.reserve(len);
-        let unused = data_buf.split_at_spare_mut_ext().1;
-        if let Err(e) = self.buf.read_full(&mut unused[..len]).await {
-            return Err(e.with_context("Could not read from wayland socket"));
-        }
-        unsafe {
-            data_buf.set_len(len);
-        }
         let object = match self.tc.objects.get(&obj_id) {
             Some(obj) => obj,
             _ => bail!(
@@ -298,11 +285,8 @@ impl Incoming {
                 obj_id
             ),
         };
-        let parser = MsgParser::new(&mut self.buf, &data_buf);
+        let parser = MsgParser::new(fds, body);
-        object.handle_request(request, parser)?;
+        object.handle_request(message, parser)?;
-        if data_buf.capacity() > 0 {
-            self.tc.bufs.push(data_buf);
-        }
         Ok(())
     }
 }

src/it/testrun.rs

@@ -58,7 +58,6 @@ impl TestRun {
             run: self.clone(),
             socket,
             client_id: Cell::new(ClientId::from_raw(0)),
-            bufs: Default::default(),
             swapchain: Default::default(),
             flush_request: Default::default(),
             incoming: Default::default(),

src/tools/tool_client.rs

@@ -10,15 +10,13 @@ use {
             asyncevent::AsyncEvent,
             bitfield::Bitfield,
             buffd::{
-                BufFdError, BufFdIn, BufFdOut, MsgFormatter, MsgParser, MsgParserError, OutBuffer,
+                BufFdError, BufFdOut, MsgFormatter, MsgParser, MsgParserError, OutBuffer,
-                OutBufferSwapchain,
+                OutBufferSwapchain, WlBufFdIn, WlMessage,
             },
             clonecell::CloneCell,
             errorfmt::ErrorFmt,
             numcell::NumCell,
             oserror::OsError,
-            stack::Stack,
-            vec_ext::VecExt,
             xrd::xrd,
         },
         wheel::{Wheel, WheelError},
@@ -59,10 +57,6 @@ pub enum ToolClientError {
     SocketPathTooLong,
     #[error("Could not connect to the compositor")]
     Connect(#[source] IoUringError),
-    #[error("The message length is smaller than 8 bytes")]
-    MsgLenTooSmall,
-    #[error("The size of the message is not a multiple of 4")]
-    UnalignedMessage,
     #[error(transparent)]
     BufFdError(#[from] BufFdError),
     #[error("Could not parse a message of type {}", .0)]
@@ -85,7 +79,6 @@ pub struct ToolClient {
             AHashMap<u32, Rc<dyn Fn(&mut MsgParser) -> Result<(), ToolClientError>>>,
         >,
     >,
-    bufs: Stack<Vec<u32>>,
     swapchain: Rc<RefCell<OutBufferSwapchain>>,
     flush_request: AsyncEvent,
     pending_futures: RefCell<AHashMap<u32, SpawnedFuture<()>>>,
@@ -186,7 +179,6 @@ impl ToolClient {
             eng,
             obj_ids: RefCell::new(obj_ids),
             handlers: Default::default(),
-            bufs: Default::default(),
             swapchain: Default::default(),
             flush_request: Default::default(),
             pending_futures: Default::default(),
@@ -209,7 +201,7 @@ impl ToolClient {
                 "tool client incoming",
                 Incoming {
                     tc: slf.clone(),
-                    buf: BufFdIn::new(&socket, &slf.ring),
+                    buf: WlBufFdIn::new(&socket, &slf.ring),
                 }
                 .run(),
             ),
@@ -528,7 +520,7 @@ impl Outgoing {
 
 struct Incoming {
     tc: Rc<ToolClient>,
-    buf: BufFdIn,
+    buf: WlBufFdIn,
 }
 
 impl Incoming {
@@ -541,44 +533,27 @@ impl Incoming {
     }
 
     async fn handle_msg(&mut self) -> Result<(), ToolClientError> {
-        let mut hdr = [0u32, 0];
+        let WlMessage {
-        if let Err(e) = self.buf.read_full(&mut hdr[..]).await {
+            obj_id,
-            return Err(ToolClientError::Read(e));
+            message,
-        }
+            body,
-        let obj_id = ObjectId::from_raw(hdr[0]);
+            fds,
-        let len = (hdr[1] >> 16) as usize;
+        } = self
-        let request = hdr[1] & 0xffff;
+            .buf
-        if len < 8 {
+            .read_message()
-            return Err(ToolClientError::MsgLenTooSmall);
+            .await
-        }
+            .map_err(ToolClientError::Read)?;
-        if len % 4 != 0 {
-            return Err(ToolClientError::UnalignedMessage);
-        }
-        let len = len / 4 - 2;
-        let mut data_buf = self.tc.bufs.pop().unwrap_or_default();
-        data_buf.clear();
-        data_buf.reserve(len);
-        let unused = data_buf.split_at_spare_mut_ext().1;
-        if let Err(e) = self.buf.read_full(&mut unused[..len]).await {
-            return Err(ToolClientError::Read(e));
-        }
-        unsafe {
-            data_buf.set_len(len);
-        }
         let mut handler = None;
         {
             let handlers = self.tc.handlers.borrow_mut();
             if let Some(handlers) = handlers.get(&obj_id) {
-                handler = handlers.get(&request).cloned();
+                handler = handlers.get(&message).cloned();
             }
         }
         if let Some(handler) = handler {
-            let mut parser = MsgParser::new(&mut self.buf, &data_buf);
+            let mut parser = MsgParser::new(fds, body);
             handler(&mut parser)?;
         }
-        if data_buf.capacity() > 0 {
-            self.tc.bufs.push(data_buf);
-        }
         Ok(())
     }
 }

src/utils/buffd.rs

@@ -6,6 +6,7 @@ pub use {
     ei_parser::{EiMsgParser, EiMsgParserError},
     formatter::MsgFormatter,
     parser::{MsgParser, MsgParserError},
+    wl_buf_in::{WlBufFdIn, WlMessage},
 };
 
 mod buf_in;
@@ -14,6 +15,7 @@ mod ei_formatter;
 mod ei_parser;
 mod formatter;
 mod parser;
+mod wl_buf_in;
 
 #[derive(Debug, Error)]
 pub enum BufFdError {
@@ -29,6 +31,12 @@ pub enum BufFdError {
     Closed,
     #[error("The connection timed out")]
     Timeout,
+    #[error("Message size is not a multiple of 4")]
+    UnalignedMessageSize,
+    #[error("Message size is larger than 4096")]
+    MessageTooLarge,
+    #[error("Message size is smaller than 8")]
+    MessageTooSmall,
 }
 
 const BUF_SIZE: usize = 4096;

src/utils/buffd/parser.rs

@@ -1,7 +1,7 @@
 use {
-    crate::{fixed::Fixed, globals::GlobalName, object::ObjectId, utils::buffd::BufFdIn},
+    crate::{fixed::Fixed, globals::GlobalName, object::ObjectId},
     bstr::{BStr, ByteSlice},
-    std::{ptr, rc::Rc},
+    std::{collections::VecDeque, ptr, rc::Rc},
     thiserror::Error,
     uapi::{OwnedFd, Pod},
 };
@@ -27,14 +27,14 @@ pub enum MsgParserError {
 }
 
 pub struct MsgParser<'a, 'b> {
-    buf: &'a mut BufFdIn,
+    fds: &'a mut VecDeque<Rc<OwnedFd>>,
     pos: usize,
     data: &'b [u32],
 }
 
 impl<'a, 'b> MsgParser<'a, 'b> {
-    pub fn new(buf: &'a mut BufFdIn, data: &'b [u32]) -> Self {
+    pub fn new(fds: &'a mut VecDeque<Rc<OwnedFd>>, data: &'b [u32]) -> Self {
-        Self { buf, pos: 0, data }
+        Self { fds, pos: 0, data }
     }
 
     #[inline(always)]
@@ -112,8 +112,8 @@ impl<'a, 'b> MsgParser<'a, 'b> {
     }
 
     pub fn fd(&mut self) -> Result<Rc<OwnedFd>, MsgParserError> {
-        match self.buf.get_fd() {
+        match self.fds.pop_front() {
-            Ok(fd) => Ok(fd),
+            Some(fd) => Ok(fd),
             _ => Err(MsgParserError::MissingFd),
         }
     }

src/utils/buffd/wl_buf_in.rs

@@ -0,0 +1,125 @@
+use {
+    crate::{
+        io_uring::IoUring,
+        object::ObjectId,
+        utils::{
+            buf::Buf,
+            buffd::{BufFdError, MAX_IN_FD},
+        },
+    },
+    std::{collections::VecDeque, ptr, rc::Rc, slice},
+    uapi::OwnedFd,
+};
+
+const WORD_SIZE: usize = 4;
+const WORD_ALIGN: usize = 4;
+const HEADER_WORDS: usize = 2;
+const HEADER_SIZE: usize = HEADER_WORDS * WORD_SIZE;
+const MAX_MESSAGE_SIZE: usize = 4096;
+const BUF_SIZE: usize = 2 * MAX_MESSAGE_SIZE;
+
+pub struct WlBufFdIn {
+    fd: Rc<OwnedFd>,
+    ring: Rc<IoUring>,
+    fds: VecDeque<Rc<OwnedFd>>,
+    buf: Buf,
+    lo: usize,
+    len: usize,
+}
+
+pub struct WlMessage<'a> {
+    pub obj_id: ObjectId,
+    pub message: u32,
+    pub body: &'a [u32],
+    pub fds: &'a mut VecDeque<Rc<OwnedFd>>,
+}
+
+impl WlBufFdIn {
+    pub fn new(fd: &Rc<OwnedFd>, ring: &Rc<IoUring>) -> Self {
+        let buf = Buf::new(BUF_SIZE);
+        assert_eq!(buf.as_ptr() as usize % WORD_ALIGN, 0);
+        Self {
+            fd: fd.clone(),
+            ring: ring.clone(),
+            fds: Default::default(),
+            buf,
+            lo: Default::default(),
+            len: Default::default(),
+        }
+    }
+
+    pub async fn read_message(&mut self) -> Result<WlMessage<'_>, BufFdError> {
+        if self.len == 0 {
+            self.lo = 0;
+        }
+        if self.len < HEADER_SIZE {
+            if self.lo > 0 {
+                self.compact();
+            }
+            while self.len < HEADER_SIZE {
+                self.recvmsg().await?;
+            }
+        }
+        let hdr: &[u32] =
+            unsafe { slice::from_raw_parts(self.buf[self.lo..].as_ptr().cast(), HEADER_WORDS) };
+        let obj_id = ObjectId::from_raw(hdr[0]);
+        let len = (hdr[1] >> 16) as usize;
+        let message = hdr[1] & 0xffff;
+        if len & 3 != 0 {
+            return Err(BufFdError::UnalignedMessageSize);
+        }
+        if len > MAX_MESSAGE_SIZE {
+            return Err(BufFdError::MessageTooLarge);
+        }
+        if len < HEADER_SIZE {
+            return Err(BufFdError::MessageTooSmall);
+        }
+        if len > self.len {
+            if self.lo + self.len >= MAX_MESSAGE_SIZE {
+                self.compact();
+            }
+            while len > self.len {
+                self.recvmsg().await?;
+            }
+        }
+        let body: &[u32] = unsafe {
+            let words = (len - HEADER_SIZE) >> 2;
+            slice::from_raw_parts(self.buf[self.lo + HEADER_SIZE..].as_ptr().cast(), words)
+        };
+        self.lo += len;
+        self.len -= len;
+        Ok(WlMessage {
+            obj_id,
+            message,
+            body,
+            fds: &mut self.fds,
+        })
+    }
+
+    #[inline(always)]
+    fn compact(&mut self) {
+        unsafe {
+            let dst = self.buf.as_mut_ptr();
+            let src = dst.add(self.lo);
+            ptr::copy(src, dst, self.len);
+            self.lo = 0;
+        }
+    }
+
+    async fn recvmsg(&mut self) -> Result<(), BufFdError> {
+        let mut buf = self.buf.slice(self.lo + self.len..);
+        match self
+            .ring
+            .recvmsg(&self.fd, slice::from_mut(&mut buf), &mut self.fds)
+            .await
+        {
+            Ok(0) => return Err(BufFdError::Closed),
+            Ok(n) => self.len += n,
+            Err(e) => return Err(BufFdError::Ring(e)),
+        }
+        if self.fds.len() > MAX_IN_FD {
+            return Err(BufFdError::TooManyFds);
+        }
+        Ok(())
+    }
+}

src/wl_usr.rs

@@ -11,15 +11,14 @@ use {
             asyncevent::AsyncEvent,
             bitfield::Bitfield,
             buffd::{
-                BufFdError, BufFdIn, BufFdOut, MsgFormatter, MsgParser, MsgParserError, OutBuffer,
+                BufFdError, BufFdOut, MsgFormatter, MsgParser, MsgParserError, OutBuffer,
-                OutBufferSwapchain,
+                OutBufferSwapchain, WlBufFdIn, WlMessage,
             },
             clonecell::CloneCell,
             copyhashmap::CopyHashMap,
             errorfmt::ErrorFmt,
             hash_map_ext::HashMapExt,
             oserror::OsError,
-            vec_ext::VecExt,
         },
         video::dmabuf::DmaBufIds,
         wheel::Wheel,
@@ -51,10 +50,6 @@ pub enum UsrConError {
     SocketPathTooLong,
     #[error("Could not connect to the compositor")]
     Connect(#[source] IoUringError),
-    #[error("The message length is smaller than 8 bytes")]
-    MsgLenTooSmall,
-    #[error("The size of the message is not a multiple of 4")]
-    UnalignedMessage,
     #[error(transparent)]
     BufFdError(#[from] BufFdError),
     #[error("Could not read from the compositor")]
@@ -168,8 +163,7 @@ impl UsrCon {
                 "wl_usr incoming",
                 Incoming {
                     con: slf.clone(),
-                    buf: BufFdIn::new(socket, &slf.ring),
+                    buf: WlBufFdIn::new(socket, &slf.ring),
-                    data: vec![],
                 }
                 .run(),
             ),
@@ -337,8 +331,7 @@ impl Outgoing {
 
 struct Incoming {
     con: Rc<UsrCon>,
-    buf: BufFdIn,
+    buf: WlBufFdIn,
-    data: Vec<u32>,
 }
 
 impl Incoming {
@@ -357,33 +350,16 @@ impl Incoming {
     }
 
     async fn handle_msg(&mut self) -> Result<(), UsrConError> {
-        let mut hdr = [0u32, 0];
+        let WlMessage {
-        if let Err(e) = self.buf.read_full(&mut hdr[..]).await {
+            obj_id,
-            return Err(UsrConError::Read(e));
+            message,
-        }
+            body,
-        let obj_id = ObjectId::from_raw(hdr[0]);
+            fds,
-        let len = (hdr[1] >> 16) as usize;
+        } = self.buf.read_message().await.map_err(UsrConError::Read)?;
-        let event = hdr[1] & 0xffff;
-        if len < 8 {
-            return Err(UsrConError::MsgLenTooSmall);
-        }
-        if len % 4 != 0 {
-            return Err(UsrConError::UnalignedMessage);
-        }
-        let len = len / 4 - 2;
-        self.data.clear();
-        self.data.reserve(len);
-        let unused = self.data.split_at_spare_mut_ext().1;
-        if let Err(e) = self.buf.read_full(&mut unused[..len]).await {
-            return Err(UsrConError::Read(e));
-        }
-        unsafe {
-            self.data.set_len(len);
-        }
         if let Some(obj) = self.con.objects.get(&obj_id) {
             if let Some(obj) = obj {
-                let parser = MsgParser::new(&mut self.buf, &self.data);
+                let parser = MsgParser::new(fds, body);
-                obj.handle_event(&self.con, event, parser)?;
+                obj.handle_event(&self.con, message, parser)?;
             }
         } else if obj_id.raw() < MIN_SERVER_ID {
             return Err(UsrConError::MissingObject(obj_id));