use log::warn;
use std::collections::HashMap;
use std::io::{Read, Write};
use std::path::Path;
const FILE_MAGIC: [u8; 4] = *b"VBC1";
const FILE_VERSION: u32 = 1;
const REC_DATA: u8 = 1;
const REC_COMMIT: u8 = 2;
const MAX_SANE_LEN: u32 = 64 * 1024 * 1024;
fn checksum(key: &[u8], vec_bytes: &[u8], tag: &[u8]) -> u64 {
let mut h = xxhash_rust::xxh3::xxh3_64(key);
h ^= xxhash_rust::xxh3::xxh3_64(vec_bytes).wrapping_mul(0x9E37_79B9_7F4A_7C15);
h ^= xxhash_rust::xxh3::xxh3_64(tag).wrapping_mul(0xBF58_476D_1CE4_E5B9);
h
}
fn write_u32(w: &mut impl Write, v: u32) -> std::io::Result<()> {
w.write_all(&v.to_le_bytes())
}
fn read_u32(r: &mut impl Read) -> std::io::Result<u32> {
let mut b = [0u8; 4];
r.read_exact(&mut b)?;
Ok(u32::from_le_bytes(b))
}
fn read_u64(r: &mut impl Read) -> std::io::Result<u64> {
let mut b = [0u8; 8];
r.read_exact(&mut b)?;
Ok(u64::from_le_bytes(b))
}
pub fn write_header(file: &mut impl std::io::Write) -> std::io::Result<()> {
file.write_all(&FILE_MAGIC)?;
write_u32(file, FILE_VERSION)?;
file.flush()?;
Ok(())
}
pub fn append_record<W: std::io::Write>(
file: &mut W,
key: &str,
vec: &[f32],
tag: &str,
nrec: u64,
) -> std::io::Result<(u64, usize)> {
let key_b = key.as_bytes();
let tag_b = tag.as_bytes();
let mut vec_b = Vec::with_capacity(vec.len() * 4);
for x in vec {
vec_b.extend_from_slice(&x.to_bits().to_le_bytes());
}
let sum = checksum(key_b, &vec_b, tag_b);
let mut rec = Vec::with_capacity(1 + 4 + key_b.len() + 4 + vec_b.len() + 4 + tag_b.len() + 8);
rec.push(REC_DATA);
rec.extend_from_slice(&(key_b.len() as u32).to_le_bytes());
rec.extend_from_slice(key_b);
rec.extend_from_slice(&(vec.len() as u32).to_le_bytes());
rec.extend_from_slice(&vec_b);
rec.extend_from_slice(&(tag_b.len() as u32).to_le_bytes());
rec.extend_from_slice(tag_b);
rec.extend_from_slice(&sum.to_le_bytes());
file.write_all(&rec)?;
let nrec = nrec + 1;
let mut commit = Vec::with_capacity(9);
commit.push(REC_COMMIT);
commit.extend_from_slice(&nrec.to_le_bytes());
file.write_all(&commit)?;
Ok((nrec, rec.len() + commit.len()))
}
pub fn replay(path: &Path, expected_tag: &str) -> (Vec<(String, Vec<f32>)>, u64) {
let mut out = Vec::new();
let mut committed: u64 = 0;
let mut valid: Vec<(String, Vec<f32>)> = Vec::new();
let mut file = match std::fs::File::open(path) {
Ok(f) => f,
Err(e) => {
warn!("persist: 无法打开 {}({}),视为空缓存", path.display(), e);
return (out, committed);
}
};
let mut magic = [0u8; 4];
if file.read_exact(&mut magic).is_err() {
return (out, committed); }
if magic != FILE_MAGIC {
warn!("persist: {} 魔数不匹配,整体弃用", path.display());
return (out, committed);
}
match read_u32(&mut file) {
Ok(v) if v == FILE_VERSION => {}
Ok(v) => {
warn!("persist: {} 版本 {} 不匹配,整体弃用", path.display(), v);
return (out, committed);
}
Err(_) => return (out, committed), }
loop {
let mut t = [0u8; 1];
match file.read_exact(&mut t) {
Ok(()) => {}
Err(_) => break, }
if t[0] == REC_COMMIT {
match read_u64(&mut file) {
Ok(n) => committed = committed.max(n),
Err(_) => break, }
continue;
}
if t[0] != REC_DATA {
warn!(
"persist: {} 未知记录类型 {},停止回放(撕裂)",
path.display(),
t[0]
);
break;
}
let key_len = match read_u32(&mut file) {
Ok(n) if n <= MAX_SANE_LEN => n as usize,
_ => break,
};
let mut key_b = vec![0u8; key_len];
if file.read_exact(&mut key_b).is_err() {
break;
}
let dim = match read_u32(&mut file) {
Ok(n) if n <= MAX_SANE_LEN / 4 => n as usize,
_ => break,
};
let mut vec_b = vec![0u8; dim * 4];
if file.read_exact(&mut vec_b).is_err() {
break;
}
let tag_len = match read_u32(&mut file) {
Ok(n) if n <= MAX_SANE_LEN => n as usize,
_ => break,
};
let mut tag_b = vec![0u8; tag_len];
if file.read_exact(&mut tag_b).is_err() {
break;
}
let mut sum_b = [0u8; 8];
if file.read_exact(&mut sum_b).is_err() {
break;
}
if u64::from_le_bytes(sum_b) != checksum(&key_b, &vec_b, &tag_b) {
warn!("persist: {} 一条记录 checksum 失败,已跳过", path.display());
continue;
}
if tag_b != expected_tag.as_bytes() {
warn!("persist: {} 一条记录指纹不匹配,已跳过", path.display());
continue;
}
let key = match String::from_utf8(key_b) {
Ok(k) => k,
Err(_) => {
warn!("persist: {} 一条记录 key 非 UTF-8,已跳过", path.display());
continue;
}
};
let mut v = Vec::with_capacity(dim);
for chunk in vec_b.as_chunks::<4>().0 {
v.push(f32::from_le_bytes(*chunk));
}
valid.push((key, v));
}
out.extend(valid.into_iter().take(committed as usize));
(out, committed)
}
pub fn compact(
path: &Path,
entries: &HashMap<String, Vec<f32>>,
tag: &str,
) -> std::io::Result<u64> {
let tmp = path.with_extension("compact.tmp");
{
let mut file = std::fs::File::create(&tmp)?;
file.write_all(&FILE_MAGIC)?;
write_u32(&mut file, FILE_VERSION)?;
let mut nrec = 0u64;
let mut keys: Vec<&String> = entries.keys().collect();
keys.sort();
for key in keys {
nrec = append_record(&mut file, key, &entries[key], tag, nrec)?.0;
}
file.flush()?;
}
std::fs::rename(&tmp, path)?;
Ok(entries.len() as u64)
}
#[cfg(test)]
pub fn record_counts(path: &Path) -> (usize, u64) {
let mut data = 0usize;
let mut last_commit = 0u64;
let mut file = match std::fs::File::open(path) {
Ok(f) => f,
Err(_) => return (0, 0),
};
let mut magic = [0u8; 4];
if file.read_exact(&mut magic).is_err() {
return (0, 0);
}
if read_u32(&mut file).is_err() {
return (0, 0);
}
loop {
let mut t = [0u8; 1];
if file.read_exact(&mut t).is_err() {
break;
}
if t[0] == REC_COMMIT {
match read_u64(&mut file) {
Ok(n) => last_commit = n,
Err(_) => break,
}
continue;
}
if t[0] != REC_DATA {
break;
}
use std::io::Seek;
let key_len = match read_u32(&mut file) {
Ok(n) => n as u64,
Err(_) => break,
};
if file
.seek(std::io::SeekFrom::Current(key_len as i64))
.is_err()
{
break;
}
let dim = match read_u32(&mut file) {
Ok(n) => n as u64,
Err(_) => break,
};
if file
.seek(std::io::SeekFrom::Current((dim * 4) as i64))
.is_err()
{
break;
}
let tag_len = match read_u32(&mut file) {
Ok(n) => n as u64,
Err(_) => break,
};
if file
.seek(std::io::SeekFrom::Current((tag_len + 8) as i64))
.is_err()
{
break;
}
data += 1;
}
(data, last_commit)
}