Skip to main content

graphyn_store/
cache.rs

1use std::collections::HashMap;
2use std::sync::atomic::{AtomicU64, Ordering};
3use std::sync::RwLock;
4
5#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
6pub struct CacheStats {
7    pub hits: u64,
8    pub misses: u64,
9    pub evictions: u64,
10    pub size: usize,
11}
12
13#[derive(Debug, Clone)]
14struct CacheEntry {
15    value: String,
16    last_access_tick: u64,
17}
18
19pub struct HotQueryCache {
20    entries: RwLock<HashMap<String, CacheEntry>>,
21    max_entries: usize,
22    tick: AtomicU64,
23    hits: AtomicU64,
24    misses: AtomicU64,
25    evictions: AtomicU64,
26}
27
28impl HotQueryCache {
29    pub fn new(max_entries: usize) -> Self {
30        Self {
31            entries: RwLock::new(HashMap::new()),
32            max_entries,
33            tick: AtomicU64::new(1),
34            hits: AtomicU64::new(0),
35            misses: AtomicU64::new(0),
36            evictions: AtomicU64::new(0),
37        }
38    }
39
40    pub fn put(&self, key: String, value: String) {
41        let current_tick = self.tick.fetch_add(1, Ordering::Relaxed);
42        if let Ok(mut entries) = self.entries.write() {
43            entries.insert(
44                key,
45                CacheEntry {
46                    value,
47                    last_access_tick: current_tick,
48                },
49            );
50            evict_if_needed(&mut entries, self.max_entries, &self.evictions);
51        }
52    }
53
54    pub fn get(&self, key: &str) -> Option<String> {
55        if let Ok(mut entries) = self.entries.write() {
56            if let Some(entry) = entries.get_mut(key) {
57                self.hits.fetch_add(1, Ordering::Relaxed);
58                entry.last_access_tick = self.tick.fetch_add(1, Ordering::Relaxed);
59                return Some(entry.value.clone());
60            }
61        }
62
63        self.misses.fetch_add(1, Ordering::Relaxed);
64        None
65    }
66
67    pub fn invalidate(&self, key: &str) {
68        if let Ok(mut entries) = self.entries.write() {
69            entries.remove(key);
70        }
71    }
72
73    pub fn clear(&self) {
74        if let Ok(mut entries) = self.entries.write() {
75            entries.clear();
76        }
77    }
78
79    pub fn stats(&self) -> CacheStats {
80        let size = self
81            .entries
82            .read()
83            .map(|entries| entries.len())
84            .unwrap_or_default();
85
86        CacheStats {
87            hits: self.hits.load(Ordering::Relaxed),
88            misses: self.misses.load(Ordering::Relaxed),
89            evictions: self.evictions.load(Ordering::Relaxed),
90            size,
91        }
92    }
93}
94
95fn evict_if_needed(
96    entries: &mut HashMap<String, CacheEntry>,
97    max_entries: usize,
98    evictions: &AtomicU64,
99) {
100    while entries.len() > max_entries {
101        let mut oldest_key = None;
102        let mut oldest_tick = u64::MAX;
103
104        for (key, entry) in entries.iter() {
105            if entry.last_access_tick < oldest_tick {
106                oldest_tick = entry.last_access_tick;
107                oldest_key = Some(key.clone());
108            }
109        }
110
111        let Some(oldest_key) = oldest_key else {
112            break;
113        };
114
115        if entries.remove(&oldest_key).is_some() {
116            evictions.fetch_add(1, Ordering::Relaxed);
117        } else {
118            break;
119        }
120    }
121}