Skip to main content

vtcode_commons/
lr_map.rs

1//! Lock-free concurrent map built on [`left_right`].
2//!
3//! [`LrMap`] keeps two copies of a `HashMap` — readers see one copy while the
4//! writer mutates the other. On publish the copies swap, giving readers a
5//! consistent, wait-free snapshot.
6//!
7//! Best for read-heavy workloads with infrequent writes (caches, registries).
8
9use hashbrown::HashMap;
10use left_right::{Absorb, ReadHandleFactory, WriteHandle};
11use std::hash::Hash;
12use std::sync::Mutex;
13
14enum MapOp<K, V> {
15    Insert(K, V),
16    Clear,
17}
18
19impl<K: Eq + Hash + Clone, V: Clone> Absorb<MapOp<K, V>> for HashMap<K, V> {
20    fn absorb_first(&mut self, operation: &mut MapOp<K, V>, _other: &Self) {
21        match operation {
22            MapOp::Insert(k, v) => {
23                self.insert(k.clone(), v.clone());
24            }
25            MapOp::Clear => self.clear(),
26        }
27    }
28
29    fn sync_with(&mut self, first: &Self) {
30        self.clone_from(first);
31    }
32}
33
34/// A concurrent map optimized for read-heavy workloads.
35///
36/// Readers never block — not even while a write is in progress. Writers are
37/// serialized through an internal [`Mutex`].
38///
39/// Trade-off: doubled memory (two copies of the map).
40pub struct LrMap<K: Eq + Hash + Clone, V: Clone> {
41    reader_factory: ReadHandleFactory<HashMap<K, V>>,
42    writer: Mutex<WriteHandle<HashMap<K, V>, MapOp<K, V>>>,
43}
44
45impl<K, V> LrMap<K, V>
46where
47    K: Eq + Hash + Clone + Send + Sync,
48    V: Clone + Send + Sync,
49{
50    pub fn new() -> Self {
51        let (writer, reader) = left_right::new_from_empty(HashMap::new());
52        let factory = reader.factory();
53        Self {
54            reader_factory: factory,
55            writer: Mutex::new(writer),
56        }
57    }
58
59    /// Lock-free lookup returning a clone of the value.
60    pub fn get<Q>(&self, key: &Q) -> Option<V>
61    where
62        K: std::borrow::Borrow<Q>,
63        Q: Hash + Eq + ?Sized,
64    {
65        let reader = self.reader_factory.handle();
66        reader.enter().and_then(|map| map.get(key).cloned())
67    }
68
69    pub fn insert(&self, key: K, value: V) {
70        match self.writer.lock() {
71            Ok(mut w) => {
72                w.append(MapOp::Insert(key, value));
73                w.publish();
74            }
75            Err(e) => {
76                // Log the lock poisoning error instead of silently dropping the write
77                eprintln!(
78                    "WARNING: LrMap::insert failed due to poisoned mutex: {}. Write dropped.",
79                    e
80                );
81            }
82        }
83    }
84
85    pub fn clear(&self) {
86        match self.writer.lock() {
87            Ok(mut w) => {
88                w.append(MapOp::Clear);
89                w.publish();
90            }
91            Err(e) => {
92                // Log the lock poisoning error instead of silently dropping the write
93                eprintln!(
94                    "WARNING: LrMap::clear failed due to poisoned mutex: {}. Operation dropped.",
95                    e
96                );
97            }
98        }
99    }
100
101    pub fn len(&self) -> usize {
102        let reader = self.reader_factory.handle();
103        reader.enter().map(|m| m.len()).unwrap_or(0)
104    }
105
106    pub fn is_empty(&self) -> bool {
107        self.len() == 0
108    }
109}
110
111impl<K, V> Default for LrMap<K, V>
112where
113    K: Eq + Hash + Clone + Send + Sync,
114    V: Clone + Send + Sync,
115{
116    fn default() -> Self {
117        Self::new()
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use std::sync::Arc;
125
126    #[test]
127    fn insert_and_get() {
128        let map: LrMap<String, i32> = LrMap::new();
129        map.insert("a".into(), 1);
130        assert_eq!(map.get("a"), Some(1));
131        assert_eq!(map.get("b"), None);
132    }
133
134    #[test]
135    fn overwrite_key() {
136        let map: LrMap<String, i32> = LrMap::new();
137        map.insert("a".into(), 1);
138        map.insert("a".into(), 2);
139        assert_eq!(map.get("a"), Some(2));
140    }
141
142    #[test]
143    fn clear_removes_all() {
144        let map: LrMap<String, i32> = LrMap::new();
145        map.insert("a".into(), 1);
146        map.insert("b".into(), 2);
147        map.clear();
148        assert!(map.is_empty());
149    }
150
151    #[test]
152    fn concurrent_reads() {
153        let map: Arc<LrMap<String, i32>> = Arc::new(LrMap::new());
154        map.insert("key".into(), 42);
155
156        let handles: Vec<_> = (0..4)
157            .map(|_| {
158                let m = Arc::clone(&map);
159                std::thread::spawn(move || {
160                    for _ in 0..100 {
161                        assert_eq!(m.get("key"), Some(42));
162                    }
163                })
164            })
165            .collect();
166
167        for h in handles {
168            h.join().expect("reader thread panicked");
169        }
170    }
171}