use std::borrow::Borrow; use std::hash::Hash; use std::num::NonZeroUsize; use lru::LruCache; use sha1::Digest; use sha1::Sha1; use tokio::sync::Mutex; use tokio::sync::MutexGuard; /// A minimal LRU cache protected by a Tokio mutex. /// Calls outside a Tokio runtime are no-ops. pub struct BlockingLruCache { inner: Mutex>, } impl BlockingLruCache where K: Eq + Hash, { /// Creates a cache with the provided non-zero capacity. #[must_use] pub fn new(capacity: NonZeroUsize) -> Self { Self { inner: Mutex::new(LruCache::new(capacity)), } } /// Returns a clone of the cached value for `key`, or computes and inserts it. pub fn get_or_insert_with(&self, key: K, value: impl FnOnce() -> V) -> V where V: Clone, { if let Some(mut guard) = lock_if_runtime(&self.inner) { if let Some(v) = guard.get(&key) { return v.clone(); } let v = value(); // Insert and return a clone to keep ownership in the cache. guard.put(key, v.clone()); return v; } value() } /// Like `get_or_insert_with`, but the value factory may fail. pub fn get_or_try_insert_with( &self, key: K, value: impl FnOnce() -> Result, ) -> Result where V: Clone, { if let Some(mut guard) = lock_if_runtime(&self.inner) { if let Some(v) = guard.get(&key) { return Ok(v.clone()); } let v = value()?; guard.put(key, v.clone()); return Ok(v); } value() } /// Builds a cache if `capacity` is non-zero, returning `None` otherwise. #[must_use] pub fn try_with_capacity(capacity: usize) -> Option { NonZeroUsize::new(capacity).map(Self::new) } /// Returns a clone of the cached value corresponding to `key`, if present. pub fn get(&self, key: &Q) -> Option where K: Borrow, Q: Hash + Eq + ?Sized, V: Clone, { let mut guard = lock_if_runtime(&self.inner)?; guard.get(key).cloned() } /// Inserts `value` for `key`, returning the previous entry if it existed. pub fn insert(&self, key: K, value: V) -> Option { let mut guard = lock_if_runtime(&self.inner)?; guard.put(key, value) } /// Removes the entry for `key` if it exists, returning it. pub fn remove(&self, key: &Q) -> Option where K: Borrow, Q: Hash + Eq + ?Sized, { let mut guard = lock_if_runtime(&self.inner)?; guard.pop(key) } /// Clears all entries from the cache. pub fn clear(&self) { if let Some(mut guard) = lock_if_runtime(&self.inner) { guard.clear(); } } /// Executes `callback` with a mutable reference to the underlying cache. pub fn with_mut(&self, callback: impl FnOnce(&mut LruCache) -> R) -> R { if let Some(mut guard) = lock_if_runtime(&self.inner) { callback(&mut guard) } else { let mut disabled = LruCache::unbounded(); callback(&mut disabled) } } /// Provides direct access to the cache guard when a Tokio runtime is available. pub fn blocking_lock(&self) -> Option>> { lock_if_runtime(&self.inner) } } fn lock_if_runtime(m: &Mutex>) -> Option>> where K: Eq + Hash, { tokio::runtime::Handle::try_current().ok()?; Some(tokio::task::block_in_place(|| m.blocking_lock())) } /// Computes the SHA-1 digest of `bytes`. /// /// Useful for content-based cache keys when you want to avoid staleness /// caused by path-only keys. #[must_use] pub fn sha1_digest(bytes: &[u8]) -> [u8; 20] { let mut hasher = Sha1::new(); hasher.update(bytes); let result = hasher.finalize(); let mut out = [0; 20]; out.copy_from_slice(&result); out } #[cfg(test)] mod tests { use super::BlockingLruCache; use std::num::NonZeroUsize; #[tokio::test(flavor = "multi_thread")] async fn stores_and_retrieves_values() { let cache = BlockingLruCache::new(NonZeroUsize::new(2).expect("capacity")); assert!(cache.get(&"first").is_none()); cache.insert("first", 1); assert_eq!(cache.get(&"first"), Some(1)); } #[tokio::test(flavor = "multi_thread")] async fn evicts_least_recently_used() { let cache = BlockingLruCache::new(NonZeroUsize::new(2).expect("capacity")); cache.insert("a", 1); cache.insert("b", 2); assert_eq!(cache.get(&"a"), Some(1)); cache.insert("c", 3); assert!(cache.get(&"b").is_none()); assert_eq!(cache.get(&"a"), Some(1)); assert_eq!(cache.get(&"c"), Some(3)); } #[test] fn disabled_without_runtime() { let cache = BlockingLruCache::new(NonZeroUsize::new(2).expect("capacity")); cache.insert("first", 1); assert!(cache.get(&"first").is_none()); assert_eq!(cache.get_or_insert_with("first", || 2), 2); assert!(cache.get(&"first").is_none()); assert!(cache.remove(&"first").is_none()); cache.clear(); let result = cache.with_mut(|inner| { inner.put("tmp", 3); inner.get(&"tmp").cloned() }); assert_eq!(result, Some(3)); assert!(cache.get(&"tmp").is_none()); assert!(cache.blocking_lock().is_none()); } }