armdb 0.6.0

sharded bitcask key-value storage optimized for NVMe
Documentation
use std::fs;
use std::path::{Path, PathBuf};

use zerocopy::FromBytes;

use crate::Key;
use crate::entry::{EntryHeader, SEQUENCE_MASK, TOMBSTONE_BIT, entry_size};
use crate::error::{DbError, DbResult};
use crate::io::direct;

type ReadFn<'a> = dyn Fn(&std::fs::File, u64, usize) -> DbResult<Vec<u8>> + 'a;

/// Size of a single hint entry: GSN(8) + Key(key_len) + Offset(8) + Len(4) + CRC(4).
#[inline]
pub const fn hint_entry_size(key_len: usize) -> usize {
    8 + key_len + 8 + 4 + 4 // GSN(8) + Key(key_len) + Offset(8) + Len(4) + CRC(4)
}

/// A parsed hint entry.
#[derive(Debug, Clone, Copy)]
pub struct HintEntry<K: Key> {
    pub gsn: u64,
    pub key: K,
    pub value_offset: u64,
    pub value_len: u32,
    pub crc32: u32,
}

impl<K: Key> HintEntry<K> {
    #[inline]
    pub fn is_tombstone(&self) -> bool {
        self.gsn & TOMBSTONE_BIT != 0
    }

    #[inline]
    pub fn sequence(&self) -> u64 {
        self.gsn & SEQUENCE_MASK
    }
}

/// Generate hint data by scanning a data file.
///
/// Reads only headers + keys (skips values), making this much cheaper
/// than a full recovery scan. Takes `key_len` at runtime so it can be
/// called from non-generic `ShardInner`.
pub fn generate_hint_data_dyn(
    data_file: &std::fs::File,
    file_len: u64,
    key_len: usize,
) -> DbResult<Vec<u8>> {
    generate_hint_data_inner(data_file, file_len, key_len, &read_plain)
}

/// Generate hint data from an encrypted data file.
#[cfg(feature = "encryption")]
pub fn generate_hint_data_dyn_encrypted(
    data_file: &std::fs::File,
    file_len: u64,
    key_len: usize,
    cipher: &crate::crypto::PageCipher,
    tag_file: &crate::io::tags::TagFile,
    file_id: u32,
) -> DbResult<Vec<u8>> {
    let reader = move |file: &std::fs::File, offset: u64, len: usize| {
        direct::pread_value_encrypted(file, tag_file, cipher, file_id, offset, len)
    };
    generate_hint_data_inner(data_file, file_len, key_len, &reader)
}

fn generate_hint_data_inner(
    data_file: &std::fs::File,
    file_len: u64,
    key_len: usize,
    read_fn: &ReadFn<'_>,
) -> DbResult<Vec<u8>> {
    let hint_sz = hint_entry_size(key_len);
    let min_data_entry = entry_size(key_len, 0);
    let estimated = file_len.checked_div(min_data_entry).unwrap_or_default() as usize;
    let mut buf = Vec::with_capacity(estimated * hint_sz);

    let header_size = size_of::<EntryHeader>() as u64;
    let mut offset: u64 = 0;

    while offset + header_size <= file_len {
        let header_bytes = match read_fn(data_file, offset, size_of::<EntryHeader>()) {
            Ok(b) => b,
            Err(_) => break,
        };
        let header: EntryHeader = match EntryHeader::read_from_bytes(&header_bytes) {
            Ok(h) => h,
            Err(_) => break,
        };

        // Page padding sentinel: force-flush pads the partial page with zeros.
        if header.gsn == 0
            && header.value_len == 0
            && header.crc32 == 0
            && crate::entry::check_page_padding_at(
                data_file,
                offset + size_of::<EntryHeader>() as u64,
                read_fn,
            )?
        {
            const PAGE_SIZE: u64 = 4096;
            let next_page =
                (offset + size_of::<EntryHeader>() as u64).div_ceil(PAGE_SIZE) * PAGE_SIZE;
            if next_page >= file_len {
                break;
            }
            offset = next_page;
            continue;
        }

        let total = entry_size(key_len, header.value_len);
        if offset + total > file_len {
            break; // partial entry at end
        }

        // Read only the key (K bytes after header)
        let key_bytes = read_fn(data_file, offset + size_of::<EntryHeader>() as u64, key_len)?;

        let value_offset = offset + size_of::<EntryHeader>() as u64 + key_len as u64;

        // Serialize: GSN(8) | Key(key_len) | Offset(8) | Len(4) | CRC(4)
        buf.extend_from_slice(&header.gsn.to_ne_bytes());
        buf.extend_from_slice(&key_bytes);
        buf.extend_from_slice(&value_offset.to_ne_bytes());
        buf.extend_from_slice(&header.value_len.to_ne_bytes());
        buf.extend_from_slice(&header.crc32.to_ne_bytes());

        offset += total;
    }

    if offset < file_len {
        return Err(DbError::CorruptedEntry { offset });
    }

    Ok(buf)
}

fn read_plain(file: &std::fs::File, offset: u64, len: usize) -> DbResult<Vec<u8>> {
    direct::pread_value(file, offset, len)
}

/// Parse hint entries from raw hint file bytes.
pub fn parse_hint_entries<K: Key>(data: &[u8]) -> impl Iterator<Item = HintEntry<K>> + '_ {
    let entry_sz = hint_entry_size(size_of::<K>());
    data.chunks_exact(entry_sz).filter_map(|chunk| {
        let gsn = u64::from_ne_bytes(chunk[..8].try_into().ok()?);
        let key: K = K::from_bytes(&chunk[8..8 + size_of::<K>()]);
        let value_offset = u64::from_ne_bytes(
            chunk[8 + size_of::<K>()..8 + size_of::<K>() + 8]
                .try_into()
                .ok()?,
        );
        let value_len = u32::from_ne_bytes(
            chunk[8 + size_of::<K>() + 8..8 + size_of::<K>() + 12]
                .try_into()
                .ok()?,
        );
        let crc32 = u32::from_ne_bytes(
            chunk[8 + size_of::<K>() + 12..8 + size_of::<K>() + 16]
                .try_into()
                .ok()?,
        );
        Some(HintEntry {
            gsn,
            key,
            value_offset,
            value_len,
            crc32,
        })
    })
}

/// Write hint data to a file. Uses standard buffered I/O (not O_DIRECT) —
/// hint files are small and written once.
pub fn write_hint_file(path: &Path, data: &[u8]) -> DbResult<()> {
    let tmp_path = path.with_extension("hint.tmp");
    fs::write(&tmp_path, data)?;
    let f = fs::File::open(&tmp_path)?;
    f.sync_data()?;
    drop(f);
    fs::rename(&tmp_path, path)?;
    Ok(())
}

/// Read a hint file into memory. Returns `None` if the file doesn't exist.
pub fn read_hint_file(path: &Path) -> DbResult<Option<Vec<u8>>> {
    match fs::read(path) {
        Ok(data) => Ok(Some(data)),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
        Err(e) => Err(DbError::Io(e)),
    }
}

/// Compute the hint file path for a data file path.
/// e.g. "000001.data" → "000001.hint"
#[inline]
pub fn hint_path_for_data(data_path: &Path) -> PathBuf {
    data_path.with_extension("hint")
}

/// Scan a directory for `.hint` files and return sorted file IDs.
pub fn scan_hint_files(dir: &Path) -> DbResult<Vec<u32>> {
    let mut ids = Vec::new();
    if !dir.exists() {
        return Ok(ids);
    }
    for entry in fs::read_dir(dir)? {
        let entry = entry?;
        let name = entry.file_name();
        let name = name.to_string_lossy();
        if name.ends_with(".hint")
            && let Ok(id) = name.trim_end_matches(".hint").parse::<u32>()
        {
            ids.push(id);
        }
    }
    ids.sort();
    Ok(ids)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn hint_sequence_strips_flag_region() {
        use crate::entry::SOFT_DELETE_BIT;
        let e = HintEntry::<[u8; 8]> {
            gsn: 7 | SOFT_DELETE_BIT,
            key: [0u8; 8],
            value_offset: 0,
            value_len: 0,
            crc32: 0,
        };
        assert_eq!(e.sequence(), 7);
        assert!(!e.is_tombstone());
    }

    #[test]
    fn hint_generation_skips_mid_file_padding() {
        use crate::entry::serialize_entry;

        let dir = std::env::temp_dir().join(format!("armdb_hintpad_{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("000001.data");

        let key_len = 8usize;
        let e1 = serialize_entry(1, &[1u8; 8], b"v1", false);
        let e2 = serialize_entry(2, &[2u8; 8], b"v2", false);

        // Layout: e1, zero-pad to next page boundary, e2.
        let mut file_bytes = Vec::new();
        file_bytes.extend_from_slice(&e1);
        let pad = 4096 - (file_bytes.len() % 4096);
        file_bytes.extend(std::iter::repeat_n(0u8, pad));
        file_bytes.extend_from_slice(&e2);

        let f = crate::io::direct::open_bulk_write(&path).unwrap();
        crate::io::direct::pwrite_at(&f, &file_bytes, 0).unwrap();
        crate::io::direct::fsync(&f).unwrap();
        drop(f);

        let rf = crate::io::direct::open_bulk_read(&path).unwrap();
        let data =
            generate_hint_data_inner(&rf, file_bytes.len() as u64, key_len, &read_plain).unwrap();
        let count = data.len() / hint_entry_size(key_len);
        assert_eq!(
            count, 2,
            "both real entries must produce hints, padding none"
        );

        std::fs::remove_dir_all(&dir).ok();
    }

    /// Encrypted variant: hint generation must skip mid-file page padding when
    /// scanning an *encrypted* data file. A page-aligned encrypted tree pads the
    /// partial page on `flush()`, so the second put lands at the next page
    /// boundary, leaving padding (ciphertext of zeros) between the two entries.
    /// `generate_hint_data_dyn_encrypted` reads through the decrypting reader, so
    /// it must recognise the decrypted-zero padding and emit hints only for the
    /// two real entries.
    #[cfg(feature = "encryption")]
    #[test]
    fn hint_generation_skips_mid_file_padding_encrypted() {
        use crate::config::Config;
        use crate::const_tree::ConstTree;
        use crate::crypto::PageCipher;
        use crate::io::tags::{TagFile, tags_path_for_data};

        let dir = tempfile::tempdir().unwrap();
        let key = [0x42u8; 32];
        let mut config = Config::test();
        config.shard_count = 1;
        config.encryption_key = Some(key);

        {
            let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config).unwrap();
            tree.put(&1u64.to_be_bytes(), &10u64.to_be_bytes()).unwrap();
            // Final flush pads the partial page; the next put starts at the next
            // page boundary, leaving mid-file padding in between.
            tree.flush().unwrap();
            tree.put(&2u64.to_be_bytes(), &20u64.to_be_bytes()).unwrap();
            tree.flush().unwrap();
            tree.close().unwrap();
        }

        // The shard's first data file is id 1 (`000001.data`). The physical file
        // is two 4096-byte pages, but production generates hints over the
        // *logical* data end (`write_offset`), not `metadata().len()`, so the
        // trailing final-page padding is excluded. Mirror that here: the two
        // real entries sit at offset 0 and at the 4096 boundary (after the
        // mid-file padding the first flush left), so the logical end is
        // `4096 + entry_size`.
        let data_path = dir.path().join("shard_000").join("000001.data");
        let data_file = crate::io::direct::open_bulk_read(&data_path).unwrap();
        let physical_len = data_file.metadata().unwrap().len();
        assert!(
            physical_len >= 8192,
            "expected at least two flushed pages, got {physical_len}"
        );
        let logical_len = 4096 + entry_size(8, 8);

        let cipher = PageCipher::new(&key).expect("create cipher");
        let tag_file = TagFile::open_read(&tags_path_for_data(&data_path)).unwrap();

        let data =
            generate_hint_data_dyn_encrypted(&data_file, logical_len, 8, &cipher, &tag_file, 1)
                .unwrap();
        let count = data.len() / hint_entry_size(8);
        assert_eq!(
            count, 2,
            "both real entries must produce hints across encrypted padding"
        );
    }
}