1use anyhow::Result;
10use rustc_hash::FxHashMap;
11use serde::{Deserialize, Serialize};
12use std::sync::{Arc, RwLock};
13use std::time::{Duration, SystemTime};
14
15pub const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(120);
17
18pub const DEFAULT_MAX_CACHE_CAPACITY: usize = 1_000;
20
21pub const MAX_CONTEXT_ITEMS: usize = 5;
23
24pub trait CacheKey: Send + Sync + std::hash::Hash + Eq + Clone + 'static {
26 fn to_cache_key(&self) -> String;
27}
28
29pub trait CacheValue: Send + Sync + Clone + 'static {}
31
32impl<T> CacheValue for T where T: Send + Sync + Clone + 'static {}
33
34#[derive(Debug, Clone, Serialize, Deserialize, Default)]
36pub struct CacheStats {
37 pub hits: u64,
38 pub misses: u64,
39 pub evictions: u64,
40 pub current_size: usize,
41 pub max_size: usize,
42 pub total_memory_bytes: u64,
43}
44
45#[derive(Debug, Clone)]
47pub struct CacheEntry<V> {
48 pub value: Arc<V>,
49 pub created_at: SystemTime,
50 pub last_accessed: SystemTime,
51 pub access_count: u64,
52 pub size_bytes: u64,
53}
54
55impl<V> CacheEntry<V> {
56 pub fn new(value: V, size_bytes: u64) -> Self {
57 let now = SystemTime::now();
58 Self {
59 value: Arc::new(value),
60 created_at: now,
61 last_accessed: now,
62 access_count: 1,
63 size_bytes,
64 }
65 }
66
67 pub fn mark_accessed(&mut self) {
68 self.last_accessed = SystemTime::now();
69 self.access_count += 1;
70 }
71
72 pub fn is_expired(&self, ttl: Duration) -> bool {
73 SystemTime::now()
74 .duration_since(self.created_at)
75 .map(|age| age > ttl)
76 .unwrap_or(true)
77 }
78}
79
80pub struct UnifiedCache<K, V> {
85 inner: RwLock<UnifiedCacheInner<K, V>>,
86}
87
88struct UnifiedCacheInner<K, V> {
90 entries: FxHashMap<K, CacheEntry<V>>,
91 max_size: usize,
92 ttl: Duration,
93 stats: CacheStats,
94 eviction_policy: EvictionPolicy,
95}
96
97#[derive(Debug, Clone, Copy)]
98pub enum EvictionPolicy {
99 Lru,
101 Lfu,
103 Fifo,
105 TtlOnly,
107}
108
109impl<K, V> UnifiedCache<K, V>
110where
111 K: CacheKey,
112 V: CacheValue,
113{
114 pub fn new(max_size: usize, ttl: Duration, eviction_policy: EvictionPolicy) -> Self {
115 Self {
116 inner: RwLock::new(UnifiedCacheInner {
117 entries: FxHashMap::with_capacity_and_hasher(max_size, Default::default()),
118 max_size,
119 ttl,
120 stats: CacheStats {
121 max_size,
122 ..Default::default()
123 },
124 eviction_policy,
125 }),
126 }
127 }
128
129 pub fn get(&self, key: &K) -> Option<Arc<V>> {
134 enum LookupState<T> {
135 Hit(Arc<T>),
136 Expired,
137 Miss,
138 }
139
140 let state = {
141 let inner = self.inner.read().ok()?;
142 let ttl = inner.ttl;
143
144 match inner.entries.get(key) {
145 Some(entry) if !entry.is_expired(ttl) => LookupState::Hit(Arc::clone(&entry.value)),
146 Some(_) => LookupState::Expired,
147 None => LookupState::Miss,
148 }
149 };
150
151 match state {
152 LookupState::Hit(value) => {
153 if let Ok(mut inner) = self.inner.try_write() {
154 if let Some(entry) = inner.entries.get_mut(key) {
155 entry.mark_accessed();
156 }
157 inner.stats.hits += 1;
158 }
159 Some(value)
160 }
161 LookupState::Expired => {
162 if let Ok(mut inner) = self.inner.try_write() {
163 let ttl = inner.ttl;
164 let should_remove = inner
165 .entries
166 .get(key)
167 .map(|entry| entry.is_expired(ttl))
168 .unwrap_or(false);
169 if should_remove {
170 Self::remove_inner(&mut inner, key);
171 }
172 inner.stats.misses += 1;
173 }
174 None
175 }
176 LookupState::Miss => {
177 if let Ok(mut inner) = self.inner.try_write() {
178 inner.stats.misses += 1;
179 }
180 None
181 }
182 }
183 }
184
185 pub fn get_owned(&self, key: &K) -> Option<V> {
187 self.get(key).map(|arc| (*arc).clone())
188 }
189
190 pub fn insert(&self, key: K, value: V, size_bytes: u64) {
192 let Ok(mut inner) = self.inner.write() else {
193 return;
194 };
195
196 Self::remove_expired_entries_inner(&mut inner);
198
199 if inner.entries.len() >= inner.max_size {
202 let to_remove = (inner.max_size / 10).max(1);
203 Self::evict_batch_inner(&mut inner, to_remove);
204 }
205
206 let entry = CacheEntry::new(value, size_bytes);
207 inner.entries.insert(key, entry);
208 inner.stats.current_size = inner.entries.len();
209 inner.stats.total_memory_bytes += size_bytes;
210 }
211
212 fn remove_expired_entries_inner(inner: &mut UnifiedCacheInner<K, V>) {
214 let expired_keys: Vec<K> = inner
215 .entries
216 .iter()
217 .filter_map(|(k, v)| {
218 if v.is_expired(inner.ttl) {
219 Some(k.clone())
220 } else {
221 None
222 }
223 })
224 .collect();
225
226 for key in expired_keys {
227 Self::remove_inner(inner, &key);
228 }
229 }
230
231 fn evict_one_inner(inner: &mut UnifiedCacheInner<K, V>) {
233 if inner.entries.is_empty() {
234 return;
235 }
236
237 let key_to_remove = match inner.eviction_policy {
238 EvictionPolicy::Lru => Self::find_lru_entry_inner(inner),
239 EvictionPolicy::Lfu => Self::find_lfu_entry_inner(inner),
240 EvictionPolicy::Fifo => Self::find_fifo_entry_inner(inner),
241 EvictionPolicy::TtlOnly => Self::find_oldest_entry_inner(inner),
242 };
243
244 if let Some(key) = key_to_remove {
245 Self::remove_inner(inner, &key);
246 inner.stats.evictions += 1;
247 }
248 }
249
250 fn evict_batch_inner(inner: &mut UnifiedCacheInner<K, V>, count: usize) {
253 if inner.entries.is_empty() {
254 return;
255 }
256
257 let keys_to_remove: Vec<K> = match inner.eviction_policy {
258 EvictionPolicy::Lru => {
259 let mut entries: Vec<_> = inner
260 .entries
261 .iter()
262 .map(|(k, e)| (k.clone(), e.last_accessed))
263 .collect();
264 entries.sort_by_key(|(_, ts)| *ts);
265 entries.into_iter().take(count).map(|(k, _)| k).collect()
266 }
267 EvictionPolicy::Lfu => {
268 let mut entries: Vec<_> = inner
269 .entries
270 .iter()
271 .map(|(k, e)| (k.clone(), e.access_count))
272 .collect();
273 entries.sort_by_key(|(_, c)| *c);
274 entries.into_iter().take(count).map(|(k, _)| k).collect()
275 }
276 EvictionPolicy::Fifo | EvictionPolicy::TtlOnly => {
277 let mut entries: Vec<_> = inner
278 .entries
279 .iter()
280 .map(|(k, e)| (k.clone(), e.created_at))
281 .collect();
282 entries.sort_by_key(|(_, ts)| *ts);
283 entries.into_iter().take(count).map(|(k, _)| k).collect()
284 }
285 };
286
287 for key in &keys_to_remove {
288 Self::remove_inner(inner, key);
289 }
290 inner.stats.evictions += keys_to_remove.len() as u64;
291 }
292
293 fn find_lru_entry_inner(inner: &UnifiedCacheInner<K, V>) -> Option<K> {
294 inner
295 .entries
296 .iter()
297 .min_by_key(|(_, entry)| entry.last_accessed)
298 .map(|(k, _)| k.clone())
299 }
300
301 fn find_lfu_entry_inner(inner: &UnifiedCacheInner<K, V>) -> Option<K> {
302 inner
303 .entries
304 .iter()
305 .min_by_key(|(_, entry)| entry.access_count)
306 .map(|(k, _)| k.clone())
307 }
308
309 fn find_fifo_entry_inner(inner: &UnifiedCacheInner<K, V>) -> Option<K> {
310 inner
311 .entries
312 .iter()
313 .min_by_key(|(_, entry)| entry.created_at)
314 .map(|(k, _)| k.clone())
315 }
316
317 fn find_oldest_entry_inner(inner: &UnifiedCacheInner<K, V>) -> Option<K> {
318 Self::find_fifo_entry_inner(inner)
319 }
320
321 fn remove_inner(inner: &mut UnifiedCacheInner<K, V>, key: &K) {
322 if let Some(entry) = inner.entries.remove(key) {
323 inner.stats.total_memory_bytes -= entry.size_bytes;
324 inner.stats.current_size = inner.entries.len();
325 }
326 }
327
328 pub fn stats(&self) -> CacheStats {
330 self.inner
331 .read()
332 .map(|inner| inner.stats.clone())
333 .unwrap_or_default()
334 }
335
336 pub fn clear(&self) {
338 if let Ok(mut inner) = self.inner.write() {
339 inner.entries.clear();
340 inner.stats.current_size = 0;
341 inner.stats.total_memory_bytes = 0;
342 }
343 }
344
345 pub fn remove(&self, key: &K) {
347 let Ok(mut inner) = self.inner.write() else {
348 return;
349 };
350 Self::remove_inner(&mut inner, key);
351 }
352
353 pub fn len(&self) -> usize {
355 self.inner
356 .read()
357 .map(|inner| inner.entries.len())
358 .unwrap_or(0)
359 }
360
361 pub fn is_empty(&self) -> bool {
362 self.len() == 0
363 }
364
365 pub fn invalidate_prefix(&self, prefix: &str) {
374 self.remove_where(|key| key.to_cache_key().starts_with(prefix));
375 }
376
377 pub fn invalidate_path(&self, path: &str) {
386 self.invalidate_prefix(&format!("{}:", path));
387 }
388
389 pub fn invalidate_suffix(&self, suffix: &str) {
397 self.remove_where(|key| key.to_cache_key().ends_with(suffix));
398 }
399
400 pub fn invalidate_containing(&self, substring: &str) {
408 self.remove_where(|key| key.to_cache_key().contains(substring));
409 }
410
411 pub fn remove_where<F>(&self, mut predicate: F) -> usize
415 where
416 F: FnMut(&K) -> bool,
417 {
418 let Ok(mut inner) = self.inner.write() else {
419 return 0;
420 };
421
422 let keys_to_remove: Vec<K> = inner
423 .entries
424 .keys()
425 .filter(|key| predicate(key))
426 .cloned()
427 .collect();
428
429 let removed_count = keys_to_remove.len();
430 for key in keys_to_remove {
431 Self::remove_inner(&mut inner, &key);
432 }
433 removed_count
434 }
435
436 pub fn total_memory_bytes(&self) -> u64 {
438 self.inner
439 .read()
440 .map(|inner| inner.stats.total_memory_bytes)
441 .unwrap_or(0)
442 }
443
444 pub fn estimate_entry_cost(entry: &CacheEntry<V>) -> u64 {
447 let base_overhead: u64 = 100; let value_size = entry.size_bytes;
450 base_overhead + value_size
451 }
452
453 pub fn reduce_ttl(&self, factor: f64) -> Duration {
456 let Ok(mut inner) = self.inner.write() else {
457 return Duration::ZERO;
458 };
459 let new_ttl = Duration::from_secs_f64(inner.ttl.as_secs_f64() * factor);
460 inner.ttl = new_ttl;
461 new_ttl
462 }
463
464 pub fn evict_under_pressure(&self, target_reduction_percent: u32) -> usize {
471 let Ok(mut inner) = self.inner.write() else {
472 return 0;
473 };
474
475 let target_percent = std::cmp::min(100, target_reduction_percent);
477
478 let expired_before = inner.entries.len();
480 Self::remove_expired_entries_inner(&mut inner);
481 let expired_removed = expired_before - inner.entries.len();
482
483 let current_size = inner.entries.len();
485 let target_size = (current_size * (100 - target_percent) as usize) / 100;
486
487 let mut evicted_count = expired_removed;
489 while inner.entries.len() > target_size && !inner.entries.is_empty() {
490 Self::evict_one_inner(&mut inner);
491 evicted_count += 1;
492 }
493
494 evicted_count
495 }
496
497 pub fn clear_least_used(&self, percent_to_clear: u32) -> usize {
500 let Ok(mut inner) = self.inner.write() else {
501 return 0;
502 };
503
504 let percent = std::cmp::min(100, percent_to_clear);
505 let entries_to_remove = (inner.entries.len() * percent as usize) / 100;
506
507 let mut removed = 0usize;
508 for _ in 0..entries_to_remove {
509 if inner.entries.is_empty() {
510 break;
511 }
512 Self::evict_one_inner(&mut inner);
513 removed += 1;
514 }
515
516 removed
517 }
518
519 pub fn entries_by_usefulness(&self) -> Vec<(K, CacheEntry<V>)> {
522 let Ok(inner) = self.inner.read() else {
523 return Vec::new();
524 };
525
526 let now = SystemTime::now();
527 let mut entries: Vec<(K, CacheEntry<V>, u64)> = inner
528 .entries
529 .iter()
530 .map(|(k, entry)| {
531 let age_secs = now
533 .duration_since(entry.last_accessed)
534 .unwrap_or_default()
535 .as_secs();
536
537 let recency_factor = std::cmp::max(1_u64, 3600 / (age_secs + 1));
539 let usefulness_score = entry.access_count * recency_factor;
540
541 (k.clone(), entry.clone(), usefulness_score)
542 })
543 .collect();
544
545 entries.sort_by_key(|(_, _, score)| std::cmp::Reverse(*score));
547 entries.into_iter().map(|(k, e, _)| (k, e)).collect()
548 }
549}
550
551pub fn estimate_json_size(value: &serde_json::Value) -> u64 {
553 match value {
554 serde_json::Value::Null => 4,
555 serde_json::Value::Bool(_) => 5,
556 serde_json::Value::Number(n) => n.to_string().len() as u64,
557 serde_json::Value::String(s) => s.len() as u64,
558 serde_json::Value::Array(arr) => arr.iter().map(estimate_json_size).sum(),
559 serde_json::Value::Object(obj) => obj
560 .iter()
561 .map(|(k, v)| k.len() as u64 + estimate_json_size(v) + 3) .sum(),
563 }
564}
565
566pub fn create_cache_key<T: Serialize>(data: &T) -> Result<String> {
568 let json_bytes = serde_json::to_vec(data)?;
569
570 let mut hash = 0u64;
572 for (i, byte) in json_bytes.iter().enumerate() {
573 hash = hash.wrapping_mul(31).wrapping_add(*byte as u64);
574 hash = hash.rotate_left((i % 64) as u32);
575 }
576
577 Ok(format!("{:016x}", hash))
578}
579
580pub struct ContextAwareCache<K, V> {
582 inner: UnifiedCache<K, V>,
583}
584
585impl<K, V> ContextAwareCache<K, V>
586where
587 K: CacheKey,
588 V: CacheValue,
589{
590 pub fn new(max_size: usize, ttl: Duration, eviction_policy: EvictionPolicy) -> Self {
591 Self {
592 inner: UnifiedCache::new(max_size, ttl, eviction_policy),
593 }
594 }
595
596 pub fn get_context_limited<F>(&self, keys: &[K], mut process_fn: F) -> Vec<V>
598 where
599 F: FnMut(&K) -> Option<V>,
600 {
601 let mut results = Vec::with_capacity(MAX_CONTEXT_ITEMS.min(keys.len()));
602 let mut overflow_count = 0;
603
604 for key in keys {
605 if results.len() >= MAX_CONTEXT_ITEMS {
606 overflow_count += 1;
607 continue;
608 }
609
610 if let Some(value) = self.inner.get(key) {
611 results.push((*value).clone());
612 } else if let Some(value) = process_fn(key) {
613 let size = size_of_val(&value) as u64;
615 self.inner.insert(key.clone(), value.clone(), size);
616 results.push(value);
617 }
618 }
619
620 if overflow_count > 0 {
622 }
625
626 results
627 }
628
629 pub fn stats(&self) -> CacheStats {
630 self.inner.stats()
631 }
632}
633
634#[cfg(test)]
635mod tests {
636 use super::*;
637
638 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
639 struct TestKey(String);
640
641 impl CacheKey for TestKey {
642 fn to_cache_key(&self) -> String {
643 self.0.clone()
644 }
645 }
646
647 #[test]
648 fn test_cache_basic_operations() {
649 let cache = UnifiedCache::new(10, DEFAULT_CACHE_TTL, EvictionPolicy::Lru);
650 let key = TestKey("test".into());
651 let value: String = "test_value".into();
652
653 cache.insert(key.clone(), value.clone(), 100);
655 assert_eq!(*cache.get(&key).unwrap(), value);
656
657 let stats = cache.stats();
659 assert_eq!(stats.hits, 1);
660 assert_eq!(stats.misses, 0);
661 assert_eq!(stats.current_size, 1);
662 }
663
664 #[test]
665 fn test_cache_expiration() {
666 let cache = UnifiedCache::new(10, Duration::from_millis(100), EvictionPolicy::Lru);
667 let key = TestKey("test".into());
668 let value: String = "test_value".into();
669
670 cache.insert(key.clone(), value, 100);
671 assert!(cache.get(&key).is_some());
672
673 std::thread::sleep(Duration::from_millis(150));
675 assert!(cache.get(&key).is_none());
676 }
677
678 #[test]
679 fn test_context_limited_cache() {
680 let cache = ContextAwareCache::new(100, DEFAULT_CACHE_TTL, EvictionPolicy::Lru);
681 let keys: Vec<TestKey> = (0..10).map(|i| TestKey(i.to_string())).collect();
682
683 let results = cache.get_context_limited(&keys, |key| Some(format!("value_{}", key.0)));
684
685 assert_eq!(results.len(), MAX_CONTEXT_ITEMS);
687 assert_eq!(results[0], "value_0");
688 assert_eq!(results[4], "value_4");
689 }
690
691 #[test]
692 fn test_pressure_aware_total_memory() {
693 let cache = UnifiedCache::new(10, DEFAULT_CACHE_TTL, EvictionPolicy::Lru);
694
695 cache.insert(TestKey("k1".into()), "v1".to_string(), 100);
697 cache.insert(TestKey("k2".into()), "v2".to_string(), 200);
698 cache.insert(TestKey("k3".into()), "v3".to_string(), 300);
699
700 assert_eq!(cache.total_memory_bytes(), 600);
702 }
703
704 #[test]
705 fn test_pressure_aware_reduce_ttl() {
706 let cache: UnifiedCache<TestKey, String> =
707 UnifiedCache::new(10, Duration::from_secs(300), EvictionPolicy::Lru);
708
709 let new_ttl = cache.reduce_ttl(0.4);
711 assert_eq!(new_ttl.as_secs(), 120); let new_ttl = cache.reduce_ttl(0.1);
715 assert_eq!(new_ttl.as_secs(), 12); }
717
718 #[test]
719 fn test_pressure_aware_evict_under_pressure() {
720 let cache: UnifiedCache<TestKey, String> =
721 UnifiedCache::new(20, Duration::from_secs(3600), EvictionPolicy::Lru);
722
723 for i in 0..10 {
725 cache.insert(TestKey(format!("key_{}", i)), format!("value_{}", i), 100);
726 }
727
728 assert_eq!(cache.len(), 10);
729
730 let removed = cache.evict_under_pressure(50);
732 assert_eq!(removed, 5);
733 assert_eq!(cache.len(), 5);
734 }
735
736 #[test]
737 fn test_pressure_aware_clear_least_used() {
738 let cache: UnifiedCache<TestKey, String> =
739 UnifiedCache::new(20, Duration::from_secs(3600), EvictionPolicy::Lru);
740
741 for i in 0..10 {
743 cache.insert(TestKey(format!("key_{}", i)), format!("value_{}", i), 100);
744 }
745
746 let _ = cache.get(&TestKey("key_0".into()));
748 let _ = cache.get(&TestKey("key_1".into()));
749
750 assert_eq!(cache.len(), 10);
751
752 let removed = cache.clear_least_used(30);
754 assert!(removed <= 4, "Should remove ~3 entries (30% of 10)");
755 assert!(cache.len() >= 6, "Should have ~7 entries left");
756 }
757
758 #[test]
759 fn test_pressure_aware_entries_by_usefulness() {
760 let cache: UnifiedCache<TestKey, String> =
761 UnifiedCache::new(20, Duration::from_secs(3600), EvictionPolicy::Lru);
762
763 cache.insert(TestKey("hot".into()), "value".to_string(), 100);
765 cache.insert(TestKey("cold".into()), "value".to_string(), 100);
766 cache.insert(TestKey("warm".into()), "value".to_string(), 100);
767
768 for _ in 0..5 {
770 let _ = cache.get(&TestKey("hot".into()));
771 }
772
773 let _ = cache.get(&TestKey("warm".into()));
775
776 let usefulness = cache.entries_by_usefulness();
779 assert_eq!(usefulness.len(), 3);
780
781 assert_eq!(usefulness[0].0.0, "hot");
783 }
784
785 #[test]
786 fn test_pressure_aware_estimate_entry_cost() {
787 let entry = CacheEntry::new("test_value".to_string(), 1000);
788 let cost = UnifiedCache::<TestKey, String>::estimate_entry_cost(&entry);
789
790 assert!(cost >= 1000);
792 assert!(cost <= 1200); }
794}