From 027f28deb93cfe70e33271eca5f4dae75e1736e9 Mon Sep 17 00:00:00 2001
From: Yuchen Liang <yuchen@neon.tech>
Date: Tue, 3 Sep 2024 11:38:00 -0400
Subject: [PATCH] remove Vec dependency

Signed-off-by: Yuchen Liang <yuchen@neon.tech>
---
 pageserver/src/virtual_file/dio.rs | 80 ++++++++++++++++++++++--------
 1 file changed, 60 insertions(+), 20 deletions(-)

diff --git a/pageserver/src/virtual_file/dio.rs b/pageserver/src/virtual_file/dio.rs
index feec21382e..4c66f82032 100644
--- a/pageserver/src/virtual_file/dio.rs
+++ b/pageserver/src/virtual_file/dio.rs
@@ -1,9 +1,11 @@
 #![allow(unused)]
 
+use core::slice;
 use std::{
     alloc::{self, Layout},
-    mem::ManuallyDrop,
+    mem::{ManuallyDrop, MaybeUninit},
     ops::{Deref, DerefMut},
+    ptr::NonNull,
 };
 
 use bytes::buf::UninitSlice;
@@ -13,7 +15,10 @@ use bytes::buf::UninitSlice;
 /// 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>>,
+    // buf: Option<Box<[MaybeUninit<u8>]>>,
+    ptr: *mut u8,
+    capacity: usize,
+    len: usize,
     align: usize,
 }
 
@@ -33,26 +38,40 @@ impl IoBufferMut {
     /// * `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`).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut buf = IoBufferMut::with_capacity_aligned(4096, 4096);
+    ///
+    /// assert_eq!(buf.len(), 0);
+    /// assert_eq!(buf.capacity(), 4096);
+    /// assert_eq!(buf.align(), 4096);
+    /// ```
     pub fn with_capacity_aligned(capacity: usize, align: usize) -> Self {
         let layout = Layout::from_size_align(capacity, align).expect("Invalid layout");
 
-        // SAFETY:  Making an allocation and construct a `Vec` from raw ptr. The memory is manually freed with the same layout.
-        let buf = unsafe {
+        // SAFETY:  Making an allocation with a sized and aligned layout. The memory is manually freed with the same layout.
+        let ptr = unsafe {
             let ptr = alloc::alloc(layout);
             if ptr.is_null() {
                 alloc::handle_alloc_error(layout);
             }
-            ManuallyDrop::new(Vec::from_raw_parts(ptr, 0, capacity))
+            ptr
         };
 
-        IoBufferMut { buf, align }
+        IoBufferMut {
+            ptr,
+            capacity,
+            len: 0,
+            align,
+        }
     }
 
-    /// Returns the total number of elements the buffer can hold without
-    /// reallocating.
+    /// Returns the total number of bytes the buffer can hold.
     #[inline]
     pub fn capacity(&self) -> usize {
-        self.buf.capacity()
+        self.capacity
     }
 
     /// Returns the alignment of the buffer.
@@ -61,16 +80,38 @@ impl IoBufferMut {
         self.align
     }
 
+    /// Returns the number of bytes in the buffer, also referred to as its 'length'.
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.len
+    }
+
     /// Force the length of the buffer to `new_len`.
     #[inline]
     unsafe fn set_len(&mut self, new_len: usize) {
-        self.buf.set_len(new_len)
+        debug_assert!(new_len <= self.capacity());
+        self.len = new_len;
+    }
+
+    /// Extracts a slice containing the entire buffer.
+    ///
+    /// Equivalent to `&s[..]`.
+    #[inline]
+    fn as_slice(&self) -> &[u8] {
+        unsafe { slice::from_raw_parts(self.ptr, self.len) }
+    }
+
+    /// Extracts a mutable slice of the entire buffer.
+    ///
+    /// Equivalent to `&mut s[..]`.
+    fn as_mut_slice(&mut self) -> &mut [u8] {
+        unsafe { slice::from_raw_parts_mut(self.ptr, self.len) }
     }
 
     /// Drops the all the contents of the buffer, setting its length to `0`.
     #[inline]
     pub fn clear(&mut self) {
-        self.buf.clear()
+        self.len = 0;
     }
 }
 
@@ -79,8 +120,8 @@ impl Drop for IoBufferMut {
         // SAFETY: memory was allocated with std::alloc::alloc with the same layout.
         unsafe {
             alloc::dealloc(
-                self.buf.as_ptr() as *mut u8,
-                Layout::from_size_align_unchecked(self.capacity(), self.align),
+                self.ptr,
+                Layout::from_size_align_unchecked(self.capacity, self.align),
             )
         }
     }
@@ -90,13 +131,13 @@ impl Deref for IoBufferMut {
     type Target = [u8];
 
     fn deref(&self) -> &Self::Target {
-        self.buf.deref()
+        self.as_slice()
     }
 }
 
 impl DerefMut for IoBufferMut {
     fn deref_mut(&mut self) -> &mut Self::Target {
-        self.buf.deref_mut()
+        self.as_mut_slice()
     }
 }
 
@@ -128,11 +169,10 @@ unsafe impl bytes::BufMut for IoBufferMut {
         let cap = self.capacity();
         let len = self.len();
 
-        let ptr = self.buf.as_mut_ptr();
-        // SAFETY: Since `ptr` is valid for `cap` bytes, `ptr.add(len)` must be
+        // SAFETY: Since `self.ptr` is valid for `cap` bytes, `self.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 { UninitSlice::from_raw_parts_mut(self.ptr.add(len), cap - len) }
     }
 }
 
@@ -149,7 +189,7 @@ fn panic_advance(idx: usize, len: usize) -> ! {
 // and the location remains stable even if [`Self`] is moved.
 unsafe impl tokio_epoll_uring::IoBuf for IoBufferMut {
     fn stable_ptr(&self) -> *const u8 {
-        self.buf.as_ptr()
+        self.ptr
     }
 
     fn bytes_init(&self) -> usize {
@@ -164,7 +204,7 @@ unsafe impl tokio_epoll_uring::IoBuf for IoBufferMut {
 // SAFETY: See above.
 unsafe impl tokio_epoll_uring::IoBufMut for IoBufferMut {
     fn stable_mut_ptr(&mut self) -> *mut u8 {
-        self.buf.as_mut_ptr()
+        self.ptr
     }
 
     unsafe fn set_init(&mut self, init_len: usize) {