feat: introduce aligned buffer for DIO

storage/Cargo.toml

@@ -27,6 +27,7 @@ lazy_static = "1.4.0"
 libc = "0.2"
 lz4 = "1.23.1"
 moka = { version = "0.7", features = ["future"] }
+nix = { version = "0.24.1", features = ["fs", "aio"] }
 parking_lot = "0.12"
 prometheus = "0.13.0"
 rand = "0.8.5"

storage/src/file_cache/buffer.rs

@@ -0,0 +1,267 @@
+use std::ops::RangeBounds;
+
+pub fn align_up(v: usize, align: usize) -> usize {
+    (v + align - 1) & !(align - 1)
+}
+
+pub fn align_down(v: usize, align: usize) -> usize {
+    v & !(align - 1)
+}
+
+pub struct AlignedBuffer<const ALIGN: usize, const SMOOTH: usize, const DEFAULT: usize> {
+    ptr: *mut u8,
+    len: usize,
+    capacity: usize,
+}
+
+impl<const ALIGN: usize, const SMOOTH: usize, const DEFAULT: usize>
+    AlignedBuffer<ALIGN, SMOOTH, DEFAULT>
+{
+    pub fn with_capacity(capacity: usize) -> Self {
+        let capacity = Self::align_up(capacity);
+        let buffer = unsafe {
+            std::alloc::alloc_zeroed(std::alloc::Layout::from_size_align_unchecked(
+                capacity, ALIGN,
+            ))
+        };
+
+        Self {
+            ptr: buffer,
+            len: 0,
+            capacity,
+        }
+    }
+
+    pub fn align(&self) -> usize {
+        ALIGN
+    }
+
+    pub fn capacity(&self) -> usize {
+        self.capacity
+    }
+
+    pub fn len(&self) -> usize {
+        self.len
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len == 0
+    }
+
+    pub fn remains(&self) -> usize {
+        self.capacity - self.len
+    }
+
+    pub fn write_at(&mut self, src: &[u8], offset: usize) {
+        let len = src.len();
+        if offset + len > self.capacity {
+            self.grow_at(offset + len)
+        }
+        unsafe {
+            let dst = std::slice::from_raw_parts_mut(self.ptr.add(offset), len);
+            dst.copy_from_slice(src);
+        }
+        self.len = std::cmp::max(self.len, offset + len);
+    }
+
+    pub fn append(&mut self, src: &[u8]) {
+        self.write_at(src, self.len)
+    }
+
+    pub fn slice(&self, range: impl RangeBounds<usize>) -> &[u8] {
+        let (start, end) = self.bounds(range);
+        let slice = unsafe { std::slice::from_raw_parts(self.ptr.add(start), end - start) };
+        slice
+    }
+
+    pub fn slice_mut(&mut self, range: impl RangeBounds<usize>) -> &mut [u8] {
+        let (start, end) = self.bounds(range);
+        let slice = unsafe { std::slice::from_raw_parts_mut(self.ptr.add(start), end - start) };
+        slice
+    }
+}
+
+impl<const ALIGN: usize, const SMOOTH: usize, const DEFAULT: usize>
+    AlignedBuffer<ALIGN, SMOOTH, DEFAULT>
+{
+    fn align_up(v: usize) -> usize {
+        align_up(v, ALIGN)
+    }
+
+    #[allow(dead_code)]
+    fn align_down(v: usize) -> usize {
+        align_down(v, ALIGN)
+    }
+
+    fn bounds(&self, range: impl RangeBounds<usize>) -> (usize, usize) {
+        let start = match range.start_bound() {
+            std::ops::Bound::Included(start) => *start,
+            std::ops::Bound::Excluded(start) => *start + 1,
+            std::ops::Bound::Unbounded => 0,
+        };
+        let end = match range.end_bound() {
+            std::ops::Bound::Included(end) => *end + 1,
+            std::ops::Bound::Excluded(end) => *end,
+            std::ops::Bound::Unbounded => self.len,
+        };
+        if end > self.len {
+            panic!(
+                "out of range: [capacity: {}] [len: {}] [given: {}]",
+                self.capacity, self.len, end,
+            );
+        }
+        (start, end)
+    }
+
+    fn grow_at(&mut self, size: usize) {
+        // smooth growth size = 1 GiB
+        let capacity = Self::grow_size_at(self.capacity, size);
+        unsafe {
+            let buffer = std::alloc::alloc_zeroed(std::alloc::Layout::from_size_align_unchecked(
+                capacity, ALIGN,
+            ));
+
+            let src = std::slice::from_raw_parts(self.ptr, self.len);
+            let dist = std::slice::from_raw_parts_mut(buffer, self.len);
+            dist.copy_from_slice(src);
+
+            std::alloc::dealloc(
+                self.ptr,
+                std::alloc::Layout::from_size_align_unchecked(self.capacity, ALIGN),
+            );
+
+            self.ptr = buffer;
+        }
+        self.capacity = capacity;
+    }
+
+    fn grow_size_at(origin: usize, size: usize) -> usize {
+        let size = Self::align_up(size);
+        let mut capacity = origin;
+        while capacity < size {
+            if capacity > SMOOTH {
+                capacity += SMOOTH;
+            } else if capacity > SMOOTH / 4 * 3 {
+                capacity = SMOOTH * 2;
+            } else if capacity > SMOOTH / 2 {
+                capacity = SMOOTH;
+            } else {
+                capacity *= 2;
+            }
+        }
+        capacity
+    }
+}
+
+impl<const ALIGN: usize, const SMOOTH: usize, const DEFAULT: usize> std::fmt::Debug
+    for AlignedBuffer<ALIGN, SMOOTH, DEFAULT>
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let data = unsafe { std::slice::from_raw_parts(self.ptr, self.len) };
+        f.debug_struct("AlignedBuffer")
+            .field("ptr", &self.ptr)
+            .field("align", &ALIGN)
+            .field("capacity", &self.capacity)
+            .field("len", &self.len)
+            .field("data", &data)
+            .finish()
+    }
+}
+
+impl<const ALIGN: usize, const SMOOTH: usize, const DEFAULT: usize, R: RangeBounds<usize>>
+    core::ops::Index<R> for AlignedBuffer<ALIGN, SMOOTH, DEFAULT>
+{
+    type Output = [u8];
+
+    fn index(&self, index: R) -> &Self::Output {
+        self.slice(index)
+    }
+}
+
+impl<const ALIGN: usize, const SMOOTH: usize, const DEFAULT: usize, R: RangeBounds<usize>>
+    core::ops::IndexMut<R> for AlignedBuffer<ALIGN, SMOOTH, DEFAULT>
+{
+    fn index_mut(&mut self, index: R) -> &mut Self::Output {
+        self.slice_mut(index)
+    }
+}
+
+impl<const ALIGN: usize, const SMOOTH: usize, const DEFAULT: usize> Default
+    for AlignedBuffer<ALIGN, SMOOTH, DEFAULT>
+{
+    fn default() -> Self {
+        Self::with_capacity(DEFAULT)
+    }
+}
+
+impl<const ALIGN: usize, const SMOOTH: usize, const DEFAULT: usize> Drop
+    for AlignedBuffer<ALIGN, SMOOTH, DEFAULT>
+{
+    fn drop(&mut self) {
+        unsafe {
+            std::alloc::dealloc(
+                self.ptr,
+                std::alloc::Layout::from_size_align_unchecked(self.capacity, ALIGN),
+            )
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use test_log::test;
+
+    use super::*;
+
+    #[test]
+    fn test_grow_capacity() {
+        assert_eq!(AlignedBuffer::<1, 64, 0>::grow_size_at(128, 192), 192);
+        assert_eq!(AlignedBuffer::<1, 64, 0>::grow_size_at(64, 128), 128);
+        assert_eq!(AlignedBuffer::<1, 64, 0>::grow_size_at(49, 128), 128);
+        assert_eq!(AlignedBuffer::<1, 64, 0>::grow_size_at(48, 64), 64);
+        assert_eq!(AlignedBuffer::<1, 64, 0>::grow_size_at(32, 64), 64);
+        assert_eq!(AlignedBuffer::<1, 64, 0>::grow_size_at(31, 62), 62);
+    }
+
+    #[test]
+    fn test_aligned_buffer() {
+        let mut buf = AlignedBuffer::<512, 1073741824 /* 1 GiB */, 0>::with_capacity(65500);
+
+        assert_eq!(buf.capacity(), 65536);
+
+        buf.append(&[b'x'; 1024]);
+        assert_eq!(&buf[0..1024], &[b'x'; 1024]);
+
+        buf.write_at(&[b'x'; 1024], 1024);
+        assert_eq!(&buf[0..2048], &[b'x'; 2048]);
+
+        buf.append(&[b'a'; 1024]);
+        (&mut buf[2048..3072]).copy_from_slice(&[b'x'; 1024]);
+        assert_eq!(&buf[0..3072], &[b'x'; 3072]);
+
+        drop(buf);
+    }
+
+    #[test]
+    fn test_growth() {
+        let mut buf = AlignedBuffer::<8, 64, 0>::with_capacity(7);
+        assert_eq!(buf.capacity(), 8);
+
+        buf.append(&[b'x'; 7]);
+        assert_eq!(buf.len(), 7);
+        assert_eq!(buf.capacity(), 8);
+
+        buf.append(&[b'x'; 8]);
+        assert_eq!(buf.len(), 15);
+        assert_eq!(buf.capacity(), 16);
+
+        buf.write_at(&[b'x'; 1], 62);
+        assert_eq!(buf.len(), 63);
+        assert_eq!(buf.capacity(), 64);
+
+        buf.write_at(&[b'x'; 1], 190);
+        assert_eq!(buf.len(), 191);
+        assert_eq!(buf.capacity(), 192);
+    }
+}

storage/src/file_cache/error.rs

@@ -0,0 +1,11 @@
+#[derive(thiserror::Error, Debug)]
+pub enum Error {
+    #[error("io error: {0}")]
+    IoError(#[from] std::io::Error),
+    #[error("nix error: {0}")]
+    NixError(#[from] nix::errno::Errno),
+    #[error("join error: {0}")]
+    JoinError(#[from] tokio::task::JoinError),
+}
+
+pub type Result<T> = core::result::Result<T, Error>;

storage/src/file_cache/file.rs

@@ -0,0 +1,55 @@
+use std::os::unix::prelude::OpenOptionsExt;
+use std::path::Path;
+
+use super::buffer::AlignedBuffer;
+use super::error::Result;
+
+// TODO: Get logical block size at `/sys/block/{device}/queue/logical_block_size`.
+const LOGICAL_BLOCK_SIZE: usize = 512;
+const SMOOTH_GROWTH_SIZE: usize = 64 * 1024 * 1024; // 64 MiB
+const DEFAULT_BUFFER_SIZE: usize = 128 * 1024; // 128 KiB
+
+pub type CacheFileBuffer =
+    AlignedBuffer<LOGICAL_BLOCK_SIZE, SMOOTH_GROWTH_SIZE, DEFAULT_BUFFER_SIZE>;
+
+pub struct CacheFile {
+    file: std::fs::File,
+}
+
+impl CacheFile {
+    pub fn open(dir: impl AsRef<Path>, id: u64) -> Result<Self> {
+        let path = dir.as_ref().join(format!("cache-{:08}", id));
+        let mut options = std::fs::OpenOptions::new();
+        options.create(true);
+        options.read(true);
+        options.write(true);
+        options.custom_flags(libc::O_DIRECT | libc::O_SYNC);
+        let file = options.open(path)?;
+        Ok(Self { file })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use std::io::{Seek, SeekFrom, Write};
+    use std::os::unix::prelude::{AsRawFd, FileExt};
+
+    use nix::sys::uio::pwrite;
+    use nix::unistd::write;
+    use test_log::test;
+
+    use super::*;
+
+    #[test]
+    fn test_dio() {
+        let tempdir = tempfile::tempdir().unwrap();
+
+        let cf = CacheFile::open(tempdir.path(), 1).unwrap();
+
+        let mut buf = CacheFileBuffer::default();
+        buf.append(&[b'x'; 4096]);
+
+        cf.file.write_all_at(&buf[0..4096], 0).unwrap();
+    }
+}

storage/src/file_cache/mod.rs

@@ -0,0 +1,6 @@
+pub mod buffer;
+pub mod error;
+#[allow(dead_code)]
+#[allow(unused_imports)]
+#[allow(unused_variables)]
+pub mod file;

storage/src/lib.rs

@@ -5,6 +5,7 @@
 #![feature(proc_macro_hygiene)]
 
 mod error;
+pub mod file_cache;
 mod lsm_tree;
 mod object_store;
 pub mod raft_log_store;

tests/integrations/test_concurrent_put_get.rs

@@ -45,7 +45,7 @@ async fn test_concurrent_put_get() {
         concurrency: 1000,
         r#loop: 3,
         raft_log_store_data_dir,
-        persist: "sync".to_string(),
+        persist: "none".to_string(),
         log_dir,
     };