Rework as using underlying bytes.len as capacity

Author: alexheretic

src/protocol/frame/init_aware_buf.rs

@@ -1,118 +1,113 @@
 use bytes::{Buf, BytesMut};
-use std::{
-    ops::{Deref, DerefMut},
-    ptr,
-};
+use std::ops::{Deref, DerefMut};
 
-/// [`BytesMut`] wrapper that tracks initialization state of its spare capacity.
+/// Buffer that provides fast & safe [`Self::resize`] & [`Self::truncate`] usage.
 ///
-/// Supports safe & efficient repeated calls to [`Self::resize`] + [`Self::truncate`].
+/// It is aware of the initialization state of its spare capacity avoiding the
+/// need to zero uninitialized bytes on resizing more than once for safe usage.
 ///
 /// 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 {
+    /// Backing buf length is used as capacity. This ensure this extra region
+    /// is always initialized (initially with zero, but otherwise the last previously
+    /// set value).
     bytes: BytesMut,
-    /// Capacity that has been initialized.
-    init_cap: usize,
+    /// Length of bytes in use (always <= bytes.len).
+    len: usize,
 }
 
 impl InitAwareBuf {
     #[inline]
     pub fn with_capacity(capacity: usize) -> Self {
-        Self { bytes: BytesMut::with_capacity(capacity), init_cap: 0 }
+        Self { bytes: BytesMut::zeroed(capacity), len: 0 }
     }
 
     #[inline]
     pub fn len(&self) -> usize {
-        self.bytes.len()
+        self.len
     }
 
+    /// Capacity that may be resized to cheaply.
     #[inline]
     pub fn capacity(&self) -> usize {
-        self.bytes.capacity()
+        self.bytes.len()
     }
 
     #[inline]
     pub fn split_to(&mut self, at: usize) -> BytesMut {
+        assert!(at <= self.len, "split_to out of bounds: {at} <= {}", self.len);
         let split = self.bytes.split_to(at);
-        self.init_cap -= at;
+        self.len -= at;
         split
     }
 
+    /// Reserve capacity for `min_additional` more bytes than the current [`Self::len`].
+    ///
+    /// `max_additional` sets the maximum number of additional bytes zeroed as extra
+    /// capacity if available after reserving in the underlying buffer. Has no effect
+    /// if `max_additional <= additional`.
     #[inline]
-    pub fn reserve(&mut self, additional: usize) {
-        // Increasing capacity doesn't change `init_cap`
-        self.bytes.reserve(additional);
+    pub fn reserve(&mut self, additional: usize, max_additional: usize) {
+        let min_len = self.len + additional;
+        let cap = self.capacity();
+        if min_len > cap {
+            self.bytes.reserve(min_len - cap);
+            let new_cap = self.bytes.capacity().min(self.len + max_additional.max(additional));
+            self.bytes.resize(new_cap, 0);
+        }
     }
 
-    /// 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.
+    /// Resizes the buffer to `new_len`.
     ///
-    /// # Panics
-    /// Panics if `len > capacity`.
+    /// If greater the new bytes will be either initialized to zero or as
+    /// they were last set to.
     #[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;
+    pub fn resize(&mut self, new_len: usize) {
+        if new_len > self.capacity() {
+            self.bytes.resize(new_len, 0);
         }
+        self.len = new_len;
     }
 
     #[inline]
     pub fn truncate(&mut self, len: usize) {
-        // truncating doesn't change `init_cap`
-        self.bytes.truncate(len);
+        if len < self.len {
+            self.len = len;
+        }
     }
 
     #[inline]
     pub fn advance(&mut self, cnt: usize) {
+        assert!(cnt <= self.len, "cannot advance past len: {cnt} <= {}", self.len);
         self.bytes.advance(cnt);
-        self.init_cap -= cnt;
+        self.len -= cnt;
     }
 }
 
 impl From<BytesMut> for InitAwareBuf {
     #[inline]
     fn from(bytes: BytesMut) -> Self {
-        let init_cap = bytes.len();
-        Self { bytes, init_cap }
+        let len = bytes.len();
+        Self { bytes, len }
     }
 }
 
 impl From<InitAwareBuf> for BytesMut {
     #[inline]
-    fn from(value: InitAwareBuf) -> Self {
-        value.bytes
+    fn from(mut zb: InitAwareBuf) -> Self {
+        zb.bytes.truncate(zb.len);
+        zb.bytes
     }
 }
 
 impl AsRef<[u8]> for InitAwareBuf {
     #[inline]
     fn as_ref(&self) -> &[u8] {
-        &self.bytes
+        &self.bytes[..self.len]
     }
 }
 
@@ -121,21 +116,21 @@ impl Deref for InitAwareBuf {
 
     #[inline]
     fn deref(&self) -> &[u8] {
-        &self.bytes
+        self.as_ref()
     }
 }
 
 impl AsMut<[u8]> for InitAwareBuf {
     #[inline]
     fn as_mut(&mut self) -> &mut [u8] {
-        &mut self.bytes
+        &mut self.bytes[..self.len]
     }
 }
 
 impl DerefMut for InitAwareBuf {
     #[inline]
     fn deref_mut(&mut self) -> &mut [u8] {
-        &mut self.bytes
+        self.as_mut()
     }
 }
 
@@ -147,18 +142,15 @@ mod 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);
+        buf.reserve(64, 0);
         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]);
 
@@ -172,14 +164,12 @@ mod test {
         }
         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]);
 
@@ -189,7 +179,6 @@ mod test {
         // 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]);
     }
@@ -198,23 +187,23 @@ mod 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);
+        assert_eq!(buf.capacity(), 5);
 
         buf.advance(2);
         assert_eq!(buf.len(), 3);
-        assert_eq!(buf.init_cap, 3);
+        assert_eq!(buf.capacity(), 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);
+        assert_eq!(buf.capacity(), 5);
 
         let split = buf.split_to(2);
         assert_eq!(buf.len(), 3);
-        assert_eq!(buf.init_cap, 3);
+        assert_eq!(buf.capacity(), 3);
         assert_eq!(&*buf, &[2, 3, 4]);
         assert_eq!(&*split, &[0, 1]);
     }

src/protocol/frame/mod.rs

@@ -135,10 +135,10 @@ impl FrameCodec {
 
     /// Create a new frame codec from partially read data.
     pub(super) fn from_partially_read(part: Vec<u8>, min_in_buf_len: usize) -> Self {
-        let mut in_buffer = BytesMut::from_iter(part);
-        in_buffer.reserve(min_in_buf_len.saturating_sub(in_buffer.len()));
+        let mut in_buffer = InitAwareBuf::from(BytesMut::from_iter(part));
+        in_buffer.reserve(min_in_buf_len.saturating_sub(in_buffer.len()), min_in_buf_len);
         Self {
-            in_buffer: in_buffer.into(),
+            in_buffer,
             in_buf_max_read: min_in_buf_len.max(FrameHeader::MAX_SIZE),
             out_buffer: <_>::default(),
             max_out_buffer_len: usize::MAX,
@@ -188,9 +188,9 @@ impl FrameCodec {
 
                     // Reserve full message length only once, even for multiple
                     // loops or if WouldBlock errors cause multiple fn calls.
-                    self.in_buffer.reserve(len);
+                    self.in_buffer.reserve(len, self.in_buf_max_read);
                 } else {
-                    self.in_buffer.reserve(FrameHeader::MAX_SIZE);
+                    self.in_buffer.reserve(FrameHeader::MAX_SIZE, self.in_buf_max_read);
                 }
             }