1use bincode::config::{BigEndian, Configuration};
37use bincode::serde::{decode_from_slice, encode_to_vec};
38use dashmap::DashMap;
39use lazy_static::lazy_static;
40use serde::{Deserialize, Serialize};
41use sha2::{Digest, Sha256};
42use std::collections::HashMap;
43use std::fs::{self, read_dir, File, OpenOptions};
44use std::io::{Read, Write};
45use std::path::{Path, PathBuf};
46use std::sync::{Arc, Mutex};
47use std::time::{Duration, SystemTime, UNIX_EPOCH};
48
49fn now() -> u128 {
50 SystemTime::now()
51 .duration_since(UNIX_EPOCH)
52 .expect("Time went backwards")
53 .as_millis()
54}
55
56#[derive(Serialize, Deserialize)]
57struct CacheEntry {
58 value: Vec<u8>,
59 expires_at: u128,
60}
61
62#[derive(Serialize, Deserialize)]
63struct PersistentCache {
64 entries: HashMap<String, CacheEntry>,
65}
66
67#[derive(Clone)]
68pub struct CacheConfig {
69 pub persistent: bool,
70 pub hash_prefix_length: usize,
71 pub cleanup_interval: Duration,
72 pub dir_path: String,
73}
74
75impl Default for CacheConfig {
76 fn default() -> Self {
77 Self {
78 persistent: true,
79 hash_prefix_length: 2,
80 cleanup_interval: Duration::from_secs(60),
81 dir_path: "cache_data".to_string(),
82 }
83 }
84}
85
86lazy_static! {
87 static ref ENTRIES: DashMap<String, CacheEntry> = DashMap::new();
88 static ref KEY_LOCKS: DashMap<String, Arc<Mutex<()>>> = DashMap::new();
89 static ref FILE_LOCKS: DashMap<String, Arc<Mutex<()>>> = DashMap::new();
90}
91
92#[derive(Clone)]
93pub struct Cache {
94 config: CacheConfig,
95}
96
97impl Cache {
98 fn config() -> Configuration<BigEndian> {
99 bincode::config::standard()
100 .with_big_endian()
101 .with_variable_int_encoding()
102 }
103
104 fn get_file_path(&self, key: &str) -> PathBuf {
105 let mut hasher = Sha256::new();
106 hasher.update(key.as_bytes());
107 let hash = hasher.finalize();
108 let prefix_len = self.config.hash_prefix_length.min(hash.len());
109 let prefix = hash[..prefix_len]
110 .iter()
111 .map(|b| format!("{:02x}", b).get(0..1).unwrap().to_string())
112 .collect::<String>();
113
114 Path::new(&self.config.dir_path).join(format!("cache_{}.bin", prefix))
115 }
116
117 pub fn new(config: CacheConfig) -> Result<Self, Box<dyn std::error::Error>> {
118 if config.persistent {
119 fs::create_dir_all(&config.dir_path)?;
120 }
121
122 let cache = Self { config };
123
124 if cache.config.persistent {
125 cache.load_persistent_data()?;
126 }
127
128 let cache_clone = cache.clone();
129 std::thread::spawn(move || loop {
130 std::thread::sleep(cache_clone.config.cleanup_interval);
131 cache_clone.cleanup();
132 });
133
134 Ok(cache)
135 }
136
137 fn load_persistent_data(&self) -> Result<(), Box<dyn std::error::Error>> {
138 for entry in read_dir(&self.config.dir_path)? {
139 let entry = entry?;
140 let path = entry.path();
141 if path.is_file() && path.extension().and_then(|s| s.to_str()) == Some("bin") {
142 let file_key = path.to_string_lossy().to_string();
143 let file_lock = FILE_LOCKS
144 .entry(file_key)
145 .or_insert_with(|| Arc::new(Mutex::new(())))
146 .clone();
147 let _guard = file_lock.lock().unwrap();
148
149 let mut file = File::open(&path)?;
150 let mut buffer = Vec::new();
151 file.read_to_end(&mut buffer)?;
152
153 let persistent_cache: HashMap<String, CacheEntry> =
154 decode_from_slice(&buffer, Self::config())?.0;
155
156 for (key, entry) in persistent_cache {
157 let key_lock = KEY_LOCKS
158 .entry(key.clone())
159 .or_insert_with(|| Arc::new(Mutex::new(())))
160 .clone();
161 let _guard = key_lock.lock().unwrap();
162
163 ENTRIES.insert(
164 key,
165 CacheEntry {
166 value: entry.value,
167 expires_at: entry.expires_at,
168 },
169 );
170 }
171 }
172 }
173 Ok(())
174 }
175
176 fn cleanup(&self) {
177 let now = now();
178 let mut rm = vec![];
179 for i in ENTRIES.iter() {
180 if i.expires_at <= now {
181 rm.push(i.key().to_string());
182 }
183 }
184 for key in rm {
185 let _ = self.remove(&key);
186 }
187 }
188
189 pub fn get_key_lock(&self, key: &str) -> Arc<Mutex<()>> {
190 KEY_LOCKS
191 .entry(key.to_string())
192 .or_insert_with(|| Arc::new(Mutex::new(())))
193 .clone()
194 }
195
196 pub fn set<V: Serialize>(
197 &self,
198 key: &str,
199 value: V,
200 ttl: Duration,
201 ) -> Result<(), Box<dyn std::error::Error>> {
202 let serialized = encode_to_vec(&value, Self::config())?;
203 let expires_at = now() + ttl.as_millis();
204
205 let key_lock = KEY_LOCKS
206 .entry(key.to_string())
207 .or_insert_with(|| Arc::new(Mutex::new(())))
208 .clone();
209 let _guard = key_lock.lock().unwrap();
210
211 ENTRIES.insert(
212 key.to_string(),
213 CacheEntry {
214 value: serialized,
215 expires_at,
216 },
217 );
218
219 if self.config.persistent {
220 self.persist_key(key)?;
221 }
222
223 Ok(())
224 }
225
226 pub fn set_without_guard<V: Serialize>(
227 &self,
228 key: &str,
229 value: V,
230 ttl: Duration,
231 ) -> Result<(), Box<dyn std::error::Error>> {
232 let serialized = encode_to_vec(&value, Self::config())?;
233 let expires_at = now() + ttl.as_millis();
234
235 ENTRIES.insert(
236 key.to_string(),
237 CacheEntry {
238 value: serialized,
239 expires_at,
240 },
241 );
242
243 if self.config.persistent {
244 self.persist_key(key)?;
245 }
246
247 Ok(())
248 }
249
250 fn persist_key(&self, key: &str) -> Result<(), Box<dyn std::error::Error>> {
251 let file_path = self.get_file_path(key);
252 let file_key = file_path.to_string_lossy().to_string();
253
254 let file_lock = FILE_LOCKS
255 .entry(file_key)
256 .or_insert_with(|| Arc::new(Mutex::new(())))
257 .clone();
258
259 let _guard = file_lock.lock().unwrap();
260
261 let mut persistent_entries = if file_path.exists() {
262 let mut file = File::open(&file_path)?;
263 let mut buffer = Vec::new();
264 file.read_to_end(&mut buffer)?;
265 let r: PersistentCache = decode_from_slice(&buffer, Self::config())?.0;
266 r.entries
267 } else {
268 HashMap::new()
269 };
270
271 if let Some(entry) = ENTRIES.get(key) {
272 persistent_entries.insert(
273 key.to_string(),
274 CacheEntry {
275 value: entry.value.clone(),
276 expires_at: entry.expires_at,
277 },
278 );
279 } else {
280 persistent_entries.remove(key);
281 }
282
283 if persistent_entries.is_empty() {
284 if file_path.exists() {
285 fs::remove_file(file_path)?;
286 }
287 return Ok(());
288 }
289
290 let persistent_cache = PersistentCache {
291 entries: persistent_entries,
292 };
293 let serialized = encode_to_vec(&persistent_cache, Self::config())?;
294
295 let mut file = OpenOptions::new()
296 .create(true)
297 .write(true)
298 .truncate(true)
299 .open(&file_path)?;
300 file.write_all(&serialized)?;
301
302 Ok(())
303 }
304
305 pub fn get<V: for<'de> Deserialize<'de>>(&self, key: &str) -> Option<V> {
306 let now = now();
307 ENTRIES.get(key).and_then(|entry| {
308 if entry.expires_at > now {
309 decode_from_slice(&entry.value, Self::config())
310 .ok()
311 .map(|(v, _)| v)
312 } else {
313 None
314 }
315 })
316 }
317
318 pub fn expire(&self, key: &str) -> Option<Duration> {
319 let now = now();
320 ENTRIES.get(key).and_then(|entry| {
321 if entry.expires_at > now {
322 let remaining = entry.expires_at - now;
323 Some(Duration::from_millis(remaining as u64))
324 } else {
325 None
326 }
327 })
328 }
329
330 pub fn remove(&self, key: &str) -> Result<(), Box<dyn std::error::Error>> {
331 {
332 let key_lock = KEY_LOCKS
333 .entry(key.to_string())
334 .or_insert_with(|| Arc::new(Mutex::new(())))
335 .clone();
336 let _guard = key_lock.lock().unwrap();
337
338 ENTRIES.remove(key);
339
340 if self.config.persistent {
341 self.persist_key(key)?;
342 }
343 }
344 KEY_LOCKS.remove(key);
345
346 Ok(())
347 }
348
349 pub fn remove_without_guard(&self, key: &str) -> Result<(), Box<dyn std::error::Error>> {
350 ENTRIES.remove(key);
351
352 if self.config.persistent {
353 self.persist_key(key)?;
354 }
355
356 Ok(())
357 }
358
359 pub fn clear(&self) -> Result<(), Box<dyn std::error::Error>> {
360 ENTRIES.clear();
361 KEY_LOCKS.clear();
362 FILE_LOCKS.clear();
363
364 if self.config.persistent {
365 for entry in read_dir(&self.config.dir_path)? {
366 let entry = entry?;
367 let path = entry.path();
368 if path.is_file() && path.extension().and_then(|s| s.to_str()) == Some("bin") {
369 fs::remove_file(path)?;
370 }
371 }
372 }
373
374 Ok(())
375 }
376
377 pub fn len(&self) -> usize {
378 ENTRIES.len()
379 }
380
381 pub fn is_empty(&self) -> bool {
382 ENTRIES.is_empty()
383 }
384}