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