Fix create_aligned_vec (#13)

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 {

src/storage/chunk_engine/src/alloc/chunk.rs

@@ -14,16 +14,16 @@ pub struct Chunk {
 pub type ChunkArc = Arc<Chunk>;
 
 lazy_static! {
-    static ref ZERO: Vec<u8> = {
-        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
     };
 }
 
 impl Chunk {
     thread_local! {
-        static BUFFER: RefCell<Vec<u8>> = RefCell::new(create_aligned_vec(CHUNK_SIZE_ULTRA));
+        static BUFFER: RefCell<AlignedBuffer> = RefCell::new(create_aligned_buf(CHUNK_SIZE_ULTRA));
     }
 
     pub fn new(meta: ChunkMeta, allocator: Arc<Allocator>) -> Self {

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);

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();

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<u8> {
-    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 {