443 lines
14 KiB
Rust
443 lines
14 KiB
Rust
use sha2::{Digest, Sha256};
|
|
use std::cmp::Ordering;
|
|
use std::collections::HashSet;
|
|
use std::fs::{self, File};
|
|
use std::io::{Read, Write};
|
|
use std::path::{Path, PathBuf};
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
|
|
|
const INDEX_MAGIC: &[u8; 8] = b"DS4KVIDX";
|
|
const INDEX_VERSION: u32 = 1;
|
|
const MAX_KEY_BYTES: u64 = 256 * 1024 * 1024;
|
|
const HIT_HALF_LIFE_SECONDS: f64 = 6.0 * 60.0 * 60.0;
|
|
|
|
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
|
pub(super) enum StoreReason {
|
|
Cold,
|
|
Continued,
|
|
}
|
|
|
|
impl StoreReason {
|
|
fn code(self) -> u8 {
|
|
match self {
|
|
Self::Cold => 1,
|
|
Self::Continued => 2,
|
|
}
|
|
}
|
|
|
|
fn from_code(value: u8) -> Option<Self> {
|
|
Some(match value {
|
|
1 => Self::Cold,
|
|
2 => Self::Continued,
|
|
_ => return None,
|
|
})
|
|
}
|
|
|
|
fn is_anchor(self) -> bool {
|
|
self == Self::Cold
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub(super) struct Entry {
|
|
pub(super) checkpoint: PathBuf,
|
|
pub(super) tag: [u8; 32],
|
|
pub(super) tokens: u32,
|
|
context: u32,
|
|
hits: u32,
|
|
created_at: u64,
|
|
last_used: u64,
|
|
reason: StoreReason,
|
|
key: Vec<u8>,
|
|
bytes: u64,
|
|
}
|
|
|
|
pub(super) struct KvStore {
|
|
directory: PathBuf,
|
|
budget_bytes: u64,
|
|
}
|
|
|
|
impl KvStore {
|
|
pub(super) fn open(directory: &Path, budget_bytes: u64) -> Result<Self, String> {
|
|
fs::create_dir_all(directory).map_err(|error| {
|
|
format!(
|
|
"Could not create transient KV cache directory {}: {error}",
|
|
directory.display()
|
|
)
|
|
})?;
|
|
let store = Self {
|
|
directory: directory.to_owned(),
|
|
budget_bytes,
|
|
};
|
|
// Reclaim first, so the budget is measured against the whole directory.
|
|
sweep_unreachable(directory);
|
|
// A lowered budget takes effect at once, like ds4's eviction on open.
|
|
store.evict(None);
|
|
Ok(store)
|
|
}
|
|
|
|
pub(super) fn find(&self, prompt_key: &[u8], context: u32) -> Option<Entry> {
|
|
self.entries()
|
|
.into_iter()
|
|
.filter(|entry| entry.context == context && prompt_key.starts_with(&entry.key))
|
|
.max_by(|left, right| {
|
|
left.key
|
|
.len()
|
|
.cmp(&right.key.len())
|
|
.then_with(|| left.tokens.cmp(&right.tokens))
|
|
})
|
|
}
|
|
|
|
pub(super) fn checkpoint_path(&self, key: &[u8]) -> PathBuf {
|
|
self.directory.join(format!("{}.bin", key_hash(key)))
|
|
}
|
|
|
|
pub(super) fn record(
|
|
&self,
|
|
checkpoint: &Path,
|
|
key: &[u8],
|
|
tag: [u8; 32],
|
|
tokens: u32,
|
|
context: u32,
|
|
reason: StoreReason,
|
|
) -> Result<bool, String> {
|
|
let checkpoint_bytes = fs::metadata(checkpoint)
|
|
.map_err(|error| error.to_string())?
|
|
.len();
|
|
let index_bytes = index_size(key)?;
|
|
if checkpoint_bytes.saturating_add(index_bytes) > self.budget_bytes {
|
|
let _ = fs::remove_file(checkpoint);
|
|
return Ok(false);
|
|
}
|
|
let now = unix_time();
|
|
let entry = Entry {
|
|
checkpoint: checkpoint.to_owned(),
|
|
tag,
|
|
tokens,
|
|
context,
|
|
hits: 0,
|
|
created_at: now,
|
|
last_used: now,
|
|
reason,
|
|
key: key.to_vec(),
|
|
bytes: checkpoint_bytes.saturating_add(index_bytes),
|
|
};
|
|
write_entry(&entry)?;
|
|
self.evict(Some(checkpoint));
|
|
Ok(checkpoint.exists())
|
|
}
|
|
|
|
pub(super) fn touch(&self, entry: &Entry) -> Result<(), String> {
|
|
let mut touched = entry.clone();
|
|
touched.hits = touched.hits.saturating_add(1);
|
|
touched.last_used = unix_time();
|
|
write_entry(&touched)
|
|
}
|
|
|
|
pub(super) fn discard(&self, entry: &Entry) {
|
|
let _ = fs::remove_file(&entry.checkpoint);
|
|
let _ = fs::remove_file(index_path(&entry.checkpoint));
|
|
}
|
|
|
|
fn entries(&self) -> Vec<Entry> {
|
|
let Ok(files) = fs::read_dir(&self.directory) else {
|
|
return Vec::new();
|
|
};
|
|
files
|
|
.filter_map(Result::ok)
|
|
.map(|file| file.path())
|
|
.filter(|path| path.extension().is_some_and(|value| value == "meta"))
|
|
.filter_map(|path| read_entry(&path).ok())
|
|
.filter(|entry| entry.checkpoint.is_file())
|
|
.collect()
|
|
}
|
|
|
|
fn evict(&self, protected: Option<&Path>) {
|
|
let mut entries = self.entries();
|
|
let mut total = entries.iter().map(|entry| entry.bytes).sum::<u64>();
|
|
let now = unix_time();
|
|
while total > self.budget_bytes {
|
|
let Some((index, _)) = entries
|
|
.iter()
|
|
.enumerate()
|
|
.filter(|(_, entry)| protected != Some(entry.checkpoint.as_path()))
|
|
.min_by(|(_, left), (_, right)| {
|
|
eviction_score(left, now)
|
|
.partial_cmp(&eviction_score(right, now))
|
|
.unwrap_or(Ordering::Equal)
|
|
.then_with(|| left.last_used.cmp(&right.last_used))
|
|
})
|
|
else {
|
|
break;
|
|
};
|
|
let entry = entries.swap_remove(index);
|
|
total = total.saturating_sub(entry.bytes);
|
|
self.discard(&entry);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Deletes what the index cannot reach: a checkpoint whose index is missing or
|
|
/// unreadable can never be found, and an index whose checkpoint is gone is
|
|
/// dead weight. Without this, an interrupted store or a checkpoint written
|
|
/// before the index existed occupies disc that eviction cannot even see, and
|
|
/// the directory grows past the budget unnoticed.
|
|
pub(crate) fn sweep_unreachable(directory: &Path) {
|
|
let Ok(files) = fs::read_dir(directory) else {
|
|
return;
|
|
};
|
|
let reachable = files
|
|
.filter_map(Result::ok)
|
|
.map(|file| file.path())
|
|
.filter(|path| path.extension().is_some_and(|value| value == "meta"))
|
|
.filter_map(|path| read_entry(&path).ok())
|
|
.filter(|entry| entry.checkpoint.is_file())
|
|
.map(|entry| entry.checkpoint)
|
|
.collect::<HashSet<_>>();
|
|
let Ok(files) = fs::read_dir(directory) else {
|
|
return;
|
|
};
|
|
for file in files.flatten() {
|
|
let path = file.path();
|
|
let dead = match path.extension().and_then(|value| value.to_str()) {
|
|
Some("bin") => !reachable.contains(&path),
|
|
Some("meta") => !reachable.contains(&path.with_extension("bin")),
|
|
// Staging files from an interrupted index write.
|
|
Some("tmp") => true,
|
|
_ => false,
|
|
};
|
|
if dead {
|
|
let _ = fs::remove_file(&path);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn eviction_score(entry: &Entry, now: u64) -> f64 {
|
|
let age = now.saturating_sub(entry.last_used.max(entry.created_at)) as f64;
|
|
let hits = f64::from(entry.hits) * f64::exp2(-age / HIT_HALF_LIFE_SECONDS);
|
|
let anchor = if entry.reason.is_anchor() { 2.0 } else { 1.0 };
|
|
anchor * (hits + 1.0) * f64::from(entry.tokens) / entry.bytes.max(1) as f64
|
|
}
|
|
|
|
fn write_entry(entry: &Entry) -> Result<(), String> {
|
|
let path = index_path(&entry.checkpoint);
|
|
let temporary = path.with_extension("meta.tmp");
|
|
let mut file = File::create(&temporary).map_err(|error| error.to_string())?;
|
|
file.write_all(INDEX_MAGIC)
|
|
.map_err(|error| error.to_string())?;
|
|
write_u32(&mut file, INDEX_VERSION)?;
|
|
file.write_all(&entry.tag)
|
|
.map_err(|error| error.to_string())?;
|
|
write_u32(&mut file, entry.tokens)?;
|
|
write_u32(&mut file, entry.context)?;
|
|
write_u32(&mut file, entry.hits)?;
|
|
file.write_all(&[entry.reason.code(), 0, 0, 0])
|
|
.map_err(|error| error.to_string())?;
|
|
write_u64(&mut file, entry.created_at)?;
|
|
write_u64(&mut file, entry.last_used)?;
|
|
write_u64(
|
|
&mut file,
|
|
u64::try_from(entry.key.len()).map_err(|_| "KV cache key is too large")?,
|
|
)?;
|
|
file.write_all(&entry.key)
|
|
.map_err(|error| error.to_string())?;
|
|
file.sync_all().map_err(|error| error.to_string())?;
|
|
fs::rename(temporary, path).map_err(|error| error.to_string())
|
|
}
|
|
|
|
fn read_entry(path: &Path) -> Result<Entry, String> {
|
|
let mut file = File::open(path).map_err(|error| error.to_string())?;
|
|
let mut magic = [0; 8];
|
|
file.read_exact(&mut magic)
|
|
.map_err(|error| error.to_string())?;
|
|
if &magic != INDEX_MAGIC || read_u32(&mut file)? != INDEX_VERSION {
|
|
return Err("KV cache index has an invalid signature".into());
|
|
}
|
|
let mut tag = [0; 32];
|
|
file.read_exact(&mut tag)
|
|
.map_err(|error| error.to_string())?;
|
|
let tokens = read_u32(&mut file)?;
|
|
let context = read_u32(&mut file)?;
|
|
let hits = read_u32(&mut file)?;
|
|
let mut reason = [0; 4];
|
|
file.read_exact(&mut reason)
|
|
.map_err(|error| error.to_string())?;
|
|
let reason = StoreReason::from_code(reason[0])
|
|
.ok_or_else(|| "KV cache index has an invalid reason".to_owned())?;
|
|
let created_at = read_u64(&mut file)?;
|
|
let last_used = read_u64(&mut file)?;
|
|
let key_bytes = read_u64(&mut file)?;
|
|
if tokens == 0 || context == 0 || key_bytes > MAX_KEY_BYTES {
|
|
return Err("KV cache index has invalid dimensions".into());
|
|
}
|
|
let mut key = vec![0; usize::try_from(key_bytes).map_err(|error| error.to_string())?];
|
|
file.read_exact(&mut key)
|
|
.map_err(|error| error.to_string())?;
|
|
let checkpoint = path.with_extension("bin");
|
|
if checkpoint.file_stem().and_then(|value| value.to_str()) != Some(&key_hash(&key)) {
|
|
return Err("KV cache index key does not match its filename".into());
|
|
}
|
|
let bytes = fs::metadata(&checkpoint)
|
|
.map_err(|error| error.to_string())?
|
|
.len()
|
|
.saturating_add(fs::metadata(path).map_err(|error| error.to_string())?.len());
|
|
Ok(Entry {
|
|
checkpoint,
|
|
tag,
|
|
tokens,
|
|
context,
|
|
hits,
|
|
created_at,
|
|
last_used,
|
|
reason,
|
|
key,
|
|
bytes,
|
|
})
|
|
}
|
|
|
|
fn index_path(checkpoint: &Path) -> PathBuf {
|
|
checkpoint.with_extension("meta")
|
|
}
|
|
|
|
fn index_size(key: &[u8]) -> Result<u64, String> {
|
|
84_u64
|
|
.checked_add(u64::try_from(key.len()).map_err(|error| error.to_string())?)
|
|
.ok_or_else(|| "KV cache index is too large".to_owned())
|
|
}
|
|
|
|
fn key_hash(key: &[u8]) -> String {
|
|
let digest: [u8; 32] = Sha256::digest(key).into();
|
|
const HEX: &[u8; 16] = b"0123456789abcdef";
|
|
let mut output = String::with_capacity(64);
|
|
for byte in digest {
|
|
output.push(HEX[(byte >> 4) as usize] as char);
|
|
output.push(HEX[(byte & 0x0f) as usize] as char);
|
|
}
|
|
output
|
|
}
|
|
|
|
fn unix_time() -> u64 {
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.map_or(0, |duration| duration.as_secs())
|
|
}
|
|
|
|
fn write_u32(file: &mut File, value: u32) -> Result<(), String> {
|
|
file.write_all(&value.to_le_bytes())
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
fn write_u64(file: &mut File, value: u64) -> Result<(), String> {
|
|
file.write_all(&value.to_le_bytes())
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
|
|
fn read_u32(file: &mut File) -> Result<u32, String> {
|
|
let mut bytes = [0; 4];
|
|
file.read_exact(&mut bytes)
|
|
.map_err(|error| error.to_string())?;
|
|
Ok(u32::from_le_bytes(bytes))
|
|
}
|
|
|
|
fn read_u64(file: &mut File) -> Result<u64, String> {
|
|
let mut bytes = [0; 8];
|
|
file.read_exact(&mut bytes)
|
|
.map_err(|error| error.to_string())?;
|
|
Ok(u64::from_le_bytes(bytes))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::settings::GIB;
|
|
|
|
fn temporary_directory(name: &str) -> PathBuf {
|
|
let path = std::env::temp_dir().join(format!(
|
|
"ds4-server-kvstore-{name}-{}-{}",
|
|
std::process::id(),
|
|
unix_time()
|
|
));
|
|
fs::create_dir_all(&path).unwrap();
|
|
path
|
|
}
|
|
|
|
fn add(store: &KvStore, key: &[u8], tokens: u32, bytes: usize) -> Entry {
|
|
let path = store.checkpoint_path(key);
|
|
fs::write(&path, vec![0; bytes]).unwrap();
|
|
assert!(
|
|
store
|
|
.record(
|
|
&path,
|
|
key,
|
|
[tokens as u8; 32],
|
|
tokens,
|
|
4096,
|
|
StoreReason::Continued,
|
|
)
|
|
.unwrap()
|
|
);
|
|
store.find(key, 4096).unwrap()
|
|
}
|
|
|
|
#[test]
|
|
fn longest_compatible_prefix_wins_and_hits_survive_reopen() {
|
|
let directory = temporary_directory("prefix");
|
|
let store = KvStore::open(&directory, GIB).unwrap();
|
|
add(&store, b"system", 100, 32);
|
|
let expected = add(&store, b"system-user-one", 200, 32);
|
|
|
|
let found = store.find(b"system-user-one-more", 4096).unwrap();
|
|
assert_eq!(found.checkpoint, expected.checkpoint);
|
|
store.touch(&found).unwrap();
|
|
assert_eq!(
|
|
KvStore::open(&directory, GIB)
|
|
.unwrap()
|
|
.find(b"system-user-one-more", 4096)
|
|
.unwrap()
|
|
.hits,
|
|
1
|
|
);
|
|
assert!(store.find(b"system-user-one-more", 2048).is_none());
|
|
|
|
fs::remove_dir_all(directory).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn opening_reclaims_checkpoints_the_index_cannot_reach() {
|
|
let directory = temporary_directory("unreachable");
|
|
let store = KvStore::open(&directory, GIB).unwrap();
|
|
let indexed = add(&store, b"indexed", 100, 32);
|
|
let stray = directory.join("deadbeef.bin");
|
|
let widow = directory.join("widow.meta");
|
|
let staging = directory.join("indexed.meta.tmp");
|
|
fs::write(&stray, vec![0; 4_096]).unwrap();
|
|
fs::write(&widow, vec![0; 64]).unwrap();
|
|
fs::write(&staging, vec![0; 64]).unwrap();
|
|
|
|
KvStore::open(&directory, GIB).unwrap();
|
|
|
|
assert!(indexed.checkpoint.is_file());
|
|
assert!(index_path(&indexed.checkpoint).is_file());
|
|
assert!(!stray.exists());
|
|
assert!(!widow.exists());
|
|
assert!(!staging.exists());
|
|
|
|
fs::remove_dir_all(directory).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn budget_evicts_the_least_valuable_unprotected_entry() {
|
|
let directory = temporary_directory("evict");
|
|
let store = KvStore::open(&directory, 300).unwrap();
|
|
let first = add(&store, b"first", 10, 32);
|
|
let second = add(&store, b"second", 100, 32);
|
|
let third = add(&store, b"third", 200, 32);
|
|
|
|
assert!(!first.checkpoint.exists());
|
|
assert!(second.checkpoint.exists());
|
|
assert!(third.checkpoint.exists());
|
|
|
|
fs::remove_dir_all(directory).unwrap();
|
|
}
|
|
}
|