k_cache/
segmented.rs

1use std::{borrow::Borrow, hash::BuildHasher};
2
3use crate::{Cache, One, Weigher};
4
5#[derive(Debug)]
6pub struct SegmentedCache<K, V, S: BuildHasher = std::hash::RandomState, W: Weigher<K, V> = One> {
7    segments: Vec<k_lock::Mutex<Cache<K, V, S, W>>>,
8    hasher: S,
9}
10
11impl<K, V, S, W> SegmentedCache<K, V, S, W>
12where
13    K: Eq + std::hash::Hash + Clone,
14    V: Clone,
15    S: BuildHasher + Default,
16    W: Weigher<K, V> + Clone,
17{
18    pub fn new(segments: usize, max_weight: usize) -> Self {
19        let weight_per_segment = max_weight / segments;
20        let segments = (0..segments)
21            .map(|_| k_lock::Mutex::new(Cache::new(S::default(), weight_per_segment)))
22            .collect();
23        Self {
24            segments,
25            hasher: S::default(),
26        }
27    }
28
29    pub fn put(&self, key: K, value: V) {
30        let slot = self.hasher.hash_one(&key) as usize % self.segments.len();
31        self.segments[slot]
32            .lock()
33            .expect("mutex must not be poisoned")
34            .put(key, value)
35    }
36
37    pub fn get<Q>(&self, key: &Q) -> Option<V>
38    where
39        K: Borrow<Q>,
40        Q: std::hash::Hash + Eq + ?Sized,
41    {
42        let slot = self.hasher.hash_one(key) as usize % self.segments.len();
43        self.segments[slot]
44            .lock()
45            .expect("mutex must not be poisoned")
46            .get(key)
47            .cloned()
48    }
49}