From 852099bc831a2775413d3699b629f53051dce626 Mon Sep 17 00:00:00 2001 From: Yuchen Liang Date: Fri, 16 Aug 2024 17:13:30 +0000 Subject: [PATCH] remove aligned-vec, use ManuallyDrop> Signed-off-by: Yuchen Liang --- Cargo.lock | 30 --- Cargo.toml | 1 - pageserver/Cargo.toml | 1 - pageserver/src/virtual_file/dio.rs | 198 ++++++++++++++---- .../owned_buffers_io/io_buf_aligned.rs | 2 + 5 files changed, 159 insertions(+), 73 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 02e2257e9e..ebef4c1d1e 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index 6b81894fee..963841e340 100644 --- a/Cargo.toml +++ b/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" diff --git a/pageserver/Cargo.toml b/pageserver/Cargo.toml index 910b8530bf..0e748ee3db 100644 --- a/pageserver/Cargo.toml +++ b/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 diff --git a/pageserver/src/virtual_file/dio.rs b/pageserver/src/virtual_file/dio.rs index 682f61dafb..7aeaf2653e 100644 --- a/pageserver/src/virtual_file/dio.rs +++ b/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); +/// An aligned buffer type used for I/O. +pub struct IoBufferMut { + /// Use `Vec` to benefit from the helper methods, but never reallocates. + buf: ManuallyDrop>, + 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); + } } diff --git a/pageserver/src/virtual_file/owned_buffers_io/io_buf_aligned.rs b/pageserver/src/virtual_file/owned_buffers_io/io_buf_aligned.rs index eedb96a4b6..19f60123c4 100644 --- a/pageserver/src/virtual_file/owned_buffers_io/io_buf_aligned.rs +++ b/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;