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;
#[inline]
pub const fn hint_entry_size(key_len: usize) -> usize {
8 + key_len + 8 + 4 + 4 }
#[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
}
}
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)
}
#[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,
};
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; }
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;
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)
}
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,
})
})
}
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(())
}
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)),
}
}
#[inline]
pub fn hint_path_for_data(data_path: &Path) -> PathBuf {
data_path.with_extension("hint")
}
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);
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();
}
#[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();
tree.flush().unwrap();
tree.put(&2u64.to_be_bytes(), &20u64.to_be_bytes()).unwrap();
tree.flush().unwrap();
tree.close().unwrap();
}
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"
);
}
}