use std::{
collections::{HashMap, VecDeque},
path::{Path, PathBuf},
sync::atomic::{AtomicU64, Ordering},
};
use bytes::Bytes;
use tokio::sync::{Mutex, RwLock};
use crate::BlakeHash;
static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);
pub(crate) enum Cache {
Mem(MemCache),
Disk(DiskCache),
}
impl Cache {
pub(crate) fn new(
cache_dir: Option<&Path>,
budget_bytes: u64,
) -> std::io::Result<Self> {
match cache_dir {
None => Ok(Self::Mem(MemCache::default())),
Some(dir) => Ok(Self::Disk(DiskCache::open(dir, budget_bytes)?)),
}
}
pub(crate) async fn get(&self, hash: &BlakeHash) -> Option<Bytes> {
match self {
Self::Mem(m) => m.get(hash).await,
Self::Disk(d) => d.get(hash).await,
}
}
pub(crate) async fn put(
&self,
hash: BlakeHash,
bytes: Bytes,
) -> std::io::Result<()> {
match self {
Self::Mem(m) => {
m.put(hash, bytes).await;
Ok(())
}
Self::Disk(d) => d.put(hash, bytes).await,
}
}
pub(crate) async fn contains(&self, hash: &BlakeHash) -> bool {
match self {
Self::Mem(m) => m.contains(hash).await,
Self::Disk(d) => d.contains(hash).await,
}
}
}
#[derive(Default)]
pub(crate) struct MemCache {
inner: RwLock<HashMap<BlakeHash, Bytes>>,
}
impl MemCache {
async fn get(&self, hash: &BlakeHash) -> Option<Bytes> {
self.inner.read().await.get(hash).cloned()
}
async fn put(&self, hash: BlakeHash, bytes: Bytes) {
self.inner.write().await.insert(hash, bytes);
}
async fn contains(&self, hash: &BlakeHash) -> bool {
self.inner.read().await.contains_key(hash)
}
}
pub(crate) struct DiskCache {
dir: PathBuf,
budget_bytes: u64,
state: Mutex<LruState>,
}
#[derive(Default)]
struct LruState {
entries: HashMap<BlakeHash, u64>,
lru: VecDeque<BlakeHash>,
total_bytes: u64,
}
impl LruState {
fn touch(&mut self, hash: &BlakeHash) {
if let Some(pos) = self.lru.iter().position(|h| h == hash) {
self.lru.remove(pos);
self.lru.push_back(*hash);
}
}
fn insert(&mut self, hash: BlakeHash, size: u64, budget: u64) -> Vec<BlakeHash> {
if let Some(old_size) = self.entries.remove(&hash) {
self.total_bytes = self.total_bytes.saturating_sub(old_size);
if let Some(pos) = self.lru.iter().position(|h| h == &hash) {
self.lru.remove(pos);
}
}
let mut evicted = Vec::new();
while self.total_bytes + size > budget {
let Some(victim) = self.lru.pop_front() else { break };
if let Some(victim_size) = self.entries.remove(&victim) {
self.total_bytes = self.total_bytes.saturating_sub(victim_size);
}
evicted.push(victim);
}
self.entries.insert(hash, size);
self.lru.push_back(hash);
self.total_bytes += size;
evicted
}
}
impl DiskCache {
pub(crate) fn open(dir: &Path, budget_bytes: u64) -> std::io::Result<Self> {
std::fs::create_dir_all(dir)?;
let mut found: Vec<(std::time::SystemTime, BlakeHash, u64)> = Vec::new();
for entry in std::fs::read_dir(dir)? {
let entry = entry?;
let path = entry.path();
let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
continue;
};
if name.ends_with(".tmp") || name.contains(".tmp.") {
let _ = std::fs::remove_file(&path);
continue;
}
let Ok(hash) = BlakeHash::from_hex(name) else { continue };
let Ok(meta) = entry.metadata() else { continue };
if !meta.is_file() {
continue;
}
let mtime = meta.modified().unwrap_or(std::time::SystemTime::UNIX_EPOCH);
found.push((mtime, hash, meta.len()));
}
found.sort_by_key(|(t, _, _)| *t);
let mut state = LruState::default();
for (_, hash, size) in found {
state.entries.insert(hash, size);
state.lru.push_back(hash);
state.total_bytes += size;
}
let dir = dir.to_path_buf();
while state.total_bytes > budget_bytes {
let Some(victim) = state.lru.pop_front() else { break };
if let Some(size) = state.entries.remove(&victim) {
state.total_bytes = state.total_bytes.saturating_sub(size);
}
let _ = std::fs::remove_file(blob_path(&dir, &victim));
}
Ok(Self {
dir,
budget_bytes,
state: Mutex::new(state),
})
}
async fn get(&self, hash: &BlakeHash) -> Option<Bytes> {
{
let state = self.state.lock().await;
if !state.entries.contains_key(hash) {
return None;
}
}
let path = blob_path(&self.dir, hash);
let data = match tokio::fs::read(&path).await {
Ok(d) => d,
Err(_) => {
let mut state = self.state.lock().await;
if let Some(size) = state.entries.remove(hash) {
state.total_bytes = state.total_bytes.saturating_sub(size);
}
if let Some(pos) = state.lru.iter().position(|h| h == hash) {
state.lru.remove(pos);
}
return None;
}
};
let bytes = Bytes::from(data);
if !hash.verify(&bytes) {
tracing::warn!(
%hash,
"disk cache BLAKE3 mismatch — evicting entry"
);
let _ = tokio::fs::remove_file(&path).await;
let mut state = self.state.lock().await;
if let Some(size) = state.entries.remove(hash) {
state.total_bytes = state.total_bytes.saturating_sub(size);
}
if let Some(pos) = state.lru.iter().position(|h| h == hash) {
state.lru.remove(pos);
}
return None;
}
self.state.lock().await.touch(hash);
Some(bytes)
}
async fn put(&self, hash: BlakeHash, bytes: Bytes) -> std::io::Result<()> {
let size = bytes.len() as u64;
if size > self.budget_bytes {
tracing::warn!(
%hash,
size,
budget = self.budget_bytes,
"disk cache: blob exceeds budget, skipping write"
);
return Ok(());
}
let final_path = blob_path(&self.dir, &hash);
let tmp_path = {
let mut p = final_path.clone();
let n = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
let pid = std::process::id();
p.set_extension(format!("tmp.{pid}.{n}"));
p
};
tokio::fs::write(&tmp_path, &bytes).await?;
tokio::fs::rename(&tmp_path, &final_path).await?;
let mut state = self.state.lock().await;
let evicted = state.insert(hash, size, self.budget_bytes);
drop(state);
for victim in evicted {
let _ = tokio::fs::remove_file(blob_path(&self.dir, &victim)).await;
}
Ok(())
}
async fn contains(&self, hash: &BlakeHash) -> bool {
self.state.lock().await.entries.contains_key(hash)
}
}
fn blob_path(dir: &Path, hash: &BlakeHash) -> PathBuf {
dir.join(hash.to_hex())
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
fn tmpdir(name: &str) -> PathBuf {
let p = std::env::temp_dir().join(format!(
"xlb-cache-test-{name}-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
std::fs::create_dir_all(&p).unwrap();
p
}
#[tokio::test]
async fn mem_cache_roundtrip() {
let cache = Cache::new(None, 1024).unwrap();
let data = Bytes::from_static(b"hello mem cache");
let hash = BlakeHash::hash(&data);
cache.put(hash, data.clone()).await.unwrap();
assert_eq!(cache.get(&hash).await.as_deref(), Some(&data[..]));
assert!(cache.contains(&hash).await);
}
#[tokio::test]
async fn disk_cache_roundtrip() {
let dir = tmpdir("roundtrip");
let cache = Cache::new(Some(&dir), 1024).unwrap();
let data = Bytes::from_static(b"hello disk cache");
let hash = BlakeHash::hash(&data);
cache.put(hash, data.clone()).await.unwrap();
assert_eq!(cache.get(&hash).await.as_deref(), Some(&data[..]));
let on_disk = std::fs::read(dir.join(hash.to_hex())).unwrap();
assert_eq!(on_disk, data.as_ref());
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn disk_cache_survives_reopen() {
let dir = tmpdir("reopen");
let data = Bytes::from_static(b"survive across restart");
let hash = BlakeHash::hash(&data);
{
let cache = Cache::new(Some(&dir), 1024).unwrap();
cache.put(hash, data.clone()).await.unwrap();
}
let cache2 = Cache::new(Some(&dir), 1024).unwrap();
assert!(cache2.contains(&hash).await);
assert_eq!(cache2.get(&hash).await.as_deref(), Some(&data[..]));
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn lru_evicts_oldest_under_budget() {
let dir = tmpdir("lru-evict");
let cache = Cache::new(Some(&dir), 100).unwrap();
let mk = |n: u8| Bytes::from(vec![n; 40]);
let a = mk(b'A');
let b = mk(b'B');
let c = mk(b'C');
let ha = BlakeHash::hash(&a);
let hb = BlakeHash::hash(&b);
let hc = BlakeHash::hash(&c);
cache.put(ha, a.clone()).await.unwrap();
cache.put(hb, b.clone()).await.unwrap();
cache.put(hc, c.clone()).await.unwrap();
assert!(!cache.contains(&ha).await, "A should be evicted (oldest)");
assert!(cache.contains(&hb).await, "B should survive");
assert!(cache.contains(&hc).await, "C should survive");
assert!(
!dir.join(ha.to_hex()).exists(),
"A's file should be gone from disk"
);
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn lru_get_touches_recency() {
let dir = tmpdir("lru-touch");
let cache = Cache::new(Some(&dir), 100).unwrap();
let a = Bytes::from(vec![b'A'; 40]);
let b = Bytes::from(vec![b'B'; 40]);
let c = Bytes::from(vec![b'C'; 40]);
let ha = BlakeHash::hash(&a);
let hb = BlakeHash::hash(&b);
let hc = BlakeHash::hash(&c);
cache.put(ha, a.clone()).await.unwrap();
cache.put(hb, b.clone()).await.unwrap();
let _ = cache.get(&ha).await;
cache.put(hc, c.clone()).await.unwrap();
assert!(cache.contains(&ha).await, "A was touched, should survive");
assert!(!cache.contains(&hb).await, "B is now oldest, should be evicted");
assert!(cache.contains(&hc).await);
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn disk_cache_verifies_on_read() {
let dir = tmpdir("verify");
let cache = Cache::new(Some(&dir), 1024).unwrap();
let data = Bytes::from_static(b"real bytes");
let hash = BlakeHash::hash(&data);
cache.put(hash, data.clone()).await.unwrap();
std::fs::write(dir.join(hash.to_hex()), b"TAMPERED").unwrap();
assert!(
cache.get(&hash).await.is_none(),
"tampered cache entry must be treated as a miss"
);
assert!(
!cache.contains(&hash).await,
"tampered entry must be evicted from the index"
);
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn concurrent_put_same_hash_atomic() {
let dir = tmpdir("concurrent");
let cache = Arc::new(Cache::new(Some(&dir), 4096).unwrap());
let data = Bytes::from(vec![0xab; 200]);
let hash = BlakeHash::hash(&data);
let mut tasks = Vec::new();
for _ in 0..8 {
let cache = cache.clone();
let data = data.clone();
tasks.push(tokio::spawn(async move {
cache.put(hash, data).await.unwrap();
}));
}
for t in tasks {
t.await.unwrap();
}
let got = cache.get(&hash).await.expect("blob must be present");
assert_eq!(got, data, "no torn writes");
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn shrunk_budget_evicts_on_open() {
let dir = tmpdir("shrink");
let a = Bytes::from(vec![b'A'; 40]);
let b = Bytes::from(vec![b'B'; 40]);
let ha = BlakeHash::hash(&a);
let hb = BlakeHash::hash(&b);
{
let cache = Cache::new(Some(&dir), 1024).unwrap();
cache.put(ha, a.clone()).await.unwrap();
std::thread::sleep(std::time::Duration::from_millis(50));
cache.put(hb, b.clone()).await.unwrap();
}
let cache = Cache::new(Some(&dir), 60).unwrap(); assert!(
!cache.contains(&ha).await,
"A (older) should have been evicted on open under shrunk budget"
);
assert!(
cache.contains(&hb).await,
"B (newer) should survive shrunk-budget open"
);
std::fs::remove_dir_all(&dir).ok();
}
#[tokio::test]
async fn temp_files_cleaned_on_open() {
let dir = tmpdir("tmp-cleanup");
std::fs::write(dir.join("abc.tmp"), b"crashed write").unwrap();
std::fs::write(dir.join("def.tmp.123.4"), b"another").unwrap();
let _cache = Cache::new(Some(&dir), 1024).unwrap();
assert!(!dir.join("abc.tmp").exists());
assert!(!dir.join("def.tmp.123.4").exists());
std::fs::remove_dir_all(&dir).ok();
}
}