near_cache/
sync.rs

1use lru::LruCache;
2use parking_lot::{Mutex, MutexGuard};
3use std::convert::Infallible;
4use std::hash::Hash;
5use std::num::NonZeroUsize;
6
7/// A wrapper around `LruCache`. This struct is thread safe, doesn't return any references to any
8/// elements inside.
9pub struct SyncLruCache<K, V> {
10    inner: Mutex<LruCache<K, V>>,
11}
12
13impl<K, V> SyncLruCache<K, V>
14where
15    K: Hash + Eq,
16    V: Clone,
17{
18    /// Creates a new `LRU` cache that holds at most `cap` items.
19    pub fn new(cap: usize) -> Self {
20        Self { inner: Mutex::new(LruCache::<K, V>::new(NonZeroUsize::new(cap).unwrap())) }
21    }
22
23    /// Returns the number of key-value pairs that are currently in the cache.
24    pub fn len(&self) -> usize {
25        self.inner.lock().len()
26    }
27
28    /// Returns true if the cache is empty and false otherwise.
29    pub fn is_empty(&self) -> bool {
30        self.inner.lock().is_empty()
31    }
32
33    /// Returns true if the cache contains the key and false otherwise.
34    pub fn contains(&self, key: &K) -> bool {
35        self.inner.lock().contains(key)
36    }
37
38    /// Pushes a key-value pair into the cache. If an entry with key `k` already exists in
39    /// the cache or another cache entry is removed (due to the lru's capacity),
40    /// then it returns the old entry's key-value pair. Otherwise, returns `None`.
41    pub fn push(&self, key: K, value: V) -> Option<(K, V)> {
42        self.inner.lock().push(key, value)
43    }
44
45    /// Return the value of they key in the cache otherwise computes the value and inserts it into
46    /// the cache. If the key is already in the cache, they get moved to the head of
47    /// the LRU list.
48    pub fn get_or_put<F>(&self, key: K, f: F) -> V
49    where
50        V: Clone,
51        F: FnOnce(&K) -> V,
52    {
53        Result::<_, Infallible>::unwrap(self.get_or_try_put(key, |k| Ok(f(k))))
54    }
55
56    /// Returns the value of they key in the cache if present, otherwise
57    /// computes the value using the provided closure.
58    ///
59    /// If the key is already in the cache, it gets moved to the head of the LRU
60    /// list.
61    ///
62    /// If the provided closure fails, the error is returned and the cache is
63    /// not updated.
64    pub fn get_or_try_put<F, E>(&self, key: K, f: F) -> Result<V, E>
65    where
66        V: Clone,
67        F: FnOnce(&K) -> Result<V, E>,
68    {
69        if let Some(result) = self.get(&key) {
70            return Ok(result);
71        }
72        let val = f(&key)?;
73        let val_clone = val.clone();
74        self.inner.lock().put(key, val_clone);
75        Ok(val)
76    }
77
78    /// Puts a key-value pair into cache. If the key already exists in the cache,
79    /// then it updates the key's value.
80    pub fn put(&self, key: K, value: V) {
81        self.inner.lock().put(key, value);
82    }
83
84    /// Returns the value of the key in the cache or None if it is not present in the cache.
85    /// Moves the key to the head of the LRU list if it exists.
86    pub fn get(&self, key: &K) -> Option<V> {
87        self.inner.lock().get(key).cloned()
88    }
89
90    /// Returns the lock over underlying LRU cache.
91    pub fn lock(&self) -> MutexGuard<LruCache<K, V>> {
92        self.inner.lock()
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    #[test]
101    fn test_cache() {
102        let cache = SyncLruCache::<u64, Vec<u64>>::new(100);
103
104        assert_eq!(cache.get(&0u64), None);
105        assert_eq!(cache.get_or_put(123u64, |key| vec![*key, 123]), vec![123u64, 123]);
106        assert_eq!(cache.get(&123u64), Some(vec![123u64, 123]));
107        assert_eq!(cache.get(&0u64), None);
108    }
109}