From 6d664788c15e2960dabe0678260f18a2717eb263 Mon Sep 17 00:00:00 2001 From: Yuchen Liang Date: Thu, 15 Aug 2024 00:34:17 +0000 Subject: [PATCH] feat(pageserver): newtype aligned-vec as aligned buffer allocation Signed-off-by: Yuchen Liang --- Cargo.lock | 32 ++++- Cargo.toml | 1 + pageserver/Cargo.toml | 1 + pageserver/src/virtual_file.rs | 1 + pageserver/src/virtual_file/dio.rs | 131 ++++++++++++++++++ .../owned_buffers_io/io_buf_ext.rs | 2 + 6 files changed, 167 insertions(+), 1 deletion(-) create mode 100644 pageserver/src/virtual_file/dio.rs diff --git a/Cargo.lock b/Cargo.lock index dee15b6aa7..02e2257e9e 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index 963841e340..6b81894fee 100644 --- a/Cargo.toml +++ b/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" diff --git a/pageserver/Cargo.toml b/pageserver/Cargo.toml index 0e748ee3db..910b8530bf 100644 --- a/pageserver/Cargo.toml +++ b/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 diff --git a/pageserver/src/virtual_file.rs b/pageserver/src/virtual_file.rs index b4695e5f40..f80aec55fc 100644 --- a/pageserver/src/virtual_file.rs +++ b/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. diff --git a/pageserver/src/virtual_file/dio.rs b/pageserver/src/virtual_file/dio.rs new file mode 100644 index 0000000000..682f61dafb --- /dev/null +++ b/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); + +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 + ); +} diff --git a/pageserver/src/virtual_file/owned_buffers_io/io_buf_ext.rs b/pageserver/src/virtual_file/owned_buffers_io/io_buf_ext.rs index 7c773b6b21..fcfe60e19b 100644 --- a/pageserver/src/virtual_file/owned_buffers_io/io_buf_ext.rs +++ b/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); +impl_io_buf_ext!(IoBufferMut);