Skip to main content

lean_ctx/core/providers/
cache.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use std::collections::HashMap;
4use std::sync::Mutex;
5use std::time::{Duration, Instant, SystemTime};
6
7const MAX_CACHE_ENTRIES: usize = 2048;
8const STALE_SERVE_GRACE: Duration = Duration::from_hours(1);
9
10/// Partition cached responses by the exact request and credential, not authorization state.
11/// No built-in provider caches a response any more (an answer is the principal's current
12/// access); the key stays for the tests that prove a warm entry under it changes nothing.
13#[cfg(test)]
14pub(super) fn request_cache_key(
15    provider_id: &str,
16    request_url: &str,
17    credential: Option<&str>,
18) -> String {
19    let mut hasher = blake3::Hasher::new();
20    for component in [provider_id.as_bytes(), request_url.as_bytes()] {
21        hasher.update(&(component.len() as u64).to_be_bytes());
22        hasher.update(component);
23    }
24    let credential_marker = [u8::from(credential.is_some())];
25    hasher.update(&(credential_marker.len() as u64).to_be_bytes());
26    hasher.update(&credential_marker);
27    if let Some(credential) = credential {
28        hasher.update(&(credential.len() as u64).to_be_bytes());
29        hasher.update(credential.as_bytes());
30    }
31    format!("{provider_id}:request:v2:{}", hasher.finalize().to_hex())
32}
33
34static PROVIDER_CACHE: std::sync::LazyLock<Mutex<ProviderCache>> =
35    std::sync::LazyLock::new(|| Mutex::new(ProviderCache::new()));
36
37struct CacheEntry {
38    data: String,
39    expires_at: Instant,
40    provider_id: String,
41}
42
43struct StaleEntry {
44    evicted_at: Instant,
45    provider_id: String,
46}
47
48/// Per-provider cache statistics.
49#[derive(Debug, Clone, Default)]
50pub struct ProviderCacheStats {
51    pub provider_id: String,
52    pub hits: u64,
53    pub misses: u64,
54    pub entry_count: usize,
55    pub last_fetch: Option<SystemTime>,
56}
57
58impl ProviderCacheStats {
59    pub fn hit_rate(&self) -> f64 {
60        let total = self.hits + self.misses;
61        if total == 0 {
62            return 0.0;
63        }
64        self.hits as f64 / total as f64
65    }
66}
67
68/// Global cache statistics across all providers.
69#[derive(Debug, Clone, Default)]
70pub struct CacheMetrics {
71    pub total_hits: u64,
72    pub total_misses: u64,
73    pub total_entries: usize,
74    pub stale_entries: usize,
75    pub max_entries: usize,
76    pub provider_stats: Vec<ProviderCacheStats>,
77}
78
79impl CacheMetrics {
80    pub fn total_hit_rate(&self) -> f64 {
81        let total = self.total_hits + self.total_misses;
82        if total == 0 {
83            return 0.0;
84        }
85        self.total_hits as f64 / total as f64
86    }
87}
88
89struct ProviderCache {
90    entries: HashMap<String, CacheEntry>,
91    access_order: Vec<String>,
92    stale_entries: HashMap<String, StaleEntry>,
93    hits: HashMap<String, u64>,
94    misses: HashMap<String, u64>,
95    last_fetch: HashMap<String, SystemTime>,
96}
97
98impl ProviderCache {
99    fn new() -> Self {
100        Self {
101            entries: HashMap::new(),
102            access_order: Vec::new(),
103            stale_entries: HashMap::new(),
104            hits: HashMap::new(),
105            misses: HashMap::new(),
106            last_fetch: HashMap::new(),
107        }
108    }
109
110    fn get(&mut self, key: &str) -> Option<&str> {
111        self.expire_entries();
112        self.purge_stale();
113        if let Some(entry) = self.entries.get(key) {
114            let provider_id = entry.provider_id.clone();
115            self.access_order.retain(|candidate| candidate != key);
116            self.access_order.push(key.to_string());
117            *self.hits.entry(provider_id).or_insert(0) += 1;
118            return self.entries.get(key).map(|entry| entry.data.as_str());
119        }
120        let provider_id = self.stale_entries.get(key).map_or_else(
121            || key.split(':').next().unwrap_or("unknown").to_string(),
122            |entry| entry.provider_id.clone(),
123        );
124        *self.misses.entry(provider_id).or_insert(0) += 1;
125        // An expired response is never an answer: serving it would keep returning
126        // a source object after the principal lost access to it (stale-allow).
127        // Evicted entries stay only for the stale-entry metric.
128        None
129    }
130
131    fn set(&mut self, key: String, data: String, ttl: Duration, provider_id: &str) {
132        self.last_fetch
133            .insert(provider_id.to_string(), SystemTime::now());
134        self.access_order.retain(|candidate| candidate != &key);
135        self.access_order.push(key.clone());
136        self.stale_entries.remove(&key);
137        self.entries.insert(
138            key,
139            CacheEntry {
140                data,
141                expires_at: Instant::now() + ttl,
142                provider_id: provider_id.to_string(),
143            },
144        );
145        self.enforce_lru_cap();
146    }
147
148    fn expire_entries(&mut self) {
149        let now = Instant::now();
150        let expired: Vec<_> = self
151            .entries
152            .iter()
153            .filter(|(_, entry)| entry.expires_at <= now)
154            .map(|(key, _)| key.clone())
155            .collect();
156        for key in expired {
157            self.move_to_stale(&key, now);
158        }
159    }
160
161    fn enforce_lru_cap(&mut self) {
162        let now = Instant::now();
163        while self.entries.len() > MAX_CACHE_ENTRIES {
164            let key = self.access_order.remove(0);
165            self.move_to_stale(&key, now);
166        }
167    }
168
169    fn move_to_stale(&mut self, key: &str, evicted_at: Instant) {
170        self.access_order.retain(|candidate| candidate != key);
171        if let Some(entry) = self.entries.remove(key) {
172            self.stale_entries.insert(
173                key.to_string(),
174                StaleEntry {
175                    evicted_at,
176                    provider_id: entry.provider_id,
177                },
178            );
179        }
180    }
181
182    fn stale_entry_count(&self) -> usize {
183        self.stale_entries.len()
184    }
185
186    fn purge_stale(&mut self) {
187        let now = Instant::now();
188        self.stale_entries
189            .retain(|_, entry| entry.evicted_at + STALE_SERVE_GRACE > now);
190    }
191
192    fn invalidate_provider(&mut self, provider_id: &str) -> usize {
193        let before = self.entries.len();
194        self.entries.retain(|_, v| v.provider_id != provider_id);
195        self.access_order
196            .retain(|key| self.entries.contains_key(key));
197        self.stale_entries
198            .retain(|_, entry| entry.provider_id != provider_id);
199        before - self.entries.len()
200    }
201
202    fn invalidate_all(&mut self) -> usize {
203        let count = self.entries.len();
204        self.entries.clear();
205        self.access_order.clear();
206        self.stale_entries.clear();
207        count
208    }
209
210    fn metrics(&mut self) -> CacheMetrics {
211        self.expire_entries();
212        self.purge_stale();
213
214        let mut by_provider: HashMap<String, ProviderCacheStats> = HashMap::new();
215
216        for entry in self.entries.values() {
217            let stats = by_provider.entry(entry.provider_id.clone()).or_default();
218            stats.provider_id.clone_from(&entry.provider_id);
219            stats.entry_count += 1;
220        }
221
222        for (pid, &count) in &self.hits {
223            let stats = by_provider.entry(pid.clone()).or_default();
224            stats.provider_id.clone_from(pid);
225            stats.hits = count;
226        }
227        for (pid, &count) in &self.misses {
228            let stats = by_provider.entry(pid.clone()).or_default();
229            stats.provider_id.clone_from(pid);
230            stats.misses = count;
231        }
232        for (pid, &ts) in &self.last_fetch {
233            let stats = by_provider.entry(pid.clone()).or_default();
234            stats.provider_id.clone_from(pid);
235            stats.last_fetch = Some(ts);
236        }
237
238        let mut provider_stats: Vec<_> = by_provider.into_values().collect();
239        provider_stats.sort_by(|a, b| a.provider_id.cmp(&b.provider_id));
240
241        CacheMetrics {
242            total_hits: self.hits.values().sum(),
243            total_misses: self.misses.values().sum(),
244            total_entries: self.entries.len(),
245            stale_entries: self.stale_entries.len(),
246            max_entries: MAX_CACHE_ENTRIES,
247            provider_stats,
248        }
249    }
250}
251
252pub fn get_cached(key: &str) -> Option<String> {
253    PROVIDER_CACHE
254        .lock()
255        .ok()
256        .and_then(|mut c| c.get(key).map(std::string::ToString::to_string))
257}
258
259pub fn set_cached(key: &str, data: &str, ttl_secs: u64) {
260    set_cached_with_provider(
261        key,
262        data,
263        ttl_secs,
264        key.split(':').next().unwrap_or("unknown"),
265    );
266}
267
268pub fn set_cached_with_provider(key: &str, data: &str, ttl_secs: u64, provider_id: &str) {
269    if let Ok(mut cache) = PROVIDER_CACHE.lock() {
270        cache.set(
271            key.to_string(),
272            data.to_string(),
273            Duration::from_secs(ttl_secs),
274            provider_id,
275        );
276    }
277}
278
279pub fn invalidate_provider(provider_id: &str) -> usize {
280    PROVIDER_CACHE
281        .lock()
282        .ok()
283        .map_or(0, |mut c| c.invalidate_provider(provider_id))
284}
285
286pub fn invalidate_all() -> usize {
287    PROVIDER_CACHE
288        .lock()
289        .ok()
290        .map_or(0, |mut c| c.invalidate_all())
291}
292
293pub fn cache_metrics() -> CacheMetrics {
294    PROVIDER_CACHE
295        .lock()
296        .ok()
297        .map_or_else(CacheMetrics::default, |mut c| c.metrics())
298}
299
300pub fn cache_entry_count() -> usize {
301    PROVIDER_CACHE.lock().ok().map_or(0, |mut cache| {
302        cache.expire_entries();
303        cache.entries.len()
304    })
305}
306
307pub fn cache_stale_count() -> usize {
308    PROVIDER_CACHE
309        .lock()
310        .ok()
311        .map_or(0, |cache| cache.stale_entry_count())
312}
313
314pub fn cache_purge_stale() -> usize {
315    PROVIDER_CACHE.lock().ok().map_or(0, |mut cache| {
316        let before = cache.stale_entry_count();
317        cache.purge_stale();
318        before - cache.stale_entry_count()
319    })
320}
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325
326    #[test]
327    fn cache_set_and_get() {
328        let mut cache = ProviderCache::new();
329        cache.set(
330            "test:key".into(),
331            "value".into(),
332            Duration::from_mins(1),
333            "test",
334        );
335        assert_eq!(cache.get("test:key"), Some("value"));
336    }
337
338    #[test]
339    fn an_expired_entry_is_a_miss_never_a_stale_answer() {
340        let mut cache = ProviderCache::new();
341        cache.set(
342            "test:key".into(),
343            "value".into(),
344            Duration::from_millis(10),
345            "test",
346        );
347        std::thread::sleep(Duration::from_millis(20));
348        assert_eq!(cache.get("test:key"), None);
349        assert_eq!(cache.metrics().stale_entries, 1);
350    }
351
352    #[test]
353    fn cache_tracks_hits_and_misses() {
354        let mut cache = ProviderCache::new();
355        cache.set(
356            "github:key".into(),
357            "data".into(),
358            Duration::from_mins(1),
359            "github",
360        );
361        cache.get("github:key"); // hit
362        cache.get("github:key"); // hit
363        cache.get("github:missing"); // miss
364
365        let metrics = cache.metrics();
366        assert_eq!(metrics.total_hits, 2);
367        assert_eq!(metrics.total_misses, 1);
368        assert!((metrics.total_hit_rate() - 0.666).abs() < 0.01);
369    }
370
371    #[test]
372    fn cache_invalidate_provider() {
373        let mut cache = ProviderCache::new();
374        cache.set(
375            "github:a".into(),
376            "1".into(),
377            Duration::from_mins(1),
378            "github",
379        );
380        cache.set(
381            "github:b".into(),
382            "2".into(),
383            Duration::from_mins(1),
384            "github",
385        );
386        cache.set(
387            "gitlab:c".into(),
388            "3".into(),
389            Duration::from_mins(1),
390            "gitlab",
391        );
392
393        let removed = cache.invalidate_provider("github");
394        assert_eq!(removed, 2);
395        assert!(cache.get("github:a").is_none());
396        assert_eq!(cache.get("gitlab:c"), Some("3"));
397    }
398
399    #[test]
400    fn cache_invalidate_all() {
401        let mut cache = ProviderCache::new();
402        cache.set("a".into(), "1".into(), Duration::from_mins(1), "x");
403        cache.set("b".into(), "2".into(), Duration::from_mins(1), "y");
404
405        let removed = cache.invalidate_all();
406        assert_eq!(removed, 2);
407        assert!(cache.get("a").is_none());
408    }
409
410    #[test]
411    fn cache_metrics_per_provider() {
412        let mut cache = ProviderCache::new();
413        cache.set(
414            "github:x".into(),
415            "a".into(),
416            Duration::from_mins(1),
417            "github",
418        );
419        cache.set(
420            "gitlab:y".into(),
421            "b".into(),
422            Duration::from_mins(1),
423            "gitlab",
424        );
425        cache.get("github:x");
426        cache.get("gitlab:miss");
427
428        let metrics = cache.metrics();
429        assert_eq!(metrics.provider_stats.len(), 2);
430
431        let gh = metrics
432            .provider_stats
433            .iter()
434            .find(|s| s.provider_id == "github")
435            .unwrap();
436        assert_eq!(gh.entry_count, 1);
437        assert_eq!(gh.hits, 1);
438
439        let gl = metrics
440            .provider_stats
441            .iter()
442            .find(|s| s.provider_id == "gitlab")
443            .unwrap();
444        assert_eq!(gl.entry_count, 1);
445        assert!(gl.last_fetch.is_some());
446    }
447
448    #[test]
449    fn test_lru_eviction_at_capacity() {
450        let mut cache = ProviderCache::new();
451        for index in 0..MAX_CACHE_ENTRIES + 10 {
452            let key = format!("test:{index}");
453            cache.set(key.clone(), key, Duration::from_mins(1), "test");
454        }
455        assert_eq!(cache.entries.len(), MAX_CACHE_ENTRIES);
456    }
457
458    #[test]
459    fn test_stale_serve_expired_grace() {
460        let mut cache = ProviderCache::new();
461        cache.stale_entries.insert(
462            "test:key".into(),
463            StaleEntry {
464                evicted_at: Instant::now()
465                    .checked_sub(STALE_SERVE_GRACE + Duration::from_millis(1))
466                    .unwrap(),
467                provider_id: "test".into(),
468            },
469        );
470        assert!(cache.get("test:key").is_none());
471    }
472
473    #[test]
474    fn test_access_order_update_on_get() {
475        let mut cache = ProviderCache::new();
476        cache.set("a".into(), "1".into(), Duration::from_mins(1), "test");
477        cache.set("b".into(), "2".into(), Duration::from_mins(1), "test");
478        cache.get("a");
479        assert_eq!(cache.access_order.last().map(String::as_str), Some("a"));
480    }
481
482    #[test]
483    fn test_metrics_include_stale_count() {
484        let mut cache = ProviderCache::new();
485        cache.set("a".into(), "1".into(), Duration::from_millis(10), "test");
486        std::thread::sleep(Duration::from_millis(20));
487        assert_eq!(cache.metrics().stale_entries, 1);
488    }
489
490    #[test]
491    fn test_purge_stale_removes_old() {
492        let mut cache = ProviderCache::new();
493        cache.stale_entries.insert(
494            "old".into(),
495            StaleEntry {
496                evicted_at: Instant::now()
497                    .checked_sub(STALE_SERVE_GRACE + Duration::from_millis(1))
498                    .unwrap(),
499                provider_id: "test".into(),
500            },
501        );
502        cache.purge_stale();
503        assert_eq!(cache.stale_entry_count(), 0);
504    }
505
506    #[test]
507    #[serial_test::serial(provider_request_cache)]
508    fn test_cache_entry_count_accuracy() {
509        invalidate_all();
510        set_cached_with_provider("test:count", "value", 60, "test");
511        assert_eq!(cache_entry_count(), 1);
512        invalidate_all();
513    }
514
515    #[test]
516    fn provider_cache_stats_hit_rate() {
517        let stats = ProviderCacheStats {
518            provider_id: "test".into(),
519            hits: 3,
520            misses: 1,
521            entry_count: 2,
522            last_fetch: None,
523        };
524        assert!((stats.hit_rate() - 0.75).abs() < f64::EPSILON);
525    }
526
527    #[test]
528    fn request_cache_key_is_deterministic_and_principal_request_bound() {
529        let request = "https://gitlab.example/api/v4/projects/group%2Fproject/issues?per_page=20&state=opened";
530        let key = request_cache_key("gitlab", request, Some("principal-token"));
531
532        assert_eq!(
533            key,
534            request_cache_key("gitlab", request, Some("principal-token"))
535        );
536        assert_ne!(
537            key,
538            request_cache_key("gitlab", request, Some("other-principal-token"))
539        );
540        assert_ne!(
541            key,
542            request_cache_key(
543                "gitlab",
544                "https://gitlab.other/api/v4/projects/group%2Fproject/issues?per_page=20&state=opened",
545                Some("principal-token")
546            )
547        );
548        assert_ne!(
549            key,
550            request_cache_key(
551                "gitlab",
552                "https://gitlab.example/api/v4/projects/group%2Fproject/issues?per_page=20&state=closed",
553                Some("principal-token")
554            )
555        );
556        assert_ne!(
557            key,
558            request_cache_key("github", request, Some("principal-token"))
559        );
560        assert_ne!(key, request_cache_key("gitlab", request, None));
561        assert_ne!(key, request_cache_key("gitlab", request, Some("")));
562        assert_ne!(
563            request_cache_key("gitlab", request, None),
564            request_cache_key("gitlab", request, Some(""))
565        );
566        assert!(!key.contains("principal-token"));
567        assert!(!key.contains("gitlab.example"));
568        assert!(!key.contains("group%2Fproject"));
569    }
570
571    #[test]
572    fn request_cache_key_does_not_reuse_legacy_raw_key() {
573        let key = request_cache_key(
574            "gitlab",
575            "https://gitlab.example/api/v4/projects/group%2Fproject/issues?per_page=20&state=opened",
576            Some("principal-token"),
577        );
578        let legacy_key = "gitlab:issues:group/project:Some(\"opened\"):None:20";
579
580        assert_ne!(key, legacy_key);
581        assert!(key.starts_with("gitlab:request:v2:"));
582    }
583
584    #[test]
585    fn request_cache_key_isolates_fresh_and_stale_responses() {
586        let request = "https://gitlab.example/api/v4/projects/42/issues?state=opened";
587        let principal_key = request_cache_key("gitlab", request, Some("principal-token"));
588        let foreign_keys = [
589            request_cache_key("gitlab", request, Some("other-principal-token")),
590            request_cache_key("gitlab", request, None),
591            request_cache_key("gitlab", request, Some("")),
592            request_cache_key("github", request, Some("principal-token")),
593            request_cache_key(
594                "gitlab",
595                "https://gitlab.other/api/v4/projects/42/issues?state=opened",
596                Some("principal-token"),
597            ),
598            request_cache_key(
599                "gitlab",
600                "https://gitlab.example/api/v4/projects/42/issues?state=closed",
601                Some("principal-token"),
602            ),
603        ];
604        let mut cache = ProviderCache::new();
605        cache.set(
606            principal_key.clone(),
607            "principal response".into(),
608            Duration::from_mins(1),
609            "gitlab",
610        );
611        assert_eq!(cache.get(&principal_key), Some("principal response"));
612        for key in &foreign_keys {
613            assert_eq!(cache.get(key), None);
614        }
615        cache.move_to_stale(&principal_key, Instant::now());
616        assert_eq!(cache.get(&principal_key), None);
617        for key in &foreign_keys {
618            assert_eq!(cache.get(key), None);
619        }
620        cache.set(
621            "gitlab:issues:42:Some(\"opened\"):None:20".into(),
622            "legacy response".into(),
623            Duration::from_mins(1),
624            "gitlab",
625        );
626        assert_eq!(cache.get(&foreign_keys[0]), None);
627    }
628}