cache_ro/
lib.rs

1//! # Persistent Cache for Rust
2//!
3//! A high-performance, thread-safe cache with:
4//! - Optional filesystem persistence
5//! - Automatic TTL-based cleanup
6//! - Key-level locking
7//! - Efficient binary serialization
8//!
9//! ## Features
10//!
11//! - **Thread-safe** - Uses DashMap for concurrent access
12//! - **Persistent** - Optional filesystem storage
13//! - **TTL Support** - Automatic expiration of entries
14//! - **Efficient** - bincode serialization
15//!
16//! ## Examples
17//!
18//! ### Basic Usage
19//!
20//! ```rust
21//! use cache_ro::{Cache, CacheConfig};
22//! use std::time::Duration;
23//!
24//! fn main() -> Result<(), Box<dyn std::error::Error>> {
25//!     let cache = Cache::new(CacheConfig::default())?;
26//!     cache.set("key", "value".to_string(), Duration::from_secs(60))?;
27//!
28//!     if let Some(value) = cache.get::<String>("key") {
29//!         println!("Retrieved: {}", value);
30//!     }
31//!     Ok(())
32//! }
33//! ```
34//!
35
36use 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)>, // key (value, expires_at)
61}
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);// 0 no run,1 running,2 in closing
87    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    /// Drops the global cache instance and clears all in-memory entries, locks, and file locks.
146    ///
147    /// After calling this, [`Cache::instance`] will return an error until [`Cache::new`] is called again.
148    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    /// Returns the globally initialized [`Cache`] instance if it exists.
161    ///
162    /// # Errors
163    /// Returns an error if [`Cache::new`] has not been called yet.
164    ///
165    /// # Example
166    /// ```
167    /// let cache = cache_ro::Cache::instance().unwrap();
168    /// ```
169    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    /// Creates and initializes a new global [`Cache`] instance.
180    ///
181    /// If persistence is enabled in the provided [`CacheConfig`], the directory will be created
182    /// and any persisted entries will be loaded into memory.
183    ///
184    /// This will also start a background thread to periodically clean up expired entries.
185    ///
186    /// # Errors
187    /// Returns an error if:
188    /// - A cache instance already exists.
189    /// - The persistence directory cannot be created.
190    /// # Example
191    /// ```
192    /// let cache = cache_ro::Cache::new(Default::default()).unwrap();
193    /// ```
194    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    /// Returns the per-key lock for the specified cache key.
279    ///
280    /// Useful when performing multiple operations atomically for a single key.
281    ///
282    /// # Arguments
283    /// * `key` - The key to lock.
284    ///
285    /// # Example
286    /// ```
287    /// let cache = cache_ro::Cache::new(Default::default()).unwrap();
288    /// let lock = cache.get_key_lock("my_key");
289    /// let _guard = lock.lock().unwrap();
290    /// // do protected work here
291    /// ```
292    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    /// Stores a value in the cache with a specified TTL (time-to-live).
301    ///
302    /// If persistence is enabled, the value will also be saved to disk.
303    ///
304    /// # Arguments
305    /// * `key` - Cache key.
306    /// * `value` - Serializable value to store.
307    /// * `ttl` - Expiration duration.
308    ///
309    /// # Errors
310    /// Returns an error if serialization or persistence fails.
311    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    /// Stores a value in the cache without acquiring a key lock.
336    ///
337    /// Intended for internal use when you already hold the lock.
338    ///
339    /// # Arguments
340    /// * `key` - Cache key.
341    /// * `value` - Serializable value to store.
342    /// * `ttl` - Expiration duration.
343    ///
344    /// # Errors
345    /// Returns an error if serialization or persistence fails.
346    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    /// Retrieves and deserializes a value from the cache if it has not expired.
414    ///
415    /// # Type Parameters
416    /// * `V` - Type to deserialize into.
417    ///
418    /// # Arguments
419    /// * `key` - Cache key.
420    ///
421    /// # Returns
422    /// `Some(value)` if the entry exists and is valid, otherwise `None`.
423    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    /// Returns the remaining TTL for a given cache key.
438    ///
439    /// # Returns
440    /// `Some(duration)` if the entry exists and is not expired, otherwise `None`.
441    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    /// Removes an entry from the cache.
454    ///
455    /// If persistence is enabled, the removal is also reflected on disk.
456    ///
457    /// # Errors
458    /// Returns an error if persistence fails.
459    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    /// Removes an entry without acquiring the key lock.
479    ///
480    /// Used internally when the lock is already held.
481    ///
482    /// # Errors
483    /// Returns an error if persistence fails.
484    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    /// Clears all entries from the cache (in memory and on disk if persistent).
495    ///
496    /// # Errors
497    /// Returns an error if persistent files cannot be deleted.
498    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    /// Returns the number of entries currently stored in memory.
517    pub fn len(&self) -> usize {
518        ENTRIES.len()
519    }
520
521    /// Checks whether the cache contains no entries.
522    pub fn is_empty(&self) -> bool {
523        ENTRIES.is_empty()
524    }
525}