diff --git a/src/storage/chunk_engine/src/alloc/allocator.rs b/src/storage/chunk_engine/src/alloc/allocator.rs index 91e5726..52b6f71 100644 --- a/src/storage/chunk_engine/src/alloc/allocator.rs +++ b/src/storage/chunk_engine/src/alloc/allocator.rs @@ -159,7 +159,7 @@ mod tests { .map(|i| { let chunks = chunks.clone(); std::thread::spawn(move || { - let mut vec = create_aligned_vec(ALIGN_SIZE); + let mut vec = create_aligned_buf(ALIGN_SIZE); vec.fill(0); for chunk in chunks.iter() { if chunk.meta().pos.index() as usize % T == i { diff --git a/src/storage/chunk_engine/src/alloc/chunk.rs b/src/storage/chunk_engine/src/alloc/chunk.rs index 7bb45d7..79d5815 100644 --- a/src/storage/chunk_engine/src/alloc/chunk.rs +++ b/src/storage/chunk_engine/src/alloc/chunk.rs @@ -14,8 +14,8 @@ pub struct Chunk { pub type ChunkArc = Arc; lazy_static! { - static ref ZERO: Vec = { - let mut vec = create_aligned_vec(CHUNK_SIZE_ULTRA); + static ref ZERO: AlignedBuffer = { + let mut vec = create_aligned_buf(CHUNK_SIZE_ULTRA); vec.fill(0); vec }; @@ -23,7 +23,7 @@ lazy_static! { impl Chunk { thread_local! { - static BUFFER: RefCell> = RefCell::new(create_aligned_vec(CHUNK_SIZE_ULTRA)); + static BUFFER: RefCell = RefCell::new(create_aligned_buf(CHUNK_SIZE_ULTRA)); } pub fn new(meta: ChunkMeta, allocator: Arc) -> Self { diff --git a/src/storage/chunk_engine/src/bin/bench.rs b/src/storage/chunk_engine/src/bin/bench.rs index 3b8c06a..cfe6606 100644 --- a/src/storage/chunk_engine/src/bin/bench.rs +++ b/src/storage/chunk_engine/src/bin/bench.rs @@ -48,7 +48,7 @@ fn main() -> Result<()> { let bytes = bytes.clone(); let running = running.clone(); - let mut vec = create_aligned_vec(CHUNK_SIZE_NORMAL); + let mut vec = create_aligned_buf(CHUNK_SIZE_NORMAL); vec.fill(i as u8); let checksum = crc32c::crc32c(&vec); running.fetch_add(1, Ordering::SeqCst); diff --git a/src/storage/chunk_engine/src/core/engine.rs b/src/storage/chunk_engine/src/core/engine.rs index 9002fee..84235a9 100644 --- a/src/storage/chunk_engine/src/core/engine.rs +++ b/src/storage/chunk_engine/src/core/engine.rs @@ -896,8 +896,8 @@ mod tests { // 2. write chunks. let engine = Engine::open(&config).unwrap(); const N: usize = 512; - for i in 0..N { - let mut data = create_aligned_vec(Size::from(i * 1024)); + for i in 1..=N { + let mut data = create_aligned_buf(Size::from(i * 1024)); data.fill(i as u8); let checksum = crc32c::crc32c(&data); let chunk = engine @@ -911,9 +911,9 @@ mod tests { let engine = Engine::open(&config).unwrap(); let chunks = engine.query_chunks([], [], u64::MAX).unwrap(); assert_eq!(chunks.len(), N); - let mut buffer = create_aligned_vec(CHUNK_SIZE_LARGE); + let mut buffer = create_aligned_buf(CHUNK_SIZE_LARGE); for (i, (chunk_id, chunk_meta)) in chunks.iter().enumerate() { - let i = N - 1 - i; // reversed order. + let i = N - i; // reversed order. assert_eq!(&i.to_be_bytes(), &chunk_id[..]); let offset = i * 512; let data_len = i * 1024; @@ -961,7 +961,7 @@ mod tests { // 2. write chunks. let engine = Engine::open(&config).unwrap(); - let mut data = create_aligned_vec(CHUNK_SIZE_SMALL); + let mut data = create_aligned_buf(CHUNK_SIZE_SMALL); for i in 0..512u32 { data.fill(i as u8); let checksum = crc32c::crc32c(&data); @@ -984,7 +984,7 @@ mod tests { drop(engine); let engine = Engine::open(&config).unwrap(); std::thread::sleep(std::time::Duration::from_millis(100)); - let mut buf = create_aligned_vec(CHUNK_SIZE_SMALL); + let mut buf = create_aligned_buf(CHUNK_SIZE_SMALL); for i in 0..512u32 { let chunk = engine.get(&i.to_le_bytes()).unwrap(); if i < 256 { @@ -1046,7 +1046,7 @@ mod tests { // 2. write chunks. let engine = Engine::open(&config).unwrap(); - let mut data = create_aligned_vec(CHUNK_SIZE_SMALL); + let mut data = create_aligned_buf(CHUNK_SIZE_SMALL); for i in 0..512u32 { data.fill(i as u8); let checksum = crc32c::crc32c(&data); @@ -1077,7 +1077,7 @@ mod tests { let engine = engine_clone; let stop = stop_clone; while !stop.load(Ordering::Acquire) { - let mut buf = create_aligned_vec(CHUNK_SIZE_SMALL); + let mut buf = create_aligned_buf(CHUNK_SIZE_SMALL); for i in 0..512u32 { let chunk = engine.get(&i.to_le_bytes()).unwrap(); if i % 2 == 0 || i == 255 || i == 511 { @@ -1128,7 +1128,7 @@ mod tests { // 2. write chunks. let engine = Engine::open(&config).unwrap(); - let mut data = create_aligned_vec(CHUNK_SIZE_SMALL); + let mut data = create_aligned_buf(CHUNK_SIZE_SMALL); for i in 0..512u32 { data.fill(i as u8); let checksum = crc32c::crc32c(&data); @@ -1148,7 +1148,7 @@ mod tests { } // 4. truncate and check. - let mut buf = create_aligned_vec(CHUNK_SIZE_NORMAL); + let mut buf = create_aligned_buf(CHUNK_SIZE_NORMAL); for i in 0..512u32 { let length = i * 137; let chunk = engine.truncate(&i.to_le_bytes(), length).unwrap(); diff --git a/src/storage/chunk_engine/src/utils/aligned.rs b/src/storage/chunk_engine/src/utils/aligned.rs index a1355dc..e47ffcf 100644 --- a/src/storage/chunk_engine/src/utils/aligned.rs +++ b/src/storage/chunk_engine/src/utils/aligned.rs @@ -1,11 +1,47 @@ use super::super::Size; +use std::alloc::Layout; pub const ALIGN_SIZE: Size = Size::new(512); -pub fn create_aligned_vec(size: Size) -> Vec { - let s: usize = size.into(); - let layout = std::alloc::Layout::from_size_align(s, ALIGN_SIZE.into()).unwrap(); - unsafe { Vec::from_raw_parts(std::alloc::alloc(layout), s, s) } +pub struct AlignedBuffer(&'static mut [u8]); + +impl AlignedBuffer { + pub fn new(size: usize) -> Self { + Self(unsafe { + let size = std::cmp::max(size, 1).next_multiple_of(ALIGN_SIZE.into()); + let layout = Layout::from_size_align_unchecked(size, ALIGN_SIZE.into()); + let ptr = std::alloc::alloc(layout); + std::slice::from_raw_parts_mut(ptr, size) + }) + } +} + +impl Drop for AlignedBuffer { + fn drop(&mut self) { + unsafe { + let layout = Layout::from_size_align_unchecked(self.0.len(), ALIGN_SIZE.into()); + std::alloc::dealloc(self.0.as_mut_ptr(), layout); + } + } +} + +impl std::ops::Deref for AlignedBuffer { + type Target = [u8]; + + #[inline(always)] + fn deref(&self) -> &Self::Target { + self.0 + } +} + +impl std::ops::DerefMut for AlignedBuffer { + fn deref_mut(&mut self) -> &mut Self::Target { + self.0 + } +} + +pub fn create_aligned_buf(size: Size) -> AlignedBuffer { + AlignedBuffer::new(size.into()) } pub fn is_aligned_buf(data: &[u8]) -> bool {