use hashlink::LinkedHashMap; use std::{ hash::Hash, time::{Duration, Instant}, }; pub struct LruTimeCache { map: LinkedHashMap, /// The time elements remain in the cache. ttl: Duration, /// The max size of the cache. capacity: usize, } impl LruTimeCache { pub fn new(ttl: Duration, capacity: Option) -> LruTimeCache { let capacity = if let Some(cap) = capacity { cap } else { usize::MAX }; LruTimeCache { map: LinkedHashMap::new(), ttl, capacity, } } /// Inserts a key-value pair into the cache. pub fn insert(&mut self, key: K, value: V) { let now = Instant::now(); self.map.insert(key, (value, now)); if self.map.len() > self.capacity { self.map.pop_front(); } } /// Retrieves a reference to the value stored under `key`, or `None` if the key doesn't exist. /// Also removes expired elements and updates the time. pub fn get(&mut self, key: &K) -> Option<&V> { self.get_mut(key).map(|value| &*value) } /// Retrieves a mutable reference to the value stored under `key`, or `None` if the key doesn't exist. /// Also removes expired elements and updates the time. pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { let now = Instant::now(); self.remove_expired_values(now); match self.map.raw_entry_mut().from_key(key) { hashlink::linked_hash_map::RawEntryMut::Occupied(mut occupied) => { occupied.get_mut().1 = now; occupied.to_back(); Some(&mut occupied.into_mut().0) } hashlink::linked_hash_map::RawEntryMut::Vacant(_) => None, } } /// Returns a reference to the value with the given `key`, if present and not expired, without /// updating the timestamp. pub fn peek(&self, key: &K) -> Option<&V> { if let Some((value, time)) = self.map.get(key) { return if *time + self.ttl >= Instant::now() { Some(value) } else { None }; } None } /// Returns the size of the cache, i.e. the number of cached non-expired key-value pairs. pub fn len(&mut self) -> usize { self.remove_expired_values(Instant::now()); self.map.len() } /// Removes a key-value pair from the cache, returning the value at the key if the key /// was previously in the map. pub fn remove(&mut self, key: &K) -> Option { self.map.remove(key).map(|v| v.0) } /// Removes expired items from the cache. fn remove_expired_values(&mut self, now: Instant) { let mut expired_keys = vec![]; for (key, (_, time)) in self.map.iter_mut() { if *time + self.ttl >= now { break; } expired_keys.push(key.clone()); } for k in expired_keys { self.map.remove(&k); } } } #[cfg(test)] mod tests { use crate::lru_time_cache::LruTimeCache; use std::time::Duration; #[test] fn insert() { let mut cache = LruTimeCache::new(Duration::from_secs(10), None); cache.insert(1, 10); cache.insert(2, 20); cache.insert(3, 30); assert_eq!(Some(&10), cache.get(&1)); assert_eq!(Some(&20), cache.get(&2)); assert_eq!(Some(&30), cache.get(&3)); } #[test] fn capacity() { let mut cache = LruTimeCache::new(Duration::from_secs(10), Some(2)); cache.insert(1, 10); cache.insert(2, 20); assert_eq!(2, cache.len()); cache.insert(3, 30); assert_eq!(2, cache.len()); assert_eq!(Some(&20), cache.get(&2)); assert_eq!(Some(&30), cache.get(&3)); } #[test] fn get() { let mut cache = LruTimeCache::new(Duration::from_secs(10), Some(2)); cache.insert(1, 10); cache.insert(2, 20); assert_eq!(Some(&10), cache.get(&1)); cache.insert(3, 30); // `1` is alive as `get()` updates the timestamp. assert_eq!(Some(&10), cache.get(&1)); // `2` is removed as `2` is oldest at the time `3` was inserted. assert_eq!(None, cache.get(&2)); } #[test] fn get_mut() { let mut cache = LruTimeCache::new(Duration::from_secs(10), None); cache.insert(1, 10); let v = cache.get_mut(&1).expect("should have value"); *v = 100; assert_eq!(Some(&100), cache.get(&1)); } #[test] fn peek() { let mut cache = LruTimeCache::new(Duration::from_secs(10), Some(2)); cache.insert(1, 10); cache.insert(2, 20); assert_eq!(Some(&10), cache.peek(&1)); cache.insert(3, 30); // `1` is removed as `peek()` does not update the time. assert_eq!(None, cache.peek(&1)); assert_eq!(Some(&20), cache.get(&2)); } #[test] fn len() { let mut cache = LruTimeCache::new(Duration::from_secs(10), None); assert_eq!(0, cache.len()); cache.insert(1, 10); cache.insert(2, 20); cache.insert(3, 30); assert_eq!(3, cache.len()); } #[test] fn remove() { let mut cache = LruTimeCache::new(Duration::from_secs(10), None); cache.insert(1, 10); assert_eq!(Some(10), cache.remove(&1)); assert_eq!(None, cache.get(&1)); assert_eq!(None, cache.remove(&1)); } mod ttl { use crate::lru_time_cache::LruTimeCache; use std::{thread::sleep, time::Duration}; const TTL: Duration = Duration::from_millis(100); #[test] fn get() { let mut cache = LruTimeCache::new(TTL, None); cache.insert(1, 10); assert_eq!(Some(&10), cache.get(&1)); sleep(TTL); assert_eq!(None, cache.get(&1)); } #[test] fn peek() { let mut cache = LruTimeCache::new(TTL, None); cache.insert(1, 10); assert_eq!(Some(&10), cache.peek(&1)); sleep(TTL); assert_eq!(None, cache.peek(&1)); } #[test] fn len() { let mut cache = LruTimeCache::new(TTL, None); cache.insert(1, 10); assert_eq!(1, cache.len()); sleep(TTL); assert_eq!(0, cache.len()); } #[test] fn ttl() { let mut cache = LruTimeCache::new(TTL, None); cache.insert(1, 10); sleep(TTL / 4); cache.insert(2, 20); sleep(TTL / 4); cache.insert(3, 30); sleep(TTL / 4); cache.insert(4, 40); sleep(TTL / 4); assert_eq!(3, cache.len()); assert_eq!(None, cache.get(&1)); assert_eq!(Some(&20), cache.get(&2)); assert_eq!(Some(&30), cache.get(&3)); assert_eq!(Some(&40), cache.get(&4)); } } }