remove aligned-vec, use ManuallyDrop<Vec<u8>>

Signed-off-by: Yuchen Liang <yuchen@neon.tech>

Cargo.lock

@@ -46,15 +46,6 @@ dependencies = [
  "memchr",
 ]
 
-[[package]]
-name = "aligned-vec"
-version = "0.6.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7e0966165eaf052580bd70eb1b32cb3d6245774c0104d1b2793e9650bf83b52a"
-dependencies = [
- "equator",
-]
-
 [[package]]
 name = "allocator-api2"
 version = "0.2.16"
@@ -1919,26 +1910,6 @@ dependencies = [
  "termcolor",
 ]
 
-[[package]]
-name = "equator"
-version = "0.2.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c35da53b5a021d2484a7cc49b2ac7f2d840f8236a286f84202369bd338d761ea"
-dependencies = [
- "equator-macro",
-]
-
-[[package]]
-name = "equator-macro"
-version = "0.2.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3bf679796c0322556351f287a51b49e48f7c4986e727b5dd78c972d30e2e16cc"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.52",
-]
-
 [[package]]
 name = "equivalent"
 version = "1.0.1"
@@ -3711,7 +3682,6 @@ dependencies = [
 name = "pageserver"
 version = "0.1.0"
 dependencies = [
- "aligned-vec",
  "anyhow",
  "arc-swap",
  "async-compression",

Cargo.toml

@@ -41,7 +41,6 @@ license = "Apache-2.0"
 
 ## All dependency versions, used in the project
 [workspace.dependencies]
-aligned-vec = "0.6.1"
 ahash = "0.8"
 anyhow = { version = "1.0", features = ["backtrace"] }
 arc-swap = "1.6"

pageserver/Cargo.toml

@@ -11,7 +11,6 @@ default = []
 testing = ["fail/failpoints"]
 
 [dependencies]
-aligned-vec.workspace = true
 anyhow.workspace = true
 arc-swap.workspace = true
 async-compression.workspace = true

pageserver/src/virtual_file/dio.rs

@@ -1,46 +1,87 @@
-use std::ops::{Deref, DerefMut};
+#![allow(unused)]
+
+use std::{
+    alloc::{self, Layout},
+    mem::ManuallyDrop,
+    ops::{Deref, DerefMut},
+};
 
-use aligned_vec::AVec;
 use bytes::buf::UninitSlice;
 
-pub const DIO_MEM_ALIGN: usize = 4096;
+// pub const DIO_MEM_ALIGN: usize = 4096;
 
-pub struct IoBufferMut(AVec<u8, aligned_vec::RuntimeAlign>);
+/// An aligned buffer type used for I/O.
+pub struct IoBufferMut {
+    /// Use `Vec` to benefit from the helper methods, but never reallocates.
+    buf: ManuallyDrop<Vec<u8>>,
+    align: usize,
+}
 
 impl IoBufferMut {
-    // pub fn new(align: usize) -> Self {
-    //     IoBufferMut(AVec::new(align))
-    // }
+    /// Constructs a new, empty `IoBufferMut` with at least the specified capacity and alignment.
+    ///
+    /// The buffer will be able to hold at most `capacity` elements and will never resize.
+    ///
+    ///
+    /// # Panics
+    ///
+    /// Panics if the new capacity exceeds `isize::MAX` _bytes_, or if the following alignment requirement is not met:
+    /// * `align` must not be zero,
+    ///
+    /// * `align` must be a power of two,
+    ///
+    /// * `capacity`, when rounded up to the nearest multiple of `align`,
+    ///    must not overflow isize (i.e., the rounded value must be
+    ///    less than or equal to `isize::MAX`).
+    pub fn with_capacity_aligned(capacity: usize, align: usize) -> Self {
+        let layout = Layout::from_size_align(capacity, align).expect("Invalid layout");
 
-    #[inline]
-    pub fn with_capacity(align: usize, capacity: usize) -> Self {
-        IoBufferMut(AVec::with_capacity(align, capacity))
+        let buf = unsafe {
+            let ptr = alloc::alloc(layout);
+            if ptr.is_null() {
+                alloc::handle_alloc_error(layout);
+            }
+            ManuallyDrop::new(Vec::from_raw_parts(ptr, 0, capacity))
+        };
+
+        IoBufferMut { buf, align }
     }
 
+    /// Returns the total number of elements the buffer can hold without
+    /// reallocating.
     #[inline]
     pub fn capacity(&self) -> usize {
-        self.0.capacity()
+        self.buf.capacity()
     }
 
+    /// Returns the alignment of the buffer.
     #[inline]
-    pub fn reserve(&mut self, additional: usize) {
-        self.0.reserve(additional)
+    pub fn align(&self) -> usize {
+        self.align
     }
 
+    /// Force the length of the buffer to `new_len`.
     #[inline]
-    pub unsafe fn set_len(&mut self, new_len: usize) {
-        self.0.set_len(new_len)
+    unsafe fn set_len(&mut self, new_len: usize) {
+        self.buf.set_len(new_len)
     }
 
-    #[inline]
-    pub fn extend_from_slice(&mut self, other: &[u8]) {
-        self.0.extend_from_slice(other)
-    }
-
-    /// Drops the all the elements of the vector, setting its length to `0`.
+    /// Drops the all the contents of the buffer, setting its length to `0`.
     #[inline]
     pub fn clear(&mut self) {
-        self.0.clear()
+        self.buf.clear()
+    }
+}
+
+impl Drop for IoBufferMut {
+    fn drop(&mut self) {
+        // SAFETY: memory was allocated with std::alloc::alloc
+        unsafe {
+            alloc::dealloc(
+                self.buf.as_ptr() as *mut u8,
+                Layout::from_size_align_unchecked(self.capacity(), self.align),
+            )
+        }
     }
 }
 
@@ -48,27 +89,28 @@ impl Deref for IoBufferMut {
     type Target = [u8];
 
     fn deref(&self) -> &Self::Target {
-        self.0.deref()
+        self.buf.deref()
     }
 }
 
 impl DerefMut for IoBufferMut {
     fn deref_mut(&mut self) -> &mut Self::Target {
-        self.0.deref_mut()
+        self.buf.deref_mut()
     }
 }
 
 unsafe impl bytes::BufMut for IoBufferMut {
     #[inline]
     fn remaining_mut(&self) -> usize {
-        // A vector can never have more than isize::MAX bytes
-        core::isize::MAX as usize - self.len()
+        // Although a `Vec` can have at most isize::MAX bytes, we never want to grow `IoBufferMut`.
+        // Thus, it can have at most `self.capacity` bytes.
+        self.capacity() - self.len()
     }
 
     #[inline]
     unsafe fn advance_mut(&mut self, cnt: usize) {
         let len = self.len();
-        let remaining = self.0.capacity() - len;
+        let remaining = self.remaining_mut();
 
         if remaining < cnt {
             panic_advance(cnt, remaining);
@@ -80,14 +122,10 @@ unsafe impl bytes::BufMut for IoBufferMut {
 
     #[inline]
     fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice {
-        if self.capacity() == self.len() {
-            self.reserve(self.0.alignment()); // Grow the vec by alignment
-        }
-
         let cap = self.capacity();
         let len = self.len();
 
-        let ptr = self.as_mut_ptr();
+        let ptr = self.buf.as_mut_ptr();
         // SAFETY: Since `ptr` is valid for `cap` bytes, `ptr.add(len)` must be
         // valid for `cap - len` bytes. The subtraction will not underflow since
         // `len <= cap`.
@@ -95,9 +133,18 @@ unsafe impl bytes::BufMut for IoBufferMut {
     }
 }
 
+/// Panic with a nice error message.
+#[cold]
+fn panic_advance(idx: usize, len: usize) -> ! {
+    panic!(
+        "advance out of bounds: the len is {} but advancing by {}",
+        len, idx
+    );
+}
+
 unsafe impl tokio_epoll_uring::IoBuf for IoBufferMut {
     fn stable_ptr(&self) -> *const u8 {
-        self.as_ptr()
+        self.buf.as_ptr()
     }
 
     fn bytes_init(&self) -> usize {
@@ -111,7 +158,7 @@ unsafe impl tokio_epoll_uring::IoBuf for IoBufferMut {
 
 unsafe impl tokio_epoll_uring::IoBufMut for IoBufferMut {
     fn stable_mut_ptr(&mut self) -> *mut u8 {
-        self.as_mut_ptr()
+        self.buf.as_mut_ptr()
     }
 
     unsafe fn set_init(&mut self, init_len: usize) {
@@ -121,11 +168,80 @@ unsafe impl tokio_epoll_uring::IoBufMut for IoBufferMut {
     }
 }
 
-/// Panic with a nice error message.
-#[cold]
-fn panic_advance(idx: usize, len: usize) -> ! {
-    panic!(
-        "advance out of bounds: the len is {} but advancing by {}",
-        len, idx
-    );
+#[cfg(test)]
+mod tests {
+
+    use super::*;
+
+    #[test]
+    fn test_with_capacity_aligned() {
+        const ALIGN: usize = 4 * 1024;
+        let v = IoBufferMut::with_capacity_aligned(ALIGN * 4, ALIGN);
+        assert_eq!(v.len(), 0);
+        assert_eq!(v.capacity(), ALIGN * 4);
+        assert_eq!(v.align(), ALIGN);
+        assert_eq!(v.as_ptr().align_offset(ALIGN), 0);
+
+        let v = IoBufferMut::with_capacity_aligned(ALIGN / 2, ALIGN);
+        assert_eq!(v.len(), 0);
+        assert_eq!(v.capacity(), ALIGN / 2);
+        assert_eq!(v.align(), ALIGN);
+        assert_eq!(v.as_ptr().align_offset(ALIGN), 0);
+    }
+
+    #[test]
+    fn test_bytes_put() {
+        use bytes::BufMut;
+        const ALIGN: usize = 4 * 1024;
+        let mut v = IoBufferMut::with_capacity_aligned(ALIGN * 4, ALIGN);
+        let x = [b'a'; ALIGN];
+
+        for _ in 0..2 {
+            for _ in 0..4 {
+                v.put(&x[..]);
+            }
+            assert_eq!(v.len(), ALIGN * 4);
+            assert_eq!(v.capacity(), ALIGN * 4);
+            assert_eq!(v.align(), ALIGN);
+            assert_eq!(v.as_ptr().align_offset(ALIGN), 0);
+            v.clear()
+        }
+        assert_eq!(v.len(), 0);
+        assert_eq!(v.capacity(), ALIGN * 4);
+        assert_eq!(v.align(), ALIGN);
+        assert_eq!(v.as_ptr().align_offset(ALIGN), 0);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_bytes_put_panic() {
+        use bytes::BufMut;
+        const ALIGN: usize = 4 * 1024;
+        let mut v = IoBufferMut::with_capacity_aligned(ALIGN * 4, ALIGN);
+        let x = [b'a'; ALIGN];
+        for _ in 0..5 {
+            v.put_slice(&x[..]);
+        }
+    }
+
+    #[test]
+    fn test_io_buf_put_slice() {
+        use tokio_epoll_uring::BoundedBufMut;
+        const ALIGN: usize = 4 * 1024;
+        let mut v = IoBufferMut::with_capacity_aligned(ALIGN, ALIGN);
+        let x = [b'a'; ALIGN];
+
+        for _ in 0..2 {
+            v.put_slice(&x[..]);
+            assert_eq!(v.len(), ALIGN);
+            assert_eq!(v.capacity(), ALIGN);
+            assert_eq!(v.align(), ALIGN);
+            assert_eq!(v.as_ptr().align_offset(ALIGN), 0);
+            v.clear()
+        }
+        assert_eq!(v.len(), 0);
+        assert_eq!(v.capacity(), ALIGN);
+        assert_eq!(v.align(), ALIGN);
+        assert_eq!(v.as_ptr().align_offset(ALIGN), 0);
+    }
 }

pageserver/src/virtual_file/owned_buffers_io/io_buf_aligned.rs

@@ -1,3 +1,5 @@
+#![allow(unused)]
+
 use tokio_epoll_uring::IoBufMut;
 
 use crate::virtual_file::dio::IoBufferMut;