subms_block_cache/features/
concurrent_shards.rs1use std::collections::hash_map::DefaultHasher;
14use std::hash::{Hash, Hasher};
15use std::sync::Mutex;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18use crate::BlockCache;
19
20pub struct ShardedCache<K, V> {
21 shards: Vec<Mutex<BlockCache<K, V>>>,
22 contention: AtomicU64,
23}
24
25impl<K, V> ShardedCache<K, V>
26where
27 K: Hash + Eq + Clone,
28 V: Clone,
29{
30 pub fn with_capacity(total_capacity: usize, num_shards: usize) -> Self {
33 let shards_n = num_shards.max(1).next_power_of_two();
34 let per_shard = total_capacity.div_ceil(shards_n).max(1);
35 let mut shards = Vec::with_capacity(shards_n);
36 for _ in 0..shards_n {
37 shards.push(Mutex::new(BlockCache::with_capacity(per_shard)));
38 }
39 Self {
40 shards,
41 contention: AtomicU64::new(0),
42 }
43 }
44
45 pub fn num_shards(&self) -> usize {
46 self.shards.len()
47 }
48 pub fn contention_events(&self) -> u64 {
49 self.contention.load(Ordering::Relaxed)
50 }
51
52 pub fn len(&self) -> usize {
55 self.shards.iter().map(|m| m.lock().unwrap().len()).sum()
56 }
57 pub fn is_empty(&self) -> bool {
58 self.len() == 0
59 }
60
61 fn shard_index(&self, key: &K) -> usize {
62 let mut h = DefaultHasher::new();
63 key.hash(&mut h);
64 (h.finish() as usize) & (self.shards.len() - 1)
65 }
66
67 pub fn get(&self, key: &K) -> Option<V> {
70 let idx = self.shard_index(key);
71 let guard = self.lock_with_contention(idx);
72 let mut g = guard;
73 g.get(key).cloned()
74 }
75
76 pub fn put(&self, key: K, value: V) -> Option<(K, V)> {
78 let idx = self.shard_index(&key);
79 let mut g = self.lock_with_contention(idx);
80 g.put(key, value)
81 }
82
83 pub fn remove(&self, key: &K) -> Option<V> {
85 let idx = self.shard_index(key);
86 let mut g = self.lock_with_contention(idx);
87 g.remove(key)
88 }
89
90 pub fn clear(&self) {
94 for shard in &self.shards {
95 shard.lock().unwrap().clear();
96 }
97 }
98
99 fn lock_with_contention(&self, idx: usize) -> std::sync::MutexGuard<'_, BlockCache<K, V>> {
100 match self.shards[idx].try_lock() {
101 Ok(g) => g,
102 Err(_) => {
103 self.contention.fetch_add(1, Ordering::Relaxed);
104 self.shards[idx].lock().unwrap()
105 }
106 }
107 }
108}
109
110#[cfg(test)]
111#[path = "concurrent_shards_tests.rs"]
112mod tests;