use bytes::Bytes as SharedBytes;
use crate::ShardCacheError;
use crate::config::EvictionPolicy;
use crate::storage::{
EmbeddedKeyRoute, EmbeddedRouteMode, PreparedPointKey, SemanticCacheError, SemanticMatch,
SharedEmbeddedConfig, SharedEmbeddedEntry, SharedEmbeddedLockPolicy, SharedEmbeddedRef,
SharedEmbeddedRefMut, SharedEmbeddedStore, SharedEmbeddedVacantEntry,
};
pub const DEFAULT_CACHE_SHARDS: usize = 64;
#[derive(Debug, Clone)]
pub struct CacheOptions {
pub total_memory_bytes: Option<usize>,
pub eviction_policy: EvictionPolicy,
pub route_mode: EmbeddedRouteMode,
pub capacity_hint: Option<usize>,
pub default_ttl_ms: Option<u64>,
pub lock_policy: SharedEmbeddedLockPolicy,
}
impl Default for CacheOptions {
fn default() -> Self {
let shared = SharedEmbeddedConfig::default();
Self {
total_memory_bytes: shared.total_memory_bytes,
eviction_policy: shared.eviction_policy,
route_mode: shared.route_mode,
capacity_hint: shared.flat_map_capacity_hint,
default_ttl_ms: None,
lock_policy: shared.lock_policy,
}
}
}
impl From<CacheOptions> for SharedEmbeddedConfig {
fn from(options: CacheOptions) -> Self {
Self {
total_memory_bytes: options.total_memory_bytes,
eviction_policy: options.eviction_policy,
route_mode: options.route_mode,
flat_map_capacity_hint: options.capacity_hint,
lock_policy: options.lock_policy,
}
}
}
impl From<SharedEmbeddedConfig> for CacheOptions {
fn from(config: SharedEmbeddedConfig) -> Self {
Self {
total_memory_bytes: config.total_memory_bytes,
eviction_policy: config.eviction_policy,
route_mode: config.route_mode,
capacity_hint: config.flat_map_capacity_hint,
default_ttl_ms: None,
lock_policy: config.lock_policy,
}
}
}
pub type CacheRef<'a> = SharedEmbeddedRef<'a>;
pub type CacheRefMut<'a> = SharedEmbeddedRefMut<'a>;
pub type CacheEntry<'a> = SharedEmbeddedEntry<'a>;
pub type CacheVacantEntry<'a> = SharedEmbeddedVacantEntry<'a>;
pub type CacheSemanticMatch = SemanticMatch;
pub type CacheSemanticError = SemanticCacheError;
#[derive(Debug, Clone)]
pub struct SharedCache<const SHARDS: usize> {
inner: SharedEmbeddedStore<SHARDS>,
default_ttl_ms: Option<u64>,
}
pub type RawShardMap = SharedCache<DEFAULT_CACHE_SHARDS>;
pub type ShardCache = RawShardMap;
pub type RawShardMapWithShards<const SHARDS: usize> = SharedCache<SHARDS>;
pub type ShardCacheWithShards<const SHARDS: usize> = SharedCache<SHARDS>;
impl<const SHARDS: usize> Default for SharedCache<SHARDS> {
fn default() -> Self {
Self::new()
}
}
impl<const SHARDS: usize> SharedCache<SHARDS> {
#[inline(always)]
pub fn new() -> Self {
Self::with_options(CacheOptions::default())
}
#[inline(always)]
pub fn with_capacity(capacity: usize) -> Self {
Self::with_options(CacheOptions {
capacity_hint: Some(capacity),
..CacheOptions::default()
})
}
#[inline(always)]
pub fn with_options(options: CacheOptions) -> Self {
let default_ttl_ms = options.default_ttl_ms;
Self {
inner: SharedEmbeddedStore::new(options.into()),
default_ttl_ms,
}
}
#[inline(always)]
pub fn from_shared_store(inner: SharedEmbeddedStore<SHARDS>) -> Self {
Self {
inner,
default_ttl_ms: None,
}
}
#[inline(always)]
pub fn as_shared_store(&self) -> &SharedEmbeddedStore<SHARDS> {
&self.inner
}
#[inline(always)]
pub fn into_shared_store(self) -> SharedEmbeddedStore<SHARDS> {
self.inner
}
#[inline(always)]
pub const fn shard_count(&self) -> usize {
SHARDS
}
#[inline(always)]
pub fn route_mode(&self) -> EmbeddedRouteMode {
self.inner.route_mode()
}
#[inline(always)]
pub const fn default_ttl_ms(&self) -> Option<u64> {
self.default_ttl_ms
}
#[inline(always)]
pub fn route_key(&self, key: &[u8]) -> EmbeddedKeyRoute {
self.inner.route_key(key)
}
#[inline(always)]
pub fn prepare_key(&self, key: &[u8]) -> PreparedPointKey {
self.inner.prepare_point_key(key)
}
#[inline(always)]
pub fn get(&self, key: &[u8]) -> Option<CacheRef<'_>> {
self.inner.get_ref(key)
}
#[inline(always)]
pub fn get_ref(&self, key: &[u8]) -> Option<CacheRef<'_>> {
self.inner.get_ref(key)
}
#[inline(always)]
pub fn get_prepared(&self, prepared: &PreparedPointKey) -> Option<CacheRef<'_>> {
self.inner.get_prepared_ref(prepared)
}
#[inline(always)]
pub fn get_owned(&self, key: &[u8]) -> Option<SharedBytes> {
self.inner.get_value_bytes(key)
}
#[inline(always)]
pub fn get_prepared_owned(&self, prepared: &PreparedPointKey) -> Option<SharedBytes> {
self.inner.get_prepared_value_bytes(prepared)
}
#[inline(always)]
pub fn contains_key(&self, key: &[u8]) -> bool {
self.inner.contains_key(key)
}
#[inline(always)]
pub fn exists(&self, key: &[u8]) -> bool {
self.contains_key(key)
}
#[inline(always)]
pub fn key_snapshot(&self) -> Vec<SharedBytes> {
self.inner.key_snapshot()
}
#[inline(always)]
pub fn visit_keys(&self, visitor: impl FnMut(&[u8]) -> bool) {
self.inner.visit_string_keys(visitor);
}
#[inline(always)]
pub fn visit_entries(&self, visitor: impl FnMut(&[u8], &[u8], Option<u64>) -> bool) {
self.inner.visit_string_entries(visitor);
}
#[inline(always)]
pub fn insert<K, V>(&self, key: K, value: V)
where
K: Into<SharedBytes>,
V: Into<SharedBytes>,
{
self.inner
.insert_with_ttl(key.into(), value.into(), self.default_ttl_ms);
}
#[inline(always)]
pub fn insert_slice(&self, key: &[u8], value: &[u8]) {
self.inner
.insert_slice_with_ttl(key, value, self.default_ttl_ms);
}
#[inline(always)]
pub fn insert_with_ttl<K, V>(&self, key: K, value: V, ttl_ms: Option<u64>)
where
K: Into<SharedBytes>,
V: Into<SharedBytes>,
{
self.inner.insert_with_ttl(key.into(), value.into(), ttl_ms);
}
#[inline(always)]
pub fn insert_slice_with_ttl(&self, key: &[u8], value: &[u8], ttl_ms: Option<u64>) {
self.inner.insert_slice_with_ttl(key, value, ttl_ms);
}
#[inline(always)]
pub fn insert_semantic_slice(
&self,
key: &[u8],
value: &[u8],
embedding: &[f32],
) -> Result<(), SemanticCacheError> {
self.inner
.insert_semantic_slice_with_ttl(key, value, embedding, self.default_ttl_ms)
}
#[inline(always)]
pub fn insert_semantic_slice_with_governance(
&self,
key: &[u8],
value: &[u8],
embedding: &[f32],
governance_metadata: &[u8],
) -> Result<(), SemanticCacheError> {
self.inner.insert_semantic_slice_with_ttl_and_governance(
key,
value,
embedding,
self.default_ttl_ms,
governance_metadata,
)
}
#[inline(always)]
pub fn insert_semantic_slice_with_ttl(
&self,
key: &[u8],
value: &[u8],
embedding: &[f32],
ttl_ms: Option<u64>,
) -> Result<(), SemanticCacheError> {
self.inner
.insert_semantic_slice_with_ttl(key, value, embedding, ttl_ms)
}
#[inline(always)]
pub fn insert_semantic_slice_with_ttl_and_governance(
&self,
key: &[u8],
value: &[u8],
embedding: &[f32],
ttl_ms: Option<u64>,
governance_metadata: &[u8],
) -> Result<(), SemanticCacheError> {
self.inner.insert_semantic_slice_with_ttl_and_governance(
key,
value,
embedding,
ttl_ms,
governance_metadata,
)
}
#[inline(always)]
pub fn semantic_search(
&self,
embedding: &[f32],
min_score: f32,
) -> Result<Option<SemanticMatch>, SemanticCacheError> {
self.inner.semantic_search(embedding, min_score)
}
#[inline(always)]
pub fn semantic_search_with_governance_filter(
&self,
embedding: &[f32],
min_score: f32,
governance_filter: impl FnMut(Option<&[u8]>) -> bool,
) -> Result<Option<SemanticMatch>, SemanticCacheError> {
self.inner
.semantic_search_with_governance_filter(embedding, min_score, governance_filter)
}
#[inline(always)]
pub fn semantic_query_cache_enabled(&self) -> bool {
self.inner.semantic_query_cache_enabled()
}
#[inline(always)]
pub fn disable_semantic_query_cache(&self) {
self.inner.disable_semantic_query_cache();
}
#[inline(always)]
pub fn try_insert<K, V>(&self, key: K, value: V) -> bool
where
K: Into<SharedBytes>,
V: Into<SharedBytes>,
{
let key = key.into();
let value = value.into();
self.inner
.insert_slice_if_absent_with_ttl(&key, &value, self.default_ttl_ms)
}
#[inline(always)]
pub fn try_insert_slice(&self, key: &[u8], value: &[u8]) -> bool {
self.inner
.insert_slice_if_absent_with_ttl(key, value, self.default_ttl_ms)
}
#[inline(always)]
pub fn try_insert_with_ttl<K, V>(&self, key: K, value: V, ttl_ms: Option<u64>) -> bool
where
K: Into<SharedBytes>,
V: Into<SharedBytes>,
{
let key = key.into();
let value = value.into();
self.inner
.insert_slice_if_absent_with_ttl(&key, &value, ttl_ms)
}
#[inline(always)]
pub fn try_insert_slice_with_ttl(&self, key: &[u8], value: &[u8], ttl_ms: Option<u64>) -> bool {
self.inner
.insert_slice_if_absent_with_ttl(key, value, ttl_ms)
}
#[inline(always)]
pub fn get_mut(&self, key: &[u8]) -> Option<CacheRefMut<'_>> {
self.inner.get_mut(key)
}
#[inline(always)]
pub fn remove(&self, key: &[u8]) -> Option<SharedBytes> {
self.inner.remove(key)
}
#[inline(always)]
pub fn entry<K>(&self, key: K) -> CacheEntry<'_>
where
K: Into<SharedBytes>,
{
self.inner.entry(key.into())
}
#[inline(always)]
pub fn try_acquire_lock(&self, key: &[u8], token: &[u8], ttl_ms: u64) -> crate::Result<bool> {
if ttl_ms == 0 {
return Err(ShardCacheError::Command(
"lock ttl_ms must be greater than zero".into(),
));
}
Ok(self
.inner
.insert_slice_if_absent_with_ttl(key, token, Some(ttl_ms)))
}
#[inline(always)]
pub fn release_lock(&self, key: &[u8], token: &[u8]) -> bool {
self.inner.remove_if_value_eq(key, token)
}
#[inline(always)]
pub fn renew_lock(&self, key: &[u8], token: &[u8], ttl_ms: u64) -> crate::Result<bool> {
if ttl_ms == 0 {
return Err(ShardCacheError::Command(
"lock ttl_ms must be greater than zero".into(),
));
}
self.inner.update_ttl_if_value_eq(key, token, ttl_ms)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn key_routed_away_from_semantic_shard<const SHARDS: usize>(
cache: &SharedCache<SHARDS>,
) -> Vec<u8> {
for index in 0..4096 {
let key = format!("semantic-cross-shard-{index}").into_bytes();
if cache.route_key(&key).shard_id != 0 {
return key;
}
}
panic!("could not find key routed away from semantic shard");
}
#[test]
fn fast_map_round_trips_like_a_shared_map() {
let cache = SharedCache::<4>::with_capacity(16);
cache.insert_slice(b"alpha", b"one");
assert_eq!(cache.get(b"alpha").unwrap().value(), b"one");
assert_eq!(cache.get_owned(b"alpha").unwrap().as_ref(), b"one");
assert!(cache.contains_key(b"alpha"));
cache.get_mut(b"alpha").unwrap().set_slice(b"two");
assert_eq!(cache.remove(b"alpha").unwrap().as_ref(), b"two");
assert!(!cache.contains_key(b"alpha"));
}
#[test]
fn fast_map_try_insert_only_writes_missing_keys() {
let cache = SharedCache::<4>::new();
assert!(cache.try_insert_slice(b"alpha", b"one"));
assert!(!cache.try_insert_slice(b"alpha", b"two"));
assert_eq!(cache.get(b"alpha").unwrap().value(), b"one");
}
#[cfg(not(feature = "no-ttl"))]
#[test]
fn default_ttl_applies_to_plain_cache_writes_and_can_be_overridden() {
let cache = SharedCache::<4>::with_options(CacheOptions {
default_ttl_ms: Some(20),
..CacheOptions::default()
});
assert_eq!(cache.default_ttl_ms(), Some(20));
cache.insert_slice(b"default", b"expires");
assert_eq!(cache.get(b"default").unwrap().value(), b"expires");
std::thread::sleep(std::time::Duration::from_millis(30));
assert!(cache.get(b"default").is_none());
cache.insert_slice_with_ttl(b"durable", b"stays", None);
std::thread::sleep(std::time::Duration::from_millis(30));
assert_eq!(cache.get(b"durable").unwrap().value(), b"stays");
cache.insert_slice_with_ttl(b"explicit", b"short", Some(10));
std::thread::sleep(std::time::Duration::from_millis(20));
assert!(cache.get(b"explicit").is_none());
}
#[test]
fn semantic_cache_returns_best_live_match() {
let cache = SharedCache::<4>::new();
cache
.insert_semantic_slice(b"cat", b"meow", &[1.0, 0.0])
.unwrap();
cache
.insert_semantic_slice(b"dog", b"woof", &[0.0, 1.0])
.unwrap();
let matched = cache.semantic_search(&[0.9, 0.1], 0.75).unwrap().unwrap();
assert_eq!(matched.key.as_slice(), b"cat");
assert_eq!(matched.value.as_ref(), b"meow");
assert!(matched.governance.is_none());
assert!(matched.score > 0.99);
}
#[test]
fn semantic_cache_returns_governance_metadata() {
let cache = SharedCache::<4>::new();
cache
.insert_semantic_slice_with_governance(
b"cat",
b"meow",
&[1.0, 0.0],
b"tenant=acme;doc=cat-faq;policy=v1",
)
.unwrap();
let matched = cache.semantic_search(&[1.0, 0.0], 0.75).unwrap().unwrap();
assert_eq!(matched.key.as_slice(), b"cat");
assert_eq!(matched.value.as_ref(), b"meow");
assert_eq!(
matched.governance.as_deref(),
Some(b"tenant=acme;doc=cat-faq;policy=v1".as_slice())
);
}
#[test]
fn semantic_governance_defaults_to_none() {
let cache = SharedCache::<4>::new();
cache
.insert_semantic_slice(b"default", b"value", &[1.0, 0.0])
.unwrap();
let matched = cache
.semantic_search_with_governance_filter(&[1.0, 0.0], 0.75, |metadata| {
metadata.is_none()
})
.unwrap()
.unwrap();
assert_eq!(matched.key.as_slice(), b"default");
assert!(matched.governance.is_none());
}
#[test]
fn semantic_governance_filter_runs_before_value_release() {
let cache = SharedCache::<4>::new();
cache
.insert_semantic_slice_with_governance(
b"restricted",
b"secret",
&[1.0, 0.0],
b"tenant=internal",
)
.unwrap();
cache
.insert_semantic_slice_with_governance(
b"allowed",
b"public",
&[0.8, 0.2],
b"tenant=acme",
)
.unwrap();
let unfiltered = cache.semantic_search(&[1.0, 0.0], 0.0).unwrap().unwrap();
assert_eq!(unfiltered.key.as_slice(), b"restricted");
let filtered = cache
.semantic_search_with_governance_filter(&[1.0, 0.0], 0.0, |metadata| {
metadata == Some(b"tenant=acme".as_slice())
})
.unwrap()
.unwrap();
assert_eq!(filtered.key.as_slice(), b"allowed");
assert_eq!(filtered.value.as_ref(), b"public");
assert_eq!(
filtered.governance.as_deref(),
Some(b"tenant=acme".as_slice())
);
assert!(
cache
.semantic_search_with_governance_filter(&[1.0, 0.0], 0.0, |metadata| {
metadata == Some(b"tenant=missing".as_slice())
})
.unwrap()
.is_none()
);
}
#[test]
fn semantic_insert_preserves_routed_point_lookup() {
let cache = SharedCache::<4>::new();
let key = key_routed_away_from_semantic_shard(&cache);
cache
.insert_semantic_slice(&key, b"value", &[1.0, 0.0])
.unwrap();
assert_eq!(cache.get(&key).unwrap().value(), b"value");
assert_eq!(cache.get_owned(&key).unwrap().as_ref(), b"value");
let matched = cache.semantic_search(&[1.0, 0.0], 0.5).unwrap().unwrap();
assert_eq!(matched.key.as_slice(), key.as_slice());
assert_eq!(matched.value.as_ref(), b"value");
}
#[test]
fn ordinary_write_invalidates_semantic_embedding() {
let cache = SharedCache::<4>::new();
let key = key_routed_away_from_semantic_shard(&cache);
cache
.insert_semantic_slice(&key, b"meow", &[1.0, 0.0])
.unwrap();
cache.insert_slice(&key, b"purr");
assert!(cache.semantic_search(&[1.0, 0.0], 0.5).unwrap().is_none());
assert_eq!(cache.get(&key).unwrap().value(), b"purr");
}
#[test]
fn mutable_guard_invalidates_semantic_embedding() {
let cache = SharedCache::<4>::new();
let key = key_routed_away_from_semantic_shard(&cache);
cache
.insert_semantic_slice(&key, b"meow", &[1.0, 0.0])
.unwrap();
cache.get_mut(&key).unwrap().set_slice(b"purr");
assert!(cache.semantic_search(&[1.0, 0.0], 0.5).unwrap().is_none());
assert_eq!(cache.get(&key).unwrap().value(), b"purr");
}
#[test]
fn ordinary_remove_invalidates_semantic_embedding() {
let cache = SharedCache::<4>::new();
let key = key_routed_away_from_semantic_shard(&cache);
cache
.insert_semantic_slice(&key, b"meow", &[1.0, 0.0])
.unwrap();
assert_eq!(cache.remove(&key).unwrap().as_ref(), b"meow");
assert!(cache.semantic_search(&[1.0, 0.0], 0.5).unwrap().is_none());
assert!(cache.get(&key).is_none());
}
#[cfg(not(feature = "no-ttl"))]
#[test]
fn semantic_cache_does_not_replay_expired_ttl_hit() {
let cache = SharedCache::<4>::new();
cache
.insert_semantic_slice_with_ttl(b"cat", b"meow", &[1.0, 0.0], Some(50))
.unwrap();
assert!(cache.semantic_search(&[1.0, 0.0], 0.5).unwrap().is_some());
std::thread::sleep(std::time::Duration::from_millis(60));
assert!(cache.semantic_search(&[1.0, 0.0], 0.5).unwrap().is_none());
}
#[test]
fn semantic_cache_rejects_invalid_embeddings() {
let cache = SharedCache::<4>::new();
assert_eq!(
cache.insert_semantic_slice(b"empty", b"value", &[]),
Err(SemanticCacheError::EmptyEmbedding)
);
assert_eq!(
cache.insert_semantic_slice(b"zero", b"value", &[0.0, 0.0]),
Err(SemanticCacheError::ZeroMagnitude)
);
assert_eq!(
cache.semantic_search(&[f32::NAN], 0.5),
Err(SemanticCacheError::NonFinite)
);
}
#[cfg(not(feature = "no-ttl"))]
#[test]
fn fast_map_token_locks_are_compare_and_delete() {
let cache = SharedCache::<4>::new();
assert!(
cache
.try_acquire_lock(b"lock:alpha", b"token-1", 60_000)
.unwrap()
);
assert!(
!cache
.try_acquire_lock(b"lock:alpha", b"token-2", 60_000)
.unwrap()
);
assert!(!cache.release_lock(b"lock:alpha", b"token-2"));
assert!(cache.renew_lock(b"lock:alpha", b"token-1", 60_000).unwrap());
assert!(cache.release_lock(b"lock:alpha", b"token-1"));
assert!(
cache
.try_acquire_lock(b"lock:alpha", b"token-2", 60_000)
.unwrap()
);
}
}