memory_cache/
lib.rs

1// Copyright 2015-2020 Parity Technologies (UK) Ltd.
2// This file is part of Parity Ethereum.
3
4// Parity Ethereum is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// Parity Ethereum is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with Parity Ethereum.  If not, see <http://www.gnu.org/licenses/>.
16
17//! Lru-cache related utilities as quick-and-dirty wrappers around the lru-cache
18//! crate.
19use lru_cache::LruCache;
20use malloc_size_of::{new_malloc_size_ops, MallocSizeOf, MallocSizeOfOps};
21
22use std::hash::Hash;
23
24const INITIAL_CAPACITY: usize = 4;
25
26/// An LRU-cache which operates on memory used.
27pub struct MemoryLruCache<K: Eq + Hash, V> {
28    inner: LruCache<K, V>,
29    malloc_size_of_ops: MallocSizeOfOps,
30    cur_size: usize,
31    max_size: usize,
32}
33
34impl<K: Eq + Hash, V: MallocSizeOf> MemoryLruCache<K, V> {
35    // amount of memory used when the item will be put on the heap.
36    fn heap_size_of(&mut self, val: &V) -> usize {
37        ::std::mem::size_of::<V>() + val.size_of(&mut self.malloc_size_of_ops)
38    }
39}
40
41impl<K: Eq + Hash, V: MallocSizeOf> MemoryLruCache<K, V> {
42    /// Create a new cache with a maximum size in bytes.
43    pub fn new(max_size: usize) -> Self {
44        MemoryLruCache {
45            inner: LruCache::new(INITIAL_CAPACITY),
46            malloc_size_of_ops: new_malloc_size_ops(),
47            max_size,
48            cur_size: 0,
49        }
50    }
51
52    /// Insert an item.
53    pub fn insert(&mut self, key: K, val: V) {
54        let cap = self.inner.capacity();
55
56        // grow the cache as necessary; it operates on amount of items
57        // but we're working based on memory usage.
58        if self.inner.len() == cap && self.cur_size < self.max_size {
59            self.inner.set_capacity(cap * 2);
60        }
61
62        self.cur_size += self.heap_size_of(&val);
63
64        // account for any element displaced from the cache.
65        if let Some(lru) = self.inner.insert(key, val) {
66            self.cur_size -= self.heap_size_of(&lru);
67        }
68
69        // remove elements until we are below the memory target.
70        while self.cur_size > self.max_size {
71            match self.inner.remove_lru() {
72                Some((_, v)) => self.cur_size -= self.heap_size_of(&v),
73                _ => break,
74            }
75        }
76    }
77
78    /// Get a reference to an item in the cache. It is a logic error for its
79    /// heap size to be altered while borrowed.
80    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
81        self.inner.get_mut(key)
82    }
83
84    /// Currently-used size of values in bytes.
85    pub fn current_size(&self) -> usize { self.cur_size }
86
87    /// Get backing LRU cache instance (read only)
88    pub fn backstore(&self) -> &LruCache<K, V> { &self.inner }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use cfx_mallocator_utils;
95    #[global_allocator]
96    static ALLOC: cfx_mallocator_utils::allocator::Allocator =
97        cfx_mallocator_utils::allocator::new_allocator();
98
99    #[test]
100    fn it_works() {
101        let mut cache = MemoryLruCache::new(256);
102        let val1 = vec![0u8; 100];
103        let size1 = cache.heap_size_of(&val1);
104        cache.insert("hello", val1);
105
106        assert_eq!(cache.current_size(), size1);
107
108        let val2 = vec![0u8; 210];
109        let size2 = cache.heap_size_of(&val2);
110        cache.insert("world", val2);
111
112        assert!(cache.get_mut(&"hello").is_none());
113        assert!(cache.get_mut(&"world").is_some());
114
115        assert_eq!(cache.current_size(), size2);
116    }
117}