pbs-tools/src/lru_cache.rs

   1 //! Least recently used (LRU) cache
   2 //!
   3 //! Implements a cache with least recently used cache replacement policy.
   4 //! A HashMap is used for fast access by a given key and a doubly linked list
   5 //! is used to keep track of the cache access order.
   6
   7 use std::collections::{HashMap, hash_map::Entry};
   8 use std::marker::PhantomData;
   9
  10 /// Interface for getting values on cache misses.
  11 pub trait Cacher<K, V> {
  12     /// Fetch a value for key on cache miss.
  13     ///
  14     /// Whenever a cache miss occurs, the fetch method provides a corresponding value.
  15     /// If no value can be obtained for the given key, None is returned, the cache is
  16     /// not updated in that case.
  17     fn fetch(&mut self, key: K) -> Result<Option<V>, anyhow::Error>;
  18 }
  19
  20 /// Node of the doubly linked list storing key and value
  21 struct CacheNode<K, V> {
  22     // We need to additionally store the key to be able to remove it
  23     // from the HashMap when removing the tail.
  24     key: K,
  25     value: V,
  26     prev: *mut CacheNode<K, V>,
  27     next: *mut CacheNode<K, V>,
  28     // Dropcheck marker. See the phantom-data section in the rustonomicon.
  29     _marker: PhantomData<Box<CacheNode<K, V>>>,
  30 }
  31
  32 impl<K, V> CacheNode<K, V> {
  33     fn new(key: K, value: V) -> Self {
  34         Self {
  35             key,
  36             value,
  37             prev: std::ptr::null_mut(),
  38             next: std::ptr::null_mut(),
  39             _marker: PhantomData,
  40         }
  41     }
  42 }
  43
  44 /// LRU cache instance.
  45 ///
  46 /// # Examples:
  47 /// ```
  48 /// # use self::proxmox_backup::tools::lru_cache::{Cacher, LruCache};
  49 /// # fn main() -> Result<(), anyhow::Error> {
  50 /// struct LruCacher {};
  51 ///
  52 /// impl Cacher<u64, u64> for LruCacher {
  53 ///     fn fetch(&mut self, key: u64) -> Result<Option<u64>, anyhow::Error> {
  54 ///         Ok(Some(key))
  55 ///     }
  56 /// }
  57 ///
  58 /// let mut cache = LruCache::new(3);
  59 ///
  60 /// assert_eq!(cache.get_mut(1), None);
  61 /// assert_eq!(cache.len(), 0);
  62 ///
  63 /// cache.insert(1, 1);
  64 /// cache.insert(2, 2);
  65 /// cache.insert(3, 3);
  66 /// cache.insert(4, 4);
  67 /// assert_eq!(cache.len(), 3);
  68 ///
  69 /// assert_eq!(cache.get_mut(1), None);
  70 /// assert_eq!(cache.get_mut(2), Some(&mut 2));
  71 /// assert_eq!(cache.get_mut(3), Some(&mut 3));
  72 /// assert_eq!(cache.get_mut(4), Some(&mut 4));
  73 ///
  74 /// cache.remove(4);
  75 /// cache.remove(3);
  76 /// cache.remove(2);
  77 /// assert_eq!(cache.len(), 0);
  78 /// assert_eq!(cache.get_mut(2), None);
  79 /// // access will fill in missing cache entry by fetching from LruCacher
  80 /// assert_eq!(cache.access(2, &mut LruCacher {}).unwrap(), Some(&mut 2));
  81 ///
  82 /// cache.insert(1, 1);
  83 /// assert_eq!(cache.get_mut(1), Some(&mut 1));
  84 ///
  85 /// cache.clear();
  86 /// assert_eq!(cache.len(), 0);
  87 /// assert_eq!(cache.get_mut(1), None);
  88 /// # Ok(())
  89 /// # }
  90 /// ```
  91 pub struct LruCache<K, V> {
  92     /// Quick access to individual nodes via the node pointer.
  93     map: HashMap<K, *mut CacheNode<K, V>>,
  94     /// Actual nodes stored in a linked list.
  95     list: LinkedList<K, V>,
  96     /// Max nodes the cache can hold, temporarily exceeded by 1 due to
  97     /// implementation details.
  98     capacity: usize,
  99     // Dropcheck marker. See the phantom-data section in the rustonomicon.
 100     _marker: PhantomData<Box<CacheNode<K, V>>>,
 101 }
 102
 103 // trivial: if our contents are Send, the whole cache is Send
 104 unsafe impl<K: Send, V: Send> Send for LruCache<K, V> {}
 105
 106 impl<K: std::cmp::Eq + std::hash::Hash + Copy, V> LruCache<K, V> {
 107     /// Create LRU cache instance which holds up to `capacity` nodes at once.
 108     pub fn new(capacity: usize) -> Self {
 109         let capacity = capacity.max(1);
 110         Self {
 111             map: HashMap::with_capacity(capacity),
 112             list: LinkedList::new(),
 113             capacity,
 114             _marker: PhantomData,
 115         }
 116     }
 117
 118     /// Clear all the entries from the cache.
 119     pub fn clear(&mut self) {
 120         // This frees only the HashMap with the node pointers.
 121         self.map.clear();
 122         // This frees the actual nodes and resets the list head and tail.
 123         self.list.clear();
 124     }
 125
 126     /// Insert or update an entry identified by `key` with the given `value`.
 127     /// This entry is placed as the most recently used node at the head.
 128     pub fn insert(&mut self, key: K, value: V) {
 129         match self.map.entry(key) {
 130             Entry::Occupied(mut o) => {
 131                 // Node present, update value
 132                 let node_ptr = *o.get_mut();
 133                 self.list.bring_to_front(node_ptr);
 134                 let mut node = unsafe { Box::from_raw(node_ptr) };
 135                 node.value = value;
 136                 let _node_ptr = Box::into_raw(node);
 137             }
 138             Entry::Vacant(v) => {
 139                 // Node not present, insert a new one
 140                 // Unfortunately we need a copy of the key here, therefore it has
 141                 // to impl the copy trait
 142                 let node = Box::new(CacheNode::new(key, value));
 143                 let node_ptr = Box::into_raw(node);
 144                 self.list.push_front(node_ptr);
 145                 v.insert(node_ptr);
 146                 // If we have more elements than capacity,
 147                 // delete the lists tail node (= oldest node).
 148                 // This needs to be executed after the insert in order to
 149                 // avoid borrow conflict. This means there are temporarily
 150                 // self.capacity + 1 cache nodes.
 151                 if self.map.len() > self.capacity {
 152                     self.pop_tail();
 153                 }
 154             }
 155         }
 156     }
 157
 158     /// Remove the given `key` and its `value` from the cache.
 159     pub fn remove(&mut self, key: K) -> Option<V> {
 160         // Remove node pointer from the HashMap and get ownership of the node
 161         let node_ptr = self.map.remove(&key)?;
 162         let node = self.list.remove(node_ptr);
 163         Some(node.value)
 164     }
 165
 166     /// Remove the least recently used node from the cache.
 167     fn pop_tail(&mut self) {
 168         if let Some(old_tail) = self.list.pop_tail() {
 169             // Remove HashMap entry for old tail
 170             self.map.remove(&old_tail.key);
 171         }
 172     }
 173
 174     /// Get a mutable reference to the value identified by `key`.
 175     /// This will update the cache entry to be the most recently used entry.
 176     /// On cache misses, None is returned.
 177     pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
 178         let node_ptr = self.map.get(&key)?;
 179         self.list.bring_to_front(*node_ptr);
 180         Some(unsafe { &mut (*self.list.head).value })
 181     }
 182
 183     /// Number of entries in the cache.
 184     pub fn len(&self) -> usize {
 185         self.map.len()
 186     }
 187
 188     /// Returns `true` when the cache is empty
 189     pub fn is_empty(&self) -> bool {
 190         self.map.is_empty()
 191     }
 192
 193     /// Get a mutable reference to the value identified by `key`.
 194     /// This will update the cache entry to be the most recently used entry.
 195     /// On cache misses, the cachers fetch method is called to get a corresponding
 196     /// value.
 197     /// If fetch returns a value, it is inserted as the most recently used entry
 198     /// in the cache.
 199     pub fn access<'a>(&'a mut self, key: K, cacher: &mut dyn Cacher<K, V>) -> Result<Option<&'a mut V>, anyhow::Error> {
 200         match self.map.entry(key) {
 201             Entry::Occupied(mut o) => {
 202                 // Cache hit, birng node to front of list
 203                 let node_ptr = *o.get_mut();
 204                 self.list.bring_to_front(node_ptr);
 205             }
 206             Entry::Vacant(v) => {
 207                 // Cache miss, try to fetch from cacher and insert at the front
 208                 match cacher.fetch(key)? {
 209                     None => return Ok(None),
 210                     Some(value) => {
 211                         // Unfortunately we need a copy of the key here, therefore it has
 212                         // to impl the copy trait
 213                         let node = Box::new(CacheNode::new(key, value));
 214                         let node_ptr = Box::into_raw(node);
 215                         self.list.push_front(node_ptr);
 216                         v.insert(node_ptr);
 217                         // If we have more elements than capacity,
 218                         // delete the lists tail node (= oldest node).
 219                         // This needs to be executed after the insert in order to
 220                         // avoid borrow conflict. This means there are temporarily
 221                         // self.capacity + 1 cache nodes.
 222                         if self.map.len() > self.capacity {
 223                             self.pop_tail();
 224                         }
 225                     }
 226                 }
 227             }
 228         }
 229
 230         Ok(Some(unsafe { &mut (*self.list.head).value }))
 231     }
 232 }
 233
 234 /// Linked list holding the nodes of the LruCache.
 235 ///
 236 /// This struct actually holds the CacheNodes via the raw linked list pointers
 237 /// and allows to define the access sequence of these via the list sequence.
 238 /// The LinkedList of the standard library unfortunately does not implement
 239 /// an efficient way to bring list entries to the front, therefore we need our own.
 240 struct LinkedList<K, V> {
 241     head: *mut CacheNode<K, V>,
 242     tail: *mut CacheNode<K, V>,
 243 }
 244
 245 impl<K, V> LinkedList<K, V> {
 246     /// Create a new empty linked list.
 247     fn new() -> Self {
 248         Self {
 249             head: std::ptr::null_mut(),
 250             tail: std::ptr::null_mut(),
 251         }
 252     }
 253
 254     /// Bring the CacheNode referenced by `node_ptr` to the front of the linked list.
 255     fn bring_to_front(&mut self, node_ptr: *mut CacheNode<K, V>) {
 256         if node_ptr == self.head {
 257             // node is already head, just return
 258             return;
 259         }
 260
 261         let mut node = unsafe { Box::from_raw(node_ptr) };
 262         // Update the prev node to point to next (or null if current node is tail)
 263         unsafe { (*node.prev).next = node.next };
 264
 265         // Update the next node or otherwise the tail
 266         if !node.next.is_null() {
 267             unsafe { (*node.next).prev = node.prev };
 268         } else {
 269             // No next node means this was the tail
 270             self.tail = node.prev;
 271         }
 272
 273         node.prev = std::ptr::null_mut();
 274         node.next = self.head;
 275         // update the head and release ownership of the node again
 276         let node_ptr = Box::into_raw(node);
 277         // Update current head
 278         unsafe { (*self.head).prev = node_ptr };
 279         // Update to new head
 280         self.head = node_ptr;
 281     }
 282
 283     /// Insert a new node at the front of the linked list.
 284     fn push_front(&mut self, node_ptr: *mut CacheNode<K, V>) {
 285         let mut node = unsafe { Box::from_raw(node_ptr) };
 286
 287         // Old head gets new heads next
 288         node.next = self.head;
 289         // Release ownership of node, rest can be handled with just the pointer.
 290         let node_ptr = Box::into_raw(node);
 291
 292         // Update the prev for the old head
 293         if !self.head.is_null() {
 294             unsafe { (*self.head).prev = node_ptr };
 295         }
 296
 297         // Update the head to the new node pointer
 298         self.head = node_ptr;
 299
 300         // If there was no old tail, this node will be the new tail too
 301         if self.tail.is_null() {
 302             self.tail = node_ptr;
 303         }
 304     }
 305
 306     /// Remove the node referenced by `node_ptr` from the linked list and return it.
 307     fn remove(&mut self, node_ptr: *mut CacheNode<K, V>) -> Box<CacheNode<K, V>> {
 308         let node = unsafe { Box::from_raw(node_ptr) };
 309
 310         // Update the previous node or otherwise the head
 311         if !node.prev.is_null() {
 312             unsafe { (*node.prev).next = node.next };
 313         } else {
 314             // No previous node means this was the head
 315             self.head = node.next;
 316         }
 317
 318         // Update the next node or otherwise the tail
 319         if !node.next.is_null() {
 320             unsafe { (*node.next).prev = node.prev };
 321         } else {
 322             // No next node means this was the tail
 323             self.tail = node.prev;
 324         }
 325         node
 326     }
 327
 328     /// Remove the tail node from the linked list and return it.
 329     fn pop_tail(&mut self) -> Option<Box<CacheNode<K, V>>> {
 330         if self.tail.is_null() {
 331             return None;
 332         }
 333
 334         let old_tail = unsafe { Box::from_raw(self.tail) };
 335         self.tail = old_tail.prev;
 336         // Update next node for new tail
 337         if !self.tail.is_null() {
 338             unsafe { (*self.tail).next = std::ptr::null_mut() };
 339         }
 340         Some(old_tail)
 341     }
 342
 343     /// Clear the linked list and free all the nodes.
 344     fn clear(&mut self) {
 345         let mut next = self.head;
 346         while !next.is_null() {
 347             // Taking ownership of node and drop it at the end of the block.
 348             let current = unsafe { Box::from_raw(next) };
 349             next = current.next;
 350         }
 351         // Reset head and tail pointers
 352         self.head = std::ptr::null_mut();
 353         self.tail = std::ptr::null_mut();
 354     }
 355 }
 356
 357 #[test]
 358 fn test_linked_list() {
 359     let mut list = LinkedList::new();
 360     for idx in 0..3 {
 361         let node = Box::new(CacheNode::new(idx, idx + 1));
 362         // Get pointer, release ownership.
 363         let node_ptr = Box::into_raw(node);
 364         list.push_front(node_ptr);
 365     }
 366     assert_eq!(unsafe { (*list.head).key }, 2);
 367     assert_eq!(unsafe { (*list.head).value }, 3);
 368     assert_eq!(unsafe { (*list.tail).key }, 0);
 369     assert_eq!(unsafe { (*list.tail).value }, 1);
 370
 371     list.bring_to_front(list.tail);
 372     assert_eq!(unsafe { (*list.head).key }, 0);
 373     assert_eq!(unsafe { (*list.head).value }, 1);
 374     assert_eq!(unsafe { (*list.tail).key }, 1);
 375     assert_eq!(unsafe { (*list.tail).value }, 2);
 376
 377     list.bring_to_front(list.tail);
 378     assert_eq!(unsafe { (*list.head).key }, 1);
 379     assert_eq!(unsafe { (*list.head).value }, 2);
 380     assert_eq!(unsafe { (*list.tail).key }, 2);
 381     assert_eq!(unsafe { (*list.tail).value }, 3);
 382
 383     let tail = list.pop_tail().unwrap();
 384     assert_eq!(tail.key, 2);
 385     assert_eq!(tail.value, 3);
 386     assert_eq!(unsafe { (*list.head).key }, 1);
 387     assert_eq!(unsafe { (*list.head).value }, 2);
 388     assert_eq!(unsafe { (*list.tail).key }, 0);
 389     assert_eq!(unsafe { (*list.tail).value }, 1);
 390
 391     list.clear();
 392     assert!(list.head.is_null());
 393     assert!(list.tail.is_null());
 394 }