refactor aligned buffer

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

pageserver/src/virtual_file.rs

@@ -18,6 +18,7 @@ use crate::page_cache::{PageWriteGuard, PAGE_SZ};
 use crate::tenant::TENANTS_SEGMENT_NAME;
 use camino::{Utf8Path, Utf8PathBuf};
 use once_cell::sync::OnceCell;
+use owned_buffers_io::aligned_buffer::{AlignedBufferMut, AlignedSlice, ConstAlign};
 use owned_buffers_io::io_buf_aligned::IoBufAlignedMut;
 use owned_buffers_io::io_buf_ext::FullSlice;
 use pageserver_api::config::defaults::DEFAULT_IO_BUFFER_ALIGNMENT;
@@ -45,7 +46,6 @@ pub(crate) use api::IoMode;
 pub(crate) use io_engine::IoEngineKind;
 pub(crate) use metadata::Metadata;
 pub(crate) use open_options::*;
-pub(crate) mod aligned_buffer;
 
 pub(crate) mod owned_buffers_io {
     //! Abstractions for IO with owned buffers.
@@ -57,6 +57,7 @@ pub(crate) mod owned_buffers_io {
     //! but for the time being we're proving out the primitives in the neon.git repo
     //! for faster iteration.
 
+    pub(crate) mod aligned_buffer;
     pub(crate) mod io_buf_aligned;
     pub(crate) mod io_buf_ext;
     pub(crate) mod slice;
@@ -1362,9 +1363,10 @@ pub(crate) const fn get_io_buffer_alignment() -> usize {
     DEFAULT_IO_BUFFER_ALIGNMENT
 }
 
-pub(crate) type IoBufferMut = aligned_buffer::AlignedBufferMut<{ get_io_buffer_alignment() }>;
+pub(crate) type IoBufferMut = AlignedBufferMut<ConstAlign<{ get_io_buffer_alignment() }>>;
+// pub(crate) type IoBuffer = AlignedBuffer<ConstAlign<{ get_io_buffer_alignment() }>>;
 pub(crate) type IoPageSlice<'a> =
-    aligned_buffer::AlignedSlice<'a, { get_io_buffer_alignment() }, PAGE_SZ>;
+    AlignedSlice<'a, PAGE_SZ, ConstAlign<{ get_io_buffer_alignment() }>>;
 
 static IO_MODE: AtomicU8 = AtomicU8::new(IoMode::preferred() as u8);
 

pageserver/src/virtual_file/owned_buffers_io/aligned_buffer.rs

@@ -0,0 +1,9 @@
+pub mod alignment;
+pub mod buffer;
+pub mod buffer_mut;
+pub mod raw;
+pub mod slice;
+
+pub use alignment::*;
+pub use buffer_mut::AlignedBufferMut;
+pub use slice::AlignedSlice;

pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/alignment.rs

@@ -0,0 +1,26 @@
+pub trait Alignment: std::marker::Unpin + 'static {
+    /// Returns the required alignments.
+    fn align(&self) -> usize;
+}
+
+/// Alignment at compile time.
+#[derive(Debug)]
+pub struct ConstAlign<const A: usize>;
+
+impl<const A: usize> Alignment for ConstAlign<A> {
+    fn align(&self) -> usize {
+        A
+    }
+}
+
+/// Alignment at run time.
+#[derive(Debug)]
+pub struct RuntimeAlign {
+    align: usize,
+}
+
+impl Alignment for RuntimeAlign {
+    fn align(&self) -> usize {
+        self.align
+    }
+}

pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/buffer.rs

@@ -0,0 +1,108 @@
+use std::{
+    ops::{Deref, Range},
+    sync::Arc,
+};
+
+use super::{alignment::Alignment, raw::RawAlignedBuffer};
+
+pub struct AlignedBuffer<A: Alignment> {
+    /// Shared raw buffer.
+    raw: Arc<RawAlignedBuffer<A>>,
+    /// Range that specifies the current slice.
+    range: Range<usize>,
+}
+
+impl<A: Alignment> AlignedBuffer<A> {
+    /// Creates an immutable `IoBuffer` from the raw buffer
+    pub(super) fn from_raw(raw: RawAlignedBuffer<A>, range: Range<usize>) -> Self {
+        AlignedBuffer {
+            raw: Arc::new(raw),
+            range,
+        }
+    }
+
+    /// Returns the number of bytes in the buffer, also referred to as its 'length'.
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.range.len()
+    }
+
+    /// Returns the alignment of the buffer.
+    #[inline]
+    pub fn align(&self) -> usize {
+        self.raw.align()
+    }
+
+    #[inline]
+    fn as_ptr(&self) -> *const u8 {
+        unsafe { self.raw.as_ptr().add(self.range.start) }
+    }
+
+    /// Extracts a slice containing the entire buffer.
+    ///
+    /// Equivalent to `&s[..]`.
+    #[inline]
+    fn as_slice(&self) -> &[u8] {
+        &self.raw.as_slice()[self.range.start..self.range.end]
+    }
+
+    /// Returns a slice of self for the index range `[begin..end)`.
+    pub fn slice(&self, begin: usize, end: usize) -> Self {
+        let len = self.len();
+
+        assert!(
+            begin <= end,
+            "range start must not be greater than end: {:?} <= {:?}",
+            begin,
+            end,
+        );
+        assert!(
+            end <= len,
+            "range end out of bounds: {:?} <= {:?}",
+            end,
+            len,
+        );
+
+        let begin = self.range.start + begin;
+        let end = self.range.start + end;
+
+        AlignedBuffer {
+            raw: Arc::clone(&self.raw),
+            range: begin..end,
+        }
+    }
+}
+
+impl<A: Alignment> Deref for AlignedBuffer<A> {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        self.as_slice()
+    }
+}
+
+impl<A: Alignment> AsRef<[u8]> for AlignedBuffer<A> {
+    fn as_ref(&self) -> &[u8] {
+        self.as_slice()
+    }
+}
+
+impl<A: Alignment> PartialEq<[u8]> for AlignedBuffer<A> {
+    fn eq(&self, other: &[u8]) -> bool {
+        self.as_slice().eq(other)
+    }
+}
+
+unsafe impl<A: Alignment> tokio_epoll_uring::IoBuf for AlignedBuffer<A> {
+    fn stable_ptr(&self) -> *const u8 {
+        self.as_ptr()
+    }
+
+    fn bytes_init(&self) -> usize {
+        self.len()
+    }
+
+    fn bytes_total(&self) -> usize {
+        self.len()
+    }
+}

pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/buffer_mut.rs

@@ -1,59 +1,17 @@
-#![allow(unused)]
+use std::ops::{Deref, DerefMut};
 
-use core::slice;
-use std::{
-    alloc::{self, Layout},
-    cmp,
-    mem::{ManuallyDrop, MaybeUninit},
-    ops::{Deref, DerefMut},
-    ptr::{addr_of_mut, NonNull},
+use super::{
+    alignment::{Alignment, ConstAlign},
+    buffer::AlignedBuffer,
+    raw::RawAlignedBuffer,
 };
 
-use bytes::buf::UninitSlice;
-
 #[derive(Debug)]
-struct IoBufferPtr(*mut u8);
-
-// SAFETY: We gurantees no one besides `IoBufferPtr` itself has the raw pointer.
-unsafe impl Send for IoBufferPtr {}
-
-/// An aligned buffer type used for I/O.
-#[derive(Debug)]
-pub struct AlignedBufferMut<const ALIGN: usize> {
-    ptr: IoBufferPtr,
-    capacity: usize,
-    len: usize,
+pub struct AlignedBufferMut<A: Alignment> {
+    raw: RawAlignedBuffer<A>,
 }
 
-pub struct AlignedSlice<'a, const ALIGN: usize, const N: usize>(&'a mut [u8; N]);
-
-impl<'a, const ALIGN: usize, const N: usize> AlignedSlice<'a, ALIGN, N> {
-    pub unsafe fn new_unchecked(buf: &'a mut [u8; N]) -> Self {
-        AlignedSlice(buf)
-    }
-}
-
-impl<'a, const ALIGN: usize, const N: usize> Deref for AlignedSlice<'a, ALIGN, N> {
-    type Target = [u8; N];
-
-    fn deref(&self) -> &Self::Target {
-        self.0
-    }
-}
-
-impl<'a, const ALIGN: usize, const N: usize> DerefMut for AlignedSlice<'a, ALIGN, N> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        self.0
-    }
-}
-
-impl<'a, const ALIGN: usize, const N: usize> AsRef<[u8; N]> for AlignedSlice<'a, ALIGN, N> {
-    fn as_ref(&self) -> &[u8; N] {
-        self.0
-    }
-}
-
-impl<const ALIGN: usize> AlignedBufferMut<ALIGN> {
+impl<const A: usize> AlignedBufferMut<ConstAlign<A>> {
     /// 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.
@@ -70,21 +28,8 @@ impl<const ALIGN: usize> AlignedBufferMut<ALIGN> {
     ///    must not overflow isize (i.e., the rounded value must be
     ///    less than or equal to `isize::MAX`).
     pub fn with_capacity(capacity: usize) -> Self {
-        let layout = Layout::from_size_align(capacity, ALIGN).expect("Invalid layout");
-
-        // 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);
-            }
-            IoBufferPtr(ptr)
-        };
-
         AlignedBufferMut {
-            ptr,
-            capacity,
-            len: 0,
+            raw: RawAlignedBuffer::with_capacity(capacity),
         }
     }
 
@@ -93,43 +38,45 @@ impl<const ALIGN: usize> AlignedBufferMut<ALIGN> {
         use bytes::BufMut;
         let mut buf = Self::with_capacity(capacity);
         buf.put_bytes(0, capacity);
-        buf.len = capacity;
+        // `put_bytes` filled the entire buffer.
+        unsafe { buf.set_len(capacity) };
         buf
     }
+}
 
+impl<A: Alignment> AlignedBufferMut<A> {
     /// Returns the total number of bytes the buffer can hold.
     #[inline]
     pub fn capacity(&self) -> usize {
-        self.capacity
+        self.raw.capacity()
     }
 
     /// Returns the alignment of the buffer.
     #[inline]
-    pub const fn align(&self) -> usize {
-        ALIGN
+    pub fn align(&self) -> usize {
+        self.raw.align()
     }
 
     /// Returns the number of bytes in the buffer, also referred to as its 'length'.
     #[inline]
     pub fn len(&self) -> usize {
-        self.len
+        self.raw.len()
     }
 
     /// Force the length of the buffer to `new_len`.
     #[inline]
     unsafe fn set_len(&mut self, new_len: usize) {
-        debug_assert!(new_len <= self.capacity());
-        self.len = new_len;
+        self.raw.set_len(new_len)
     }
 
     #[inline]
     fn as_ptr(&self) -> *const u8 {
-        self.ptr.0
+        self.raw.as_ptr()
     }
 
     #[inline]
     fn as_mut_ptr(&mut self) -> *mut u8 {
-        self.ptr.0
+        self.raw.as_mut_ptr()
     }
 
     /// Extracts a slice containing the entire buffer.
@@ -137,22 +84,20 @@ impl<const ALIGN: usize> AlignedBufferMut<ALIGN> {
     /// Equivalent to `&s[..]`.
     #[inline]
     fn as_slice(&self) -> &[u8] {
-        // SAFETY: The pointer is valid and `len` bytes are initialized.
-        unsafe { slice::from_raw_parts(self.as_ptr(), self.len) }
+        self.raw.as_slice()
     }
 
     /// Extracts a mutable slice of the entire buffer.
     ///
     /// Equivalent to `&mut s[..]`.
     fn as_mut_slice(&mut self) -> &mut [u8] {
-        // SAFETY: The pointer is valid and `len` bytes are initialized.
-        unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
+        self.raw.as_mut_slice()
     }
 
     /// Drops the all the contents of the buffer, setting its length to `0`.
     #[inline]
     pub fn clear(&mut self) {
-        self.len = 0;
+        self.raw.clear()
     }
 
     /// Reserves capacity for at least `additional` more bytes to be inserted
@@ -165,90 +110,26 @@ impl<const ALIGN: usize> AlignedBufferMut<ALIGN> {
     ///
     /// Panics if the new capacity exceeds `isize::MAX` _bytes_.
     pub fn reserve(&mut self, additional: usize) {
-        if additional > self.capacity() - self.len() {
-            self.reserve_inner(additional);
-        }
-    }
-
-    fn reserve_inner(&mut self, additional: usize) {
-        let Some(required_cap) = self.len().checked_add(additional) else {
-            capacity_overflow()
-        };
-
-        let old_capacity = self.capacity();
-        let align = self.align();
-        // This guarantees exponential growth. The doubling cannot overflow
-        // because `cap <= isize::MAX` and the type of `cap` is `usize`.
-        let cap = cmp::max(old_capacity * 2, required_cap);
-
-        if !is_valid_alloc(cap) {
-            capacity_overflow()
-        }
-        let new_layout = Layout::from_size_align(cap, self.align()).expect("Invalid layout");
-
-        let old_ptr = self.as_mut_ptr();
-
-        // SAFETY: old allocation was allocated with std::alloc::alloc with the same layout,
-        // and we panics on null pointer.
-        let (ptr, cap) = unsafe {
-            let old_layout = Layout::from_size_align_unchecked(old_capacity, align);
-            let ptr = alloc::realloc(old_ptr, old_layout, new_layout.size());
-            if ptr.is_null() {
-                alloc::handle_alloc_error(new_layout);
-            }
-            (IoBufferPtr(ptr), cap)
-        };
-
-        self.ptr = ptr;
-        self.capacity = cap;
+        self.raw.reserve(additional);
     }
 
     /// Shortens the buffer, keeping the first len bytes.
     pub fn truncate(&mut self, len: usize) {
-        if len > self.len {
-            return;
-        }
-        self.len = len;
+        self.raw.truncate(len);
     }
 
     /// Consumes and leaks the `IoBufferMut`, returning a mutable reference to the contents, &'a mut [u8].
     pub fn leak<'a>(self) -> &'a mut [u8] {
-        let mut buf = ManuallyDrop::new(self);
-        // SAFETY: leaking the buffer as intended.
-        unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr(), buf.len) }
+        self.raw.leak()
+    }
+
+    pub fn freeze(self) -> AlignedBuffer<A> {
+        let len = self.len();
+        AlignedBuffer::from_raw(self.raw, 0..len)
     }
 }
 
-fn capacity_overflow() -> ! {
-    panic!("capacity overflow")
-}
-
-// We need to guarantee the following:
-// * We don't ever allocate `> isize::MAX` byte-size objects.
-// * We don't overflow `usize::MAX` and actually allocate too little.
-//
-// On 64-bit we just need to check for overflow since trying to allocate
-// `> isize::MAX` bytes will surely fail. On 32-bit and 16-bit we need to add
-// an extra guard for this in case we're running on a platform which can use
-// all 4GB in user-space, e.g., PAE or x32.
-#[inline]
-fn is_valid_alloc(alloc_size: usize) -> bool {
-    !(usize::BITS < 64 && alloc_size > isize::MAX as usize)
-}
-
-impl<const ALIGN: usize> Drop for AlignedBufferMut<ALIGN> {
-    fn drop(&mut self) {
-        // SAFETY: memory was allocated with std::alloc::alloc with the same layout.
-        unsafe {
-            alloc::dealloc(
-                self.as_mut_ptr(),
-                Layout::from_size_align_unchecked(self.capacity, ALIGN),
-            )
-        }
-    }
-}
-
-impl<const ALIGN: usize> Deref for AlignedBufferMut<ALIGN> {
+impl<A: Alignment> Deref for AlignedBufferMut<A> {
     type Target = [u8];
 
     fn deref(&self) -> &Self::Target {
@@ -256,32 +137,32 @@ impl<const ALIGN: usize> Deref for AlignedBufferMut<ALIGN> {
     }
 }
 
-impl<const ALIGN: usize> DerefMut for AlignedBufferMut<ALIGN> {
+impl<A: Alignment> DerefMut for AlignedBufferMut<A> {
     fn deref_mut(&mut self) -> &mut Self::Target {
         self.as_mut_slice()
     }
 }
 
-impl<const ALIGN: usize> AsRef<[u8]> for AlignedBufferMut<ALIGN> {
+impl<A: Alignment> AsRef<[u8]> for AlignedBufferMut<A> {
     fn as_ref(&self) -> &[u8] {
         self.as_slice()
     }
 }
 
-impl<const ALIGN: usize> AsMut<[u8]> for AlignedBufferMut<ALIGN> {
+impl<A: Alignment> AsMut<[u8]> for AlignedBufferMut<A> {
     fn as_mut(&mut self) -> &mut [u8] {
         self.as_mut_slice()
     }
 }
 
-impl<const ALIGN: usize> PartialEq<[u8]> for AlignedBufferMut<ALIGN> {
+impl<A: Alignment> PartialEq<[u8]> for AlignedBufferMut<A> {
     fn eq(&self, other: &[u8]) -> bool {
         self.as_slice().eq(other)
     }
 }
 
 /// SAFETY: When advancing the internal cursor, the caller needs to make sure the bytes advcanced past have been initialized.
-unsafe impl<const ALIGN: usize> bytes::BufMut for AlignedBufferMut<ALIGN> {
+unsafe impl<A: Alignment> bytes::BufMut for AlignedBufferMut<A> {
     #[inline]
     fn remaining_mut(&self) -> usize {
         // Although a `Vec` can have at most isize::MAX bytes, we never want to grow `IoBufferMut`.
@@ -311,7 +192,9 @@ unsafe impl<const ALIGN: usize> bytes::BufMut for AlignedBufferMut<ALIGN> {
         // 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(self.as_mut_ptr().add(len), cap - len) }
+        unsafe {
+            bytes::buf::UninitSlice::from_raw_parts_mut(self.as_mut_ptr().add(len), cap - len)
+        }
     }
 }
 
@@ -326,7 +209,7 @@ fn panic_advance(idx: usize, len: usize) -> ! {
 
 /// Safety: [`IoBufferMut`] has exclusive ownership of the io buffer,
 /// and the location remains stable even if [`Self`] is moved.
-unsafe impl<const ALIGN: usize> tokio_epoll_uring::IoBuf for AlignedBufferMut<ALIGN> {
+unsafe impl<A: Alignment> tokio_epoll_uring::IoBuf for AlignedBufferMut<A> {
     fn stable_ptr(&self) -> *const u8 {
         self.as_ptr()
     }
@@ -341,7 +224,7 @@ unsafe impl<const ALIGN: usize> tokio_epoll_uring::IoBuf for AlignedBufferMut<AL
 }
 
 // SAFETY: See above.
-unsafe impl<const ALIGN: usize> tokio_epoll_uring::IoBufMut for AlignedBufferMut<ALIGN> {
+unsafe impl<A: Alignment> tokio_epoll_uring::IoBufMut for AlignedBufferMut<A> {
     fn stable_mut_ptr(&mut self) -> *mut u8 {
         self.as_mut_ptr()
     }
@@ -359,7 +242,7 @@ mod tests {
     use super::*;
 
     const ALIGN: usize = 4 * 1024;
-    type TestIoBufferMut = AlignedBufferMut<ALIGN>;
+    type TestIoBufferMut = AlignedBufferMut<ConstAlign<ALIGN>>;
 
     #[test]
     fn test_with_capacity() {

pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/raw.rs

@@ -0,0 +1,213 @@
+use core::slice;
+use std::{
+    alloc::{self, Layout},
+    cmp,
+    mem::ManuallyDrop,
+};
+
+use super::alignment::{Alignment, ConstAlign};
+
+#[derive(Debug)]
+struct AlignedBufferPtr(*mut u8);
+
+// SAFETY: We gurantees no one besides `IoBufferPtr` itself has the raw pointer.
+unsafe impl Send for AlignedBufferPtr {}
+
+/// An aligned buffer type used for I/O.
+#[derive(Debug)]
+pub struct RawAlignedBuffer<A: Alignment> {
+    ptr: AlignedBufferPtr,
+    capacity: usize,
+    len: usize,
+    align: A,
+}
+
+impl<const A: usize> RawAlignedBuffer<ConstAlign<A>> {
+    /// 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(capacity: usize) -> Self {
+        let align = ConstAlign::<A>;
+        let layout = Layout::from_size_align(capacity, align.align()).expect("Invalid layout");
+
+        // 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);
+            }
+            AlignedBufferPtr(ptr)
+        };
+
+        RawAlignedBuffer {
+            ptr,
+            capacity,
+            len: 0,
+            align,
+        }
+    }
+}
+
+impl<A: Alignment> RawAlignedBuffer<A> {
+    /// Returns the total number of bytes the buffer can hold.
+    #[inline]
+    pub fn capacity(&self) -> usize {
+        self.capacity
+    }
+
+    /// Returns the alignment of the buffer.
+    #[inline]
+    pub fn align(&self) -> usize {
+        self.align.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]
+    pub unsafe fn set_len(&mut self, new_len: usize) {
+        debug_assert!(new_len <= self.capacity());
+        self.len = new_len;
+    }
+
+    #[inline]
+    pub fn as_ptr(&self) -> *const u8 {
+        self.ptr.0
+    }
+
+    #[inline]
+    pub fn as_mut_ptr(&mut self) -> *mut u8 {
+        self.ptr.0
+    }
+
+    /// Extracts a slice containing the entire buffer.
+    ///
+    /// Equivalent to `&s[..]`.
+    #[inline]
+    pub fn as_slice(&self) -> &[u8] {
+        // SAFETY: The pointer is valid and `len` bytes are initialized.
+        unsafe { slice::from_raw_parts(self.as_ptr(), self.len) }
+    }
+
+    /// Extracts a mutable slice of the entire buffer.
+    ///
+    /// Equivalent to `&mut s[..]`.
+    pub fn as_mut_slice(&mut self) -> &mut [u8] {
+        // SAFETY: The pointer is valid and `len` bytes are initialized.
+        unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
+    }
+
+    /// Drops the all the contents of the buffer, setting its length to `0`.
+    #[inline]
+    pub fn clear(&mut self) {
+        self.len = 0;
+    }
+
+    /// Reserves capacity for at least `additional` more bytes to be inserted
+    /// in the given `IoBufferMut`. The collection may reserve more space to
+    /// speculatively avoid frequent reallocations. After calling `reserve`,
+    /// capacity will be greater than or equal to `self.len() + additional`.
+    /// Does nothing if capacity is already sufficient.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the new capacity exceeds `isize::MAX` _bytes_.
+    pub fn reserve(&mut self, additional: usize) {
+        if additional > self.capacity() - self.len() {
+            self.reserve_inner(additional);
+        }
+    }
+
+    fn reserve_inner(&mut self, additional: usize) {
+        let Some(required_cap) = self.len().checked_add(additional) else {
+            capacity_overflow()
+        };
+
+        let old_capacity = self.capacity();
+        let align = self.align();
+        // This guarantees exponential growth. The doubling cannot overflow
+        // because `cap <= isize::MAX` and the type of `cap` is `usize`.
+        let cap = cmp::max(old_capacity * 2, required_cap);
+
+        if !is_valid_alloc(cap) {
+            capacity_overflow()
+        }
+        let new_layout = Layout::from_size_align(cap, self.align()).expect("Invalid layout");
+
+        let old_ptr = self.as_mut_ptr();
+
+        // SAFETY: old allocation was allocated with std::alloc::alloc with the same layout,
+        // and we panics on null pointer.
+        let (ptr, cap) = unsafe {
+            let old_layout = Layout::from_size_align_unchecked(old_capacity, align);
+            let ptr = alloc::realloc(old_ptr, old_layout, new_layout.size());
+            if ptr.is_null() {
+                alloc::handle_alloc_error(new_layout);
+            }
+            (AlignedBufferPtr(ptr), cap)
+        };
+
+        self.ptr = ptr;
+        self.capacity = cap;
+    }
+
+    /// Shortens the buffer, keeping the first len bytes.
+    pub fn truncate(&mut self, len: usize) {
+        if len > self.len {
+            return;
+        }
+        self.len = len;
+    }
+
+    /// Consumes and leaks the `IoBufferMut`, returning a mutable reference to the contents, &'a mut [u8].
+    pub fn leak<'a>(self) -> &'a mut [u8] {
+        let mut buf = ManuallyDrop::new(self);
+        // SAFETY: leaking the buffer as intended.
+        unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr(), buf.len) }
+    }
+}
+
+fn capacity_overflow() -> ! {
+    panic!("capacity overflow")
+}
+
+// We need to guarantee the following:
+// * We don't ever allocate `> isize::MAX` byte-size objects.
+// * We don't overflow `usize::MAX` and actually allocate too little.
+//
+// On 64-bit we just need to check for overflow since trying to allocate
+// `> isize::MAX` bytes will surely fail. On 32-bit and 16-bit we need to add
+// an extra guard for this in case we're running on a platform which can use
+// all 4GB in user-space, e.g., PAE or x32.
+#[inline]
+fn is_valid_alloc(alloc_size: usize) -> bool {
+    !(usize::BITS < 64 && alloc_size > isize::MAX as usize)
+}
+
+impl<A: Alignment> Drop for RawAlignedBuffer<A> {
+    fn drop(&mut self) {
+        // SAFETY: memory was allocated with std::alloc::alloc with the same layout.
+        unsafe {
+            alloc::dealloc(
+                self.as_mut_ptr(),
+                Layout::from_size_align_unchecked(self.capacity, self.align.align()),
+            )
+        }
+    }
+}

pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/slice.rs

@@ -0,0 +1,40 @@
+use std::ops::{Deref, DerefMut};
+
+use super::alignment::{Alignment, ConstAlign};
+
+/// Newtype for an aligned slice.
+pub struct AlignedSlice<'a, const N: usize, A: Alignment> {
+    /// underlying byte slice
+    buf: &'a mut [u8; N],
+    /// alignment marker
+    _align: A,
+}
+
+impl<'a, const N: usize, const A: usize> AlignedSlice<'a, N, ConstAlign<A>> {
+    /// Create a new aligned slice from a mutable byte slice. The input must already satisify the alignment.
+    pub unsafe fn new_unchecked(buf: &'a mut [u8; N]) -> Self {
+        let _align = ConstAlign::<A>;
+        assert_eq!(buf.as_ptr().align_offset(_align.align()), 0);
+        AlignedSlice { buf, _align }
+    }
+}
+
+impl<'a, const N: usize, A: Alignment> Deref for AlignedSlice<'a, N, A> {
+    type Target = [u8; N];
+
+    fn deref(&self) -> &Self::Target {
+        self.buf
+    }
+}
+
+impl<'a, const N: usize, A: Alignment> DerefMut for AlignedSlice<'a, N, A> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.buf
+    }
+}
+
+impl<'a, const N: usize, A: Alignment> AsRef<[u8; N]> for AlignedSlice<'a, N, A> {
+    fn as_ref(&self) -> &[u8; N] {
+        &self.buf
+    }
+}