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::{ErrorKind, Read, Write};
45use std::path::{Path, PathBuf};
46use std::sync::atomic::{AtomicU8, Ordering};
47use std::sync::{Arc, Mutex, RwLock};
48use std::thread::{sleep, 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 pub cleanup_interval: Duration,
69}
70
71impl Default for CacheConfig {
72 fn default() -> Self {
73 Self {
74 persistent: true,
75 hash_prefix_length: 2,
76 dir_path: "cache_data".to_string(),
77 cleanup_interval: Duration::from_secs(10),
78 }
79 }
80}
81
82lazy_static! {
83 static ref ENTRIES: DashMap<String, (Vec<u8>,u128)> = DashMap::new();
84 static ref FILE_LOCKS: DashMap<String, Arc<Mutex<()>>> = DashMap::new();
85 static ref CACHE: RwLock<Option<Cache>> = RwLock::new(None);
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 loop {
93 if CACHESTATE.load(Ordering::SeqCst) != 1 {
94 CACHESTATE.store(0, Ordering::SeqCst);
95 break;
96 }
97 let now = now();
98 let expired_keys: Vec<String> = ENTRIES
99 .iter()
100 .filter(|entry| entry.1 <= now)
101 .map(|entry| entry.key().clone())
102 .collect();
103
104 for key in expired_keys {
105 let _ = cache.remove(&key);
106 }
107 sleep(cache.config.cleanup_interval);
108 }
109 })
110}
111
112#[derive(Clone)]
113pub struct Cache {
114 config: CacheConfig,
115}
116
117impl Cache {
118 pub fn drop() {
122 CACHESTATE.store(2, Ordering::SeqCst);
123 if let Some(handle) = CLEANUP_THREAD_HANDLE.lock().unwrap().take() {
124 let _ = handle.join();
125 }
126
127 ENTRIES.clear();
128 FILE_LOCKS.clear();
129 let mut conf = CACHE.write().unwrap();
130 *conf = None;
131 }
132
133 pub fn instance() -> Result<Self, Box<dyn std::error::Error>> {
143 if let Some(cache) = CACHE.read().unwrap().as_ref() {
144 return Ok(cache.clone());
145 }
146 Err(Box::new(std::io::Error::new(
147 ErrorKind::Other,
148 "Cache::new not running",
149 )))
150 }
151
152 pub fn new(config: CacheConfig) -> Result<Self, Box<dyn std::error::Error>> {
168 if CACHESTATE.load(Ordering::SeqCst) != 0 {
169 return Err(Box::new(std::io::Error::new(
170 ErrorKind::Other,
171 "Cache is running",
172 )));
173 }
174 CACHESTATE.store(1, Ordering::SeqCst);
175
176 if config.persistent {
177 fs::create_dir_all(&config.dir_path)?;
178 }
179
180 let cache = Self { config };
181
182 if cache.config.persistent {
183 cache.load_persistent_data()?;
184 }
185
186 {
187 let mut conf = CACHE.write().unwrap();
188 *conf = Some(cache.clone());
189 }
190
191 let handle = start_cleanup_thread(cache.clone());
192 *CLEANUP_THREAD_HANDLE.lock().unwrap() = Some(handle);
193
194 Ok(cache)
195 }
196
197 fn config() -> Configuration<BigEndian> {
198 bincode::config::standard()
199 .with_big_endian()
200 .with_variable_int_encoding()
201 }
202
203 fn get_file_path(&self, key: &str) -> PathBuf {
204 let mut hasher = Sha256::new();
205 hasher.update(key.as_bytes());
206 let hash = hasher.finalize();
207 let prefix_len = self.config.hash_prefix_length.min(hash.len());
208 let prefix = hash[..prefix_len]
209 .iter()
210 .map(|b| format!("{:02x}", b).get(0..1).unwrap().to_string())
211 .collect::<String>();
212
213 Path::new(&self.config.dir_path).join(format!("cache_{}.bin", prefix))
214 }
215
216 fn load_persistent_data(&self) -> Result<(), Box<dyn std::error::Error>> {
217 for entry in read_dir(&self.config.dir_path)? {
218 let entry = entry?;
219 let path = entry.path();
220 if path.is_file() && path.extension().and_then(|s| s.to_str()) == Some("bin") {
221 let file_key = path.to_string_lossy().to_string();
222 let file_lock = FILE_LOCKS
223 .entry(file_key)
224 .or_insert_with(|| Arc::new(Mutex::new(())))
225 .clone();
226 let _guard = file_lock.lock().unwrap();
227
228 let mut file = File::open(&path)?;
229 let mut buffer = Vec::new();
230 file.read_to_end(&mut buffer)?;
231
232 let persistent_cache: HashMap<String, (Vec<u8>, u128)> =
233 decode_from_slice(&buffer, Self::config())?.0;
234
235 for (key, (value, expires_at)) in persistent_cache {
236 ENTRIES.insert(key.clone(), (value, expires_at));
237 }
238 }
239 }
240 Ok(())
241 }
242
243 pub fn set<V: Serialize>(
255 &self,
256 key: &str,
257 value: V,
258 ttl: Duration,
259 ) -> Result<(), Box<dyn std::error::Error>> {
260 let serialized = encode_to_vec(&value, Self::config())?;
261 let expires_at = now() + ttl.as_millis();
262
263 ENTRIES.insert(key.to_string(), (serialized, expires_at));
264
265 if self.config.persistent {
266 self.persist_key(key)?;
267 }
268
269 Ok(())
270 }
271
272 fn persist_key(&self, key: &str) -> Result<(), Box<dyn std::error::Error>> {
273 let file_path = self.get_file_path(key);
274 let file_key = file_path.to_string_lossy().to_string();
275
276 let file_lock = FILE_LOCKS
277 .entry(file_key)
278 .or_insert_with(|| Arc::new(Mutex::new(())))
279 .clone();
280
281 let _guard = file_lock.lock().unwrap();
282
283 let mut persistent_entries = if file_path.exists() {
284 let mut file = File::open(&file_path)?;
285 let mut buffer = Vec::new();
286 file.read_to_end(&mut buffer)?;
287 let r: PersistentCache = decode_from_slice(&buffer, Self::config())?.0;
288 r.entries
289 } else {
290 HashMap::new()
291 };
292
293 if let Some(v) = ENTRIES.get(key) {
294 persistent_entries.insert(key.to_string(), v.clone());
295 } else {
296 persistent_entries.remove(key);
297 }
298
299 if persistent_entries.is_empty() {
300 if file_path.exists() {
301 fs::remove_file(file_path)?;
302 }
303 return Ok(());
304 }
305
306 let persistent_cache = PersistentCache {
307 entries: persistent_entries,
308 };
309 let serialized = encode_to_vec(&persistent_cache, Self::config())?;
310
311 let mut file = OpenOptions::new()
312 .create(true)
313 .write(true)
314 .truncate(true)
315 .open(&file_path)?;
316 file.write_all(&serialized)?;
317
318 Ok(())
319 }
320
321 pub fn get<V: for<'de> Deserialize<'de>>(&self, key: &str) -> Option<V> {
332 let now = now();
333
334 ENTRIES.get(key).and_then(|entry| {
335 if entry.1 > now {
336 decode_from_slice(&entry.0, Self::config())
337 .ok()
338 .map(|(v, _)| v)
339 } else {
340 None
341 }
342 })
343 }
344
345 pub fn expire(&self, key: &str) -> Option<Duration> {
350 let now = now();
351
352 ENTRIES.get(key).and_then(|entry| {
353 if entry.1 > now {
354 let remaining = entry.1 - now;
355 Some(Duration::from_millis(remaining as u64))
356 } else {
357 None
358 }
359 })
360 }
361
362 pub fn remove(&self, key: &str) -> Result<(), Box<dyn std::error::Error>> {
369 ENTRIES.remove(key);
370
371 if self.config.persistent {
372 self.persist_key(key)?;
373 }
374
375 Ok(())
376 }
377
378 pub fn clear(&self) -> Result<(), Box<dyn std::error::Error>> {
383 ENTRIES.clear();
384 FILE_LOCKS.clear();
385
386 if self.config.persistent {
387 for entry in read_dir(&self.config.dir_path)? {
388 let entry = entry?;
389 let path = entry.path();
390 if path.is_file() && path.extension().and_then(|s| s.to_str()) == Some("bin") {
391 fs::remove_file(path)?;
392 }
393 }
394 }
395
396 Ok(())
397 }
398
399 pub fn len(&self) -> usize {
401 ENTRIES.len()
402 }
403
404 pub fn is_empty(&self) -> bool {
406 ENTRIES.is_empty()
407 }
408
409 pub fn memory_usage(&self) -> usize {
410 let entries_size = ENTRIES
412 .iter()
413 .fold(0, |acc, entry| {
414 acc + size_of_val(entry.key())
415 + size_of_val(&entry.value().0) + size_of::<u128>() + size_of::<String>() });
419
420 let file_locks_size = FILE_LOCKS
422 .iter()
423 .fold(0, |acc, entry| {
424 acc + size_of_val(entry.key())
425 + size_of::<Arc<Mutex<()>>>() + size_of::<String>() });
428
429 entries_size+file_locks_size
430 }
431}