armdb 0.1.14

sharded bitcask key-value storage optimized for NVMe
Documentation
use std::alloc::{Layout, alloc, dealloc};
use std::ops::{Deref, DerefMut};

const ALIGNMENT: usize = 4096;

/// A byte buffer aligned to 4096 bytes, suitable for O_DIRECT I/O.
pub struct AlignedBuf {
    ptr: *mut u8,
    len: usize,
    capacity: usize,
}

impl AlignedBuf {
    /// Allocate a zeroed buffer of `len` bytes, aligned to 4096.
    pub fn zeroed(len: usize) -> Self {
        if len == 0 {
            return Self {
                ptr: std::ptr::NonNull::dangling().as_ptr(),
                len: 0,
                capacity: 0,
            };
        }
        // Round up to alignment
        let capacity = (len + ALIGNMENT - 1) & !(ALIGNMENT - 1);
        let layout = Layout::from_size_align(capacity, ALIGNMENT).expect("invalid layout");
        let ptr = unsafe { alloc(layout) };
        if ptr.is_null() {
            std::alloc::handle_alloc_error(layout);
        }
        unsafe {
            std::ptr::write_bytes(ptr, 0, capacity);
        }
        Self { ptr, len, capacity }
    }

    pub fn len(&self) -> usize {
        self.len
    }
}

impl Deref for AlignedBuf {
    type Target = [u8];

    fn deref(&self) -> &[u8] {
        if self.capacity == 0 {
            return &[];
        }
        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
    }
}

impl DerefMut for AlignedBuf {
    fn deref_mut(&mut self) -> &mut [u8] {
        if self.capacity == 0 {
            return &mut [];
        }
        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
    }
}

impl Drop for AlignedBuf {
    fn drop(&mut self) {
        if self.capacity == 0 {
            return;
        }
        let layout = Layout::from_size_align(self.capacity, ALIGNMENT).expect("invalid layout");
        unsafe {
            dealloc(self.ptr, layout);
        }
    }
}

// SAFETY: AlignedBuf owns its allocation exclusively.
unsafe impl Send for AlignedBuf {}
unsafe impl Sync for AlignedBuf {}