1use std::collections::HashMap;
10use std::sync::atomic::{AtomicUsize, Ordering};
11use std::sync::{Arc, RwLock};
12use std::time::{Duration, Instant};
13
14#[derive(Clone, Debug)]
16struct CachedResponse {
17 body: Arc<[u8]>,
19 etag: String,
21 expires_at: Instant,
23}
24
25#[derive(Clone, Debug, Default)]
27pub struct CacheStats {
28 pub total_requests: usize,
29 pub cache_hits: usize,
30 pub cache_misses: usize,
31 pub etag_hits: usize,
32}
33
34impl CacheStats {
35 pub fn hit_ratio(&self) -> f64 {
36 if self.total_requests == 0 {
37 0.0
38 } else {
39 (self.cache_hits + self.etag_hits) as f64 / self.total_requests as f64
40 }
41 }
42}
43
44#[derive(Debug, Default)]
46struct AtomicCacheStats {
47 total_requests: AtomicUsize,
48 cache_hits: AtomicUsize,
49 cache_misses: AtomicUsize,
50 etag_hits: AtomicUsize,
51}
52
53#[derive(Debug)]
55pub struct ApiResponseCache {
56 cache: Arc<RwLock<HashMap<String, CachedResponse>>>,
57 stats: Arc<AtomicCacheStats>,
58}
59
60impl ApiResponseCache {
61 pub fn new() -> Self {
63 Self {
64 cache: Arc::new(RwLock::new(HashMap::new())),
65 stats: Arc::new(AtomicCacheStats::default()),
66 }
67 }
68
69 pub fn get_or_compute<F>(
77 &self,
78 key: &str,
79 ttl: Duration,
80 compute_fn: F,
81 ) -> Result<(Arc<[u8]>, String), String>
82 where
83 F: FnOnce() -> Result<Vec<u8>, String>,
84 {
85 self.stats.total_requests.fetch_add(1, Ordering::Relaxed);
86
87 {
89 let cache = self.cache.read().unwrap();
90 if let Some(cached) = cache.get(key) {
91 if cached.expires_at > Instant::now() {
92 self.stats.cache_hits.fetch_add(1, Ordering::Relaxed);
93 return Ok((Arc::clone(&cached.body), cached.etag.clone()));
94 }
95 }
96 }
97
98 self.stats.cache_misses.fetch_add(1, Ordering::Relaxed);
100 let body = compute_fn()?;
101 let etag = compute_etag(&body);
102 let body_arc: Arc<[u8]> = Arc::from(body);
103
104 let mut cache = self.cache.write().unwrap();
105 cache.insert(
106 key.to_string(),
107 CachedResponse {
108 body: Arc::clone(&body_arc),
109 etag: etag.clone(),
110 expires_at: Instant::now() + ttl,
111 },
112 );
113
114 Ok((body_arc, etag))
115 }
116
117 pub fn etag_matches(&self, key: &str, client_etag: &str) -> bool {
119 let cache = self.cache.read().unwrap();
120 if let Some(cached) = cache.get(key) {
121 if cached.expires_at > Instant::now() {
122 return cached.etag == client_etag;
123 }
124 }
125 false
126 }
127
128 pub fn record_etag_hit(&self) {
130 self.stats.etag_hits.fetch_add(1, Ordering::Relaxed);
131 }
132
133 pub fn stats(&self) -> CacheStats {
135 CacheStats {
136 total_requests: self.stats.total_requests.load(Ordering::Relaxed),
137 cache_hits: self.stats.cache_hits.load(Ordering::Relaxed),
138 cache_misses: self.stats.cache_misses.load(Ordering::Relaxed),
139 etag_hits: self.stats.etag_hits.load(Ordering::Relaxed),
140 }
141 }
142
143 pub fn clear(&self) {
145 self.cache.write().unwrap().clear();
146 self.stats.total_requests.store(0, Ordering::Relaxed);
147 self.stats.cache_hits.store(0, Ordering::Relaxed);
148 self.stats.cache_misses.store(0, Ordering::Relaxed);
149 self.stats.etag_hits.store(0, Ordering::Relaxed);
150 }
151
152 pub fn pre_warm<F>(&self, endpoints: Vec<(&str, Duration)>, compute_fn: F) -> Result<(), String>
154 where
155 F: Fn(&str) -> Result<Vec<u8>, String>,
156 {
157 for (key, ttl) in endpoints {
158 self.get_or_compute(key, ttl, || compute_fn(key))?;
159 }
160 Ok(())
161 }
162
163 pub fn size(&self) -> usize {
165 self.cache.read().unwrap().len()
166 }
167}
168
169impl Default for ApiResponseCache {
170 fn default() -> Self {
171 Self::new()
172 }
173}
174
175impl Clone for ApiResponseCache {
176 fn clone(&self) -> Self {
177 Self {
178 cache: Arc::clone(&self.cache),
179 stats: Arc::clone(&self.stats),
180 }
181 }
182}
183
184fn compute_etag(body: &[u8]) -> String {
194 let crc = crate::protocol::crc32(body);
198 format!("W/\"{:x}-{:x}\"", crc, body.len())
199}
200
201#[cfg(test)]
202mod tests {
203 use super::*;
204
205 #[test]
206 fn cache_miss_computes_value() {
207 let cache = ApiResponseCache::new();
208 let body = b"response data".to_vec();
209 let body_clone = body.clone();
210
211 let (result, etag) = cache
212 .get_or_compute("test", Duration::from_secs(60), || Ok(body_clone))
213 .unwrap();
214
215 assert_eq!(result.as_ref(), body.as_slice());
216 assert!(!etag.is_empty());
217 }
218
219 #[test]
220 fn cache_hit_reuses() {
221 let cache = ApiResponseCache::new();
222 let body1 = b"response1".to_vec();
223 let body1_clone = body1.clone();
224
225 let (result1, etag1) = cache
226 .get_or_compute("test", Duration::from_secs(60), || Ok(body1_clone))
227 .unwrap();
228
229 let (result2, etag2) = cache
230 .get_or_compute(
231 "test",
232 Duration::from_secs(60),
233 || Ok(b"response2".to_vec()),
234 )
235 .unwrap();
236
237 assert_eq!(result1, result2);
238 assert_eq!(etag1, etag2);
239 assert_eq!(cache.stats().cache_hits, 1);
240 }
241
242 #[test]
243 fn etag_match_detection() {
244 let cache = ApiResponseCache::new();
245 let body = b"data".to_vec();
246 let body_clone = body.clone();
247
248 let (_result, etag) = cache
249 .get_or_compute("test", Duration::from_secs(60), || Ok(body_clone))
250 .unwrap();
251
252 assert!(cache.etag_matches("test", &etag));
253 assert!(!cache.etag_matches("test", "wrong-etag"));
254 }
255
256 #[test]
257 fn expiration_and_recompute() {
258 let cache = ApiResponseCache::new();
259
260 cache
262 .get_or_compute("test", Duration::from_millis(5), || Ok(b"data1".to_vec()))
263 .unwrap();
264
265 std::thread::sleep(Duration::from_millis(50));
267
268 let (result2, _etag2) = cache
270 .get_or_compute("test", Duration::from_secs(60), || Ok(b"data2".to_vec()))
271 .unwrap();
272
273 assert_eq!(result2.as_ref(), b"data2");
274 }
275
276 #[test]
277 fn statistics_tracking() {
278 let cache = ApiResponseCache::new();
279
280 for i in 0..5 {
281 let body = format!("response{}", i).into_bytes();
282 cache
283 .get_or_compute(&format!("endpoint{}", i), Duration::from_secs(60), || {
284 Ok(body)
285 })
286 .unwrap();
287 }
288
289 let stats = cache.stats();
290 assert_eq!(stats.total_requests, 5);
291 assert_eq!(stats.cache_misses, 5);
292 assert_eq!(stats.cache_hits, 0);
293 }
294
295 #[test]
296 fn hit_ratio_calculation() {
297 let cache = ApiResponseCache::new();
298
299 for _ in 0..3 {
300 cache
301 .get_or_compute("test", Duration::from_secs(60), || Ok(b"data".to_vec()))
302 .unwrap();
303 cache
304 .get_or_compute("test", Duration::from_secs(60), || Ok(b"data".to_vec()))
305 .unwrap();
306 }
307
308 let stats = cache.stats();
309 let ratio = stats.hit_ratio();
310 assert!(ratio > 0.5);
311 }
312
313 #[test]
314 fn pre_warm_populates() {
315 let cache = ApiResponseCache::new();
316
317 cache
318 .pre_warm(
319 vec![
320 ("models", Duration::from_secs(60)),
321 ("runtime", Duration::from_secs(300)),
322 ],
323 |_key| Ok(b"warm data".to_vec()),
324 )
325 .unwrap();
326
327 assert_eq!(cache.size(), 2);
328 }
329
330 #[test]
331 fn clone_shares_state() {
332 let cache1 = ApiResponseCache::new();
333 cache1
334 .get_or_compute("test", Duration::from_secs(60), || Ok(b"data".to_vec()))
335 .unwrap();
336
337 let cache2 = cache1.clone();
338 assert_eq!(cache1.size(), cache2.size());
339 }
340}