memory_cache/
lib.rs

1// Copyright 2015-2020 Parity Technologies (UK) Ltd.
2// This file is part of Tetsy Vapory.
3
4// Tetsy Vapory 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// Tetsy Vapory 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 Tetsy Vapory.  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.
19// TODO: push changes upstream in a clean way.
20
21extern crate tetsy_util_mem;
22extern crate lru_cache;
23
24use tetsy_util_mem::{MallocSizeOf, MallocSizeOfExt};
25use lru_cache::LruCache;
26
27use std::hash::Hash;
28
29const INITIAL_CAPACITY: usize = 4;
30
31/// An LRU-cache which operates on memory used.
32pub struct MemoryLruCache<K: Eq + Hash, V> {
33	inner: LruCache<K, V>,
34	cur_size: usize,
35	max_size: usize,
36}
37
38// amount of memory used when the item will be put on the heap.
39fn heap_size_of<T: MallocSizeOf>(val: &T) -> usize {
40	::std::mem::size_of::<T>() + val.malloc_size_of()
41}
42
43impl<K: Eq + Hash, V: MallocSizeOf> MemoryLruCache<K, V> {
44	/// Create a new cache with a maximum size in bytes.
45	pub fn new(max_size: usize) -> Self {
46		MemoryLruCache {
47			inner: LruCache::new(INITIAL_CAPACITY),
48			max_size: max_size,
49			cur_size: 0,
50		}
51	}
52
53	/// Insert an item.
54	pub fn insert(&mut self, key: K, val: V) {
55		let cap = self.inner.capacity();
56
57		// grow the cache as necessary; it operates on amount of items
58		// but we're working based on memory usage.
59		if self.inner.len() == cap && self.cur_size < self.max_size {
60			self.inner.set_capacity(cap * 2);
61		}
62
63		self.cur_size += heap_size_of(&val);
64
65		// account for any element displaced from the cache.
66		if let Some(lru) = self.inner.insert(key, val) {
67			self.cur_size -= heap_size_of(&lru);
68		}
69
70		// remove elements until we are below the memory target.
71		while self.cur_size > self.max_size {
72			match self.inner.remove_lru() {
73				Some((_, v)) => self.cur_size -= heap_size_of(&v),
74				_ => break,
75			}
76		}
77	}
78
79	/// Get a reference to an item in the cache. It is a logic error for its
80	/// heap size to be altered while borrowed.
81	pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
82		self.inner.get_mut(key)
83	}
84
85	/// Currently-used size of values in bytes.
86	pub fn current_size(&self) -> usize {
87		self.cur_size
88	}
89
90	/// Get backing LRU cache instance (read only)
91	pub fn backstore(&self) -> &LruCache<K, V> {
92		&self.inner
93	}
94}
95
96#[cfg(test)]
97mod tests {
98	use super::*;
99
100	#[test]
101	fn it_works() {
102		let mut cache = MemoryLruCache::new(256);
103		let val1 = vec![0u8; 100];
104		let size1 = heap_size_of(&val1);
105		cache.insert("hello", val1);
106
107		assert_eq!(cache.current_size(), size1);
108
109		let val2 = vec![0u8; 210];
110		let size2 = heap_size_of(&val2);
111		cache.insert("world", val2);
112
113		assert!(cache.get_mut(&"hello").is_none());
114		assert!(cache.get_mut(&"world").is_some());
115
116		assert_eq!(cache.current_size(), size2);
117	}
118}