matecito/
cache.rs

1use crate::bloom_filter::BloomFilter;
2use crate::matecito_internal::MatecitoInternal;
3
4use parking_lot::Mutex;
5use std::hash::{BuildHasher, Hasher};
6use std::sync::Arc;
7
8const NUM_SHARDS: usize = 256;
9pub(crate) struct Cache<K, T> {
10    hash_builder: twox_hash::RandomXxHashBuilder64,
11    sharded_matecitos: Arc<Vec<Arc<Mutex<MatecitoInternal<K, T>>>>>,
12    bloom_filter: Arc<Mutex<BloomFilter>>,
13    put_threshold: usize,
14}
15
16impl<K: Clone + std::hash::Hash + Ord, T: std::fmt::Debug + Clone> Cache<K, T> {
17    pub fn new(max_size: usize, put_threshold: usize) -> Self {
18        let mut v = Vec::with_capacity(NUM_SHARDS);
19        let hash_builder = twox_hash::RandomXxHashBuilder64::default();
20        for _ in 0..NUM_SHARDS {
21            v.push(Arc::new(Mutex::new(MatecitoInternal::<K, T>::new(
22                max_size / NUM_SHARDS,
23            ))))
24        }
25        let sharded_matecitos = Arc::new(v);
26        let bloom_filter = BloomFilter::new(max_size as u64, 1000);
27        Self {
28            hash_builder,
29            sharded_matecitos,
30            bloom_filter: Arc::new(Mutex::new(bloom_filter)),
31            put_threshold,
32        }
33    }
34
35    pub(crate) fn put(&self, key: K, value: T) {
36        // TODO: Check whether it makes sense to put it or not... smart stats missing.
37        let (count, max_count) = {
38            let mut bloom_filter = (*self.bloom_filter).lock();
39            bloom_filter.increment(&key);
40            (bloom_filter.count_present(&key), bloom_filter.max_count)
41        };
42        let slot = self.get_shard(key.clone());
43        let mut matecito = (*self.sharded_matecitos)[slot].lock();
44        // It makes sense to add it either:
45        // - there is space in the cache
46        // - we've seen it a lot!
47        if !matecito.is_full() || count > self.put_threshold as u64 {
48            if let Some(another_key) = matecito.put(&key, value) {
49                let mut bloom_filter = (*self.bloom_filter).lock();
50                bloom_filter.decrement(another_key);
51            }
52            if max_count > (self.put_threshold * 10) as u64 {
53                let mut bloom_filter = (*self.bloom_filter).lock();
54                bloom_filter.halve();
55            }
56        }
57    }
58
59    pub(crate) fn get(&self, key: K) -> Option<T> {
60        let mut bloom_filter = (*self.bloom_filter).lock();
61        bloom_filter.increment(&key);
62
63        let slot = self.get_shard(key.clone());
64        let mut matecito = (*self.sharded_matecitos)[slot].lock();
65        matecito.get(key).cloned()
66    }
67
68    #[inline]
69    fn get_shard(&self, key: K) -> usize {
70        let mut state = self.hash_builder.build_hasher();
71        key.hash(&mut state);
72        state.finish() as usize % NUM_SHARDS
73    }
74}
75
76impl<K, T> Clone for Cache<K, T> {
77    fn clone(&self) -> Self {
78        let sharded_matecitos = self.sharded_matecitos.clone();
79        let hash_builder = self.hash_builder.clone();
80        let bloom_filter = self.bloom_filter.clone();
81        let put_threshold = self.put_threshold;
82        Self {
83            sharded_matecitos,
84            hash_builder,
85            bloom_filter,
86            put_threshold,
87        }
88    }
89}
90
91unsafe impl<K, T> Send for Cache<K, T> {}
92unsafe impl<K, T> Sync for Cache<K, T> {}