Designing a Custom Write-Ahead Log (WAL) with Zero-Copy Direct I/O in Rust — code snippets

snippet-1.txt

use std::alloc::{alloc, dealloc, Layout};
use std::ptr::NonNull;

pub struct AlignedBuffer {
    ptr: NonNull<u8>,
    layout: Layout,
    capacity: usize,
}

impl AlignedBuffer {
    pub fn new(capacity: usize, alignment: usize) -> Self {
        assert!(alignment.is_power_of_two(), "Alignment must be a power of two");
        assert!(capacity % alignment == 0, "Capacity must be a multiple of alignment");
        
        let layout = Layout::from_size_align(capacity, alignment)
            .expect("Invalid layout definition");
        
        let raw_ptr = unsafe { alloc(layout) };
        if raw_ptr.is_null() {
            panic!("Memory allocation failed for capacity {}", capacity);
        }
        let ptr = NonNull::new(raw_ptr).expect("Acquired pointer is null");
        
        Self { ptr, layout, capacity }
    }

    pub fn as_slice(&self) -> &[u8] {
        unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.capacity) }
    }

    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.capacity) }
    }

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

impl Drop for AlignedBuffer {
    fn drop(&mut self) {
        unsafe {
            dealloc(self.ptr.as_ptr(), self.layout);
        }
    }
}

snippet-2.txt

use std::fs::{File, OpenOptions};
use std::os::unix::fs::OpenOptionsExt;
use std::path::Path;

pub fn open_wal_file_direct<P: AsRef<Path>>(path: P) -> std::io::Result<File> {
    let mut options = OpenOptions::new();
    options.read(true).write(true).create(true);
    
    #[cfg(target_os = "linux")]
    {
        // Inject O_DIRECT to bypass page cache and O_DSYNC for synchronous data integrity
        options.custom_flags(libc::O_DIRECT | libc::O_DSYNC);
    }
    
    let file = options.open(path)?;
    Ok(file)
}

snippet-3.txt

use std::convert::TryInto;

pub const PAGE_SIZE: usize = 4096;
pub const SECTOR_SIZE: usize = 512;
pub const HEADER_SIZE: usize = 16; // 8-byte payload length + 8-byte CRC32C checksum

#[repr(C, align(4096))]
pub struct AlignedPage {
    pub data: [u8; PAGE_SIZE],
}

pub struct WalRecordHeader {
    pub length: u64,
    pub checksum: u64,
}

impl WalRecordHeader {
    pub fn to_bytes(&self) -> [u8; HEADER_SIZE] {
        let mut bytes = [0u8; HEADER_SIZE];
        bytes[0..8].copy_from_slice(&self.length.to_le_bytes());
        bytes[8..16].copy_from_slice(&self.checksum.to_le_bytes());
        bytes
    }
    
    pub fn from_bytes(bytes: &[u8; HEADER_SIZE]) -> Self {
        let length = u64::from_le_bytes(bytes[0..8].try_into().unwrap());
        let checksum = u64::from_le_bytes(bytes[8..16].try_into().unwrap());
        Self { length, checksum }
    }
}

snippet-4.txt

use crc32fast::Hasher;

pub struct WalWriter {
    active_page: AlignedBuffer,
    write_offset: usize,
    file_offset: u64,
}

impl WalWriter {
    pub fn new() -> Self {
        Self {
            active_page: AlignedBuffer::new(PAGE_SIZE, PAGE_SIZE),
            write_offset: 0,
            file_offset: 0,
        }
    }

    pub fn append(&mut self, file: &File, payload: &[u8]) -> Result<usize, String> {
        let record_len = payload.len();
        let total_write_len = HEADER_SIZE + record_len;
        
        // If the record cannot fit in the remaining page space, flush the current page first
        if self.write_offset + total_write_len > PAGE_SIZE {
            self.flush_active_page(file)?;
        }
        
        let header = WalRecordHeader {
            length: record_len as u64,
            checksum: self.calculate_crc(payload),
        };
        
        let dest = self.active_page.as_mut_slice();
        
        // Copy Header
        dest[self.write_offset..self.write_offset + HEADER_SIZE].copy_from_slice(&header.to_bytes());
        self.write_offset += HEADER_SIZE;
        
        // Copy Payload
        dest[self.write_offset..self.write_offset + record_len].copy_from_slice(payload);
        self.write_offset += record_len;
        
        Ok(self.write_offset)
    }

    fn calculate_crc(&self, data: &[u8]) -> u64 {
        let mut hasher = Hasher::new();
        hasher.update(data);
        hasher.finalize() as u64
    }
}

snippet-5.txt

use std::os::unix::io::AsRawFd;

impl WalWriter {
    pub fn flush_active_page(&mut self, file: &File) -> Result<(), String> {
        if self.write_offset == 0 {
            return Ok(());
        }
        
        // Round up the write length to the nearest 512-byte sector boundary
        let aligned_write_len = (self.write_offset + SECTOR_SIZE - 1) & !(SECTOR_SIZE - 1);
        
        // Pad the tail of the buffer to the sector boundary with zeros
        let active_slice = self.active_page.as_mut_slice();
        if aligned_write_len > self.write_offset {
            active_slice[self.write_offset..aligned_write_len].fill(0);
        }
        
        let fd = file.as_raw_fd();
        let buf_ptr = active_slice.as_ptr() as *const libc::c_void;
        
        let n = unsafe {
            libc::pwrite(fd, buf_ptr, aligned_write_len, self.file_offset as libc::off_t)
        };
        
        if n < 0 {
            return Err(format!(
                "Direct pwrite failed: {} (Offset: {}, Size: {})",
                std::io::Error::last_os_error(),
                self.file_offset,
                aligned_write_len
            ));
        }
        
        if n as usize != aligned_write_len {
            return Err(format!(
                "Short write: requested {}, wrote {}",
                aligned_write_len, n
            ));
        }
        
        // Increment the file offset. The next write starts at this sector boundary.
        self.file_offset += aligned_write_len as u64;
        self.write_offset = 0; // Reset offset inside the buffer
        Ok(())
    }
}

snippet-6.txt

use std::os::unix::io::AsRawFd;

pub struct WalRecoveryParser {
    file: File,
    read_buffer: AlignedBuffer,
    file_offset: u64,
}

impl WalRecoveryParser {
    pub fn new(file: File) -> Self {
        Self {
            file,
            read_buffer: AlignedBuffer::new(PAGE_SIZE, PAGE_SIZE),
            file_offset: 0,
        }
    }

    pub fn recover_records(&mut self) -> Result<Vec<Vec<u8>>, String> {
        let mut records = Vec::new();
        let fd = self.file.as_raw_fd();
        let buf_slice = self.read_buffer.as_mut_slice();
        
        loop {
            let buf_ptr = buf_slice.as_mut_ptr() as *mut libc::c_void;
            let n = unsafe {
                libc::pread(fd, buf_ptr, PAGE_SIZE, self.file_offset as libc::off_t)
            };
            
            if n < 0 {
                return Err(format!("Read failed: {}", std::io::Error::last_os_error()));
            }
            if n == 0 {
                break; // Physical End of File (EOF)
            }
            
            let mut offset = 0;
            while offset + HEADER_SIZE <= n as usize {
                let header_bytes: &[u8; HEADER_SIZE] = buf_slice[offset..offset + HEADER_SIZE]
                    .try_into()
                    .unwrap();
                let header = WalRecordHeader::from_bytes(header_bytes);
                
                // A zero length indicates the start of zero-padding at the end of a page
                if header.length == 0 {
                    break;
                }
                
                let next_offset = offset + HEADER_SIZE + header.length as usize;
                if next_offset > n as usize {
                    break; // Record spans past the current buffer block (or is corrupt)
                }
                
                let payload = &buf_slice[offset + HEADER_SIZE..next_offset];
                
                // Validate payload integrity via CRC32C
                let mut hasher = Hasher::new();
                hasher.update(payload);
                if hasher.finalize() as u64 != header.checksum {
                    return Err(format!(
                        "Integrity validation failed at offset {}: Checksum mismatch",
                        self.file_offset + offset as u64
                    ));
                }
                
                records.push(payload.to_vec());
                offset = next_offset;
            }
            
            self.file_offset += n as u64;
        }
        
        Ok(records)
    }
}