Skip to main content

subms_block_cache/features/
concurrent_shards.rs

1//! Sharded block cache: split keyspace into N independent shards by
2//! `hash(key) % N`. Each shard wraps the base `BlockCache` behind its
3//! own `Mutex`, so concurrent readers/writers on different shards
4//! don't contend.
5//!
6//! Per-shard contention counter (`try_lock` failures) is exposed via
7//! `contention_events()` for the `metrics` feature integration. The
8//! ShardedCache itself only tracks the counter when a put/get path
9//! actually backs off; it does NOT change correctness behaviour - if
10//! `try_lock` would have failed, we fall through to the blocking lock
11//! and still complete the op.
12
13use 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    /// Create a sharded cache with `total_capacity` distributed across
31    /// `num_shards` shards (rounded up). Each shard gets at least 1 slot.
32    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    /// Aggregate length across shards. Snapshot; under concurrent
53    /// access the returned number may be slightly stale.
54    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    /// Get a cloned value for `key`. Cloning matters because returning
68    /// a borrow would extend the MutexGuard across the call site.
69    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    /// Insert or update. Returns the evicted entry if eviction occurred.
77    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    /// Invalidate `key` in its own shard. Only that shard is locked.
84    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    /// Drop every entry. Shards are cleared one at a time, so a concurrent
91    /// writer can land in an already-cleared shard - this is a bulk
92    /// invalidation, not a global barrier.
93    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;