1use crate::epoch::{Epoch, Snapshot};
19use crate::page::CachedPage;
20use std::collections::HashMap;
21use std::path::PathBuf;
22use std::sync::atomic::{AtomicU64, Ordering};
23
24#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
29pub struct CacheStats {
30 pub hits: u64,
31 pub misses: u64,
32 pub try_lock_misses: u64,
36}
37
38impl CacheStats {
39 pub fn hit_rate(&self) -> f64 {
41 let total = self.hits + self.misses;
42 if total == 0 {
43 0.0
44 } else {
45 self.hits as f64 / total as f64
46 }
47 }
48}
49
50pub struct PageCache {
53 map: HashMap<[u8; 32], Entry>,
54 ring: Vec<[u8; 32]>,
56 hand: usize,
57 capacity_bytes: u64,
58 used_bytes: u64,
59 dir: Option<PathBuf>,
62 persistent: bool,
63 hits: AtomicU64,
65 misses: AtomicU64,
67}
68
69struct Entry {
70 bytes: bytes::Bytes,
71 epoch: Epoch,
72 freq: u8,
73}
74
75const FREQ_MAX: u8 = 3;
76
77impl PageCache {
78 pub fn new(capacity_bytes: u64) -> Self {
79 Self {
80 map: HashMap::new(),
81 ring: Vec::new(),
82 hand: 0,
83 capacity_bytes,
84 used_bytes: 0,
85 dir: None,
86 persistent: false,
87 hits: AtomicU64::new(0),
88 misses: AtomicU64::new(0),
89 }
90 }
91
92 pub fn stats(&self) -> CacheStats {
95 CacheStats {
96 hits: self.hits.load(Ordering::Relaxed),
97 misses: self.misses.load(Ordering::Relaxed),
98 try_lock_misses: 0,
99 }
100 }
101
102 pub fn reset_stats(&self) {
104 self.hits.store(0, Ordering::Relaxed);
105 self.misses.store(0, Ordering::Relaxed);
106 }
107
108 pub fn with_persistence(mut self, dir: PathBuf) -> Self {
112 self.persistent = true;
113 self.dir = Some(dir.clone());
114 self.load_from_disk();
115 self
116 }
117
118 pub fn insert(&mut self, page: CachedPage) {
121 let key = page.content_hash;
122 let size = page.bytes.len() as u64;
123 let epoch = page.committed_epoch;
124 let was_new = !self.map.contains_key(&key);
125 if let Some(old) = self.map.insert(
126 key,
127 Entry {
128 bytes: page.bytes.clone(),
129 epoch,
130 freq: 1,
131 },
132 ) {
133 self.used_bytes = self.used_bytes.saturating_sub(old.bytes.len() as u64);
134 } else if was_new {
135 self.ring.push(key);
136 }
137 self.used_bytes += size;
138 self.evict_if_needed();
139 }
140
141 pub fn get(&mut self, content_hash: &[u8; 32], snapshot: Snapshot) -> Option<bytes::Bytes> {
144 let hit = self
145 .map
146 .get_mut(content_hash)
147 .filter(|e| e.epoch <= snapshot.epoch)
148 .map(|entry| {
149 if entry.freq < FREQ_MAX {
150 entry.freq += 1;
151 }
152 entry.bytes.clone()
153 });
154 let counter = if hit.is_some() {
155 &self.hits
156 } else {
157 &self.misses
158 };
159 counter.fetch_add(1, Ordering::Relaxed);
160 hit
161 }
162
163 pub fn try_get(&self, content_hash: &[u8; 32], snapshot: Snapshot) -> Option<bytes::Bytes> {
168 let hit = self.map.get(content_hash).and_then(|e| {
169 if e.epoch <= snapshot.epoch {
170 Some(e.bytes.clone())
171 } else {
172 None
173 }
174 });
175 let counter = if hit.is_some() {
176 &self.hits
177 } else {
178 &self.misses
179 };
180 counter.fetch_add(1, Ordering::Relaxed);
181 hit
182 }
183
184 pub fn used_bytes(&self) -> u64 {
185 self.used_bytes
186 }
187
188 pub fn len(&self) -> usize {
189 self.map.len()
190 }
191
192 pub fn is_empty(&self) -> bool {
193 self.map.is_empty()
194 }
195
196 pub fn flush_to_disk(&self) {
198 if !self.persistent {
199 return;
200 }
201 for (key, entry) in &self.map {
202 self.spill(key, entry.epoch, &entry.bytes);
203 }
204 }
205
206 fn spill(&self, key: &[u8; 32], epoch: Epoch, bytes: &[u8]) {
207 if !self.persistent {
208 return;
209 }
210 let Some(dir) = &self.dir else { return };
211 let _ = std::fs::create_dir_all(dir);
212 let hex = hex_key(key);
213 let path = dir.join(format!("{hex}.{}", epoch.0));
219 let tmp = dir.join(format!("{hex}.{}.tmp", epoch.0));
220 if std::fs::write(&tmp, bytes).is_ok() {
221 let _ = std::fs::rename(&tmp, &path);
222 }
223 }
224
225 fn load_from_disk(&mut self) {
226 let Some(dir) = &self.dir else { return };
227 let entries = match std::fs::read_dir(dir) {
228 Ok(e) => e,
229 Err(_) => return,
230 };
231 for entry in entries.flatten() {
232 let name = entry.file_name();
233 let Some(s) = name.to_str() else { continue };
234 if s.ends_with(".tmp") {
236 continue;
237 }
238 let (key, epoch) = match s.rsplit_once('.') {
241 Some((hex, suffix)) => {
242 let key = match decode_hex_key(hex) {
243 Some(k) => k,
244 None => continue,
245 };
246 let epoch = match suffix.parse::<u64>() {
247 Ok(e) => Epoch(e),
248 Err(_) => continue,
249 };
250 (key, epoch)
251 }
252 None => {
253 let key = match decode_hex_key(s) {
254 Some(k) => k,
255 None => continue,
256 };
257 (key, Epoch(u64::MAX))
258 }
259 };
260 if self.map.contains_key(&key) {
261 continue;
262 }
263 let Ok(bytes) = std::fs::read(entry.path()) else {
264 continue;
265 };
266 let size = bytes.len() as u64;
267 if self.used_bytes.saturating_add(size) > self.capacity_bytes {
268 break; }
270 self.map.insert(
271 key,
272 Entry {
273 bytes: bytes::Bytes::from(bytes),
274 epoch,
275 freq: 0,
276 },
277 );
278 self.ring.push(key);
279 self.used_bytes += size;
280 }
281 }
282
283 fn evict_if_needed(&mut self) {
286 if self.ring.is_empty() {
287 return;
288 }
289 while self.used_bytes > self.capacity_bytes {
290 let mut scanned = 0;
292 let len = self.ring.len();
293 loop {
294 let key = self.ring[self.hand];
295 let Some(entry) = self.map.get_mut(&key) else {
296 self.ring.swap_remove(self.hand);
298 if self.hand >= self.ring.len() && !self.ring.is_empty() {
299 self.hand %= self.ring.len();
300 }
301 break;
302 };
303 if entry.freq == 0 {
304 let removed = self.map.remove(&key).expect("entry present");
306 self.spill(&key, removed.epoch, &removed.bytes);
307 self.used_bytes = self.used_bytes.saturating_sub(removed.bytes.len() as u64);
308 self.ring.swap_remove(self.hand);
309 if !self.ring.is_empty() {
310 self.hand %= self.ring.len();
311 }
312 break;
313 } else {
314 entry.freq -= 1;
315 scanned += 1;
316 if scanned > len * 2 {
317 let removed = self.map.remove(&key).expect("entry present");
321 self.spill(&key, removed.epoch, &removed.bytes);
322 self.used_bytes =
323 self.used_bytes.saturating_sub(removed.bytes.len() as u64);
324 self.ring.swap_remove(self.hand);
325 if !self.ring.is_empty() {
326 self.hand %= self.ring.len();
327 }
328 break;
329 }
330 self.hand = (self.hand + 1) % len;
331 }
332 }
333 if self.ring.is_empty() {
334 break;
335 }
336 }
337 }
338}
339
340fn hex_key(key: &[u8; 32]) -> String {
341 let mut s = String::with_capacity(64);
342 for b in key {
343 s.push_str(&format!("{b:02x}"));
344 }
345 s
346}
347
348fn decode_hex_key(hex: &str) -> Option<[u8; 32]> {
349 if hex.len() != 64 {
350 return None;
351 }
352 let mut out = [0u8; 32];
353 for i in 0..32 {
354 out[i] = u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16).ok()?;
355 }
356 Some(out)
357}
358
359pub struct DecodedPageCache {
367 map: HashMap<[u8; 32], std::sync::Arc<crate::columnar::NativeColumn>>,
368 ring: Vec<[u8; 32]>,
369 hand: usize,
370 capacity_bytes: u64,
371 used_bytes: u64,
372 hits: AtomicU64,
374 misses: AtomicU64,
376}
377
378impl DecodedPageCache {
379 pub fn new(capacity_bytes: u64) -> Self {
380 Self {
381 map: HashMap::new(),
382 ring: Vec::new(),
383 hand: 0,
384 capacity_bytes,
385 used_bytes: 0,
386 hits: AtomicU64::new(0),
387 misses: AtomicU64::new(0),
388 }
389 }
390
391 pub fn stats(&self) -> CacheStats {
394 CacheStats {
395 hits: self.hits.load(Ordering::Relaxed),
396 misses: self.misses.load(Ordering::Relaxed),
397 try_lock_misses: 0,
398 }
399 }
400
401 pub fn reset_stats(&self) {
403 self.hits.store(0, Ordering::Relaxed);
404 self.misses.store(0, Ordering::Relaxed);
405 }
406
407 pub fn try_get(&self, key: &[u8; 32]) -> Option<std::sync::Arc<crate::columnar::NativeColumn>> {
410 match self.map.get(key).cloned() {
411 Some(v) => {
412 self.hits.fetch_add(1, Ordering::Relaxed);
413 Some(v)
414 }
415 None => {
416 self.misses.fetch_add(1, Ordering::Relaxed);
417 None
418 }
419 }
420 }
421
422 pub fn insert(&mut self, key: [u8; 32], col: std::sync::Arc<crate::columnar::NativeColumn>) {
425 let size = col.approx_bytes();
426 let was_new = !self.map.contains_key(&key);
427 if let Some(old) = self.map.insert(key, col) {
428 self.used_bytes = self.used_bytes.saturating_sub(old.approx_bytes());
429 } else if was_new {
430 self.ring.push(key);
431 }
432 self.used_bytes += size;
433 self.evict_if_needed();
434 }
435
436 fn evict_if_needed(&mut self) {
437 while self.used_bytes > self.capacity_bytes && !self.ring.is_empty() {
442 let idx = self.hand.min(self.ring.len() - 1);
443 let key = self.ring.swap_remove(idx);
444 if let Some(col) = self.map.remove(&key) {
445 self.used_bytes = self.used_bytes.saturating_sub(col.approx_bytes());
446 }
447 if self.ring.is_empty() {
448 self.hand = 0;
449 } else {
450 self.hand %= self.ring.len();
451 }
452 }
453 }
454
455 pub fn len(&self) -> usize {
456 self.map.len()
457 }
458
459 pub fn is_empty(&self) -> bool {
460 self.map.is_empty()
461 }
462
463 pub fn used_bytes(&self) -> u64 {
464 self.used_bytes
465 }
466}
467
468pub struct Sharded<T> {
473 shards: Vec<parking_lot::Mutex<T>>,
474 try_lock_misses: AtomicU64,
475}
476
477pub const CACHE_SHARDS: usize = 16;
480
481impl<T> Sharded<T> {
482 pub fn new(n_shards: usize, make: impl FnMut() -> T) -> Self {
483 let mut make = make;
484 let shards = (0..n_shards)
485 .map(|_| parking_lot::Mutex::new(make()))
486 .collect();
487 Self {
488 shards,
489 try_lock_misses: AtomicU64::new(0),
490 }
491 }
492
493 fn idx(&self, key: &[u8; 32]) -> usize {
494 let h = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
495 (h as usize) % self.shards.len()
496 }
497
498 pub fn try_lock(&self, key: &[u8; 32]) -> Option<parking_lot::MutexGuard<'_, T>> {
499 match self.shards[self.idx(key)].try_lock() {
500 Some(g) => Some(g),
501 None => {
502 self.try_lock_misses.fetch_add(1, Ordering::Relaxed);
503 None
504 }
505 }
506 }
507
508 pub fn lock(&self, key: &[u8; 32]) -> parking_lot::MutexGuard<'_, T> {
509 self.shards[self.idx(key)].lock()
510 }
511
512 pub fn try_lock_misses(&self) -> u64 {
513 self.try_lock_misses.load(Ordering::Relaxed)
514 }
515}
516
517impl Sharded<PageCache> {
518 pub fn stats(&self) -> CacheStats {
519 let mut total = CacheStats::default();
520 for s in &self.shards {
521 let st = s.lock().stats();
522 total.hits += st.hits;
523 total.misses += st.misses;
524 }
525 total.try_lock_misses = self.try_lock_misses();
526 total
527 }
528
529 pub fn reset_stats(&self) {
530 for s in &self.shards {
531 s.lock().reset_stats();
532 }
533 self.try_lock_misses.store(0, Ordering::Relaxed);
534 }
535
536 pub fn flush_to_disk(&self) {
537 for s in &self.shards {
538 s.lock().flush_to_disk();
539 }
540 }
541
542 pub fn len(&self) -> usize {
543 self.shards.iter().map(|s| s.lock().len()).sum()
544 }
545
546 #[allow(clippy::len_without_is_empty)]
547 pub fn is_empty(&self) -> bool {
548 self.len() == 0
549 }
550}
551
552impl Sharded<DecodedPageCache> {
553 pub fn stats(&self) -> CacheStats {
554 let mut total = CacheStats::default();
555 for s in &self.shards {
556 let st = s.lock().stats();
557 total.hits += st.hits;
558 total.misses += st.misses;
559 }
560 total.try_lock_misses = self.try_lock_misses();
561 total
562 }
563
564 pub fn reset_stats(&self) {
565 for s in &self.shards {
566 s.lock().reset_stats();
567 }
568 self.try_lock_misses.store(0, Ordering::Relaxed);
569 }
570
571 pub fn len(&self) -> usize {
572 self.shards.iter().map(|s| s.lock().len()).sum()
573 }
574}
575
576#[cfg(test)]
577mod tests {
578 use super::*;
579 use tempfile::tempdir;
580
581 fn page(hash: [u8; 32], epoch: u64, data: &[u8]) -> CachedPage {
582 CachedPage {
583 committed_epoch: Epoch(epoch),
584 content_hash: hash,
585 bytes: bytes::Bytes::copy_from_slice(data),
586 }
587 }
588
589 #[test]
590 fn mvcc_visibility() {
591 let mut cache = PageCache::new(1 << 20);
592 let hash = [1u8; 32];
593 cache.insert(page(hash, 3, b"v3"));
594 assert!(cache.get(&hash, Snapshot::at(Epoch(2))).is_none());
596 assert_eq!(
598 cache.get(&hash, Snapshot::at(Epoch(3))),
599 Some(bytes::Bytes::copy_from_slice(b"v3"))
600 );
601 }
602
603 #[test]
604 fn hit_miss_counters_track_lookups() {
605 let mut cache = PageCache::new(1 << 20);
606 let hash = [7u8; 32];
607 cache.insert(page(hash, 1, b"v"));
608 assert_eq!(
609 cache.stats(),
610 CacheStats {
611 hits: 0,
612 misses: 0,
613 try_lock_misses: 0
614 }
615 );
616
617 assert!(cache.get(&hash, Snapshot::at(Epoch(1))).is_some());
619 assert!(cache.try_get(&hash, Snapshot::at(Epoch(1))).is_some());
620 assert!(cache.get(&[9u8; 32], Snapshot::at(Epoch(1))).is_none());
622 assert!(cache.get(&hash, Snapshot::at(Epoch(0))).is_none());
623
624 let s = cache.stats();
625 assert_eq!(
626 s,
627 CacheStats {
628 hits: 2,
629 misses: 2,
630 try_lock_misses: 0
631 }
632 );
633 assert!((s.hit_rate() - 0.5).abs() < 1e-9);
634
635 cache.reset_stats();
636 assert_eq!(
637 cache.stats(),
638 CacheStats {
639 hits: 0,
640 misses: 0,
641 try_lock_misses: 0
642 }
643 );
644 assert_eq!(cache.stats().hit_rate(), 0.0);
645 }
646
647 #[test]
648 fn content_addressed_replacement_ages_out_old() {
649 let mut cache = PageCache::new(1 << 20);
650 let hash_a = [1u8; 32];
651 let hash_b = [2u8; 32];
652 cache.insert(page(hash_a, 1, b"old"));
653 cache.insert(page(hash_b, 1, b"new"));
654 assert_eq!(cache.len(), 2);
655 assert!(cache.get(&hash_a, Snapshot::at(Epoch(1))).is_some());
656 }
657
658 #[test]
659 fn capacity_eviction_removes_cold_first() {
660 let mut cache = PageCache::new(10);
661 cache.insert(page([1u8; 32], 1, b"0123456789")); assert_eq!(cache.used_bytes(), 10);
663 cache.insert(page([2u8; 32], 1, b"x")); assert_eq!(cache.used_bytes(), 1);
665 assert!(cache.get(&[1u8; 32], Snapshot::at(Epoch(1))).is_none());
666 assert!(cache.get(&[2u8; 32], Snapshot::at(Epoch(1))).is_some());
667 }
668
669 #[test]
670 fn frequency_keeps_hot_pages_alive() {
671 let mut cache = PageCache::new(2);
674 let a = [10u8; 32];
675 cache.insert(page(a, 1, b"A"));
676 for _ in 0..5 {
678 let _ = cache.get(&a, Snapshot::at(Epoch(1)));
679 }
680 cache.insert(page([20u8; 32], 1, b"B"));
681 cache.insert(page([30u8; 32], 1, b"C"));
682 assert!(
684 cache.get(&a, Snapshot::at(Epoch(1))).is_some(),
685 "hot page should survive eviction"
686 );
687 }
688
689 #[test]
690 fn persistent_cache_survives_restart() {
691 let dir = tempdir().unwrap();
692 let data = b"page-payload-1234";
693 let hash = [42u8; 32];
694
695 {
697 let mut cache = PageCache::new(1 << 20).with_persistence(dir.path().to_path_buf());
698 cache.insert(page(hash, 5, data));
699 cache.flush_to_disk();
700 }
701 let backing = dir.path().join(format!("{}.{}", hex_key(&hash), 5));
703 assert!(backing.exists(), "cache file should be spilled");
704
705 let mut cache = PageCache::new(1 << 20).with_persistence(dir.path().to_path_buf());
707 let got = cache.get(&hash, Snapshot::at(Epoch(u64::MAX)));
708 assert_eq!(got, Some(bytes::Bytes::copy_from_slice(data)));
709 }
710
711 #[test]
712 fn try_get_does_not_block() {
713 let mut cache = PageCache::new(1 << 20);
714 let hash = [7u8; 32];
715 cache.insert(page(hash, 1, b"hi"));
716 let got = cache.try_get(&hash, Snapshot::at(Epoch(1)));
717 assert!(got.is_some());
718 let miss = cache.try_get(&[8u8; 32], Snapshot::at(Epoch(1)));
719 assert!(miss.is_none());
720 }
721
722 #[test]
723 fn hot_eviction_storm_does_not_orphan_entries() {
724 let mut cache = PageCache::new(6); let mut next = 1u8;
729 for _ in 0..200 {
730 let key = [next; 32];
731 cache.insert(page(key, 1, b"aa")); for _ in 0..4 {
735 let _ = cache.get(&key, Snapshot::at(Epoch(1)));
736 }
737 next = next.wrapping_add(1);
738 }
739 assert!(
742 cache.used_bytes() <= cache.capacity_bytes + 2, "used_bytes {} leaked past capacity {}",
744 cache.used_bytes(),
745 cache.capacity_bytes
746 );
747 let mut ok = true;
749 for k in &cache.ring {
750 if !cache.map.contains_key(k) {
751 ok = false;
752 }
753 }
754 assert!(
755 ok,
756 "ring references a key absent from the map (orphan slot)"
757 );
758 assert_eq!(
759 cache.ring.len(),
760 cache.map.len(),
761 "ring/map size mismatch (orphaned entries or slots)"
762 );
763 }
764
765 #[test]
766 fn decoded_cache_hit_miss_counters() {
767 use crate::columnar::NativeColumn;
770 let mut cache = DecodedPageCache::new(1 << 20);
771 let key = [9u8; 32];
772 cache.insert(
773 key,
774 std::sync::Arc::new(NativeColumn::Int64 {
775 data: vec![1, 2, 3],
776 validity: vec![],
777 }),
778 );
779 assert!(cache.try_get(&key).is_some());
781 assert!(cache.try_get(&key).is_some());
782 assert!(cache.try_get(&[0u8; 32]).is_none());
783 let s = cache.stats();
784 assert_eq!(s.hits, 2);
785 assert_eq!(s.misses, 1);
786 assert!(s.hit_rate() > 0.66 && s.hit_rate() < 0.67);
787 cache.reset_stats();
788 assert_eq!(cache.stats(), CacheStats::default());
789 }
790}