moka/sync_base/
key_lock.rs

Help
1use std::{
2    hash::{BuildHasher, Hash},
3    sync::Arc,
4};
5
6use crate::{cht::SegmentedHashMap, common::concurrent::arc::MiniArc};
7
8use parking_lot::{Mutex, MutexGuard};
9
10const LOCK_MAP_NUM_SEGMENTS: usize = 64;
11
12type LockMap<K, S> = SegmentedHashMap<Arc<K>, MiniArc<Mutex<()>>, S>;
13
14// We need the `where` clause here because of the Drop impl.
15pub(crate) struct KeyLock<'a, K, S>
16where
17    K: Eq + Hash,
18    S: BuildHasher,
19{
20    map: &'a LockMap<K, S>,
21    key: Arc<K>,
22    hash: u64,
23    lock: MiniArc<Mutex<()>>,
24}
25
26impl<K, S> Drop for KeyLock<'_, K, S>
27where
28    K: Eq + Hash,
29    S: BuildHasher,
30{
31    fn drop(&mut self) {
32        if MiniArc::count(&self.lock) <= 2 {
33            self.map.remove_if(
34                self.hash,
35                |k| k == &self.key,
36                |_k, v| MiniArc::count(v) <= 2,
37            );
38        }
39    }
40}
41
42impl<'a, K, S> KeyLock<'a, K, S>
43where
44    K: Eq + Hash,
45    S: BuildHasher,
46{
47    fn new(map: &'a LockMap<K, S>, key: &Arc<K>, hash: u64, lock: MiniArc<Mutex<()>>) -> Self {
48        Self {
49            map,
50            key: Arc::clone(key),
51            hash,
52            lock,
53        }
54    }
55
56    pub(crate) fn lock(&self) -> MutexGuard<'_, ()> {
57        self.lock.lock()
58    }
59}
60
61pub(crate) struct KeyLockMap<K, S> {
62    locks: LockMap<K, S>,
63}
64
65impl<K, S> KeyLockMap<K, S>
66where
67    K: Eq + Hash,
68    S: BuildHasher,
69{
70    pub(crate) fn with_hasher(hasher: S) -> Self {
71        Self {
72            locks: SegmentedHashMap::with_num_segments_and_hasher(LOCK_MAP_NUM_SEGMENTS, hasher),
73        }
74    }
75
76    pub(crate) fn key_lock(&self, key: &Arc<K>) -> KeyLock<'_, K, S> {
77        let hash = self.locks.hash(key);
78        let kl = MiniArc::new(Mutex::new(()));
79        match self
80            .locks
81            .insert_if_not_present(Arc::clone(key), hash, kl.clone())
82        {
83            None => KeyLock::new(&self.locks, key, hash, kl),
84            Some(existing_kl) => KeyLock::new(&self.locks, key, hash, existing_kl),
85        }
86    }
87}
88
89#[cfg(test)]
90impl<K, S> KeyLockMap<K, S> {
91    pub(crate) fn is_empty(&self) -> bool {
92        self.locks.len() == 0
93    }
94}
moka/sync_base/key_lock.rs

moka/sync_base/
key_lock.rs
