Skip to main content

lean_ctx/core/providers/
cache.rs

1use std::collections::HashMap;
2use std::sync::Mutex;
3use std::time::{Duration, Instant, SystemTime};
4
5const MAX_CACHE_ENTRIES: usize = 2048;
6const STALE_SERVE_GRACE: Duration = Duration::from_hours(1);
7
8static PROVIDER_CACHE: std::sync::LazyLock<Mutex<ProviderCache>> =
9    std::sync::LazyLock::new(|| Mutex::new(ProviderCache::new()));
10
11struct CacheEntry {
12    data: String,
13    expires_at: Instant,
14    provider_id: String,
15}
16
17struct StaleEntry {
18    data: String,
19    evicted_at: Instant,
20    provider_id: String,
21}
22
23/// Per-provider cache statistics.
24#[derive(Debug, Clone, Default)]
25pub struct ProviderCacheStats {
26    pub provider_id: String,
27    pub hits: u64,
28    pub misses: u64,
29    pub entry_count: usize,
30    pub last_fetch: Option<SystemTime>,
31}
32
33impl ProviderCacheStats {
34    pub fn hit_rate(&self) -> f64 {
35        let total = self.hits + self.misses;
36        if total == 0 {
37            return 0.0;
38        }
39        self.hits as f64 / total as f64
40    }
41}
42
43/// Global cache statistics across all providers.
44#[derive(Debug, Clone, Default)]
45pub struct CacheMetrics {
46    pub total_hits: u64,
47    pub total_misses: u64,
48    pub total_entries: usize,
49    pub stale_entries: usize,
50    pub max_entries: usize,
51    pub provider_stats: Vec<ProviderCacheStats>,
52}
53
54impl CacheMetrics {
55    pub fn total_hit_rate(&self) -> f64 {
56        let total = self.total_hits + self.total_misses;
57        if total == 0 {
58            return 0.0;
59        }
60        self.total_hits as f64 / total as f64
61    }
62}
63
64struct ProviderCache {
65    entries: HashMap<String, CacheEntry>,
66    access_order: Vec<String>,
67    stale_entries: HashMap<String, StaleEntry>,
68    hits: HashMap<String, u64>,
69    misses: HashMap<String, u64>,
70    last_fetch: HashMap<String, SystemTime>,
71}
72
73impl ProviderCache {
74    fn new() -> Self {
75        Self {
76            entries: HashMap::new(),
77            access_order: Vec::new(),
78            stale_entries: HashMap::new(),
79            hits: HashMap::new(),
80            misses: HashMap::new(),
81            last_fetch: HashMap::new(),
82        }
83    }
84
85    fn get(&mut self, key: &str) -> Option<&str> {
86        self.expire_entries();
87        self.purge_stale();
88        if let Some(entry) = self.entries.get(key) {
89            let provider_id = entry.provider_id.clone();
90            self.access_order.retain(|candidate| candidate != key);
91            self.access_order.push(key.to_string());
92            *self.hits.entry(provider_id).or_insert(0) += 1;
93            return self.entries.get(key).map(|entry| entry.data.as_str());
94        }
95        let provider_id = self.stale_entries.get(key).map_or_else(
96            || key.split(':').next().unwrap_or("unknown").to_string(),
97            |entry| entry.provider_id.clone(),
98        );
99        *self.misses.entry(provider_id.clone()).or_insert(0) += 1;
100        let stale = self.stale_entries.get(key)?;
101        tracing::warn!(key, provider_id, "serving stale provider cache entry");
102        Some(stale.data.as_str())
103    }
104
105    fn set(&mut self, key: String, data: String, ttl: Duration, provider_id: &str) {
106        self.last_fetch
107            .insert(provider_id.to_string(), SystemTime::now());
108        self.access_order.retain(|candidate| candidate != &key);
109        self.access_order.push(key.clone());
110        self.stale_entries.remove(&key);
111        self.entries.insert(
112            key,
113            CacheEntry {
114                data,
115                expires_at: Instant::now() + ttl,
116                provider_id: provider_id.to_string(),
117            },
118        );
119        self.enforce_lru_cap();
120    }
121
122    fn expire_entries(&mut self) {
123        let now = Instant::now();
124        let expired: Vec<_> = self
125            .entries
126            .iter()
127            .filter(|(_, entry)| entry.expires_at <= now)
128            .map(|(key, _)| key.clone())
129            .collect();
130        for key in expired {
131            self.move_to_stale(&key, now);
132        }
133    }
134
135    fn enforce_lru_cap(&mut self) {
136        let now = Instant::now();
137        while self.entries.len() > MAX_CACHE_ENTRIES {
138            let key = self.access_order.remove(0);
139            self.move_to_stale(&key, now);
140        }
141    }
142
143    fn move_to_stale(&mut self, key: &str, evicted_at: Instant) {
144        self.access_order.retain(|candidate| candidate != key);
145        if let Some(entry) = self.entries.remove(key) {
146            self.stale_entries.insert(
147                key.to_string(),
148                StaleEntry {
149                    data: entry.data,
150                    evicted_at,
151                    provider_id: entry.provider_id,
152                },
153            );
154        }
155    }
156
157    fn stale_entry_count(&self) -> usize {
158        self.stale_entries.len()
159    }
160
161    fn purge_stale(&mut self) {
162        let now = Instant::now();
163        self.stale_entries
164            .retain(|_, entry| entry.evicted_at + STALE_SERVE_GRACE > now);
165    }
166
167    fn invalidate_provider(&mut self, provider_id: &str) -> usize {
168        let before = self.entries.len();
169        self.entries.retain(|_, v| v.provider_id != provider_id);
170        self.access_order
171            .retain(|key| self.entries.contains_key(key));
172        self.stale_entries
173            .retain(|_, entry| entry.provider_id != provider_id);
174        before - self.entries.len()
175    }
176
177    fn invalidate_all(&mut self) -> usize {
178        let count = self.entries.len();
179        self.entries.clear();
180        self.access_order.clear();
181        self.stale_entries.clear();
182        count
183    }
184
185    fn metrics(&mut self) -> CacheMetrics {
186        self.expire_entries();
187        self.purge_stale();
188
189        let mut by_provider: HashMap<String, ProviderCacheStats> = HashMap::new();
190
191        for entry in self.entries.values() {
192            let stats = by_provider.entry(entry.provider_id.clone()).or_default();
193            stats.provider_id.clone_from(&entry.provider_id);
194            stats.entry_count += 1;
195        }
196
197        for (pid, &count) in &self.hits {
198            let stats = by_provider.entry(pid.clone()).or_default();
199            stats.provider_id.clone_from(pid);
200            stats.hits = count;
201        }
202        for (pid, &count) in &self.misses {
203            let stats = by_provider.entry(pid.clone()).or_default();
204            stats.provider_id.clone_from(pid);
205            stats.misses = count;
206        }
207        for (pid, &ts) in &self.last_fetch {
208            let stats = by_provider.entry(pid.clone()).or_default();
209            stats.provider_id.clone_from(pid);
210            stats.last_fetch = Some(ts);
211        }
212
213        let mut provider_stats: Vec<_> = by_provider.into_values().collect();
214        provider_stats.sort_by(|a, b| a.provider_id.cmp(&b.provider_id));
215
216        CacheMetrics {
217            total_hits: self.hits.values().sum(),
218            total_misses: self.misses.values().sum(),
219            total_entries: self.entries.len(),
220            stale_entries: self.stale_entries.len(),
221            max_entries: MAX_CACHE_ENTRIES,
222            provider_stats,
223        }
224    }
225}
226
227pub fn get_cached(key: &str) -> Option<String> {
228    PROVIDER_CACHE
229        .lock()
230        .ok()
231        .and_then(|mut c| c.get(key).map(std::string::ToString::to_string))
232}
233
234pub fn set_cached(key: &str, data: &str, ttl_secs: u64) {
235    set_cached_with_provider(
236        key,
237        data,
238        ttl_secs,
239        key.split(':').next().unwrap_or("unknown"),
240    );
241}
242
243pub fn set_cached_with_provider(key: &str, data: &str, ttl_secs: u64, provider_id: &str) {
244    if let Ok(mut cache) = PROVIDER_CACHE.lock() {
245        cache.set(
246            key.to_string(),
247            data.to_string(),
248            Duration::from_secs(ttl_secs),
249            provider_id,
250        );
251    }
252}
253
254pub fn invalidate_provider(provider_id: &str) -> usize {
255    PROVIDER_CACHE
256        .lock()
257        .ok()
258        .map_or(0, |mut c| c.invalidate_provider(provider_id))
259}
260
261pub fn invalidate_all() -> usize {
262    PROVIDER_CACHE
263        .lock()
264        .ok()
265        .map_or(0, |mut c| c.invalidate_all())
266}
267
268pub fn cache_metrics() -> CacheMetrics {
269    PROVIDER_CACHE
270        .lock()
271        .ok()
272        .map_or_else(CacheMetrics::default, |mut c| c.metrics())
273}
274
275pub fn cache_entry_count() -> usize {
276    PROVIDER_CACHE.lock().ok().map_or(0, |mut cache| {
277        cache.expire_entries();
278        cache.entries.len()
279    })
280}
281
282pub fn cache_stale_count() -> usize {
283    PROVIDER_CACHE
284        .lock()
285        .ok()
286        .map_or(0, |cache| cache.stale_entry_count())
287}
288
289pub fn cache_purge_stale() -> usize {
290    PROVIDER_CACHE.lock().ok().map_or(0, |mut cache| {
291        let before = cache.stale_entry_count();
292        cache.purge_stale();
293        before - cache.stale_entry_count()
294    })
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300
301    #[test]
302    fn cache_set_and_get() {
303        let mut cache = ProviderCache::new();
304        cache.set(
305            "test:key".into(),
306            "value".into(),
307            Duration::from_mins(1),
308            "test",
309        );
310        assert_eq!(cache.get("test:key"), Some("value"));
311    }
312
313    #[test]
314    fn test_stale_serve_on_expired_entry() {
315        let mut cache = ProviderCache::new();
316        cache.set(
317            "test:key".into(),
318            "value".into(),
319            Duration::from_millis(10),
320            "test",
321        );
322        std::thread::sleep(Duration::from_millis(20));
323        assert_eq!(cache.get("test:key"), Some("value"));
324    }
325
326    #[test]
327    fn cache_tracks_hits_and_misses() {
328        let mut cache = ProviderCache::new();
329        cache.set(
330            "github:key".into(),
331            "data".into(),
332            Duration::from_mins(1),
333            "github",
334        );
335        cache.get("github:key"); // hit
336        cache.get("github:key"); // hit
337        cache.get("github:missing"); // miss
338
339        let metrics = cache.metrics();
340        assert_eq!(metrics.total_hits, 2);
341        assert_eq!(metrics.total_misses, 1);
342        assert!((metrics.total_hit_rate() - 0.666).abs() < 0.01);
343    }
344
345    #[test]
346    fn cache_invalidate_provider() {
347        let mut cache = ProviderCache::new();
348        cache.set(
349            "github:a".into(),
350            "1".into(),
351            Duration::from_mins(1),
352            "github",
353        );
354        cache.set(
355            "github:b".into(),
356            "2".into(),
357            Duration::from_mins(1),
358            "github",
359        );
360        cache.set(
361            "gitlab:c".into(),
362            "3".into(),
363            Duration::from_mins(1),
364            "gitlab",
365        );
366
367        let removed = cache.invalidate_provider("github");
368        assert_eq!(removed, 2);
369        assert!(cache.get("github:a").is_none());
370        assert_eq!(cache.get("gitlab:c"), Some("3"));
371    }
372
373    #[test]
374    fn cache_invalidate_all() {
375        let mut cache = ProviderCache::new();
376        cache.set("a".into(), "1".into(), Duration::from_mins(1), "x");
377        cache.set("b".into(), "2".into(), Duration::from_mins(1), "y");
378
379        let removed = cache.invalidate_all();
380        assert_eq!(removed, 2);
381        assert!(cache.get("a").is_none());
382    }
383
384    #[test]
385    fn cache_metrics_per_provider() {
386        let mut cache = ProviderCache::new();
387        cache.set(
388            "github:x".into(),
389            "a".into(),
390            Duration::from_mins(1),
391            "github",
392        );
393        cache.set(
394            "gitlab:y".into(),
395            "b".into(),
396            Duration::from_mins(1),
397            "gitlab",
398        );
399        cache.get("github:x");
400        cache.get("gitlab:miss");
401
402        let metrics = cache.metrics();
403        assert_eq!(metrics.provider_stats.len(), 2);
404
405        let gh = metrics
406            .provider_stats
407            .iter()
408            .find(|s| s.provider_id == "github")
409            .unwrap();
410        assert_eq!(gh.entry_count, 1);
411        assert_eq!(gh.hits, 1);
412
413        let gl = metrics
414            .provider_stats
415            .iter()
416            .find(|s| s.provider_id == "gitlab")
417            .unwrap();
418        assert_eq!(gl.entry_count, 1);
419        assert!(gl.last_fetch.is_some());
420    }
421
422    #[test]
423    fn test_lru_eviction_at_capacity() {
424        let mut cache = ProviderCache::new();
425        for index in 0..MAX_CACHE_ENTRIES + 10 {
426            let key = format!("test:{index}");
427            cache.set(key.clone(), key, Duration::from_mins(1), "test");
428        }
429        assert_eq!(cache.entries.len(), MAX_CACHE_ENTRIES);
430    }
431
432    #[test]
433    fn test_stale_serve_expired_grace() {
434        let mut cache = ProviderCache::new();
435        cache.stale_entries.insert(
436            "test:key".into(),
437            StaleEntry {
438                data: "value".into(),
439                evicted_at: Instant::now()
440                    .checked_sub(STALE_SERVE_GRACE + Duration::from_millis(1))
441                    .unwrap(),
442                provider_id: "test".into(),
443            },
444        );
445        assert!(cache.get("test:key").is_none());
446    }
447
448    #[test]
449    fn test_access_order_update_on_get() {
450        let mut cache = ProviderCache::new();
451        cache.set("a".into(), "1".into(), Duration::from_mins(1), "test");
452        cache.set("b".into(), "2".into(), Duration::from_mins(1), "test");
453        cache.get("a");
454        assert_eq!(cache.access_order.last().map(String::as_str), Some("a"));
455    }
456
457    #[test]
458    fn test_metrics_include_stale_count() {
459        let mut cache = ProviderCache::new();
460        cache.set("a".into(), "1".into(), Duration::from_millis(10), "test");
461        std::thread::sleep(Duration::from_millis(20));
462        assert_eq!(cache.metrics().stale_entries, 1);
463    }
464
465    #[test]
466    fn test_purge_stale_removes_old() {
467        let mut cache = ProviderCache::new();
468        cache.stale_entries.insert(
469            "old".into(),
470            StaleEntry {
471                data: "value".into(),
472                evicted_at: Instant::now()
473                    .checked_sub(STALE_SERVE_GRACE + Duration::from_millis(1))
474                    .unwrap(),
475                provider_id: "test".into(),
476            },
477        );
478        cache.purge_stale();
479        assert_eq!(cache.stale_entry_count(), 0);
480    }
481
482    #[test]
483    fn test_cache_entry_count_accuracy() {
484        invalidate_all();
485        set_cached_with_provider("test:count", "value", 60, "test");
486        assert_eq!(cache_entry_count(), 1);
487        invalidate_all();
488    }
489
490    #[test]
491    fn provider_cache_stats_hit_rate() {
492        let stats = ProviderCacheStats {
493            provider_id: "test".into(),
494            hits: 3,
495            misses: 1,
496            entry_count: 2,
497            last_fetch: None,
498        };
499        assert!((stats.hit_rate() - 0.75).abs() < f64::EPSILON);
500    }
501}