starnix_rcu/
rcu_hash_map.rs

1// Copyright 2025 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use fuchsia_rcu::{RcuDroppable, RcuReadScope};
6use fuchsia_rcu_collections::rcu_raw_hash_map::{InsertionResult, RcuRawHashMap};
7use starnix_sync::Mutex;
8use std::borrow::Borrow;
9use std::hash::{BuildHasher, Hash};
10
11/// A concurrent hash map that uses RCU for read synchronization and a mutex for write synchronization.
12///
13/// This map allows concurrent readers to access entries without blocking, while writers are
14/// synchronized via a `Mutex`.
15///
16/// By default, this map uses `rapidhash::RapidBuildHasher`, which provides high performance.
17/// However, if this map holds keys which may be attacker-controlled, consider using
18/// `std::collections::hash_map::RandomState` instead.
19#[derive(RcuDroppable)]
20pub struct RcuHashMap<K, V, S = rapidhash::RapidBuildHasher>
21where
22    K: Eq + Hash + Clone + RcuDroppable + Sync,
23    V: Clone + RcuDroppable + Sync,
24    S: BuildHasher + Send + Sync + 'static,
25{
26    map: RcuRawHashMap<K, V, S>,
27    mutex: Mutex<()>,
28}
29
30impl<K, V> Default for RcuHashMap<K, V, rapidhash::RapidBuildHasher>
31where
32    K: Eq + Hash + Clone + RcuDroppable + Sync,
33    V: Clone + RcuDroppable + Sync,
34{
35    fn default() -> Self {
36        Self { map: Default::default(), mutex: Mutex::new(()) }
37    }
38}
39
40impl<K, V, S> RcuHashMap<K, V, S>
41where
42    K: Eq + Hash + Clone + RcuDroppable + Sync,
43    V: Clone + RcuDroppable + Sync,
44    S: BuildHasher + Send + Sync + 'static,
45{
46    /// Creates a new hash map with the given capacity and hasher.
47    pub fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> Self {
48        Self {
49            map: RcuRawHashMap::with_capacity_and_hasher(capacity, hash_builder),
50            mutex: Mutex::new(()),
51        }
52    }
53
54    /// Creates a new hash map with the given hasher.
55    pub fn with_hasher(hash_builder: S) -> Self {
56        Self { map: RcuRawHashMap::with_hasher(hash_builder), mutex: Mutex::new(()) }
57    }
58
59    /// Returns a reference to the value associated with the given key, if it exists.
60    ///
61    /// The returned reference is bound to the lifetime of the `RcuReadScope`.
62    pub fn get<'a, Q>(&self, scope: &'a RcuReadScope, key: &Q) -> Option<&'a V>
63    where
64        K: Borrow<Q>,
65        Q: ?Sized + Hash + Eq,
66    {
67        self.map.get(scope, key)
68    }
69
70    /// Locks the map for exclusive access, returning a guard that allows mutation.
71    pub fn lock(&self) -> RcuHashMapGuard<'_, K, V, S> {
72        RcuHashMapGuard { map: &self.map, _guard: self.mutex.lock() }
73    }
74
75    /// Inserts a key-value pair into the map, returning the old value if the key was already present.
76    pub fn insert(&self, key: K, value: V) -> Option<V> {
77        self.lock().insert(key, value)
78    }
79
80    /// Removes a key from the map, returning the value if the key was present.
81    pub fn remove<Q>(&self, key: &Q) -> Option<V>
82    where
83        K: Borrow<Q>,
84        Q: ?Sized + Hash + Eq,
85    {
86        self.lock().remove(key)
87    }
88
89    /// Returns an iterator over the map's entries.
90    pub fn iter<'a>(&'a self, scope: &'a RcuReadScope) -> impl Iterator<Item = (&'a K, &'a V)> {
91        let mut cursor = self.map.cursor(scope);
92        std::iter::from_fn(move || {
93            let current = cursor.current();
94            if current.is_some() {
95                cursor.advance();
96            }
97            current
98        })
99    }
100
101    /// Returns an iterator over the map's keys.
102    pub fn keys<'a>(&'a self, scope: &'a RcuReadScope) -> impl Iterator<Item = &'a K> {
103        self.iter(scope).map(|(k, _)| k)
104    }
105
106    /// Returns the number of entries in the map.
107    pub fn len(&self) -> usize {
108        self.map.len()
109    }
110}
111
112// TODO(https://fxbug.dev/482462174): switch back to #[derive(Debug)]
113impl<K, V, S> std::fmt::Debug for RcuHashMap<K, V, S>
114where
115    K: Eq + Hash + std::fmt::Debug + Clone + RcuDroppable + Sync,
116    V: std::fmt::Debug + Clone + RcuDroppable + Sync,
117    S: BuildHasher + Send + Sync + 'static,
118{
119    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120        f.debug_struct("RcuHashMap").field("map", &self.map).finish()
121    }
122}
123
124/// A guard that provides exclusive access to the `RcuHashMap`.
125pub struct RcuHashMapGuard<'a, K, V, S = rapidhash::RapidBuildHasher>
126where
127    K: Eq + Hash + Clone + RcuDroppable + Sync,
128    V: Clone + RcuDroppable + Sync,
129    S: BuildHasher + Send + Sync + 'static,
130{
131    map: &'a RcuRawHashMap<K, V, S>,
132    _guard: starnix_sync::MutexGuard<'a, ()>,
133}
134
135impl<'a, K, V, S> RcuHashMapGuard<'a, K, V, S>
136where
137    K: Eq + Hash + Clone + RcuDroppable + Sync,
138    V: Clone + RcuDroppable + Sync,
139    S: BuildHasher + Send + Sync + 'static,
140{
141    /// Returns a copy (clone) of the value associated with the given key, if it exists.
142    pub fn get<Q>(&self, key: &Q) -> Option<V>
143    where
144        K: Borrow<Q>,
145        Q: ?Sized + Hash + Eq,
146    {
147        let scope = RcuReadScope::new();
148        self.map.get(&scope, key).cloned()
149    }
150
151    /// Inserts a key-value pair into the map.
152    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
153        let scope = RcuReadScope::new();
154        // SAFETY: We have exclusive access to the map because we have exclusive access to the mutex.
155        match unsafe { self.map.insert(&scope, key, value) } {
156            InsertionResult::Inserted(_) => None,
157            InsertionResult::Updated(old_value) => Some(old_value),
158        }
159    }
160
161    /// Removes a key from the map.
162    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
163    where
164        K: Borrow<Q>,
165        Q: ?Sized + Hash + Eq,
166    {
167        // SAFETY: We have exclusive access to the map because we have exclusive access to the mutex.
168        unsafe { self.map.remove(key) }
169    }
170
171    /// Removes all values from the map and returns them.
172    pub fn drain<'b>(&'b mut self) -> impl Iterator<Item = (K, V)> + 'b {
173        let scope = RcuReadScope::new();
174        // We collect the keys first because we cannot iterate and modify the map at the same time.
175        #[allow(clippy::needless_collect)]
176        let keys: Vec<_> = self.map.keys(&scope).map(Clone::clone).collect();
177        keys.into_iter().filter_map(move |k| self.remove(&k).map(|v| (k, v)))
178    }
179
180    /// Returns true if the map contains a value for the specified key.
181    pub fn contains_key<Q>(&self, key: &Q) -> bool
182    where
183        K: Borrow<Q>,
184        Q: ?Sized + Hash + Eq,
185    {
186        self.get(key).is_some()
187    }
188
189    /// Gets the given key's corresponding entry in the map for in-place manipulation.
190    pub fn entry<'b>(&'b mut self, key: K) -> Entry<'b, 'a, K, V, S> {
191        if self.get(&key).is_some() {
192            Entry::Occupied(OccupiedEntry { guard: self, key })
193        } else {
194            Entry::Vacant(VacantEntry { guard: self, key })
195        }
196    }
197}
198
199/// A view into a single entry in the map, which may either be vacant or occupied.
200pub enum Entry<'b, 'a, K, V, S = rapidhash::RapidBuildHasher>
201where
202    K: Eq + Hash + Clone + RcuDroppable + Sync,
203    V: Clone + RcuDroppable + Sync,
204    S: BuildHasher + Send + Sync + 'static,
205{
206    /// An occupied entry.
207    Occupied(OccupiedEntry<'b, 'a, K, V, S>),
208    /// A vacant entry.
209    Vacant(VacantEntry<'b, 'a, K, V, S>),
210}
211
212impl<'b, 'a, K, V, S> Entry<'b, 'a, K, V, S>
213where
214    K: Eq + Hash + Clone + RcuDroppable + Sync,
215    V: Clone + RcuDroppable + Sync,
216    S: BuildHasher + Send + Sync + 'static,
217{
218    /// Ensures a value is in the entry by inserting the result of the default function if empty,
219    /// and returns an occupied entry.
220    pub fn or_insert_with<F: FnOnce() -> V>(self, default: F) -> OccupiedEntry<'b, 'a, K, V, S> {
221        match self {
222            Entry::Occupied(entry) => entry,
223            Entry::Vacant(entry) => entry.insert_entry(default()),
224        }
225    }
226}
227
228/// A view into an occupied entry in a `RcuHashMap`.
229pub struct OccupiedEntry<'b, 'a, K, V, S = rapidhash::RapidBuildHasher>
230where
231    K: Eq + Hash + Clone + RcuDroppable + Sync,
232    V: Clone + RcuDroppable + Sync,
233    S: BuildHasher + Send + Sync + 'static,
234{
235    guard: &'b mut RcuHashMapGuard<'a, K, V, S>,
236    key: K,
237}
238
239impl<K, V, S> OccupiedEntry<'_, '_, K, V, S>
240where
241    K: Eq + Hash + Clone + RcuDroppable + Sync,
242    V: Clone + RcuDroppable + Sync,
243    S: BuildHasher + Send + Sync + 'static,
244{
245    /// Gets a copy (clone) of the value in the entry.
246    pub fn get(&self) -> V {
247        self.guard.get(&self.key).unwrap()
248    }
249
250    /// Sets the value of the entry, returning the old value.
251    pub fn insert(&mut self, value: V) -> V {
252        self.guard.insert(self.key.clone(), value).unwrap()
253    }
254
255    /// Removes the entry from the map, returning the value.
256    pub fn remove(self) -> V {
257        self.guard.remove(&self.key).unwrap()
258    }
259}
260
261/// A view into a vacant entry in a `RcuHashMap`.
262pub struct VacantEntry<'b, 'a, K, V, S = rapidhash::RapidBuildHasher>
263where
264    K: Eq + Hash + Clone + RcuDroppable + Sync,
265    V: Clone + RcuDroppable + Sync,
266    S: BuildHasher + Send + Sync + 'static,
267{
268    guard: &'b mut RcuHashMapGuard<'a, K, V, S>,
269    key: K,
270}
271
272impl<'b, 'a, K, V, S> VacantEntry<'b, 'a, K, V, S>
273where
274    K: Eq + Hash + Clone + RcuDroppable + Sync,
275    V: Clone + RcuDroppable + Sync,
276    S: BuildHasher + Send + Sync + 'static,
277{
278    /// Sets the value of the entry with the VacantEntry's key.
279    pub fn insert(self, value: V) {
280        self.guard.insert(self.key, value);
281    }
282
283    /// Sets the value of the entry with the VacantEntry's key, and returns an occupied entry.
284    pub fn insert_entry(self, value: V) -> OccupiedEntry<'b, 'a, K, V, S> {
285        self.guard.insert(self.key.clone(), value);
286        OccupiedEntry { guard: self.guard, key: self.key }
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293    use fuchsia_rcu::rcu_run_callbacks;
294
295    #[test]
296    fn test_rcu_hash_map_custom_hasher() {
297        use std::collections::hash_map::DefaultHasher;
298        use std::hash::BuildHasherDefault;
299        let hasher = BuildHasherDefault::<DefaultHasher>::default();
300        let map = RcuHashMap::with_capacity_and_hasher(10, hasher);
301        let mut guard = map.lock();
302        guard.insert(1, 10);
303        assert_eq!(guard.get(&1), Some(10));
304    }
305
306    #[test]
307    fn test_rcu_hash_map_insert_and_get() {
308        let map = RcuHashMap::<i32, i32>::default();
309        let mut guard = map.lock();
310        let scope = RcuReadScope::new();
311
312        guard.insert(1, 10);
313        guard.insert(2, 20);
314
315        assert_eq!(guard.get(&1), Some(10));
316        assert_eq!(guard.get(&2), Some(20));
317        assert_eq!(guard.get(&3), None);
318
319        // Verify we can read without the lock too
320        drop(guard);
321        assert_eq!(map.get(&scope, &1), Some(&10));
322        assert_eq!(map.get(&scope, &2), Some(&20));
323
324        drop(scope);
325        rcu_run_callbacks();
326    }
327
328    #[test]
329    fn test_rcu_hash_map_update() {
330        let map = RcuHashMap::<i32, i32>::default();
331        let mut guard = map.lock();
332        let scope = RcuReadScope::new();
333
334        guard.insert(1, 10);
335        assert_eq!(guard.get(&1), Some(10));
336
337        guard.insert(1, 20);
338        assert_eq!(guard.get(&1), Some(20));
339
340        drop(guard);
341        assert_eq!(map.get(&scope, &1), Some(&20));
342
343        drop(scope);
344        rcu_run_callbacks();
345    }
346
347    #[test]
348    fn test_rcu_hash_map_remove() {
349        let map = RcuHashMap::<i32, i32>::default();
350        let mut guard = map.lock();
351        let scope = RcuReadScope::new();
352
353        guard.insert(1, 10);
354        assert_eq!(guard.get(&1), Some(10));
355
356        guard.remove(&1);
357        assert_eq!(guard.get(&1), None);
358
359        drop(guard);
360        assert_eq!(map.get(&scope, &1), None);
361
362        drop(scope);
363        rcu_run_callbacks();
364    }
365
366    #[test]
367    fn test_rcu_hash_map_entry_api() {
368        let map = RcuHashMap::<i32, i32>::default();
369        let mut guard = map.lock();
370
371        // Vacant entry
372        match guard.entry(1) {
373            Entry::Vacant(e) => e.insert(10),
374            Entry::Occupied(_) => panic!("Should be vacant"),
375        }
376        assert_eq!(guard.get(&1), Some(10));
377
378        // Occupied entry
379        match guard.entry(1) {
380            Entry::Occupied(mut e) => {
381                assert_eq!(e.get(), 10);
382                e.insert(20);
383            }
384            Entry::Vacant(_) => panic!("Should be occupied"),
385        }
386        assert_eq!(guard.get(&1), Some(20));
387
388        drop(guard);
389        rcu_run_callbacks();
390    }
391
392    #[test]
393    fn test_rcu_hash_map_iter() {
394        let map = RcuHashMap::<i32, i32>::default();
395        let scope = RcuReadScope::new();
396        map.insert(1, 10);
397        map.insert(2, 20);
398        map.insert(3, 30);
399
400        let mut items: Vec<_> = map.iter(&scope).collect();
401        items.sort_by_key(|(k, _)| **k);
402        assert_eq!(items, vec![(&1, &10), (&2, &20), (&3, &30)]);
403    }
404
405    #[test]
406    fn test_rcu_hash_map_keys() {
407        let map = RcuHashMap::<i32, i32>::default();
408        let scope = RcuReadScope::new();
409        map.insert(1, 10);
410        map.insert(2, 20);
411        map.insert(3, 30);
412
413        let mut keys: Vec<_> = map.keys(&scope).collect();
414        keys.sort();
415        assert_eq!(keys, vec![&1, &2, &3]);
416    }
417
418    #[test]
419    fn test_rcu_hash_map_len() {
420        let map = RcuHashMap::<i32, i32>::default();
421        map.insert(1, 10);
422        map.insert(2, 20);
423        map.insert(3, 30);
424
425        assert_eq!(map.len(), 3);
426    }
427
428    #[test]
429    fn test_rcu_hash_map_or_insert_with() {
430        let map = RcuHashMap::<i32, i32>::default();
431        let mut guard = map.lock();
432
433        // test or_insert_with
434        guard.entry(1).or_insert_with(|| 10);
435        assert!(guard.contains_key(&1));
436        assert_eq!(guard.get(&1), Some(10));
437
438        // test or_insert_with existing
439        guard.entry(1).or_insert_with(|| 20);
440        assert_eq!(guard.get(&1), Some(10));
441
442        // test OccupiedEntry::remove
443        match guard.entry(1) {
444            Entry::Occupied(e) => {
445                assert_eq!(e.remove(), 10);
446            }
447            Entry::Vacant(_) => panic!("Should be occupied"),
448        }
449        assert!(!guard.contains_key(&1));
450    }
451
452    #[test]
453    fn test_rcu_hash_map_drain() {
454        let map = RcuHashMap::<i32, i32>::default();
455        let mut guard = map.lock();
456
457        guard.insert(1, 10);
458        guard.insert(2, 20);
459        guard.insert(3, 30);
460
461        let mut items: Vec<_> = guard.drain().collect();
462        items.sort_by_key(|(k, _)| *k);
463        assert_eq!(items, vec![(1, 10), (2, 20), (3, 30)]);
464
465        assert!(!guard.contains_key(&1));
466        assert!(!guard.contains_key(&2));
467        assert!(!guard.contains_key(&3));
468    }
469
470    #[test]
471    fn test_rcu_hash_map_capacity_zero() {
472        use std::collections::hash_map::RandomState;
473        let map =
474            RcuHashMap::<i32, i32, RandomState>::with_capacity_and_hasher(0, RandomState::new());
475        let mut guard = map.lock();
476
477        assert_eq!(guard.get(&1), None);
478
479        guard.insert(1, 10);
480        assert_eq!(guard.get(&1), Some(10));
481
482        guard.remove(&1);
483        assert_eq!(guard.get(&1), None);
484    }
485}
starnix_rcu/rcu_hash_map.rs

starnix_rcu/
rcu_hash_map.rs
