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::{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)>, // 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    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);// 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        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    /// Drops the global cache instance and clears all in-memory entries, locks, and file locks.
119    ///
120    /// After calling this, [`Cache::instance`] will return an error until [`Cache::new`] is called again.
121    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    /// Returns the globally initialized [`Cache`] instance if it exists.
134    ///
135    /// # Errors
136    /// Returns an error if [`Cache::new`] has not been called yet.
137    ///
138    /// # Example
139    /// ```
140    /// let cache = cache_ro::Cache::instance().unwrap();
141    /// ```
142    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    /// Creates and initializes a new global [`Cache`] instance.
153    ///
154    /// If persistence is enabled in the provided [`CacheConfig`], the directory will be created
155    /// and any persisted entries will be loaded into memory.
156    ///
157    /// This will also start a background thread to periodically clean up expired entries.
158    ///
159    /// # Errors
160    /// Returns an error if:
161    /// - A cache instance already exists.
162    /// - The persistence directory cannot be created.
163    /// # Example
164    /// ```
165    /// let cache = cache_ro::Cache::new(Default::default()).unwrap();
166    /// ```
167    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    /// Stores a value in the cache with a specified TTL (time-to-live).
244    ///
245    /// If persistence is enabled, the value will also be saved to disk.
246    ///
247    /// # Arguments
248    /// * `key` - Cache key.
249    /// * `value` - Serializable value to store.
250    /// * `ttl` - Expiration duration.
251    ///
252    /// # Errors
253    /// Returns an error if serialization or persistence fails.
254    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    /// Retrieves and deserializes a value from the cache if it has not expired.
322    ///
323    /// # Type Parameters
324    /// * `V` - Type to deserialize into.
325    ///
326    /// # Arguments
327    /// * `key` - Cache key.
328    ///
329    /// # Returns
330    /// `Some(value)` if the entry exists and is valid, otherwise `None`.
331    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    /// Returns the remaining TTL for a given cache key.
346    ///
347    /// # Returns
348    /// `Some(duration)` if the entry exists and is not expired, otherwise `None`.
349    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    /// Removes an entry from the cache.
363    ///
364    /// If persistence is enabled, the removal is also reflected on disk.
365    ///
366    /// # Errors
367    /// Returns an error if persistence fails.
368    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    /// Clears all entries from the cache (in memory and on disk if persistent).
379    ///
380    /// # Errors
381    /// Returns an error if persistent files cannot be deleted.
382    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    /// Returns the number of entries currently stored in memory.
400    pub fn len(&self) -> usize {
401        ENTRIES.len()
402    }
403
404    /// Checks whether the cache contains no entries.
405    pub fn is_empty(&self) -> bool {
406        ENTRIES.is_empty()
407    }
408
409    pub fn memory_usage(&self) -> usize {
410        // Estimate ENTRIES memory usage
411        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)  // Serialized value size
416                    + size_of::<u128>()             // expires_at size
417                    + size_of::<String>()           // Overhead for String in DashMap
418            });
419
420        // Estimate FILE_LOCKS memory usage
421        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 the lock structure
426                    + size_of::<String>()           // Overhead for String in DashMap
427            });
428
429        entries_size+file_locks_size
430    }
431}