key_mutex/
std.rs

1//! Implementations of [`KeyMutex`] and [`KeyRwLock`] using `std::sync::{Mutex, RwLock}`.
2
3use crate::{
4    inner::{KeyRef, LockMap},
5    Empty,
6};
7use guardian::{ArcMutexGuardian, ArcRwLockReadGuardian, ArcRwLockWriteGuardian};
8use std::{
9    hash::Hash,
10    mem::ManuallyDrop,
11    ops::{Deref, DerefMut},
12    sync::{self, LockResult, PoisonError},
13};
14
15fn map_guard<G, R>(result: LockResult<G>, f: impl FnOnce(G) -> R) -> LockResult<R> {
16    match result {
17        Ok(guard) => Ok(f(guard)),
18        Err(poisoned) => Err(PoisonError::new(f(poisoned.into_inner()))),
19    }
20}
21
22decl_mutex_guard!(ArcMutexGuardian);
23#[derive(Clone, Default)]
24pub struct KeyMutex<K: Eq + Hash, V>(LockMap<K, sync::Mutex<V>>);
25impl<K: Eq + Hash + Clone, V: Empty + Default> KeyMutex<K, V> {
26    pub fn new() -> Self {
27        Self(LockMap::new())
28    }
29
30    pub fn lock(&self, key: K) -> LockResult<OwnedMutexGuard<K, V>> {
31        let lock = self.0.obtain(key.clone());
32        map_guard(ArcMutexGuardian::take(lock), |guard| OwnedMutexGuard {
33            key_ref: KeyRef::new(&self.0, key),
34            guard: ManuallyDrop::new(guard),
35        })
36    }
37
38    pub fn len(&self) -> usize {
39        self.0.len()
40    }
41
42    pub fn is_empty(&self) -> bool {
43        self.0.is_empty()
44    }
45}
46
47decl_rwlock_guard!(ArcRwLockReadGuardian, ArcRwLockWriteGuardian);
48#[derive(Clone, Default)]
49pub struct KeyRwLock<K: Eq + Hash, V>(LockMap<K, sync::RwLock<V>>);
50impl<K: Eq + Hash + Clone, V: Empty + Default> KeyRwLock<K, V> {
51    pub fn new() -> Self {
52        Self(LockMap::new())
53    }
54
55    pub fn read(&self, key: K) -> LockResult<OwnedReadGuard<K, V>> {
56        let lock = self.0.obtain(key.clone());
57        map_guard(ArcRwLockReadGuardian::take(lock), |guard| OwnedReadGuard {
58            key: KeyRef::new(&self.0, key),
59            guard: ManuallyDrop::new(guard),
60        })
61    }
62
63    pub fn write(&self, key: K) -> LockResult<OwnedWriteGuard<K, V>> {
64        let lock = self.0.obtain(key.clone());
65        map_guard(ArcRwLockWriteGuardian::take(lock), |guard| {
66            OwnedWriteGuard {
67                key: KeyRef::new(&self.0, key),
68                guard: ManuallyDrop::new(guard),
69            }
70        })
71    }
72
73    pub fn len(&self) -> usize {
74        self.0.len()
75    }
76
77    pub fn is_empty(&self) -> bool {
78        self.0.is_empty()
79    }
80}
81
82impl<K: Eq + Hash, V> Empty for KeyMutex<K, V> {
83    fn is_empty(&self) -> bool {
84        self.0.is_empty()
85    }
86}
87
88impl<K: Eq + Hash, V> Empty for KeyRwLock<K, V> {
89    fn is_empty(&self) -> bool {
90        self.0.is_empty()
91    }
92}
93
94#[cfg(test)]
95mod test {
96    use crate::KeyMutex;
97    use std::collections::BTreeSet;
98
99    #[test]
100    fn drop_only_if_empty() {
101        let locks = KeyMutex::<u32, BTreeSet<String>>::new();
102
103        let mut lock = locks.lock(1).unwrap();
104        lock.insert("Hello".to_owned());
105        lock.insert("World".to_owned());
106        drop(lock);
107
108        // Value is not empty and thus is not dropped
109        assert_eq!(locks.len(), 1);
110
111        let mut lock = locks.lock(1).unwrap();
112        assert_eq!(lock.len(), 2);
113        lock.clear();
114        drop(lock);
115
116        // Should be dropped now
117        assert_eq!(locks.len(), 0);
118    }
119}
key_mutex/std.rs

key_mutex/
std.rs
