1use lru::LruCache;
6use std::num::NonZeroUsize;
7use std::sync::Mutex;
8use std::sync::atomic::{AtomicU64, Ordering};
9
10pub struct RowCache {
12 inner: Mutex<LruCache<CacheKey, Vec<u8>>>,
13 hit_count: AtomicU64,
14 miss_count: AtomicU64,
15}
16
17#[derive(Debug, Clone, PartialEq, Eq, Hash)]
19struct CacheKey {
20 table: String,
21 key: Vec<u8>,
22}
23
24impl RowCache {
25 pub fn new(capacity: usize) -> Self {
27 let cap = NonZeroUsize::new(capacity).expect("capacity must be > 0");
28 Self {
29 inner: Mutex::new(LruCache::new(cap)),
30 hit_count: AtomicU64::new(0),
31 miss_count: AtomicU64::new(0),
32 }
33 }
34
35 pub fn with_default_capacity() -> Self {
37 Self::new(10_000)
38 }
39
40 pub fn auto_tune(&self) -> Option<usize> {
44 let ratio = self.hit_ratio();
45 let current_size = {
46 let cache = self.inner.lock().unwrap();
47 cache.len()
48 };
49
50 if ratio < 0.95 && current_size > 0 {
52 Some((current_size as f64 * 1.5) as usize)
54 } else {
55 None }
57 }
58
59 pub fn get(&self, table: &str, key: &[u8]) -> Option<Vec<u8>> {
61 let cache_key = CacheKey {
62 table: table.to_string(),
63 key: key.to_vec(),
64 };
65
66 let mut cache = self.inner.lock().unwrap();
67 if let Some(value) = cache.get(&cache_key) {
68 self.hit_count.fetch_add(1, Ordering::Relaxed);
69 Some(value.clone())
70 } else {
71 self.miss_count.fetch_add(1, Ordering::Relaxed);
72 None
73 }
74 }
75
76 pub fn insert(&self, table: &str, key: &[u8], value: &[u8]) {
78 let cache_key = CacheKey {
79 table: table.to_string(),
80 key: key.to_vec(),
81 };
82
83 let mut cache = self.inner.lock().unwrap();
84 cache.put(cache_key, value.to_vec());
85 }
86
87 pub fn invalidate(&self, table: &str, key: &[u8]) {
89 let cache_key = CacheKey {
90 table: table.to_string(),
91 key: key.to_vec(),
92 };
93
94 let mut cache = self.inner.lock().unwrap();
95 cache.pop(&cache_key);
96 }
97
98 pub fn invalidate_table(&self, table: &str) {
100 let mut cache = self.inner.lock().unwrap();
101 let keys_to_remove: Vec<CacheKey> = cache
103 .iter()
104 .filter(|(k, _)| k.table == table)
105 .map(|(k, _)| k.clone())
106 .collect();
107 for key in keys_to_remove {
108 cache.pop(&key);
109 }
110 }
111
112 pub fn hit_ratio(&self) -> f64 {
114 let hits = self.hit_count.load(Ordering::Relaxed);
115 let misses = self.miss_count.load(Ordering::Relaxed);
116 let total = hits + misses;
117
118 if total == 0 {
119 0.0
120 } else {
121 hits as f64 / total as f64
122 }
123 }
124
125 pub fn stats(&self) -> CacheStats {
127 CacheStats {
128 hits: self.hit_count.load(Ordering::Relaxed),
129 misses: self.miss_count.load(Ordering::Relaxed),
130 hit_ratio: self.hit_ratio(),
131 }
132 }
133
134 pub fn clear(&self) {
136 let mut cache = self.inner.lock().unwrap();
137 cache.clear();
138 self.hit_count.store(0, Ordering::Relaxed);
139 self.miss_count.store(0, Ordering::Relaxed);
140 }
141}
142
143#[derive(Debug, Clone)]
145pub struct CacheStats {
146 pub hits: u64,
147 pub misses: u64,
148 pub hit_ratio: f64,
149}
150
151#[cfg(test)]
152mod tests {
153 use super::*;
154
155 #[test]
156 fn test_cache_basic() {
157 let cache = RowCache::new(3);
158
159 cache.insert("users", b"key1", b"value1");
161 cache.insert("users", b"key2", b"value2");
162
163 assert_eq!(cache.get("users", b"key1"), Some(b"value1".to_vec()));
165 assert_eq!(cache.get("users", b"key2"), Some(b"value2".to_vec()));
166 assert_eq!(cache.get("users", b"key3"), None);
167 }
168
169 #[test]
170 fn test_cache_lru_eviction() {
171 let cache = RowCache::new(2);
172
173 cache.insert("users", b"key1", b"value1");
174 cache.insert("users", b"key2", b"value2");
175 cache.insert("users", b"key3", b"value3"); assert_eq!(cache.get("users", b"key1"), None);
178 assert_eq!(cache.get("users", b"key2"), Some(b"value2".to_vec()));
179 assert_eq!(cache.get("users", b"key3"), Some(b"value3".to_vec()));
180 }
181
182 #[test]
183 fn test_cache_invalidate() {
184 let cache = RowCache::new(3);
185
186 cache.insert("users", b"key1", b"value1");
187 cache.insert("users", b"key2", b"value2");
188
189 cache.invalidate("users", b"key1");
190
191 assert_eq!(cache.get("users", b"key1"), None);
192 assert_eq!(cache.get("users", b"key2"), Some(b"value2".to_vec()));
193 }
194
195 #[test]
196 fn test_cache_hit_ratio() {
197 let cache = RowCache::new(3);
198
199 cache.insert("users", b"key1", b"value1");
200
201 cache.get("users", b"key1"); cache.get("users", b"key1"); cache.get("users", b"key2"); let stats = cache.stats();
206 assert_eq!(stats.hits, 2);
207 assert_eq!(stats.misses, 1);
208 assert!((stats.hit_ratio - 0.666).abs() < 0.01);
209 }
210
211 #[test]
212 fn test_cache_clear() {
213 let cache = RowCache::new(3);
214
215 cache.insert("users", b"key1", b"value1");
216 cache.get("users", b"key1");
217
218 cache.clear();
219
220 assert_eq!(cache.get("users", b"key1"), None);
221 assert_eq!(cache.stats().hits, 0);
222 assert_eq!(cache.stats().misses, 1);
223 }
224}