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::{BTreeMap, HashMap};
43use std::fs::{self, read_dir, File, OpenOptions};
44use std::io::{ErrorKind, Read, Write};
45use std::path::{Path, PathBuf};
46use std::sync::{Arc, Mutex, RwLock};
47use std::sync::atomic::{AtomicU8, Ordering};
48use std::thread::JoinHandle;
49use std::time::{Duration, SystemTime, UNIX_EPOCH};
50
51fn now() -> u128 {
52 SystemTime::now()
53 .duration_since(UNIX_EPOCH)
54 .expect("Time went backwards")
55 .as_millis()
56}
57
58#[derive(Serialize, Deserialize)]
59struct PersistentCache {
60 entries: HashMap<String, (Vec<u8>, u128)>, }
62
63#[derive(Clone)]
64pub struct CacheConfig {
65 pub persistent: bool,
66 pub hash_prefix_length: usize,
67 pub dir_path: String,
68}
69
70impl Default for CacheConfig {
71 fn default() -> Self {
72 Self {
73 persistent: true,
74 hash_prefix_length: 2,
75 dir_path: "cache_data".to_string(),
76 }
77 }
78}
79
80lazy_static! {
81 static ref ENTRIES: DashMap<String, (Vec<u8>,u128)> = DashMap::new();
82 static ref KEY_LOCKS: DashMap<String, Arc<Mutex<()>>> = DashMap::new();
83 static ref FILE_LOCKS: DashMap<String, Arc<Mutex<()>>> = DashMap::new();
84 static ref CACHE: RwLock<Option<Cache>> = RwLock::new(None);
85 static ref EXPIRATION_QUEUE: DashMap<String,u128> = DashMap::new();
86 static ref CACHESTATE: AtomicU8 = AtomicU8::new(0);static ref CLEANUP_THREAD_HANDLE: Mutex<Option<JoinHandle<()>>> = Mutex::new(None);
88}
89
90fn start_cleanup_thread(cache: Cache) -> JoinHandle<()>{
91 std::thread::spawn(move || {
92 let mut heap:BTreeMap<u128, Vec<String>>=BTreeMap::new();
93 loop {
94 if CACHESTATE.load(Ordering::SeqCst) != 1 {
95 CACHESTATE.store(0, Ordering::SeqCst);
96 break;
97 }
98 let mut removes =vec![];
99 {
100 for item in EXPIRATION_QUEUE.iter() {
101 heap.entry(*item.value()).or_default().push(item.key().clone());
102 removes.push(item.key().clone());
103 }
104 }
105 for rm in removes{
106 EXPIRATION_QUEUE.remove(&rm).unwrap();
107 }
108
109 if let Some((next_ts, _)) = heap.iter().next() {
110 let now_ms = now();
111 if next_ts > &now_ms {
112 let sleep_ms = (next_ts - now_ms).min(5000);
113 std::thread::sleep(Duration::from_millis(sleep_ms as u64));
114 continue;
115 }
116 let now_ms = now();
117 let expired_ts: Vec<u128> = heap
118 .iter()
119 .take_while(|&(&ts, _)| ts <= now_ms)
120 .map(|(&ts, _)| ts)
121 .collect();
122 if expired_ts.is_empty() {
123 break;
124 }
125 for ts in expired_ts {
126 if let Some(keys) = heap.remove(&ts) {
127 for key in keys {
128 let _ = cache.remove(&key);
129 }
130 }
131 }
132 } else {
133 std::thread::sleep(Duration::from_secs(1));
134 }
135 }
136 })}
137
138#[derive(Clone)]
139pub struct Cache {
140 config: CacheConfig,
141}
142
143impl Cache {
144
145 pub fn drop() {
149 CACHESTATE.store(2, Ordering::SeqCst);
150 if let Some(handle) = CLEANUP_THREAD_HANDLE.lock().unwrap().take() {
151 let _ = handle.join();
152 }
153 ENTRIES.clear();
154 KEY_LOCKS.clear();
155 FILE_LOCKS.clear();
156 let mut conf = CACHE.write().unwrap();
157 *conf = None;
158 }
159
160 pub fn instance() -> Result<Self, Box<dyn std::error::Error>> {
170 if let Some(cache) = CACHE.read().unwrap().as_ref() {
171 return Ok(cache.clone());
172 }
173 Err(Box::new(std::io::Error::new(
174 ErrorKind::Other,
175 "Cache::new not running",
176 )))
177 }
178
179 pub fn new(config: CacheConfig) -> Result<Self, Box<dyn std::error::Error>> {
195 if CACHESTATE.load(Ordering::SeqCst) != 0 {
196 return Err(Box::new(std::io::Error::new(
197 ErrorKind::Other,
198 "Cache is running",
199 )));
200 }
201 CACHESTATE.store(1, Ordering::SeqCst);
202
203 if config.persistent {
204 fs::create_dir_all(&config.dir_path)?;
205 }
206
207 let cache = Self { config };
208
209 if cache.config.persistent {
210 cache.load_persistent_data()?;
211 }
212
213 {
214 let mut conf = CACHE.write().unwrap();
215 *conf = Some(cache.clone());
216 }
217
218 let handle = start_cleanup_thread(cache.clone());
219 *CLEANUP_THREAD_HANDLE.lock().unwrap() = Some(handle);
220
221 Ok(cache)
222 }
223
224 fn config() -> Configuration<BigEndian> {
225 bincode::config::standard()
226 .with_big_endian()
227 .with_variable_int_encoding()
228 }
229
230 fn get_file_path(&self, key: &str) -> PathBuf {
231 let mut hasher = Sha256::new();
232 hasher.update(key.as_bytes());
233 let hash = hasher.finalize();
234 let prefix_len = self.config.hash_prefix_length.min(hash.len());
235 let prefix = hash[..prefix_len]
236 .iter()
237 .map(|b| format!("{:02x}", b).get(0..1).unwrap().to_string())
238 .collect::<String>();
239
240 Path::new(&self.config.dir_path).join(format!("cache_{}.bin", prefix))
241 }
242
243 fn load_persistent_data(&self) -> Result<(), Box<dyn std::error::Error>> {
244 for entry in read_dir(&self.config.dir_path)? {
245 let entry = entry?;
246 let path = entry.path();
247 if path.is_file() && path.extension().and_then(|s| s.to_str()) == Some("bin") {
248 let file_key = path.to_string_lossy().to_string();
249 let file_lock = FILE_LOCKS
250 .entry(file_key)
251 .or_insert_with(|| Arc::new(Mutex::new(())))
252 .clone();
253 let _guard = file_lock.lock().unwrap();
254
255 let mut file = File::open(&path)?;
256 let mut buffer = Vec::new();
257 file.read_to_end(&mut buffer)?;
258
259 let persistent_cache: HashMap<String, (Vec<u8>,u128)> =
260 decode_from_slice(&buffer, Self::config())?.0;
261
262 for (key, (value,expires_at)) in persistent_cache {
263 let key_lock = KEY_LOCKS
264 .entry(key.clone())
265 .or_insert_with(|| Arc::new(Mutex::new(())))
266 .clone();
267 let _guard = key_lock.lock().unwrap();
268
269 ENTRIES.insert(key.clone(), (value,expires_at));
270 EXPIRATION_QUEUE.insert( key,expires_at);
271 }
272 }
273 }
274 Ok(())
275 }
276
277
278 pub fn get_key_lock(&self, key: &str) -> Arc<Mutex<()>> {
293 KEY_LOCKS
294 .entry(key.to_string())
295 .or_insert_with(|| Arc::new(Mutex::new(())))
296 .clone()
297 }
298
299
300 pub fn set<V: Serialize>(
312 &self,
313 key: &str,
314 value: V,
315 ttl: Duration,
316 ) -> Result<(), Box<dyn std::error::Error>> {
317 let serialized = encode_to_vec(&value, Self::config())?;
318 let expires_at = now() + ttl.as_millis();
319
320 let key_lock = KEY_LOCKS
321 .entry(key.to_string())
322 .or_insert_with(|| Arc::new(Mutex::new(())))
323 .clone();
324 let _guard = key_lock.lock().unwrap();
325
326 ENTRIES.insert(key.to_string(), (serialized,expires_at));
327 EXPIRATION_QUEUE.insert( key.to_string(),expires_at);
328 if self.config.persistent {
329 self.persist_key(key)?;
330 }
331
332 Ok(())
333 }
334
335 pub fn set_without_guard<V: Serialize>(
347 &self,
348 key: &str,
349 value: V,
350 ttl: Duration,
351 ) -> Result<(), Box<dyn std::error::Error>> {
352 let serialized = encode_to_vec(&value, Self::config())?;
353 let expires_at = now() + ttl.as_millis();
354
355 ENTRIES.insert(key.to_string(), (serialized,expires_at), );
356 EXPIRATION_QUEUE.insert( key.to_string(),expires_at);
357 if self.config.persistent {
358 self.persist_key(key)?;
359 }
360
361 Ok(())
362 }
363
364 fn persist_key(&self, key: &str) -> Result<(), Box<dyn std::error::Error>> {
365 let file_path = self.get_file_path(key);
366 let file_key = file_path.to_string_lossy().to_string();
367
368 let file_lock = FILE_LOCKS
369 .entry(file_key)
370 .or_insert_with(|| Arc::new(Mutex::new(())))
371 .clone();
372
373 let _guard = file_lock.lock().unwrap();
374
375 let mut persistent_entries = if file_path.exists() {
376 let mut file = File::open(&file_path)?;
377 let mut buffer = Vec::new();
378 file.read_to_end(&mut buffer)?;
379 let r: PersistentCache = decode_from_slice(&buffer, Self::config())?.0;
380 r.entries
381 } else {
382 HashMap::new()
383 };
384
385 if let Some(v) = ENTRIES.get(key) {
386 persistent_entries.insert(key.to_string(), v.clone());
387 } else {
388 persistent_entries.remove(key);
389 }
390
391 if persistent_entries.is_empty() {
392 if file_path.exists() {
393 fs::remove_file(file_path)?;
394 }
395 return Ok(());
396 }
397
398 let persistent_cache = PersistentCache {
399 entries: persistent_entries,
400 };
401 let serialized = encode_to_vec(&persistent_cache, Self::config())?;
402
403 let mut file = OpenOptions::new()
404 .create(true)
405 .write(true)
406 .truncate(true)
407 .open(&file_path)?;
408 file.write_all(&serialized)?;
409
410 Ok(())
411 }
412
413 pub fn get<V: for<'de> Deserialize<'de>>(&self, key: &str) -> Option<V> {
424 let now = now();
425 ENTRIES.get(key).and_then(|entry| {
426 if entry.1 > now {
427 decode_from_slice(&entry.0, Self::config())
428 .ok()
429 .map(|(v, _)| v)
430 } else {
431 None
432 }
433 })
434 }
435
436
437 pub fn expire(&self, key: &str) -> Option<Duration> {
442 let now = now();
443 ENTRIES.get(key).and_then(|entry| {
444 if entry.1 > now {
445 let remaining = entry.1 - now;
446 Some(Duration::from_millis(remaining as u64))
447 } else {
448 None
449 }
450 })
451 }
452
453 pub fn remove(&self, key: &str) -> Result<(), Box<dyn std::error::Error>> {
460 {
461 let key_lock = KEY_LOCKS
462 .entry(key.to_string())
463 .or_insert_with(|| Arc::new(Mutex::new(())))
464 .clone();
465 let _guard = key_lock.lock().unwrap();
466
467 ENTRIES.remove(key);
468
469 if self.config.persistent {
470 self.persist_key(key)?;
471 }
472 }
473 KEY_LOCKS.remove(key);
474
475 Ok(())
476 }
477
478 pub fn remove_without_guard(&self, key: &str) -> Result<(), Box<dyn std::error::Error>> {
485 ENTRIES.remove(key);
486
487 if self.config.persistent {
488 self.persist_key(key)?;
489 }
490
491 Ok(())
492 }
493
494 pub fn clear(&self) -> Result<(), Box<dyn std::error::Error>> {
499 ENTRIES.clear();
500 KEY_LOCKS.clear();
501 FILE_LOCKS.clear();
502
503 if self.config.persistent {
504 for entry in read_dir(&self.config.dir_path)? {
505 let entry = entry?;
506 let path = entry.path();
507 if path.is_file() && path.extension().and_then(|s| s.to_str()) == Some("bin") {
508 fs::remove_file(path)?;
509 }
510 }
511 }
512
513 Ok(())
514 }
515
516 pub fn len(&self) -> usize {
518 ENTRIES.len()
519 }
520
521 pub fn is_empty(&self) -> bool {
523 ENTRIES.is_empty()
524 }
525}