use std::collections::HashMap;
use std::sync::Arc;
use oxcache::backend::MokaMemoryBackend;
use oxcache::cache::Cache;
pub(crate) struct OxCacheBackend {
cache: Option<Cache<String, Vec<f32>>>,
enabled: bool,
}
impl OxCacheBackend {
pub fn new(capacity: usize) -> Self {
let moka = MokaMemoryBackend::builder()
.capacity(capacity.max(1) as u64)
.build();
let cache = Cache::with_dependencies(Arc::new(moka));
Self {
cache: Some(cache),
enabled: true,
}
}
pub fn disabled() -> Self {
Self {
cache: None,
enabled: false,
}
}
pub fn is_enabled(&self) -> bool {
self.enabled
}
pub async fn get(&self, key: &str) -> Option<Vec<f32>> {
if !self.enabled {
return None;
}
let cache = self.cache.as_ref()?;
cache.get(&key.to_string()).await.ok().flatten()
}
pub async fn put(&self, key: &str, value: Vec<f32>) {
if !self.enabled {
return;
}
if let Some(cache) = &self.cache {
let _ = cache.set(&key.to_string(), &value).await;
}
}
pub async fn get_or_insert<F, Fut, E>(&self, key: &str, f: F) -> Result<Vec<f32>, E>
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = Result<Vec<f32>, E>>,
{
if let Some(cached) = self.get(key).await {
return Ok(cached);
}
let embedding = f().await?;
self.put(key, embedding.clone()).await;
Ok(embedding)
}
pub async fn remove(&self, key: &str) -> bool {
if !self.enabled {
return false;
}
match &self.cache {
Some(cache) => cache.delete(&key.to_string()).await.is_ok(),
None => false,
}
}
pub async fn clear(&self) {
if let Some(cache) = &self.cache {
let _ = cache.clear().await;
}
}
pub async fn len(&self) -> usize {
match &self.cache {
Some(cache) => cache.len().await.map(|n| n as usize).unwrap_or(0),
None => 0,
}
}
pub async fn is_empty(&self) -> bool {
self.len().await == 0
}
pub async fn warm_up(&self, entries: HashMap<String, Vec<f32>>) {
if !self.enabled {
return;
}
for (key, value) in entries {
self.put(&key, value).await;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_put_and_get_hit() {
let cache = OxCacheBackend::new(16);
assert!(cache.is_enabled());
cache.put("text:hello", vec![0.1, 0.2, 0.3]).await;
let got = cache.get("text:hello").await;
assert_eq!(got, Some(vec![0.1, 0.2, 0.3]));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_get_miss_returns_none() {
let cache = OxCacheBackend::new(16);
assert!(cache.get("text:missing").await.is_none());
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_eviction_under_capacity_pressure() {
let cache = OxCacheBackend::new(10);
for i in 0..30 {
cache.put(&format!("k{}", i), vec![i as f32]).await;
}
tokio::time::sleep(Duration::from_millis(300)).await;
let mut remaining = 0;
for i in 0..30 {
if cache.get(&format!("k{}", i)).await.is_some() {
remaining += 1;
}
}
assert!(
remaining < 30,
"capacity pressure should evict some entries, {} still present",
remaining
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_ttl_expiry() {
let moka = MokaMemoryBackend::builder()
.capacity(16)
.ttl(Duration::from_millis(50))
.build();
let cache = Cache::with_dependencies(Arc::new(moka));
cache
.set(&"k".to_string(), &vec![1.0_f32])
.await
.expect("set");
assert_eq!(
cache.get(&"k".to_string()).await.unwrap(),
Some(vec![1.0]),
"should hit before TTL"
);
tokio::time::sleep(Duration::from_millis(120)).await;
assert_eq!(
cache.get(&"k".to_string()).await.unwrap(),
None,
"should expire after TTL"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_clear_empties_cache() {
let cache = OxCacheBackend::new(16);
cache.put("a", vec![1.0]).await;
cache.put("b", vec![2.0]).await;
cache.clear().await;
tokio::time::sleep(Duration::from_millis(20)).await;
assert!(cache.get("a").await.is_none(), "a should be cleared");
assert!(cache.get("b").await.is_none(), "b should be cleared");
assert!(cache.is_empty().await, "cache should be empty");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_disabled_backend_is_noop() {
let cache = OxCacheBackend::disabled();
assert!(!cache.is_enabled());
cache.put("k", vec![1.0]).await;
assert!(cache.get("k").await.is_none());
assert_eq!(cache.len().await, 0);
assert!(cache.is_empty().await);
assert!(!cache.remove("k").await);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_get_or_insert_first_call_invokes_fallback() {
let cache = OxCacheBackend::new(16);
let called = Arc::new(std::sync::atomic::AtomicBool::new(false));
let called_clone = called.clone();
let val = cache
.get_or_insert::<_, _, std::convert::Infallible>("key", || async move {
called_clone.store(true, std::sync::atomic::Ordering::SeqCst);
Ok(vec![9.9])
})
.await
.expect("get_or_insert");
assert_eq!(val, vec![9.9]);
assert!(called.load(std::sync::atomic::Ordering::SeqCst));
let called2 = Arc::new(std::sync::atomic::AtomicBool::new(false));
let called2_clone = called2.clone();
let val = cache
.get_or_insert::<_, _, std::convert::Infallible>("key", || async move {
called2_clone.store(true, std::sync::atomic::Ordering::SeqCst);
Ok(vec![0.0])
})
.await
.expect("get_or_insert second");
assert_eq!(val, vec![9.9], "should return cached value");
assert!(
!called2.load(std::sync::atomic::Ordering::SeqCst),
"fallback should not be called on cache hit"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_warm_up_inserts_all_entries() {
let cache = OxCacheBackend::new(32);
let mut entries = HashMap::new();
entries.insert("a".to_string(), vec![1.0]);
entries.insert("b".to_string(), vec![2.0]);
entries.insert("c".to_string(), vec![3.0]);
cache.warm_up(entries).await;
tokio::time::sleep(Duration::from_millis(20)).await;
assert_eq!(cache.get("a").await, Some(vec![1.0]));
assert_eq!(cache.get("b").await, Some(vec![2.0]));
assert_eq!(cache.get("c").await, Some(vec![3.0]));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_remove_returns_false_for_missing() {
let cache = OxCacheBackend::new(16);
let _ = cache.remove("missing").await;
}
}