Skip to main content

starnix_rcu/
rcu_cache.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::{LockDepGuard, LockDepMutex, RcuCacheLock};
8use std::hash::Hash;
9
10pub enum RcuCacheInsertionResult<V> {
11    /// The entry was inserted.
12    Inserted,
13
14    /// The entry was updated.
15    ///
16    /// The old value is returned.
17    Updated(V),
18
19    /// The entry was inserted and caused another entry to be evicted.
20    ///
21    /// The evicted value is returned.
22    Evicted(V),
23}
24
25/// A cache that uses RCU to provide thread-safe access to a hash map.
26///
27/// This is similar to `RcuHashMap`, but it also evicts items when the cache
28/// exceeds a specified capacity.
29///
30/// Entries are evicted in a FIFO manner.
31///
32/// By default, this map uses `rapidhash::RapidBuildHasher`, which provides high performance.
33/// However, if this map holds keys which may be attacker-controlled, consider using
34/// `std::collections::hash_map::RandomState` instead.
35#[derive(Debug, RcuDroppable)]
36pub struct RcuCache<K, V, S = rapidhash::RapidBuildHasher>
37where
38    K: Eq + Hash + Clone + RcuDroppable + Sync,
39    V: Clone + RcuDroppable + Sync,
40    S: std::hash::BuildHasher + Send + Sync + 'static,
41{
42    /// The maximum number of entries in the cache.
43    capacity: usize,
44
45    /// The underlying hash map.
46    map: RcuRawHashMap<K, V, S>,
47
48    /// A mutex to provide synchronization for writing to the map.
49    mutex: LockDepMutex<(), RcuCacheLock>,
50}
51
52impl<K, V> RcuCache<K, V, rapidhash::RapidBuildHasher>
53where
54    K: Eq + Hash + Clone + RcuDroppable + Sync,
55    V: Clone + RcuDroppable + Sync,
56{
57    /// Creates a new `RcuCache` with the specified capacity.
58    pub fn new(capacity: usize) -> Self {
59        Self {
60            capacity,
61            map: RcuRawHashMap::with_capacity(capacity + 1),
62            mutex: LockDepMutex::new(()),
63        }
64    }
65}
66
67impl<K, V, S> RcuCache<K, V, S>
68where
69    K: Eq + Hash + Clone + RcuDroppable + Sync,
70    V: Clone + RcuDroppable + Sync,
71    S: std::hash::BuildHasher + Send + Sync + 'static,
72{
73    /// Creates a new `RcuCache` with the specified capacity and hasher.
74    pub fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> Self {
75        Self {
76            capacity,
77            map: RcuRawHashMap::with_capacity_and_hasher(capacity + 1, hash_builder),
78            mutex: LockDepMutex::new(()),
79        }
80    }
81}
82
83impl<K, V, S> RcuCache<K, V, S>
84where
85    K: Eq + Hash + Clone + RcuDroppable + Sync,
86    V: Clone + RcuDroppable + Sync,
87    S: std::hash::BuildHasher + Send + Sync + 'static,
88{
89    /// Returns the capacity with which this instance was initialized.
90    pub fn capacity(&self) -> usize {
91        self.capacity
92    }
93
94    /// Returns the number of entries in the cache.
95    pub fn len(&self) -> usize {
96        self.map.len()
97    }
98
99    /// Returns a reference to the value associated with the key.
100    pub fn get<'a>(&self, scope: &'a RcuReadScope, key: &K) -> Option<&'a V> {
101        self.map.get(scope, key)
102    }
103
104    pub fn lock(&self) -> RcuCacheGuard<'_, K, V, S> {
105        let guard = self.mutex.lock();
106        RcuCacheGuard { cache: self, _guard: guard }
107    }
108
109    /// Removes all entries from the cache.
110    pub fn clear(&self) {
111        let _guard = self.mutex.lock();
112        let scope = RcuReadScope::new();
113        let mut cursor = self.map.cursor(&scope);
114        loop {
115            // SAFETY: We have exclusive access to the map because we have exclusive access to the
116            // mutex.
117            let removed = unsafe { cursor.remove() };
118            if removed.is_none() {
119                break;
120            }
121        }
122    }
123}
124
125pub struct RcuCacheGuard<'a, K, V, S = rapidhash::RapidBuildHasher>
126where
127    K: Eq + Hash + Clone + RcuDroppable + Sync,
128    V: Clone + RcuDroppable + Sync,
129    S: std::hash::BuildHasher + Send + Sync + 'static,
130{
131    cache: &'a RcuCache<K, V, S>,
132    _guard: LockDepGuard<'a, ()>,
133}
134
135impl<'a, K, V, S> RcuCacheGuard<'a, K, V, S>
136where
137    K: Eq + Hash + Clone + RcuDroppable + Sync,
138    V: Clone + RcuDroppable + Sync,
139    S: std::hash::BuildHasher + Send + Sync + 'static,
140{
141    pub fn get<'rcu>(&self, scope: &'rcu RcuReadScope, key: &K) -> Option<&'rcu V> {
142        self.cache.map.get(scope, key)
143    }
144
145    /// Inserts a key-value pair into the cache.
146    ///
147    /// If the cache exceeds its capacity, entries are evicted in a FIFO manner.
148    pub fn insert(&self, scope: &RcuReadScope, key: K, value: V) -> RcuCacheInsertionResult<V> {
149        // SAFETY: We have exclusive access to the map because we have exclusive access to the mutex.
150        match unsafe { self.cache.map.insert(scope, key, value) } {
151            InsertionResult::Inserted(count) => {
152                if count > self.cache.capacity {
153                    // The mutex should prevent any other modifications to the map while the insert
154                    // operation is in progress.
155                    assert!(count == self.cache.capacity + 1);
156                    let mut cursor = self.cache.map.cursor(&scope);
157                    // SAFETY: We have exclusive access to the map because we have exclusive access
158                    // to the mutex.
159                    if let Some(old_value) = unsafe { cursor.remove() } {
160                        RcuCacheInsertionResult::Evicted(old_value)
161                    } else {
162                        unreachable!("cache is full but no entries to evict")
163                    }
164                } else {
165                    RcuCacheInsertionResult::Inserted
166                }
167            }
168            InsertionResult::Updated(old_value) => RcuCacheInsertionResult::Updated(old_value),
169        }
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use fuchsia_rcu::rcu_run_callbacks;
177
178    #[test]
179    fn test_rcu_cache_fifo_eviction() {
180        let capacity = 3;
181        let cache = RcuCache::new(capacity);
182        let guard = cache.lock();
183        let scope = RcuReadScope::new();
184
185        // Insert items up to capacity
186        guard.insert(&scope, 1, 10);
187        guard.insert(&scope, 2, 20);
188        guard.insert(&scope, 3, 30);
189
190        assert_eq!(guard.get(&scope, &1), Some(&10));
191        assert_eq!(guard.get(&scope, &2), Some(&20));
192        assert_eq!(guard.get(&scope, &3), Some(&30));
193
194        // Insert an item beyond capacity, should evict 1
195        guard.insert(&scope, 4, 40);
196
197        assert_eq!(cache.get(&scope, &1), None);
198        assert_eq!(cache.get(&scope, &2), Some(&20));
199        assert_eq!(cache.get(&scope, &3), Some(&30));
200        assert_eq!(cache.get(&scope, &4), Some(&40));
201
202        // Insert another item, should evict 2
203        guard.insert(&scope, 5, 50);
204
205        assert_eq!(cache.get(&scope, &1), None);
206        assert_eq!(cache.get(&scope, &2), None);
207        assert_eq!(cache.get(&scope, &3), Some(&30));
208        assert_eq!(cache.get(&scope, &4), Some(&40));
209        assert_eq!(cache.get(&scope, &5), Some(&50));
210
211        // Update an existing item, should not evict and not change order for eviction
212        guard.insert(&scope, 3, 300);
213
214        assert_eq!(cache.get(&scope, &1), None);
215        assert_eq!(cache.get(&scope, &2), None);
216        assert_eq!(cache.get(&scope, &3), Some(&300));
217        assert_eq!(cache.get(&scope, &4), Some(&40));
218        assert_eq!(cache.get(&scope, &5), Some(&50));
219
220        // Insert another item, should evict 4 (because 3 was updated, not re-inserted)
221        guard.insert(&scope, 6, 60);
222
223        assert_eq!(cache.get(&scope, &1), None);
224        assert_eq!(cache.get(&scope, &2), None);
225        assert_eq!(cache.get(&scope, &3), Some(&300));
226        assert_eq!(cache.get(&scope, &4), None);
227        assert_eq!(cache.get(&scope, &5), Some(&50));
228        assert_eq!(cache.get(&scope, &6), Some(&60));
229
230        std::mem::drop(guard);
231        std::mem::drop(scope);
232        rcu_run_callbacks();
233    }
234}