Add prefix-aware KV session cache
This commit is contained in:
Georg Bauer
385
src/engine/kvstore.rs
Normal file
385
src/engine/kvstore.rs
Normal file
@@ -0,0 +1,385 @@
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::cmp::Ordering;
|
||||
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 DEFAULT_BUDGET_BYTES: u64 = 4 * 1024 * 1024 * 1024;
|
||||
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) -> Result<Self, String> {
|
||||
fs::create_dir_all(directory).map_err(|error| {
|
||||
format!(
|
||||
"Could not create transient KV cache directory {}: {error}",
|
||||
directory.display()
|
||||
)
|
||||
})?;
|
||||
Ok(Self {
|
||||
directory: directory.to_owned(),
|
||||
budget_bytes: DEFAULT_BUDGET_BYTES,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn with_budget(directory: &Path, budget_bytes: u64) -> Result<Self, String> {
|
||||
let mut store = Self::open(directory)?;
|
||||
store.budget_bytes = budget_bytes;
|
||||
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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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::*;
|
||||
|
||||
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).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)
|
||||
.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 budget_evicts_the_least_valuable_unprotected_entry() {
|
||||
let directory = temporary_directory("evict");
|
||||
let store = KvStore::with_budget(&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();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user