k_cache/
segmented.rs

1use std::{borrow::Borrow, hash::BuildHasher};
2
3use crate::{
4    cache::{DefaultLifecycle, Lifecycle},
5    Cache, One, Weigher,
6};
7
8#[derive(Debug)]
9pub struct SegmentedCache<
10    K,
11    V,
12    S: BuildHasher = std::hash::RandomState,
13    W: Weigher<K, V> = One,
14    L: Lifecycle<K, V> = DefaultLifecycle,
15> {
16    #[allow(clippy::type_complexity)]
17    segments: Vec<k_lock::Mutex<Cache<K, V, S, W, L>>>,
18    hasher: S,
19}
20
21impl<K, V, S, W, L> SegmentedCache<K, V, S, W, L>
22where
23    K: Eq + std::hash::Hash + Clone,
24    V: Clone,
25    S: BuildHasher + Default,
26    W: Weigher<K, V> + Clone,
27    L: Lifecycle<K, V> + Default,
28{
29    pub fn new(segments: usize, max_weight: usize) -> Self {
30        let weight_per_segment = max_weight / segments;
31        let segments = (0..segments)
32            .map(|_| {
33                k_lock::Mutex::new(Cache::new_with_lifecycle(
34                    S::default(),
35                    weight_per_segment,
36                    L::default(),
37                ))
38            })
39            .collect();
40        Self {
41            segments,
42            hasher: S::default(),
43        }
44    }
45}
46
47impl<K, V, S, W, L> SegmentedCache<K, V, S, W, L>
48where
49    K: Eq + std::hash::Hash + Clone,
50    V: Clone,
51    S: BuildHasher + Default,
52    W: Weigher<K, V> + Clone,
53    L: Lifecycle<K, V> + Clone,
54{
55    pub fn new_with_lifecycle(segments: usize, max_weight: usize, lifecycle: L) -> Self {
56        let weight_per_segment = max_weight / segments;
57        let segments = (0..segments)
58            .map(|_| {
59                k_lock::Mutex::new(Cache::new_with_lifecycle(
60                    S::default(),
61                    weight_per_segment,
62                    lifecycle.clone(),
63                ))
64            })
65            .collect();
66        Self {
67            segments,
68            hasher: S::default(),
69        }
70    }
71
72    pub fn put(&self, key: K, value: V) {
73        let slot = self.hasher.hash_one(&key) as usize % self.segments.len();
74        self.segments[slot]
75            .lock()
76            .expect("mutex must not be poisoned")
77            .put(key, value)
78    }
79
80    pub fn remove(&self, key: &K) -> Option<V> {
81        let slot = self.hasher.hash_one(key) as usize % self.segments.len();
82        self.segments[slot]
83            .lock()
84            .expect("mutex must not be poisoned")
85            .remove(key)
86    }
87
88    pub fn get<Q>(&self, key: &Q) -> Option<V>
89    where
90        K: Borrow<Q>,
91        Q: std::hash::Hash + Eq + ?Sized,
92    {
93        let slot = self.hasher.hash_one(key) as usize % self.segments.len();
94        self.segments[slot]
95            .lock()
96            .expect("mutex must not be poisoned")
97            .get(key)
98            .cloned()
99    }
100    /// Clears the cache, via eviction.
101    ///
102    /// If you want to handle the evicted entries, use the Lifecycle trait.
103    /// The order of eviction is unspecified.
104    ///
105    /// This function does not guarantee that the cache is empty after calling it
106    /// in the presence of concurrent writes. If you want to empty the cache,
107    /// you need to stop writes first, and then call this function.
108    /// If you just want to reset the cache, then you can call this function with
109    /// concurrent writes. It's safe, it's just not going to be empty.
110    ///
111    /// It will block `1/segments` of the cache at a time for segment size time, so
112    /// it's potentially fairly intrusive with large segments, especially with an
113    /// expensive on_eviction Lifecycle if you have one.
114    pub fn evict_all(&self) {
115        for segment in &self.segments {
116            segment
117                .lock()
118                .expect("mutex must not be poisoned")
119                .evict_all();
120        }
121    }
122}
k_cache/segmented.rs

k_cache/
segmented.rs
