From 8f9679ce9527f61c42d5548cf46fffd00066f201 Mon Sep 17 00:00:00 2001 From: Yuchen Liang Date: Wed, 9 Oct 2024 10:30:25 -0400 Subject: [PATCH] refactor aligned buffer Signed-off-by: Yuchen Liang --- pageserver/src/virtual_file.rs | 8 +- .../owned_buffers_io/aligned_buffer.rs | 9 + .../aligned_buffer/alignment.rs | 26 +++ .../owned_buffers_io/aligned_buffer/buffer.rs | 108 +++++++++ .../aligned_buffer/buffer_mut.rs} | 203 ++++------------- .../owned_buffers_io/aligned_buffer/raw.rs | 213 ++++++++++++++++++ .../owned_buffers_io/aligned_buffer/slice.rs | 40 ++++ 7 files changed, 444 insertions(+), 163 deletions(-) create mode 100644 pageserver/src/virtual_file/owned_buffers_io/aligned_buffer.rs create mode 100644 pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/alignment.rs create mode 100644 pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/buffer.rs rename pageserver/src/virtual_file/{aligned_buffer.rs => owned_buffers_io/aligned_buffer/buffer_mut.rs} (59%) create mode 100644 pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/raw.rs create mode 100644 pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/slice.rs diff --git a/pageserver/src/virtual_file.rs b/pageserver/src/virtual_file.rs index 9243b19054..2d9f4218e7 100644 --- a/pageserver/src/virtual_file.rs +++ b/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>; +// pub(crate) type IoBuffer = AlignedBuffer>; 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); diff --git a/pageserver/src/virtual_file/owned_buffers_io/aligned_buffer.rs b/pageserver/src/virtual_file/owned_buffers_io/aligned_buffer.rs new file mode 100644 index 0000000000..8ffc29b93d --- /dev/null +++ b/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; diff --git a/pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/alignment.rs b/pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/alignment.rs new file mode 100644 index 0000000000..933b78a13b --- /dev/null +++ b/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; + +impl Alignment for ConstAlign { + 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 + } +} diff --git a/pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/buffer.rs b/pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/buffer.rs new file mode 100644 index 0000000000..4669a1b353 --- /dev/null +++ b/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 { + /// Shared raw buffer. + raw: Arc>, + /// Range that specifies the current slice. + range: Range, +} + +impl AlignedBuffer { + /// Creates an immutable `IoBuffer` from the raw buffer + pub(super) fn from_raw(raw: RawAlignedBuffer, range: Range) -> 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 Deref for AlignedBuffer { + type Target = [u8]; + + fn deref(&self) -> &Self::Target { + self.as_slice() + } +} + +impl AsRef<[u8]> for AlignedBuffer { + fn as_ref(&self) -> &[u8] { + self.as_slice() + } +} + +impl PartialEq<[u8]> for AlignedBuffer { + fn eq(&self, other: &[u8]) -> bool { + self.as_slice().eq(other) + } +} + +unsafe impl tokio_epoll_uring::IoBuf for AlignedBuffer { + fn stable_ptr(&self) -> *const u8 { + self.as_ptr() + } + + fn bytes_init(&self) -> usize { + self.len() + } + + fn bytes_total(&self) -> usize { + self.len() + } +} diff --git a/pageserver/src/virtual_file/aligned_buffer.rs b/pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/buffer_mut.rs similarity index 59% rename from pageserver/src/virtual_file/aligned_buffer.rs rename to pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/buffer_mut.rs index 99cb4a4c8d..ab617ad322 100644 --- a/pageserver/src/virtual_file/aligned_buffer.rs +++ b/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 { - ptr: IoBufferPtr, - capacity: usize, - len: usize, +pub struct AlignedBufferMut { + raw: RawAlignedBuffer, } -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 AlignedBufferMut { +impl AlignedBufferMut> { /// 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 AlignedBufferMut { /// 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 AlignedBufferMut { 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 AlignedBufferMut { /// 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 AlignedBufferMut { /// 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 AlignedBufferMut { /// /// 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 { + 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 Drop for AlignedBufferMut { - 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 Deref for AlignedBufferMut { +impl Deref for AlignedBufferMut { type Target = [u8]; fn deref(&self) -> &Self::Target { @@ -256,32 +137,32 @@ impl Deref for AlignedBufferMut { } } -impl DerefMut for AlignedBufferMut { +impl DerefMut for AlignedBufferMut { fn deref_mut(&mut self) -> &mut Self::Target { self.as_mut_slice() } } -impl AsRef<[u8]> for AlignedBufferMut { +impl AsRef<[u8]> for AlignedBufferMut { fn as_ref(&self) -> &[u8] { self.as_slice() } } -impl AsMut<[u8]> for AlignedBufferMut { +impl AsMut<[u8]> for AlignedBufferMut { fn as_mut(&mut self) -> &mut [u8] { self.as_mut_slice() } } -impl PartialEq<[u8]> for AlignedBufferMut { +impl PartialEq<[u8]> for AlignedBufferMut { 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 bytes::BufMut for AlignedBufferMut { +unsafe impl bytes::BufMut for AlignedBufferMut { #[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 bytes::BufMut for AlignedBufferMut { // 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 tokio_epoll_uring::IoBuf for AlignedBufferMut { +unsafe impl tokio_epoll_uring::IoBuf for AlignedBufferMut { fn stable_ptr(&self) -> *const u8 { self.as_ptr() } @@ -341,7 +224,7 @@ unsafe impl tokio_epoll_uring::IoBuf for AlignedBufferMut tokio_epoll_uring::IoBufMut for AlignedBufferMut { +unsafe impl tokio_epoll_uring::IoBufMut for AlignedBufferMut { 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; + type TestIoBufferMut = AlignedBufferMut>; #[test] fn test_with_capacity() { diff --git a/pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/raw.rs b/pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/raw.rs new file mode 100644 index 0000000000..b979606424 --- /dev/null +++ b/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 { + ptr: AlignedBufferPtr, + capacity: usize, + len: usize, + align: A, +} + +impl RawAlignedBuffer> { + /// 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::; + 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 RawAlignedBuffer { + /// 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 Drop for RawAlignedBuffer { + 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()), + ) + } + } +} diff --git a/pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/slice.rs b/pageserver/src/virtual_file/owned_buffers_io/aligned_buffer/slice.rs new file mode 100644 index 0000000000..ee90746f87 --- /dev/null +++ b/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> { + /// 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::; + 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 + } +}