armdb 0.1.13

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, 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).
#[inline]
pub const fn hint_entry_size(key_len: usize) -> usize {
    8 + key_len + 8 + 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,
}

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 & !TOMBSTONE_BIT
    }
}

/// 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 = if min_data_entry > 0 {
        (file_len / min_data_entry) as usize
    } else {
        0
    };
    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,
        };

        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)
        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());

        offset += total;
    }

    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()?,
        );
        Some(HintEntry {
            gsn,
            key,
            value_offset,
            value_len,
        })
    })
}

/// 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<()> {
    fs::write(path, data)?;
    let f = fs::File::open(path)?;
    f.sync_data()?;
    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)
}