lean_ctx/core/ocla/
response_cache.rs

1//! Bounded in-memory cache for identical model responses.
2
3use std::collections::{HashMap, VecDeque};
4use std::sync::{Mutex, OnceLock, PoisonError};
5use std::time::{Duration, Instant};
6
7/// Maximum number of responses retained by a response cache.
8pub const MAX_ENTRIES: usize = 512;
9/// Default lifetime for a cached response.
10pub const DEFAULT_TTL: Duration = Duration::from_mins(5);
11
12/// Selects how cached responses determine their lifetime.
13#[derive(Clone, Debug)]
14pub enum CachePolicy {
15    /// Global TTL for all models.
16    Uniform,
17    /// Per-model TTL overrides selected by model-prefix match.
18    ModelAware {
19        overrides: HashMap<String, Duration>,
20    },
21}
22
23static GLOBAL_RESPONSE_CACHE: OnceLock<ResponseCache> = OnceLock::new();
24
25pub(crate) fn global_response_cache() -> &'static ResponseCache {
26    GLOBAL_RESPONSE_CACHE.get_or_init(ResponseCache::default)
27}
28
29/// Stable cache key derived from response-defining request fields.
30#[derive(Clone, Debug, Eq, Hash, PartialEq)]
31pub struct ResponseCacheKey {
32    /// First 64 bits of the digest of the key components.
33    pub hash: u64,
34    /// Model name used for model-aware cache expiry.
35    pub model: String,
36}
37
38impl ResponseCacheKey {
39    /// Hashes model, prompt hash, temperature, and maximum output tokens.
40    pub fn new(
41        model: impl AsRef<str>,
42        prompt_hash: u64,
43        temperature: f32,
44        max_tokens: u64,
45    ) -> Self {
46        let mut hasher = blake3::Hasher::new();
47        let model = model.as_ref();
48        hasher.update(&(model.len() as u64).to_be_bytes());
49        hasher.update(model.as_bytes());
50        hasher.update(&prompt_hash.to_be_bytes());
51        hasher.update(&temperature.to_bits().to_be_bytes());
52        hasher.update(&max_tokens.to_be_bytes());
53
54        let digest = hasher.finalize();
55        let mut bytes = [0; 8];
56        bytes.copy_from_slice(&digest.as_bytes()[..8]);
57        Self {
58            hash: u64::from_be_bytes(bytes),
59            model: model.to_owned(),
60        }
61    }
62}
63
64/// A response stored in the cache.
65#[derive(Clone, Debug, PartialEq, Eq)]
66pub struct CachedResponse {
67    pub body: Vec<u8>,
68    pub status: u16,
69    pub tokens: u64,
70    pub created_at: Instant,
71    pub ttl: Duration,
72}
73
74#[derive(Debug, Default)]
75struct CacheState {
76    entries: VecDeque<(ResponseCacheKey, CachedResponse)>,
77    hits: u64,
78    misses: u64,
79    evictions: u64,
80    evictions_by_reason: HashMap<String, u64>,
81}
82
83/// Thread-safe bounded LRU response cache.
84#[derive(Debug)]
85pub struct ResponseCache {
86    capacity: usize,
87    ttl: Duration,
88    policy: CachePolicy,
89    state: Mutex<CacheState>,
90}
91
92impl Default for ResponseCache {
93    fn default() -> Self {
94        Self::new(MAX_ENTRIES, DEFAULT_TTL)
95    }
96}
97
98impl ResponseCache {
99    /// Creates a cache with the requested capacity and defaulting zero TTLs.
100    ///
101    /// Capacity is clamped to the inclusive range 1..=512.
102    pub fn new(capacity: usize, ttl: Duration) -> Self {
103        Self {
104            capacity: capacity.clamp(1, MAX_ENTRIES),
105            ttl,
106            policy: CachePolicy::Uniform,
107            state: Mutex::new(CacheState::default()),
108        }
109    }
110
111    /// Creates a cache with the requested TTL and the maximum capacity.
112    pub fn with_ttl(ttl: Duration) -> Self {
113        Self::new(MAX_ENTRIES, ttl)
114    }
115
116    /// Replaces the cache expiry policy for existing and future entries.
117    pub fn set_policy(&mut self, policy: CachePolicy) {
118        self.policy = policy;
119    }
120
121    /// Looks up a response, refreshing its LRU position on a live hit.
122    pub fn get(&self, key: &ResponseCacheKey) -> Option<CachedResponse> {
123        let mut state = self.lock_state();
124        let now = Instant::now();
125        let position = state
126            .entries
127            .iter()
128            .position(|(entry_key, _)| entry_key == key);
129
130        let Some(position) = position else {
131            state.misses += 1;
132            return None;
133        };
134
135        let expiration_reason = {
136            let (entry_key, response) = &state.entries[position];
137            self.expiration_reason(entry_key, response, now)
138        };
139        if let Some(reason) = expiration_reason {
140            state.entries.remove(position);
141            record_eviction(&mut state, reason);
142            state.misses += 1;
143            return None;
144        }
145
146        let entry = state.entries.remove(position)?;
147        state.entries.push_back(entry.clone());
148        state.hits += 1;
149        Some(entry.1)
150    }
151
152    /// Inserts or replaces a response, evicting the least recently used entry
153    /// when the bounded capacity is reached.
154    pub fn put(&self, key: ResponseCacheKey, mut response: CachedResponse) {
155        let mut state = self.lock_state();
156
157        if response.ttl.is_zero() {
158            response.ttl = self.ttl;
159        }
160        self.remove_expired(&mut state, Instant::now());
161
162        if let Some(position) = state
163            .entries
164            .iter()
165            .position(|(entry_key, _)| entry_key == &key)
166        {
167            state.entries.remove(position);
168        } else if state.entries.len() >= self.capacity {
169            state.entries.pop_front();
170            record_eviction(&mut state, "lru_capacity");
171        }
172
173        state.entries.push_back((key, response));
174    }
175
176    /// Returns cumulative hit, miss, and eviction counters.
177    pub fn stats(&self) -> CacheStats {
178        let state = self.lock_state();
179        let total = state.hits + state.misses;
180        CacheStats {
181            entries: state.entries.len(),
182            hits: state.hits,
183            misses: state.misses,
184            evictions: state.evictions,
185            evictions_by_reason: state.evictions_by_reason.clone(),
186            hit_rate: if total == 0 {
187                0.0
188            } else {
189                state.hits as f64 / total as f64
190            },
191        }
192    }
193
194    fn lock_state(&self) -> std::sync::MutexGuard<'_, CacheState> {
195        self.state.lock().unwrap_or_else(PoisonError::into_inner)
196    }
197
198    fn effective_ttl(&self, model: &str) -> Duration {
199        match &self.policy {
200            CachePolicy::Uniform => self.ttl,
201            CachePolicy::ModelAware { overrides } => overrides
202                .iter()
203                .filter(|(prefix, _)| model.starts_with(prefix.as_str()))
204                .max_by_key(|(prefix, _)| prefix.len())
205                .map_or(self.ttl, |(_, ttl)| *ttl),
206        }
207    }
208
209    fn expiration_reason(
210        &self,
211        key: &ResponseCacheKey,
212        response: &CachedResponse,
213        now: Instant,
214    ) -> Option<&'static str> {
215        let ttl = self.effective_ttl(&key.model);
216        let expired = now
217            .checked_duration_since(response.created_at)
218            .unwrap_or_default()
219            >= ttl;
220        expired.then_some(match &self.policy {
221            CachePolicy::Uniform => "ttl_expired",
222            CachePolicy::ModelAware { .. } => "model_ttl_expired",
223        })
224    }
225
226    fn remove_expired(&self, state: &mut CacheState, now: Instant) {
227        let mut retained = VecDeque::with_capacity(state.entries.len());
228        while let Some((key, response)) = state.entries.pop_front() {
229            if let Some(reason) = self.expiration_reason(&key, &response, now) {
230                record_eviction(state, reason);
231            } else {
232                retained.push_back((key, response));
233            }
234        }
235        state.entries = retained;
236    }
237}
238
239fn record_eviction(state: &mut CacheState, reason: &str) {
240    state.evictions += 1;
241    *state
242        .evictions_by_reason
243        .entry(reason.to_owned())
244        .or_default() += 1;
245}
246
247/// Snapshot of cache activity.
248#[derive(Clone, Debug, PartialEq)]
249pub struct CacheStats {
250    pub entries: usize,
251    pub hits: u64,
252    pub misses: u64,
253    pub evictions: u64,
254    pub evictions_by_reason: HashMap<String, u64>,
255    pub hit_rate: f64,
256}
257
258#[cfg(test)]
259mod tests {
260    use super::*;
261
262    fn cache_key(model: &str) -> ResponseCacheKey {
263        ResponseCacheKey::new(model, 42, 0.2, 128)
264    }
265
266    fn response(body: &[u8], created_at: Instant, ttl: Duration) -> CachedResponse {
267        CachedResponse {
268            body: body.to_vec(),
269            status: 200,
270            tokens: body.len() as u64,
271            created_at,
272            ttl,
273        }
274    }
275
276    #[test]
277    fn key_hash_changes_when_request_fields_change() {
278        let base = cache_key("model-a");
279        assert_ne!(base, cache_key("model-b"));
280        assert_ne!(base, ResponseCacheKey::new("model-a", 43, 0.2, 128));
281        assert_ne!(base, ResponseCacheKey::new("model-a", 42, 0.3, 128));
282        assert_ne!(base, ResponseCacheKey::new("model-a", 42, 0.2, 129));
283    }
284
285    #[test]
286    fn cache_hit_and_miss_update_stats() {
287        let cache = ResponseCache::new(4, Duration::from_mins(1));
288        let key = cache_key("model-a");
289        cache.put(
290            key.clone(),
291            response(b"answer", Instant::now(), Duration::ZERO),
292        );
293
294        assert_eq!(cache.get(&key).unwrap().body, b"answer");
295        assert!(cache.get(&cache_key("model-b")).is_none());
296
297        let stats = cache.stats();
298        assert_eq!(stats.hits, 1);
299        assert_eq!(stats.misses, 1);
300        assert_eq!(stats.evictions, 0);
301        assert!(stats.evictions_by_reason.is_empty());
302        assert!((stats.hit_rate - 0.5).abs() < f64::EPSILON);
303    }
304
305    #[test]
306    fn expired_entries_count_as_misses() {
307        let cache = ResponseCache::with_ttl(Duration::from_mins(1));
308        let key = cache_key("model-a");
309        cache.put(
310            key.clone(),
311            response(
312                b"old",
313                Instant::now().checked_sub(Duration::from_secs(61)).unwrap(),
314                Duration::from_mins(1),
315            ),
316        );
317
318        assert!(cache.get(&key).is_none());
319        let stats = cache.stats();
320        assert_eq!(stats.hits, 0);
321        assert_eq!(stats.misses, 1);
322        assert_eq!(stats.evictions, 1);
323        assert_eq!(stats.evictions_by_reason["ttl_expired"], 1);
324    }
325
326    #[test]
327    fn lru_eviction_removes_least_recently_used_entry() {
328        let cache = ResponseCache::new(2, Duration::from_mins(1));
329        let first = cache_key("first");
330        let second = cache_key("second");
331        let third = cache_key("third");
332
333        cache.put(
334            first.clone(),
335            response(b"1", Instant::now(), Duration::ZERO),
336        );
337        cache.put(
338            second.clone(),
339            response(b"2", Instant::now(), Duration::ZERO),
340        );
341        assert!(cache.get(&first).is_some());
342        cache.put(
343            third.clone(),
344            response(b"3", Instant::now(), Duration::ZERO),
345        );
346
347        assert!(cache.get(&second).is_none());
348        assert!(cache.get(&first).is_some());
349        assert!(cache.get(&third).is_some());
350
351        let stats = cache.stats();
352        assert_eq!(stats.evictions, 1);
353        assert_eq!(stats.evictions_by_reason["lru_capacity"], 1);
354        assert_eq!(stats.hits, 3);
355        assert_eq!(stats.misses, 1);
356    }
357
358    #[test]
359    fn capacity_is_hard_capped() {
360        let cache = ResponseCache::new(MAX_ENTRIES + 1, Duration::from_mins(1));
361        for index in 0..=MAX_ENTRIES {
362            cache.put(
363                ResponseCacheKey {
364                    hash: index as u64,
365                    model: String::new(),
366                },
367                response(b"x", Instant::now(), Duration::ZERO),
368            );
369        }
370
371        assert_eq!(cache.stats().evictions, 1);
372    }
373
374    #[test]
375    fn default_policy_uses_five_minute_ttl() {
376        let cache = ResponseCache::default();
377        let key = cache_key("default");
378        cache.put(
379            key.clone(),
380            response(
381                b"answer",
382                Instant::now()
383                    .checked_sub(Duration::from_secs(301))
384                    .unwrap(),
385                Duration::ZERO,
386            ),
387        );
388
389        assert!(cache.get(&key).is_none());
390    }
391
392    #[test]
393    fn model_aware_ttl_overrides_default() {
394        let mut cache = ResponseCache::new(4, Duration::from_mins(1));
395        let gpt = cache_key("gpt-4o");
396        let other = cache_key("other-model");
397        let created_at = Instant::now().checked_sub(Duration::from_secs(5)).unwrap();
398
399        cache.put(gpt.clone(), response(b"gpt", created_at, Duration::ZERO));
400        cache.put(
401            other.clone(),
402            response(b"other", created_at, Duration::ZERO),
403        );
404        cache.set_policy(CachePolicy::ModelAware {
405            overrides: HashMap::from([(String::from("gpt"), Duration::from_secs(1))]),
406        });
407
408        assert!(cache.get(&gpt).is_none());
409        assert_eq!(cache.get(&other).unwrap().body, b"other");
410        assert_eq!(cache.stats().evictions_by_reason["model_ttl_expired"], 1);
411    }
412
413    #[test]
414    fn uniform_policy_uses_global_ttl() {
415        let cache = ResponseCache::new(4, Duration::from_secs(1));
416        let created_at = Instant::now().checked_sub(Duration::from_secs(2)).unwrap();
417        let first = cache_key("gpt-4o");
418        let second = cache_key("other-model");
419
420        cache.put(
421            first.clone(),
422            response(b"first", created_at, Duration::from_mins(1)),
423        );
424        cache.put(
425            second.clone(),
426            response(b"second", created_at, Duration::from_mins(1)),
427        );
428
429        assert!(cache.get(&first).is_none());
430        assert!(cache.get(&second).is_none());
431        assert_eq!(cache.stats().evictions_by_reason["ttl_expired"], 2);
432    }
433
434    #[test]
435    fn set_policy_changes_behavior() {
436        let mut cache = ResponseCache::new(4, Duration::from_mins(1));
437        let gpt = cache_key("gpt-4o");
438        let created_at = Instant::now().checked_sub(Duration::from_secs(5)).unwrap();
439
440        cache.put(gpt.clone(), response(b"answer", created_at, Duration::ZERO));
441        assert!(cache.get(&gpt).is_some());
442
443        cache.put(gpt.clone(), response(b"answer", created_at, Duration::ZERO));
444        cache.set_policy(CachePolicy::ModelAware {
445            overrides: HashMap::from([(String::from("gpt"), Duration::from_secs(1))]),
446        });
447        assert!(cache.get(&gpt).is_none());
448    }
449
450    #[test]
451    fn eviction_stats_track_reason() {
452        let cache = ResponseCache::new(1, Duration::from_secs(1));
453        let expired = cache_key("expired");
454        let first = cache_key("first");
455        let second = cache_key("second");
456
457        cache.put(
458            expired,
459            response(
460                b"expired",
461                Instant::now().checked_sub(Duration::from_secs(2)).unwrap(),
462                Duration::ZERO,
463            ),
464        );
465        cache.put(first, response(b"first", Instant::now(), Duration::ZERO));
466        cache.put(second, response(b"second", Instant::now(), Duration::ZERO));
467
468        let stats = cache.stats();
469        assert_eq!(stats.evictions, 2);
470        assert_eq!(stats.evictions_by_reason["ttl_expired"], 1);
471        assert_eq!(stats.evictions_by_reason["lru_capacity"], 1);
472    }
473}
lean_ctx/core/ocla/response_cache.rs

lean_ctx/core/ocla/
response_cache.rs
