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::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}