feat(pageserver): newtype aligned-vec as aligned buffer allocation

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

Cargo.lock

@@ -46,6 +46,15 @@ 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"
@@ -1910,6 +1919,26 @@ 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"
@@ -3682,6 +3711,7 @@ dependencies = [
 name = "pageserver"
 version = "0.1.0"
 dependencies = [
+ "aligned-vec",
  "anyhow",
  "arc-swap",
  "async-compression",
@@ -4147,7 +4177,7 @@ dependencies = [
 [[package]]
 name = "postgres-protocol"
 version = "0.6.4"
-source = "git+https://github.com/neondatabase/rust-postgres.git?branch=neon#20031d7a9ee1addeae6e0968e3899ae6bf01cee2"
+source = "git+https://github.com/neondatabase/rust-postgres.git?branch=neon#cff6927e4f58b1af6ecc2ee7279df1f2ff537295"
 dependencies = [
  "base64 0.20.0",
  "byteorder",

Cargo.toml

@@ -41,6 +41,7 @@ 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,6 +11,7 @@ default = []
 testing = ["fail/failpoints"]
 
 [dependencies]
+aligned-vec.workspace = true
 anyhow.workspace = true
 arc-swap.workspace = true
 async-compression.workspace = true

pageserver/src/virtual_file.rs

@@ -40,6 +40,7 @@ pub(crate) use api::DirectIoMode;
 pub(crate) use io_engine::IoEngineKind;
 pub(crate) use metadata::Metadata;
 pub(crate) use open_options::*;
+pub(crate) mod dio;
 
 pub(crate) mod owned_buffers_io {
     //! Abstractions for IO with owned buffers.

pageserver/src/virtual_file/dio.rs

@@ -0,0 +1,131 @@
+use std::ops::{Deref, DerefMut};
+
+use aligned_vec::AVec;
+use bytes::buf::UninitSlice;
+
+pub const DIO_MEM_ALIGN: usize = 4096;
+
+pub struct IoBufferMut(AVec<u8, aligned_vec::RuntimeAlign>);
+
+impl IoBufferMut {
+    // pub fn new(align: usize) -> Self {
+    //     IoBufferMut(AVec::new(align))
+    // }
+
+    #[inline]
+    pub fn with_capacity(align: usize, capacity: usize) -> Self {
+        IoBufferMut(AVec::with_capacity(align, capacity))
+    }
+
+    #[inline]
+    pub fn capacity(&self) -> usize {
+        self.0.capacity()
+    }
+
+    #[inline]
+    pub fn reserve(&mut self, additional: usize) {
+        self.0.reserve(additional)
+    }
+
+    #[inline]
+    pub unsafe fn set_len(&mut self, new_len: usize) {
+        self.0.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`.
+    #[inline]
+    pub fn clear(&mut self) {
+        self.0.clear()
+    }
+}
+
+impl Deref for IoBufferMut {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        self.0.deref()
+    }
+}
+
+impl DerefMut for IoBufferMut {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.0.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()
+    }
+
+    #[inline]
+    unsafe fn advance_mut(&mut self, cnt: usize) {
+        let len = self.len();
+        let remaining = self.0.capacity() - len;
+
+        if remaining < cnt {
+            panic_advance(cnt, remaining);
+        }
+
+        // Addition will not overflow since the sum is at most the capacity.
+        self.set_len(len + cnt);
+    }
+
+    #[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();
+        // 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`.
+        unsafe { UninitSlice::from_raw_parts_mut(ptr.add(len), cap - len) }
+    }
+}
+
+unsafe impl tokio_epoll_uring::IoBuf for IoBufferMut {
+    fn stable_ptr(&self) -> *const u8 {
+        self.as_ptr()
+    }
+
+    fn bytes_init(&self) -> usize {
+        self.len()
+    }
+
+    fn bytes_total(&self) -> usize {
+        self.capacity()
+    }
+}
+
+unsafe impl tokio_epoll_uring::IoBufMut for IoBufferMut {
+    fn stable_mut_ptr(&mut self) -> *mut u8 {
+        self.as_mut_ptr()
+    }
+
+    unsafe fn set_init(&mut self, init_len: usize) {
+        if self.len() < init_len {
+            self.set_len(init_len);
+        }
+    }
+}
+
+/// 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
+    );
+}

pageserver/src/virtual_file/owned_buffers_io/io_buf_ext.rs

@@ -1,5 +1,6 @@
 //! See [`FullSlice`].
 
+use crate::virtual_file::dio::IoBufferMut;
 use bytes::{Bytes, BytesMut};
 use std::ops::{Deref, Range};
 use tokio_epoll_uring::{BoundedBuf, IoBuf, Slice};
@@ -76,3 +77,4 @@ macro_rules! impl_io_buf_ext {
 impl_io_buf_ext!(Bytes);
 impl_io_buf_ext!(BytesMut);
 impl_io_buf_ext!(Vec<u8>);
+impl_io_buf_ext!(IoBufferMut);