use somatize_core::cache::{CacheKey, CacheStore, CacheTier, EntryMeta};
use somatize_core::error::Result;
use somatize_core::value::Value;
pub struct TieredCache {
tiers: Vec<(CacheTier, Box<dyn CacheStore>)>,
}
impl TieredCache {
pub fn new(tiers: Vec<(CacheTier, Box<dyn CacheStore>)>) -> Self {
Self { tiers }
}
pub fn memory_and_local(memory: Box<dyn CacheStore>, local: Box<dyn CacheStore>) -> Self {
Self {
tiers: vec![(CacheTier::Memory, memory), (CacheTier::Local, local)],
}
}
}
impl CacheStore for TieredCache {
fn tier(&self) -> CacheTier {
self.tiers
.first()
.map(|(t, _)| *t)
.unwrap_or(CacheTier::Memory)
}
fn get(&self, key: &CacheKey) -> Result<Option<Value>> {
Ok(self.get_located(key)?.map(|(value, _)| value))
}
fn get_located(&self, key: &CacheKey) -> Result<Option<(Value, CacheTier)>> {
for (i, (tier, store)) in self.tiers.iter().enumerate() {
if let Some(value) = store.get(key)? {
for (_, faster_store) in &self.tiers[..i] {
let _ = faster_store.put(key, &value);
}
return Ok(Some((value, *tier)));
}
}
Ok(None)
}
fn put(&self, key: &CacheKey, value: &Value) -> Result<()> {
for (_, store) in &self.tiers {
store.put(key, value)?;
}
Ok(())
}
fn put_with_origin(
&self,
key: &CacheKey,
value: &Value,
origin: &somatize_core::cache::Origin,
) -> Result<()> {
for (_, store) in &self.tiers {
store.put_with_origin(key, value, origin)?;
}
Ok(())
}
fn exists(&self, key: &CacheKey) -> Result<bool> {
for (_, store) in &self.tiers {
if store.exists(key)? {
return Ok(true);
}
}
Ok(false)
}
fn remove(&self, key: &CacheKey) -> Result<()> {
for (_, store) in &self.tiers {
store.remove(key)?;
}
Ok(())
}
fn metadata(&self, key: &CacheKey) -> Result<Option<EntryMeta>> {
for (_, store) in &self.tiers {
if let Some(meta) = store.metadata(key)? {
return Ok(Some(meta));
}
}
Ok(None)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cache::local::LocalCache;
use crate::cache::memory::MemoryCache;
use std::env;
use std::fs;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
static COUNTER: AtomicU64 = AtomicU64::new(0);
fn temp_dir() -> PathBuf {
let id = COUNTER.fetch_add(1, Ordering::Relaxed);
let dir = env::temp_dir().join(format!("soma_tiered_test_{}_{id}", std::process::id()));
let _ = fs::remove_dir_all(&dir);
dir
}
fn make_tiered() -> (TieredCache, PathBuf) {
let dir = temp_dir();
let memory = Box::new(MemoryCache::default());
let local = Box::new(LocalCache::new(&dir).unwrap());
(TieredCache::memory_and_local(memory, local), dir)
}
#[test]
fn get_located_names_the_tier_that_answered() {
let (cache, dir) = make_tiered();
let key = CacheKey::hash_data(b"only-on-disk");
let value = Value::tensor(vec![7.0], vec![1]);
let local = LocalCache::new(&dir).unwrap();
local.put(&key, &value).unwrap();
let (got, tier) = cache.get_located(&key).unwrap().unwrap();
assert_eq!(got, value);
assert_eq!(tier, CacheTier::Local);
let (_, tier) = cache.get_located(&key).unwrap().unwrap();
assert_eq!(tier, CacheTier::Memory);
let _ = fs::remove_dir_all(&dir);
}
#[test]
fn put_writes_to_all_tiers() {
let (cache, dir) = make_tiered();
let key = CacheKey::hash_data(b"test");
let value = Value::tensor(vec![1.0, 2.0], vec![2]);
cache.put(&key, &value).unwrap();
assert!(cache.tiers[0].1.exists(&key).unwrap()); assert!(cache.tiers[1].1.exists(&key).unwrap());
let _ = fs::remove_dir_all(&dir);
}
#[test]
fn get_from_memory_first() {
let (cache, dir) = make_tiered();
let key = CacheKey::hash_data(b"test");
let value = Value::tensor(vec![1.0], vec![1]);
cache.put(&key, &value).unwrap();
let result = cache.get(&key).unwrap().unwrap();
assert_eq!(result, value);
let _ = fs::remove_dir_all(&dir);
}
#[test]
fn promotes_from_local_to_memory() {
let dir = temp_dir();
let memory = Box::new(MemoryCache::default());
let local = Box::new(LocalCache::new(&dir).unwrap());
let key = CacheKey::hash_data(b"local_only");
let value = Value::tensor(vec![42.0], vec![1]);
local.put(&key, &value).unwrap();
let tiered = TieredCache::memory_and_local(memory, local);
assert!(!tiered.tiers[0].1.exists(&key).unwrap());
let result = tiered.get(&key).unwrap().unwrap();
assert_eq!(result, value);
assert!(tiered.tiers[0].1.exists(&key).unwrap());
let _ = fs::remove_dir_all(&dir);
}
#[test]
fn miss_returns_none() {
let (cache, dir) = make_tiered();
assert!(cache.get(&CacheKey::hash_data(b"nope")).unwrap().is_none());
let _ = fs::remove_dir_all(&dir);
}
#[test]
fn remove_from_all_tiers() {
let (cache, dir) = make_tiered();
let key = CacheKey::hash_data(b"test");
cache.put(&key, &Value::Empty).unwrap();
cache.remove(&key).unwrap();
assert!(!cache.tiers[0].1.exists(&key).unwrap());
assert!(!cache.tiers[1].1.exists(&key).unwrap());
let _ = fs::remove_dir_all(&dir);
}
}