1use anyhow::Result;
42use serde::{Deserialize, Serialize};
43use std::collections::{HashMap, VecDeque};
44use std::sync::{Arc, RwLock};
45use std::time::{Duration, SystemTime};
46
47use crate::cache::CacheCoordinator;
49
50pub struct QueryResultCache {
52 entries: Arc<RwLock<HashMap<String, CacheEntry>>>,
54 lru_queue: Arc<RwLock<VecDeque<String>>>,
56 config: CacheConfig,
58 stats: Arc<RwLock<CacheStatistics>>,
60 invalidation_coordinator: Option<Arc<CacheCoordinator>>,
62 invalidated_entries: Arc<RwLock<std::collections::HashSet<String>>>,
64}
65
66#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct CacheConfig {
69 pub max_entries: usize,
71 pub ttl: Duration,
73 pub enable_compression: bool,
75 pub max_result_size: usize,
77 pub enable_stats: bool,
79 pub eviction_batch_size: usize,
81}
82
83impl Default for CacheConfig {
84 fn default() -> Self {
85 Self {
86 max_entries: 10_000,
87 ttl: Duration::from_secs(3600), enable_compression: false,
89 max_result_size: 10 * 1024 * 1024, enable_stats: true,
91 eviction_batch_size: 100,
92 }
93 }
94}
95
96impl CacheConfig {
97 pub fn with_max_entries(mut self, max: usize) -> Self {
99 self.max_entries = max;
100 self
101 }
102
103 pub fn with_ttl(mut self, ttl: Duration) -> Self {
105 self.ttl = ttl;
106 self
107 }
108
109 pub fn with_compression(mut self, enabled: bool) -> Self {
111 self.enable_compression = enabled;
112 self
113 }
114
115 pub fn with_max_result_size(mut self, size: usize) -> Self {
117 self.max_result_size = size;
118 self
119 }
120}
121
122#[derive(Debug, Clone, Serialize, Deserialize)]
124struct CacheEntry {
125 fingerprint_hash: String,
127 results: Vec<u8>,
129 original_size: usize,
131 created_at: SystemTime,
133 last_accessed: SystemTime,
135 access_count: u64,
137 is_compressed: bool,
139}
140
141#[derive(Debug, Clone, Default, Serialize, Deserialize)]
143pub struct CacheStatistics {
144 pub hits: u64,
146 pub misses: u64,
148 pub puts: u64,
150 pub evictions: u64,
152 pub invalidations: u64,
154 pub size_bytes: usize,
156 pub entry_count: usize,
158 pub hit_rate: f64,
160 pub avg_result_size: usize,
162 pub compression_ratio: f64,
164}
165
166impl CacheStatistics {
167 fn calculate_hit_rate(&mut self) {
169 let total = self.hits + self.misses;
170 self.hit_rate = if total > 0 {
171 self.hits as f64 / total as f64
172 } else {
173 0.0
174 };
175 }
176}
177
178impl QueryResultCache {
179 pub fn new(config: CacheConfig) -> Self {
181 Self {
182 entries: Arc::new(RwLock::new(HashMap::new())),
183 lru_queue: Arc::new(RwLock::new(VecDeque::new())),
184 config,
185 stats: Arc::new(RwLock::new(CacheStatistics::default())),
186 invalidation_coordinator: None,
187 invalidated_entries: Arc::new(RwLock::new(std::collections::HashSet::new())),
188 }
189 }
190
191 pub fn with_invalidation_coordinator(
193 config: CacheConfig,
194 coordinator: Arc<CacheCoordinator>,
195 ) -> Self {
196 Self {
197 entries: Arc::new(RwLock::new(HashMap::new())),
198 lru_queue: Arc::new(RwLock::new(VecDeque::new())),
199 config,
200 stats: Arc::new(RwLock::new(CacheStatistics::default())),
201 invalidation_coordinator: Some(coordinator),
202 invalidated_entries: Arc::new(RwLock::new(std::collections::HashSet::new())),
203 }
204 }
205
206 pub fn attach_coordinator(&mut self, coordinator: Arc<CacheCoordinator>) {
208 self.invalidation_coordinator = Some(coordinator);
209 }
210
211 pub fn put(&self, fingerprint_hash: String, results: Vec<u8>) -> Result<()> {
213 if results.len() > self.config.max_result_size {
215 return Ok(()); }
217
218 let mut entries = self.entries.write().expect("lock poisoned");
219 let mut lru = self.lru_queue.write().expect("lock poisoned");
220
221 if entries.len() >= self.config.max_entries {
223 self.evict_lru(&mut entries, &mut lru)?;
224 }
225
226 let (stored_results, is_compressed) = if self.config.enable_compression {
228 match self.compress_results(&results) {
229 Ok(compressed) => (compressed, true),
230 Err(_) => (results.clone(), false),
231 }
232 } else {
233 (results.clone(), false)
234 };
235
236 let entry = CacheEntry {
237 fingerprint_hash: fingerprint_hash.clone(),
238 results: stored_results.clone(),
239 original_size: results.len(),
240 created_at: SystemTime::now(),
241 last_accessed: SystemTime::now(),
242 access_count: 0,
243 is_compressed,
244 };
245
246 entries.insert(fingerprint_hash.clone(), entry);
248 lru.push_back(fingerprint_hash);
249
250 if self.config.enable_stats {
252 let mut stats = self.stats.write().expect("lock poisoned");
253 stats.puts += 1;
254 stats.entry_count = entries.len();
255 stats.size_bytes += stored_results.len();
256 stats.avg_result_size = stats.size_bytes.checked_div(stats.entry_count).unwrap_or(0);
257 }
258
259 Ok(())
260 }
261
262 pub fn get(&self, fingerprint_hash: &str) -> Option<Vec<u8>> {
264 {
266 let invalidated = self.invalidated_entries.read().expect("lock poisoned");
267 if invalidated.contains(fingerprint_hash) {
268 if self.config.enable_stats {
270 let mut stats = self.stats.write().expect("lock poisoned");
271 stats.misses += 1;
272 stats.invalidations += 1;
273 stats.calculate_hit_rate();
274 }
275 return None;
276 }
277 }
278
279 let mut entries = self.entries.write().expect("lock poisoned");
280 let mut lru = self.lru_queue.write().expect("lock poisoned");
281
282 if let Some(entry) = entries.get_mut(fingerprint_hash) {
283 if let Ok(elapsed) = entry.created_at.elapsed() {
285 if elapsed > self.config.ttl {
286 entries.remove(fingerprint_hash);
288 lru.retain(|k| k != fingerprint_hash);
289
290 if self.config.enable_stats {
292 let mut stats = self.stats.write().expect("lock poisoned");
293 stats.misses += 1;
294 stats.evictions += 1;
295 stats.calculate_hit_rate();
296 }
297 return None;
298 }
299 }
300
301 entry.last_accessed = SystemTime::now();
303 entry.access_count += 1;
304
305 lru.retain(|k| k != fingerprint_hash);
307 lru.push_back(fingerprint_hash.to_string());
308
309 let results = if entry.is_compressed {
311 self.decompress_results(&entry.results).ok()?
312 } else {
313 entry.results.clone()
314 };
315
316 if self.config.enable_stats {
318 let mut stats = self.stats.write().expect("lock poisoned");
319 stats.hits += 1;
320 stats.calculate_hit_rate();
321 }
322
323 Some(results)
324 } else {
325 if self.config.enable_stats {
327 let mut stats = self.stats.write().expect("lock poisoned");
328 stats.misses += 1;
329 stats.calculate_hit_rate();
330 }
331 None
332 }
333 }
334
335 pub fn invalidate(&self, fingerprint_hash: &str) -> Result<()> {
337 {
339 let mut invalidated = self.invalidated_entries.write().expect("lock poisoned");
340 invalidated.insert(fingerprint_hash.to_string());
341 }
342
343 let mut entries = self.entries.write().expect("lock poisoned");
344 let mut lru = self.lru_queue.write().expect("lock poisoned");
345
346 if entries.remove(fingerprint_hash).is_some() {
347 lru.retain(|k| k != fingerprint_hash);
348
349 if self.config.enable_stats {
350 let mut stats = self.stats.write().expect("lock poisoned");
351 stats.invalidations += 1;
352 stats.entry_count = entries.len();
353 }
354 }
355
356 Ok(())
357 }
358
359 pub fn mark_invalidated(&self, fingerprint_hash: &str) -> Result<()> {
361 let mut invalidated = self.invalidated_entries.write().expect("lock poisoned");
362 invalidated.insert(fingerprint_hash.to_string());
363
364 if self.config.enable_stats {
365 let mut stats = self.stats.write().expect("lock poisoned");
366 stats.invalidations += 1;
367 }
368
369 Ok(())
370 }
371
372 pub fn invalidate_all(&self) -> Result<()> {
374 let mut entries = self.entries.write().expect("lock poisoned");
375 let mut lru = self.lru_queue.write().expect("lock poisoned");
376 let mut invalidated = self.invalidated_entries.write().expect("lock poisoned");
377
378 let count = entries.len();
379
380 for key in entries.keys() {
382 invalidated.insert(key.clone());
383 }
384
385 entries.clear();
386 lru.clear();
387
388 if self.config.enable_stats {
389 let mut stats = self.stats.write().expect("lock poisoned");
390 stats.invalidations += count as u64;
391 stats.entry_count = 0;
392 stats.size_bytes = 0;
393 }
394
395 Ok(())
396 }
397
398 pub fn statistics(&self) -> CacheStatistics {
400 self.stats.read().expect("lock poisoned").clone()
401 }
402
403 pub fn size(&self) -> usize {
405 self.entries.read().expect("lock poisoned").len()
406 }
407
408 pub fn contains(&self, fingerprint_hash: &str) -> bool {
410 self.entries
411 .read()
412 .expect("lock poisoned")
413 .contains_key(fingerprint_hash)
414 }
415
416 fn evict_lru(
418 &self,
419 entries: &mut HashMap<String, CacheEntry>,
420 lru: &mut VecDeque<String>,
421 ) -> Result<()> {
422 let batch_size = self.config.eviction_batch_size.min(entries.len() / 10 + 1);
423
424 for _ in 0..batch_size {
425 if let Some(oldest) = lru.pop_front() {
426 if let Some(entry) = entries.remove(&oldest) {
427 if self.config.enable_stats {
428 let mut stats = self.stats.write().expect("lock poisoned");
429 stats.evictions += 1;
430 stats.size_bytes = stats.size_bytes.saturating_sub(entry.results.len());
431 stats.entry_count = entries.len();
432 }
433 }
434 }
435 }
436
437 Ok(())
438 }
439
440 fn compress_results(&self, results: &[u8]) -> Result<Vec<u8>> {
442 Ok(oxiarc_deflate::gzip_compress(results, 1)?)
444 }
445
446 fn decompress_results(&self, compressed: &[u8]) -> Result<Vec<u8>> {
448 Ok(oxiarc_deflate::gzip_decompress(compressed)?)
449 }
450}
451
452pub struct QueryResultCacheBuilder {
454 config: CacheConfig,
455}
456
457impl QueryResultCacheBuilder {
458 pub fn new() -> Self {
460 Self {
461 config: CacheConfig::default(),
462 }
463 }
464
465 pub fn max_entries(mut self, max: usize) -> Self {
467 self.config.max_entries = max;
468 self
469 }
470
471 pub fn ttl(mut self, ttl: Duration) -> Self {
473 self.config.ttl = ttl;
474 self
475 }
476
477 pub fn compression(mut self, enabled: bool) -> Self {
479 self.config.enable_compression = enabled;
480 self
481 }
482
483 pub fn build(self) -> QueryResultCache {
485 QueryResultCache::new(self.config)
486 }
487}
488
489impl Default for QueryResultCacheBuilder {
490 fn default() -> Self {
491 Self::new()
492 }
493}
494
495#[cfg(test)]
496mod tests {
497 use super::*;
498
499 #[test]
500 fn test_cache_basic_operations() {
501 let cache = QueryResultCache::new(CacheConfig::default());
502
503 let hash = "test_hash_123".to_string();
504 let results = vec![1, 2, 3, 4, 5];
505
506 cache.put(hash.clone(), results.clone()).unwrap();
508 let retrieved = cache.get(&hash).unwrap();
509 assert_eq!(results, retrieved);
510
511 let stats = cache.statistics();
513 assert_eq!(stats.puts, 1);
514 assert_eq!(stats.hits, 1);
515 assert_eq!(stats.misses, 0);
516 }
517
518 #[test]
519 fn test_gzip_compress_decompress_roundtrip() {
520 let cache = QueryResultCache::new(CacheConfig::default());
523
524 let original: Vec<u8> = b"oxirs-arq result cache gzip payload "
526 .iter()
527 .cycle()
528 .take(4096)
529 .copied()
530 .collect();
531
532 let compressed = cache
533 .compress_results(&original)
534 .expect("gzip compression failed");
535 assert_eq!(&compressed[..2], &[0x1f, 0x8b]);
537 assert!(compressed.len() < original.len());
538
539 let decompressed = cache
540 .decompress_results(&compressed)
541 .expect("gzip decompression failed");
542 assert_eq!(decompressed, original);
543 }
544
545 #[test]
546 fn test_gzip_roundtrip_empty_and_small() {
547 let cache = QueryResultCache::new(CacheConfig::default());
548 for original in [Vec::new(), vec![42u8], b"hello".to_vec()] {
549 let compressed = cache
550 .compress_results(&original)
551 .expect("gzip compression failed");
552 let decompressed = cache
553 .decompress_results(&compressed)
554 .expect("gzip decompression failed");
555 assert_eq!(decompressed, original);
556 }
557 }
558
559 #[test]
560 fn test_cache_miss() {
561 let cache = QueryResultCache::new(CacheConfig::default());
562
563 let result = cache.get("nonexistent");
564 assert!(result.is_none());
565
566 let stats = cache.statistics();
567 assert_eq!(stats.misses, 1);
568 }
569
570 #[test]
571 fn test_cache_invalidation() {
572 let cache = QueryResultCache::new(CacheConfig::default());
573
574 let hash = "test_hash".to_string();
575 let results = vec![1, 2, 3];
576
577 cache.put(hash.clone(), results).unwrap();
578 assert!(cache.contains(&hash));
579
580 cache.invalidate(&hash).unwrap();
581 assert!(!cache.contains(&hash));
582
583 let stats = cache.statistics();
584 assert_eq!(stats.invalidations, 1);
585 }
586
587 #[test]
588 fn test_lru_eviction() {
589 let config = CacheConfig::default().with_max_entries(3);
590 let cache = QueryResultCache::new(config);
591
592 cache.put("hash1".to_string(), vec![1]).unwrap();
594 cache.put("hash2".to_string(), vec![2]).unwrap();
595 cache.put("hash3".to_string(), vec![3]).unwrap();
596
597 cache.put("hash4".to_string(), vec![4]).unwrap();
599
600 assert!(!cache.contains("hash1"));
602 assert!(cache.contains("hash4"));
603 }
604
605 #[test]
606 fn test_cache_compression() {
607 let config = CacheConfig::default().with_compression(true);
608 let cache = QueryResultCache::new(config);
609
610 let hash = "compressed_hash".to_string();
611 let large_results = vec![0u8; 10_000]; cache.put(hash.clone(), large_results.clone()).unwrap();
614 let retrieved = cache.get(&hash).unwrap();
615 assert_eq!(large_results, retrieved);
616
617 let stats = cache.statistics();
618 assert!(stats.compression_ratio > 1.0 || stats.size_bytes < large_results.len());
619 }
620
621 #[test]
622 fn test_cache_ttl_expiration() {
623 use std::thread;
624
625 let config = CacheConfig::default().with_ttl(Duration::from_millis(100));
626 let cache = QueryResultCache::new(config);
627
628 let hash = "expiring_hash".to_string();
629 cache.put(hash.clone(), vec![1, 2, 3]).unwrap();
630
631 assert!(cache.get(&hash).is_some());
633
634 thread::sleep(Duration::from_millis(150));
636
637 assert!(cache.get(&hash).is_none());
639 }
640
641 #[test]
642 fn test_cache_builder() {
643 let cache = QueryResultCacheBuilder::new()
644 .max_entries(5000)
645 .ttl(Duration::from_secs(1800))
646 .compression(true)
647 .build();
648
649 assert_eq!(cache.config.max_entries, 5000);
650 assert_eq!(cache.config.ttl, Duration::from_secs(1800));
651 assert!(cache.config.enable_compression);
652 }
653
654 #[test]
655 fn test_cache_statistics_accuracy() {
656 let cache = QueryResultCache::new(CacheConfig::default());
657
658 cache.put("h1".to_string(), vec![1]).unwrap();
660 cache.put("h2".to_string(), vec![2]).unwrap();
661 cache.get("h1"); cache.get("h3"); cache.invalidate("h1").unwrap();
664
665 let stats = cache.statistics();
666 assert_eq!(stats.puts, 2);
667 assert_eq!(stats.hits, 1);
668 assert_eq!(stats.misses, 1);
669 assert_eq!(stats.invalidations, 1);
670 assert_eq!(stats.hit_rate, 0.5);
671 }
672
673 #[test]
674 fn test_cache_max_result_size() {
675 let config = CacheConfig::default().with_max_result_size(100);
676 let cache = QueryResultCache::new(config);
677
678 cache.put("small".to_string(), vec![1; 50]).unwrap();
680 assert!(cache.contains("small"));
681
682 cache.put("large".to_string(), vec![1; 200]).unwrap();
684 assert!(!cache.contains("large"));
685 }
686
687 #[test]
688 fn test_cache_access_tracking() {
689 let cache = QueryResultCache::new(CacheConfig::default());
690
691 let hash = "tracked".to_string();
692 cache.put(hash.clone(), vec![1, 2, 3]).unwrap();
693
694 for _ in 0..5 {
696 cache.get(&hash);
697 }
698
699 let entries = cache.entries.read().unwrap_or_else(|e| e.into_inner());
700 let entry = entries.get(&hash).unwrap();
701 assert_eq!(entry.access_count, 5);
702 }
703}