Implementing a Zero-Copy LSM-Tree Storage Engine with io_uring in Rust — code snippets

snippet-1.txt

use io_uring::{opcode, squeue, types, IoUring};
use std::fs::File;
use std::os::unix::io::AsRawFd;
use std::alloc::{alloc, dealloc, Layout};

pub struct StorageRing {
    ring: IoUring,
    registered_buffers: Vec<*mut u8>,
    buffer_layout: Layout,
}

impl StorageRing {
    pub fn new(entries: u32, buf_count: usize, buf_size: usize) -> Result<Self, std::io::Error> {
        // Build io_uring with SQPOLL to bypass syscalls in the hot path.
        let ring = IoUring::builder()
            .setup_sqpoll(2000) // 2000ms idle timeout for kernel poll thread
            .build(entries)?;

        // Ensure 4KB page alignment for O_DIRECT read/write compliance
        let buffer_layout = Layout::from_size_align(buf_size, 4096)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidInput, e))?;

        let mut registered_buffers = Vec::with_capacity(buf_count);
        let mut io_vectors = Vec::with_capacity(buf_count);

        for _ in 0..buf_count {
            let ptr = unsafe { alloc(buffer_layout) };
            if ptr.is_null() {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::OutOfMemory,
                    "Failed to allocate aligned block buffer",
                ));
            }
            registered_buffers.push(ptr);

            // Populate the standard POSIX I/O vector for registration
            let iov = libc::iovec {
                iov_base: ptr as *mut libc::c_void,
                iov_len: buf_size,
            };
            io_vectors.push(iov);
        }

        // Register buffers with the kernel to enable IORING_OP_READ_FIXED/WRITE_FIXED
        unsafe {
            ring.submitter().register_buffers(&io_vectors)?;
        }

        Ok(Self {
            ring,
            registered_buffers,
            buffer_layout,
        })
    }
}

impl Drop for StorageRing {
    fn drop(&mut self) {
        // Safely deallocate the page-aligned buffers on teardown
        for &ptr in &self.registered_buffers {
            unsafe {
                dealloc(ptr, self.buffer_layout);
            }
        }
    }
}

snippet-2.txt

use std::collections::VecDeque;
use std::sync::{Arc, Mutex};

pub struct BufferLease {
    pub index: u32,
    pub ptr: *mut u8,
    pool: Arc<InnerPool>,
}

impl Drop for BufferLease {
    fn drop(&mut self) {
        self.pool.release(self.index);
    }
}

// Safety: The raw pointers are backed by locked memory managed by StorageRing.
unsafe impl Send for BufferLease {}
unsafe impl Sync for BufferLease {}

struct InnerPool {
    free_indices: Mutex<VecDeque<u32>>,
    ptrs: Vec<*mut u8>,
}

impl InnerPool {
    fn release(&self, index: u32) {
        let mut guard = self.free_indices.lock().unwrap();
        guard.push_back(index);
    }
}

pub struct FixedBufferPool {
    inner: Arc<InnerPool>,
}

impl FixedBufferPool {
    pub fn new(ptrs: Vec<*mut u8>) -> Self {
        let count = ptrs.len() as u32;
        let free_indices = Mutex::new((0..count).collect());
        Self {
            inner: Arc::new(InnerPool { free_indices, ptrs }),
        }
    }

    pub fn acquire(&self) -> Option<BufferLease> {
        let mut guard = self.inner.free_indices.lock().unwrap();
        guard.pop_front().map(|index| BufferLease {
            index,
            ptr: self.inner.ptrs[index as usize],
            pool: Arc::clone(&self.inner),
        })
    }
}

snippet-3.txt

use io_uring::{opcode, squeue, types, IoUring};
use std::os::unix::io::RawFd;

pub struct WalWriter {
    fd: RawFd,
    offset: u64,
}

impl WalWriter {
    pub fn new(fd: RawFd) -> Self {
        Self { fd, offset: 0 }
    }

    pub fn write_record(
        &mut self,
        ring: &mut IoUring,
        lease: &BufferLease,
        len: usize,
        user_data: u64,
    ) -> Result<(), std::io::Error> {
        // Enforce block-size alignment for direct writes to prevent EINVAL.
        assert!(
            len % 512 == 0,
            "WAL write length must be aligned to 512-byte sector boundary for O_DIRECT"
        );
        assert!(
            self.offset % 512 == 0,
            "WAL file offset must be aligned to sector boundary"
        );

        // Prepare the write fixed operation using the pre-registered buffer index.
        let write_sqe = opcode::WriteFixed::new(
            types::Fd(self.fd),
            lease.ptr,
            len as u32,
            lease.index,
        )
        .offset(self.offset)
        .build()
        .user_data(user_data);

        // Push directly to the submission queue
        unsafe {
            ring.submission()
                .push(&write_sqe)
                .map_err(|_| {
                    std::io::Error::new(
                        std::io::ErrorKind::Other,
                        "io_uring submission queue is full",
                    )
                })?;
        }

        // Notify the kernel SQ poll thread of new work
        ring.submit()?;
        self.offset += len as u64;
        Ok(())
    }
}

snippet-4.txt

use io_uring::{opcode, squeue, types, IoUring};
use std::os::unix::io::RawFd;

pub struct SstableReader {
    fd: RawFd,
}

impl SstableReader {
    pub fn new(fd: RawFd) -> Self {
        Self { fd }
    }

    pub fn submit_block_read(
        &self,
        ring: &mut IoUring,
        lease: &BufferLease,
        block_offset: u64,
        block_size: u32,
        user_data: u64,
    ) -> Result<(), std::io::Error> {
        // SSTable blocks must align strictly with page size boundaries.
        assert!(
            block_offset % 4096 == 0,
            "File block offset must be 4KB aligned for O_DIRECT reads"
        );
        assert!(
            block_size % 4096 == 0,
            "Read block size must be 4KB aligned for O_DIRECT reads"
        );

        let read_sqe = opcode::ReadFixed::new(
            types::Fd(self.fd),
            lease.ptr,
            block_size,
            lease.index,
        )
        .offset(block_offset)
        .build()
        .user_data(user_data);

        unsafe {
            ring.submission()
                .push(&read_sqe)
                .map_err(|_| {
                    std::io::Error::new(
                        std::io::ErrorKind::Other,
                        "io_uring submission queue overflow",
                    )
                })?;
        }

        ring.submit()?;
        Ok(())
    }
}

snippet-5.txt

use io_uring::IoUring;
use std::sync::mpsc::Sender;

pub struct CompletionEvent {
    pub user_data: u64,
    pub result: i32, // Contains the number of bytes read/written or negative errno
}

pub struct EventReaper {
    ring: IoUring,
    tx: Sender<CompletionEvent>,
}

impl EventReaper {
    pub fn new(ring: IoUring, tx: Sender<CompletionEvent>) -> Self {
        Self { ring, tx }
    }

    pub fn run_loop(&mut self) -> Result<(), std::io::Error> {
        loop {
            // Block until at least 1 CQE is available.
            // Under SQPOLL, this avoids system calls when active but yields cleanly when idle.
            self.ring.submit_and_wait(1)?;

            let mut cq = self.ring.completion();
            for cqe in &mut cq {
                let res = cqe.result();
                let user_data = cqe.user_data();

                // If result is negative, it represents a standard POSIX error code (e.g., -EINVAL, -EIO)
                if res < 0 {
                    let errno = -res;
                    // Handle specific errors like EAGAIN (35), EINVAL (22), EIO (5)
                    eprintln!(
                        "io_uring operation failed. SQE UserData: {}, Linux errno: {}",
                        user_data, errno
                    );
                }

                let event = CompletionEvent {
                    user_data,
                    result: res,
                };

                // Propagate completion status to database futures
                if self.tx.send(event).is_err() {
                    // Receiver channel was closed, meaning system shutdown is underway
                    return Ok(());
                }
            }
            // Sync reaped state back to kernel to release queue space
            cq.sync();
        }
    }
}

snippet-6.txt

use io_uring::{opcode, squeue, types, IoUring};
use std::os::unix::io::RawFd;

pub struct CompactionBatch {
    sqes: Vec<squeue::Entry>,
}

impl CompactionBatch {
    pub fn new() -> Self {
        Self { sqes: Vec::new() }
    }

    pub fn add_block_read(
        &mut self,
        fd: RawFd,
        lease: &BufferLease,
        offset: u64,
        size: u32,
        user_data: u64,
    ) {
        let sqe = opcode::ReadFixed::new(
            types::Fd(fd),
            lease.ptr,
            size,
            lease.index,
        )
        .offset(offset)
        .build()
        .user_data(user_data);
        self.sqes.push(sqe);
    }

    pub fn submit_all(&mut self, ring: &mut IoUring) -> Result<usize, std::io::Error> {
        let batch_size = self.sqes.len();
        if batch_size == 0 {
            return Ok(0);
        }

        let mut sq = ring.submission();
        let mut submitted = 0;

        for sqe in self.sqes.drain(..) {
            if unsafe { sq.push(&sqe).is_err() } {
                // If the ring fills up, flush it to disk and reopen the submission queue guard
                drop(sq);
                ring.submit()?;
                sq = ring.submission();
                
                // Retry push. If it fails again, abort batching.
                if unsafe { sq.push(&sqe).is_err() } {
                    return Err(std::io::Error::new(
                        std::io::ErrorKind::Other,
                        "io_uring SQ ring overflow during compaction batch submission",
                    ));
                }
            }
            submitted += 1;
        }

        drop(sq);
        ring.submit()?;
        Ok(submitted)
    }
}