use rustc_hash::FxHashMap;
use std::hash::Hash;
use std::io;
use std::path::{Path, PathBuf};
use std::sync::{Arc, OnceLock, RwLock};
use std::time::SystemTime;
pub struct ProcessCache<K, V> {
inner: OnceLock<RwLock<FxHashMap<K, Arc<V>>>>,
}
impl<K, V> ProcessCache<K, V> {
pub const fn new() -> Self {
Self {
inner: OnceLock::new(),
}
}
fn map(&self) -> &RwLock<FxHashMap<K, Arc<V>>> {
self.inner.get_or_init(|| RwLock::new(FxHashMap::default()))
}
}
impl<K, V> Default for ProcessCache<K, V> {
fn default() -> Self {
Self::new()
}
}
impl<K, V> ProcessCache<K, V>
where
K: Eq + Hash + Clone,
{
pub fn get_or_compute<F>(&self, key: K, f: F) -> Arc<V>
where
F: FnOnce() -> V,
{
if let Some(v) = self
.map()
.read()
.expect("ProcessCache lock poisoned")
.get(&key)
{
return Arc::clone(v);
}
let value = Arc::new(f());
let mut w = self.map().write().expect("ProcessCache lock poisoned");
let stored = w.entry(key).or_insert_with(|| Arc::clone(&value));
Arc::clone(stored)
}
pub fn get(&self, key: &K) -> Option<Arc<V>> {
self.map()
.read()
.expect("ProcessCache lock poisoned")
.get(key)
.map(Arc::clone)
}
pub fn insert(&self, key: K, value: Arc<V>) {
self.map()
.write()
.expect("ProcessCache lock poisoned")
.insert(key, value);
}
pub fn invalidate(&self, key: &K) -> Option<Arc<V>> {
self.map()
.write()
.expect("ProcessCache lock poisoned")
.remove(key)
}
}
pub struct DiskCache {
root: PathBuf,
}
impl DiskCache {
pub fn new(root: impl Into<PathBuf>) -> Self {
Self { root: root.into() }
}
pub fn root(&self) -> &Path {
&self.root
}
fn path_for(&self, key: &[u8]) -> PathBuf {
let hash = blake3::hash(key).to_hex();
let hex = hash.as_str();
self.root.join(&hex[..2]).join(hex)
}
pub fn read_bytes(&self, key: &[u8]) -> io::Result<Option<Vec<u8>>> {
let path = self.path_for(key);
match std::fs::read(&path) {
Ok(b) => Ok(Some(b)),
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
Err(e) => Err(e),
}
}
pub fn write_bytes(&self, key: &[u8], bytes: &[u8]) -> io::Result<()> {
let path = self.path_for(key);
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
crate::fs_atomic::atomic_write(&path, bytes)
}
pub fn remove(&self, key: &[u8]) -> io::Result<()> {
let path = self.path_for(key);
match std::fs::remove_file(&path) {
Ok(()) => Ok(()),
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(e),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FreshnessSnapshot {
pub mtime: SystemTime,
pub size: u64,
pub hash: [u8; 32],
}
impl FreshnessSnapshot {
pub fn capture(path: &Path) -> io::Result<Self> {
let meta = std::fs::metadata(path)?;
let mtime = meta.modified()?;
let size = meta.len();
let bytes = std::fs::read(path)?;
let hash = *blake3::hash(&bytes).as_bytes();
Ok(Self { mtime, size, hash })
}
pub fn is_fresh(&self, path: &Path) -> io::Result<bool> {
let meta = std::fs::metadata(path)?;
if meta.len() != self.size {
return Ok(false);
}
if let Ok(mtime) = meta.modified()
&& mtime == self.mtime
{
return Ok(true);
}
let bytes = std::fs::read(path)?;
let hash = *blake3::hash(&bytes).as_bytes();
Ok(hash == self.hash)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn tempdir() -> PathBuf {
let base = std::env::temp_dir().join(format!(
"aube-cache-test-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0)
));
std::fs::create_dir_all(&base).unwrap();
base
}
#[test]
fn process_cache_computes_once() {
let cache: ProcessCache<&'static str, u32> = ProcessCache::new();
let n = std::sync::atomic::AtomicU32::new(0);
let _a = cache.get_or_compute("k", || {
n.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
42
});
let _b = cache.get_or_compute("k", || {
n.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
42
});
assert_eq!(n.load(std::sync::atomic::Ordering::SeqCst), 1);
}
#[test]
fn process_cache_returns_arc_clone() {
let cache: ProcessCache<u32, String> = ProcessCache::new();
let a = cache.get_or_compute(1, || "hello".to_string());
let b = cache.get_or_compute(1, || "world".to_string());
assert!(Arc::ptr_eq(&a, &b));
assert_eq!(*a, "hello");
}
#[test]
fn disk_cache_roundtrip() {
let dir = tempdir();
let cache = DiskCache::new(dir.join("dc"));
assert!(cache.read_bytes(b"key1").unwrap().is_none());
cache.write_bytes(b"key1", b"value-bytes").unwrap();
assert_eq!(
cache.read_bytes(b"key1").unwrap().as_deref(),
Some(b"value-bytes".as_ref())
);
cache.remove(b"key1").unwrap();
assert!(cache.read_bytes(b"key1").unwrap().is_none());
}
#[test]
fn freshness_detects_size_change() {
let dir = tempdir();
let path = dir.join("file");
std::fs::write(&path, b"hello").unwrap();
let snap = FreshnessSnapshot::capture(&path).unwrap();
assert!(snap.is_fresh(&path).unwrap());
std::fs::write(&path, b"hello world").unwrap();
assert!(!snap.is_fresh(&path).unwrap());
}
#[test]
fn freshness_handles_touch_with_same_content() {
let dir = tempdir();
let path = dir.join("file");
std::fs::write(&path, b"hello").unwrap();
let snap = FreshnessSnapshot::capture(&path).unwrap();
std::thread::sleep(std::time::Duration::from_millis(20));
std::fs::write(&path, b"hello").unwrap();
assert!(snap.is_fresh(&path).unwrap());
}
}