Optimise read buf initialization performance

Particularly improves large read_buffer_size performance

Author: alexheretic

CHANGELOG.md

@@ -1,3 +1,6 @@
+# Unreleased 
+* Optimise read performance particularly for larger than default `read_buffer_size`s.
+
 # 0.28.0
 * Reduce `Error` size 136 -> **32** by boxing internals of `Error::Http`, `Error::WriteBufferFull`,
   `ProtocolError::InvalidHeader`, `TlsError::Native`, `TlsError::Rustls`.

src/protocol/frame/init_aware_buf.rs

@@ -0,0 +1,221 @@
+use bytes::{Buf, BytesMut};
+use std::{
+    ops::{Deref, DerefMut},
+    ptr,
+};
+
+/// [`BytesMut`] wrapper that tracks initialization state of its spare capacity.
+///
+/// Supports safe & efficient repeated calls to [`Self::resize`] + [`Self::truncate`].
+///
+/// This optimisation is useful for [`std::io::Read`] to safely provide spare
+/// capacity as an initialized slice.
+///
+/// Related, may be obsoleted by: <https://github.com/rust-lang/rust/issues/78485>
+#[derive(Debug, Default)]
+pub struct InitAwareBuf {
+    bytes: BytesMut,
+    /// Capacity that has been initialized.
+    init_cap: usize,
+}
+
+impl InitAwareBuf {
+    #[inline]
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self { bytes: BytesMut::with_capacity(capacity), init_cap: 0 }
+    }
+
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.bytes.len()
+    }
+
+    #[inline]
+    pub fn capacity(&self) -> usize {
+        self.bytes.capacity()
+    }
+
+    #[inline]
+    pub fn split_to(&mut self, at: usize) -> BytesMut {
+        let split = self.bytes.split_to(at);
+        self.init_cap -= at;
+        split
+    }
+
+    #[inline]
+    pub fn reserve(&mut self, additional: usize) {
+        // Increasing capacity doesn't change `init_cap`
+        self.bytes.reserve(additional);
+    }
+
+    /// Sets the length of the buffer to `len`. If above the current
+    /// initialized capacity any uninitialized bytes will be zeroed.
+    ///
+    /// This is more efficient that [`BytesMut::resize`] as spare capacity
+    /// is only initialized **once** past the initialized_capacity. This
+    /// allow the method to be efficiently called after truncating.
+    ///
+    /// # Panics
+    /// Panics if `len > capacity`.
+    #[inline]
+    pub fn resize(&mut self, len: usize) {
+        if len <= self.init_cap {
+            // SAFETY: init_cap tracks initialised bytes.
+            unsafe {
+                self.bytes.set_len(len);
+            }
+        } else {
+            assert!(len <= self.capacity());
+            let cur_len = self.bytes.len();
+            let spare = self.bytes.spare_capacity_mut();
+            let already_init = self.init_cap - cur_len;
+            let zeroes = len - self.init_cap;
+            debug_assert!(already_init + zeroes <= spare.len());
+            unsafe {
+                // SAFETY: spare capacity is sufficient for `zeroes` extra bytes
+                ptr::write_bytes(spare[already_init..].as_mut_ptr().cast::<u8>(), 0, zeroes);
+                // SAFETY: len has been initialized
+                self.bytes.set_len(len);
+            }
+            self.init_cap = len;
+        }
+    }
+
+    #[inline]
+    pub fn truncate(&mut self, len: usize) {
+        // truncating doesn't change `init_cap`
+        self.bytes.truncate(len);
+    }
+
+    #[inline]
+    pub fn advance(&mut self, cnt: usize) {
+        self.bytes.advance(cnt);
+        self.init_cap -= cnt;
+    }
+}
+
+impl From<BytesMut> for InitAwareBuf {
+    #[inline]
+    fn from(bytes: BytesMut) -> Self {
+        let init_cap = bytes.len();
+        Self { bytes, init_cap }
+    }
+}
+
+impl From<InitAwareBuf> for BytesMut {
+    #[inline]
+    fn from(value: InitAwareBuf) -> Self {
+        value.bytes
+    }
+}
+
+impl AsRef<[u8]> for InitAwareBuf {
+    #[inline]
+    fn as_ref(&self) -> &[u8] {
+        &self.bytes
+    }
+}
+
+impl Deref for InitAwareBuf {
+    type Target = [u8];
+
+    #[inline]
+    fn deref(&self) -> &[u8] {
+        &self.bytes
+    }
+}
+
+impl AsMut<[u8]> for InitAwareBuf {
+    #[inline]
+    fn as_mut(&mut self) -> &mut [u8] {
+        &mut self.bytes
+    }
+}
+
+impl DerefMut for InitAwareBuf {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut [u8] {
+        &mut self.bytes
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn reserve_resize_truncate() {
+        let mut buf = InitAwareBuf::default();
+        assert_eq!(buf.len(), 0);
+        assert_eq!(buf.init_cap, 0);
+        assert_eq!(buf.capacity(), 0);
+
+        buf.reserve(64);
+        assert_eq!(buf.len(), 0);
+        assert_eq!(buf.init_cap, 0);
+        let new_capacity = buf.capacity();
+        assert!(new_capacity >= 64);
+
+        buf.resize(10);
+        assert_eq!(buf.len(), 10);
+        assert_eq!(buf.init_cap, 10);
+        assert_eq!(buf.capacity(), new_capacity);
+        assert_eq!(&*buf, &[0; 10]);
+
+        // write 3 bytes =8
+        buf[0] = 8;
+        buf[1] = 8;
+        buf[2] = 8;
+        // mark the other bytes as =44
+        for i in 3..10 {
+            buf[i] = 44;
+        }
+        buf.truncate(3);
+        assert_eq!(buf.len(), 3);
+        assert_eq!(buf.init_cap, 10);
+        assert_eq!(buf.capacity(), new_capacity);
+        assert_eq!(&*buf, &[8; 3]);
+
+        // resizing should need do nothing now since this has already been initialized once
+        buf.resize(10);
+        assert_eq!(buf.len(), 10);
+        assert_eq!(buf.init_cap, 10);
+        assert_eq!(buf.capacity(), new_capacity);
+        assert_eq!(&*buf, &[8, 8, 8, 44, 44, 44, 44, 44, 44, 44]);
+
+        buf.truncate(3);
+        assert_eq!(&*buf, &[8; 3]);
+
+        // resizing should only init to zero the 3 bytes that hadn't previously been
+        buf.resize(13);
+        assert_eq!(buf.len(), 13);
+        assert_eq!(buf.init_cap, 13);
+        assert_eq!(buf.capacity(), new_capacity);
+        assert_eq!(&*buf, &[8, 8, 8, 44, 44, 44, 44, 44, 44, 44, 0, 0, 0]);
+    }
+
+    #[test]
+    fn advance() {
+        let mut buf = InitAwareBuf::from(BytesMut::from(&[0, 1, 2, 3, 4][..]));
+        assert_eq!(buf.len(), 5);
+        assert_eq!(buf.init_cap, 5);
+
+        buf.advance(2);
+        assert_eq!(buf.len(), 3);
+        assert_eq!(buf.init_cap, 3);
+        assert_eq!(&*buf, &[2, 3, 4]);
+    }
+
+    #[test]
+    fn split_to() {
+        let mut buf = InitAwareBuf::from(BytesMut::from(&[0, 1, 2, 3, 4][..]));
+        assert_eq!(buf.len(), 5);
+        assert_eq!(buf.init_cap, 5);
+
+        let split = buf.split_to(2);
+        assert_eq!(buf.len(), 3);
+        assert_eq!(buf.init_cap, 3);
+        assert_eq!(&*buf, &[2, 3, 4]);
+        assert_eq!(&*split, &[0, 1]);
+    }
+}

src/protocol/frame/mod.rs

@@ -4,6 +4,7 @@ pub mod coding;
 
 #[allow(clippy::module_inception)]
 mod frame;
+mod init_aware_buf;
 mod mask;
 mod utf8;
 
@@ -14,7 +15,7 @@ pub use self::{
 
 use crate::{
     error::{CapacityError, Error, ProtocolError, Result},
-    protocol::frame::mask::apply_mask,
+    protocol::frame::{init_aware_buf::InitAwareBuf, mask::apply_mask},
     Message,
 };
 use bytes::BytesMut;
@@ -46,7 +47,7 @@ impl<Stream> FrameSocket<Stream> {
 
     /// Extract a stream from the socket.
     pub fn into_inner(self) -> (Stream, BytesMut) {
-        (self.stream, self.codec.in_buffer)
+        (self.stream, self.codec.in_buffer.into())
     }
 
     /// Returns a shared reference to the inner stream.
@@ -103,7 +104,7 @@ where
 #[derive(Debug)]
 pub(super) struct FrameCodec {
     /// Buffer to read data from the stream.
-    in_buffer: BytesMut,
+    in_buffer: InitAwareBuf,
     in_buf_max_read: usize,
     /// Buffer to send packets to the network.
     out_buffer: Vec<u8>,
@@ -123,7 +124,7 @@ impl FrameCodec {
     /// Create a new frame codec.
     pub(super) fn new(in_buf_len: usize) -> Self {
         Self {
-            in_buffer: BytesMut::with_capacity(in_buf_len),
+            in_buffer: InitAwareBuf::with_capacity(in_buf_len),
             in_buf_max_read: in_buf_len.max(FrameHeader::MAX_SIZE),
             out_buffer: <_>::default(),
             max_out_buffer_len: usize::MAX,
@@ -137,7 +138,7 @@ impl FrameCodec {
         let mut in_buffer = BytesMut::from_iter(part);
         in_buffer.reserve(min_in_buf_len.saturating_sub(in_buffer.len()));
         Self {
-            in_buffer,
+            in_buffer: in_buffer.into(),
             in_buf_max_read: min_in_buf_len.max(FrameHeader::MAX_SIZE),
             out_buffer: <_>::default(),
             max_out_buffer_len: usize::MAX,
@@ -172,7 +173,7 @@ impl FrameCodec {
                 let mut cursor = Cursor::new(&mut self.in_buffer);
                 self.header = FrameHeader::parse(&mut cursor)?;
                 let advanced = cursor.position();
-                bytes::Buf::advance(&mut self.in_buffer, advanced as _);
+                self.in_buffer.advance(advanced as _);
 
                 if let Some((_, len)) = &self.header {
                     let len = *len as usize;
@@ -233,7 +234,7 @@ impl FrameCodec {
     fn read_in(&mut self, stream: &mut impl Read) -> io::Result<usize> {
         let len = self.in_buffer.len();
         debug_assert!(self.in_buffer.capacity() > len);
-        self.in_buffer.resize(self.in_buffer.capacity().min(len + self.in_buf_max_read), 0);
+        self.in_buffer.resize(self.in_buffer.capacity().min(len + self.in_buf_max_read));
         let size = stream.read(&mut self.in_buffer[len..]);
         self.in_buffer.truncate(len + size.as_ref().copied().unwrap_or(0));
         size