use crate::error::CacheResult;
use crate::traits::CacheStore;
use async_trait::async_trait;
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap, VecDeque};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use tokio::sync::RwLock;
pub struct TieredCache<L1, L2>
where
L1: CacheStore,
L2: CacheStore,
{
l1: Arc<L1>,
l2: Arc<L2>,
config: TieredCacheConfig,
metrics: Arc<TieredMetrics>,
}
#[derive(Debug, Default)]
struct TieredMetrics {
l1_hits: AtomicU64,
l2_hits: AtomicU64,
misses: AtomicU64,
promotions: AtomicU64,
}
#[derive(Debug, Clone)]
pub struct TieredCacheConfig {
pub enable_l1: bool,
pub enable_l2: bool,
pub write_through: bool,
pub promote_to_l1: bool,
pub l1_ttl_fraction: f64,
pub l1_promote_ttl: Option<Duration>,
}
impl Default for TieredCacheConfig {
fn default() -> Self {
Self {
enable_l1: true,
enable_l2: true,
write_through: true,
promote_to_l1: true,
l1_ttl_fraction: 0.25, l1_promote_ttl: Some(Duration::from_secs(60)),
}
}
}
impl<L1, L2> TieredCache<L1, L2>
where
L1: CacheStore,
L2: CacheStore,
{
pub fn new(l1: Arc<L1>, l2: Arc<L2>) -> Self {
Self::with_config(l1, l2, TieredCacheConfig::default())
}
pub fn with_config(l1: Arc<L1>, l2: Arc<L2>, config: TieredCacheConfig) -> Self {
Self {
l1,
l2,
config,
metrics: Arc::new(TieredMetrics::default()),
}
}
pub async fn get(&self, key: &str) -> CacheResult<Option<String>> {
if self.config.enable_l1
&& let Some(value) = self.l1.get_json(key).await?
{
self.metrics.l1_hits.fetch_add(1, Ordering::Relaxed);
return Ok(Some(value));
}
if self.config.enable_l2
&& let Some(value) = self.l2.get_json(key).await?
{
self.metrics.l2_hits.fetch_add(1, Ordering::Relaxed);
if self.config.enable_l1 && self.config.promote_to_l1 {
let l1_ttl = self.config.l1_promote_ttl;
if self.l1.set_json(key, value.clone(), l1_ttl).await.is_ok() {
self.metrics.promotions.fetch_add(1, Ordering::Relaxed);
}
}
return Ok(Some(value));
}
self.metrics.misses.fetch_add(1, Ordering::Relaxed);
Ok(None)
}
pub async fn set(&self, key: &str, value: String, ttl: Option<Duration>) -> CacheResult<()> {
if self.config.enable_l2 {
self.l2.set_json(key, value.clone(), ttl).await?;
}
if self.config.enable_l1 && (self.config.write_through || !self.config.enable_l2) {
let l1_ttl = ttl.map(|ttl| {
Duration::from_secs_f64(ttl.as_secs_f64() * self.config.l1_ttl_fraction)
});
self.l1.set_json(key, value, l1_ttl).await?;
}
Ok(())
}
pub async fn delete(&self, key: &str) -> CacheResult<()> {
if self.config.enable_l1 {
self.l1.delete(key).await?;
}
if self.config.enable_l2 {
self.l2.delete(key).await?;
}
Ok(())
}
pub async fn exists(&self, key: &str) -> CacheResult<bool> {
if self.config.enable_l1 && self.l1.exists(key).await? {
return Ok(true);
}
if self.config.enable_l2 {
return self.l2.exists(key).await;
}
Ok(false)
}
pub async fn clear(&self) -> CacheResult<()> {
if self.config.enable_l1 {
self.l1.clear().await?;
}
if self.config.enable_l2 {
self.l2.clear().await?;
}
Ok(())
}
pub async fn stats(&self) -> CacheStats {
CacheStats {
l1_enabled: self.config.enable_l1,
l2_enabled: self.config.enable_l2,
write_through: self.config.write_through,
promote_to_l1: self.config.promote_to_l1,
l1_hits: self.metrics.l1_hits.load(Ordering::Relaxed),
l2_hits: self.metrics.l2_hits.load(Ordering::Relaxed),
misses: self.metrics.misses.load(Ordering::Relaxed),
promotions: self.metrics.promotions.load(Ordering::Relaxed),
}
}
}
impl<L1, L2> Clone for TieredCache<L1, L2>
where
L1: CacheStore,
L2: CacheStore,
{
fn clone(&self) -> Self {
Self {
l1: self.l1.clone(),
l2: self.l2.clone(),
config: self.config.clone(),
metrics: self.metrics.clone(),
}
}
}
#[derive(Debug, Clone)]
pub struct CacheStats {
pub l1_enabled: bool,
pub l2_enabled: bool,
pub write_through: bool,
pub promote_to_l1: bool,
pub l1_hits: u64,
pub l2_hits: u64,
pub misses: u64,
pub promotions: u64,
}
pub const DEFAULT_MAX_ENTRIES: usize = 10_000;
pub struct InMemoryCache {
data: Arc<RwLock<CacheState>>,
max_entries: usize,
}
#[derive(Default)]
struct CacheState {
entries: HashMap<String, CacheEntry>,
by_expiry: BinaryHeap<Reverse<(tokio::time::Instant, String)>>,
without_expiry: VecDeque<String>,
}
#[derive(Clone)]
struct CacheEntry {
value: String,
expires_at: Option<tokio::time::Instant>,
}
impl CacheState {
fn insert(&mut self, key: String, entry: CacheEntry) {
match entry.expires_at {
Some(expires_at) => self.by_expiry.push(Reverse((expires_at, key.clone()))),
None => self.without_expiry.push_back(key.clone()),
}
self.entries.insert(key, entry);
self.compact_if_slack();
}
fn is_current(&self, key: &str, expires_at: Option<tokio::time::Instant>) -> bool {
self.entries
.get(key)
.is_some_and(|entry| entry.expires_at == expires_at)
}
fn prune_expired(&mut self, now: tokio::time::Instant) {
while matches!(self.by_expiry.peek(), Some(Reverse((exp, _))) if *exp <= now) {
let Some(Reverse((expires_at, key))) = self.by_expiry.pop() else {
break;
};
if self.is_current(&key, Some(expires_at)) {
self.entries.remove(&key);
}
}
}
fn evict_one(&mut self) {
while let Some(Reverse((expires_at, key))) = self.by_expiry.pop() {
if self.is_current(&key, Some(expires_at)) {
self.entries.remove(&key);
return;
}
}
while let Some(key) = self.without_expiry.pop_front() {
if self.is_current(&key, None) {
self.entries.remove(&key);
return;
}
}
}
fn compact_if_slack(&mut self) {
let tracked = self.by_expiry.len() + self.without_expiry.len();
if tracked > 2 * self.entries.len().max(16) {
self.compact();
}
}
fn compact(&mut self) {
let mut by_expiry = BinaryHeap::with_capacity(self.entries.len());
let mut without_expiry = VecDeque::with_capacity(self.entries.len());
for (key, entry) in &self.entries {
match entry.expires_at {
Some(expires_at) => by_expiry.push(Reverse((expires_at, key.clone()))),
None => without_expiry.push_back(key.clone()),
}
}
self.by_expiry = by_expiry;
self.without_expiry = without_expiry;
}
fn clear(&mut self) {
self.entries.clear();
self.by_expiry.clear();
self.without_expiry.clear();
}
}
impl InMemoryCache {
pub fn new() -> Self {
Self::with_capacity(DEFAULT_MAX_ENTRIES)
}
pub fn with_capacity(max_entries: usize) -> Self {
Self {
data: Arc::new(RwLock::new(CacheState::default())),
max_entries,
}
}
pub async fn len(&self) -> usize {
self.data.read().await.entries.len()
}
pub async fn is_empty(&self) -> bool {
self.data.read().await.entries.is_empty()
}
pub async fn cleanup_expired(&self) {
let mut data = self.data.write().await;
let now = tokio::time::Instant::now();
data.prune_expired(now);
}
}
impl Default for InMemoryCache {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl CacheStore for InMemoryCache {
async fn get_json(&self, key: &str) -> CacheResult<Option<String>> {
{
let data = self.data.read().await;
match data.entries.get(key) {
None => return Ok(None),
Some(entry) => match entry.expires_at {
Some(expires_at) if tokio::time::Instant::now() > expires_at => {
}
_ => return Ok(Some(entry.value.clone())),
},
}
}
let mut data = self.data.write().await;
if let Some(entry) = data.entries.get(key)
&& entry
.expires_at
.is_some_and(|exp| tokio::time::Instant::now() > exp)
{
data.entries.remove(key);
}
Ok(None)
}
async fn set_json(&self, key: &str, value: String, ttl: Option<Duration>) -> CacheResult<()> {
let now = tokio::time::Instant::now();
let expires_at = ttl.map(|d| now + d);
let entry = CacheEntry { value, expires_at };
let mut data = self.data.write().await;
if self.max_entries != 0
&& data.entries.len() >= self.max_entries
&& !data.entries.contains_key(key)
{
data.prune_expired(now);
if data.entries.len() >= self.max_entries {
data.evict_one();
}
}
data.insert(key.to_string(), entry);
Ok(())
}
async fn delete(&self, key: &str) -> CacheResult<()> {
self.data.write().await.entries.remove(key);
Ok(())
}
async fn exists(&self, key: &str) -> CacheResult<bool> {
self.get_json(key).await.map(|v| v.is_some())
}
async fn clear(&self) -> CacheResult<()> {
self.data.write().await.clear();
Ok(())
}
async fn ttl(&self, key: &str) -> CacheResult<Option<Duration>> {
let data = self.data.read().await;
let now = tokio::time::Instant::now();
Ok(data
.entries
.get(key)
.and_then(|e| e.expires_at)
.filter(|&x| x > now)
.map(|x| x - now))
}
async fn expire(&self, key: &str, ttl: Duration) -> CacheResult<()> {
let mut data = self.data.write().await;
let expires_at = tokio::time::Instant::now() + ttl;
let updated = match data.entries.get_mut(key) {
Some(entry) => {
entry.expires_at = Some(expires_at);
true
}
None => false,
};
if updated {
data.by_expiry.push(Reverse((expires_at, key.to_string())));
data.compact_if_slack();
}
Ok(())
}
async fn increment(&self, key: &str, delta: i64) -> CacheResult<i64> {
let mut data = self.data.write().await;
let new_value = match data.entries.get_mut(key) {
Some(entry) => {
let current: i64 = entry.value.parse().unwrap_or(0);
let new_value = current + delta;
entry.value = new_value.to_string();
new_value
}
None => {
data.insert(
key.to_string(),
CacheEntry {
value: delta.to_string(),
expires_at: None,
},
);
delta
}
};
Ok(new_value)
}
async fn decrement(&self, key: &str, delta: i64) -> CacheResult<i64> {
self.increment(key, -delta).await
}
}
#[cfg(test)]
mod tests_tiered {
use super::*;
#[tokio::test]
async fn test_tiered_cache() {
let l1 = Arc::new(InMemoryCache::new());
let l2 = Arc::new(InMemoryCache::new());
let cache = TieredCache::new(l1.clone(), l2.clone());
cache.set("test", "value".to_string(), None).await.unwrap();
let value = l1.get_json("test").await.unwrap();
assert!(value.is_some());
let value = cache.get("test").await.unwrap();
assert_eq!(value, Some("value".to_string()));
cache.delete("test").await.unwrap();
let value = cache.get("test").await.unwrap();
assert_eq!(value, None);
}
#[tokio::test]
async fn test_l2_promotion() {
let l1 = Arc::new(InMemoryCache::new());
let l2 = Arc::new(InMemoryCache::new());
let cache = TieredCache::new(l1.clone(), l2.clone());
l2.set_json("key", "value".to_string(), None).await.unwrap();
let value = cache.get("key").await.unwrap();
assert_eq!(value, Some("value".to_string()));
let l1_value = l1.get_json("key").await.unwrap();
assert!(l1_value.is_some());
}
#[tokio::test]
async fn test_promotion_uses_fixed_l1_ttl_no_l2_roundtrip() {
let l1 = Arc::new(InMemoryCache::new());
let l2 = Arc::new(InMemoryCache::new());
let config = TieredCacheConfig {
l1_promote_ttl: Some(Duration::from_secs(30)),
..TieredCacheConfig::default()
};
let cache = TieredCache::with_config(l1.clone(), l2.clone(), config);
l2.set_json("key", "value".to_string(), None).await.unwrap();
let value = cache.get("key").await.unwrap();
assert_eq!(value, Some("value".to_string()));
let l1_ttl = l1.ttl("key").await.unwrap();
let l1_ttl = l1_ttl.expect("promoted L1 entry should have a TTL");
assert!(l1_ttl > Duration::from_secs(0));
assert!(l1_ttl <= Duration::from_secs(30));
}
#[tokio::test]
async fn test_promotion_with_no_l1_ttl_stores_without_expiry() {
let l1 = Arc::new(InMemoryCache::new());
let l2 = Arc::new(InMemoryCache::new());
let config = TieredCacheConfig {
l1_promote_ttl: None,
..TieredCacheConfig::default()
};
let cache = TieredCache::with_config(l1.clone(), l2.clone(), config);
l2.set_json("key", "value".to_string(), None).await.unwrap();
let _ = cache.get("key").await.unwrap();
assert_eq!(l1.ttl("key").await.unwrap(), None);
assert!(l1.get_json("key").await.unwrap().is_some());
}
#[tokio::test]
async fn test_stats_track_hits_misses_promotions() {
let l1 = Arc::new(InMemoryCache::new());
let l2 = Arc::new(InMemoryCache::new());
let cache = TieredCache::new(l1.clone(), l2.clone());
assert_eq!(cache.get("absent").await.unwrap(), None);
cache.set("k", "v".to_string(), None).await.unwrap();
assert_eq!(cache.get("k").await.unwrap(), Some("v".to_string()));
l2.set_json("only2", "v2".to_string(), None).await.unwrap();
assert_eq!(cache.get("only2").await.unwrap(), Some("v2".to_string()));
let stats = cache.stats().await;
assert_eq!(stats.misses, 1, "one miss expected");
assert_eq!(stats.l1_hits, 1, "one L1 hit expected");
assert_eq!(stats.l2_hits, 1, "one L2 hit expected");
assert_eq!(stats.promotions, 1, "one promotion expected");
}
#[tokio::test(start_paused = true)]
async fn test_l1_expired_entries_are_evicted_on_read() {
let cache = InMemoryCache::new();
cache
.set_json("k", "v".to_string(), Some(Duration::from_secs(1)))
.await
.unwrap();
assert_eq!(cache.len().await, 1);
tokio::time::advance(Duration::from_secs(2)).await;
assert_eq!(cache.get_json("k").await.unwrap(), None);
assert_eq!(
cache.len().await,
0,
"expired entry must be evicted from the map, not retained"
);
}
#[tokio::test]
async fn test_l1_capacity_bound_is_enforced() {
let cache = InMemoryCache::with_capacity(2);
cache.set_json("a", "1".to_string(), None).await.unwrap();
cache.set_json("b", "2".to_string(), None).await.unwrap();
cache.set_json("c", "3".to_string(), None).await.unwrap();
assert!(
cache.len().await <= 2,
"cache must not exceed its configured capacity of 2, got {}",
cache.len().await
);
assert_eq!(cache.get_json("c").await.unwrap(), Some("3".to_string()));
}
#[tokio::test(start_paused = true)]
async fn test_capacity_prefers_reclaiming_expired() {
let cache = InMemoryCache::with_capacity(2);
cache
.set_json("short", "x".to_string(), Some(Duration::from_secs(1)))
.await
.unwrap();
cache.set_json("keep", "y".to_string(), None).await.unwrap();
tokio::time::advance(Duration::from_secs(2)).await;
cache.set_json("new", "z".to_string(), None).await.unwrap();
assert!(cache.len().await <= 2);
assert_eq!(cache.get_json("keep").await.unwrap(), Some("y".to_string()));
assert_eq!(cache.get_json("new").await.unwrap(), Some("z".to_string()));
}
#[tokio::test(start_paused = true)]
async fn test_evicts_the_entry_nearest_to_expiry() {
let cache = InMemoryCache::with_capacity(3);
cache
.set_json("far", "1".to_string(), Some(Duration::from_secs(300)))
.await
.unwrap();
cache
.set_json("soon", "2".to_string(), Some(Duration::from_secs(10)))
.await
.unwrap();
cache
.set_json("mid", "3".to_string(), Some(Duration::from_secs(60)))
.await
.unwrap();
cache.set_json("new", "4".to_string(), None).await.unwrap();
assert_eq!(cache.len().await, 3);
assert_eq!(
cache.get_json("soon").await.unwrap(),
None,
"the soonest-to-expire entry must be the victim"
);
assert_eq!(cache.get_json("far").await.unwrap(), Some("1".to_string()));
assert_eq!(cache.get_json("mid").await.unwrap(), Some("3".to_string()));
assert_eq!(cache.get_json("new").await.unwrap(), Some("4".to_string()));
}
#[tokio::test(start_paused = true)]
async fn test_entries_without_ttl_are_evicted_last() {
let cache = InMemoryCache::with_capacity(2);
cache
.set_json("nottl", "1".to_string(), None)
.await
.unwrap();
cache
.set_json("ttl", "2".to_string(), Some(Duration::from_secs(300)))
.await
.unwrap();
cache.set_json("new", "3".to_string(), None).await.unwrap();
assert_eq!(
cache.get_json("ttl").await.unwrap(),
None,
"a TTL-carrying entry must be evicted before an unexpiring one"
);
assert_eq!(
cache.get_json("nottl").await.unwrap(),
Some("1".to_string())
);
cache
.set_json("newest", "4".to_string(), None)
.await
.unwrap();
assert_eq!(cache.get_json("nottl").await.unwrap(), None);
assert_eq!(cache.get_json("new").await.unwrap(), Some("3".to_string()));
assert_eq!(
cache.get_json("newest").await.unwrap(),
Some("4".to_string())
);
}
#[tokio::test(start_paused = true)]
async fn test_expire_updates_eviction_order() {
let cache = InMemoryCache::with_capacity(2);
cache.set_json("a", "1".to_string(), None).await.unwrap();
cache.set_json("b", "2".to_string(), None).await.unwrap();
cache.expire("b", Duration::from_secs(300)).await.unwrap();
cache.set_json("c", "3".to_string(), None).await.unwrap();
assert_eq!(cache.get_json("b").await.unwrap(), None);
assert_eq!(cache.get_json("a").await.unwrap(), Some("1".to_string()));
assert_eq!(cache.get_json("c").await.unwrap(), Some("3".to_string()));
}
#[tokio::test(start_paused = true)]
async fn test_repeated_overwrites_do_not_grow_ordering_structures() {
let cache = InMemoryCache::with_capacity(8);
for round in 0..500 {
for key in ["a", "b", "c", "d"] {
cache
.set_json(key, round.to_string(), Some(Duration::from_secs(300)))
.await
.unwrap();
}
}
let state = cache.data.read().await;
assert_eq!(state.entries.len(), 4);
assert!(
state.by_expiry.len() + state.without_expiry.len() <= 2 * 16,
"stale ordering records must be compacted away, got {} tracked for {} entries",
state.by_expiry.len() + state.without_expiry.len(),
state.entries.len()
);
}
#[tokio::test(start_paused = true)]
async fn test_admission_churn_keeps_cache_bounded() {
let cache = InMemoryCache::with_capacity(64);
for i in 0..2_000 {
cache
.set_json(
&format!("k{i}"),
i.to_string(),
if i % 2 == 0 {
Some(Duration::from_secs(300 + i as u64))
} else {
None
},
)
.await
.unwrap();
}
assert_eq!(cache.len().await, 64);
assert_eq!(
cache.get_json("k1999").await.unwrap(),
Some("1999".to_string()),
"the most recent write must survive"
);
}
}