s3_fifo/
lib.rs

1//! A non-thread-safe implementation of an S3-FIFO
2//! Paper here: https://jasony.me/publication/sosp23-s3fifo.pdf
3
4use std::{
5    collections::VecDeque,
6    sync::atomic::{AtomicI8, Ordering},
7};
8
9mod key;
10pub use key::S3FIFOKey;
11
12/// S3FIFO is a non-thread-safe implementation of an S3-FIFO
13///
14/// Paper here: https://jasony.me/publication/sosp23-s3fifo.pdf
15///
16/// S3FIFO is a cache that is split into three parts:
17/// 1. A small cache that holds the most recently used items
18/// 2. A main cache that holds the most frequently used items
19/// 3. A ghost cache that holds keys that have been evicted from the main cache
20///
21/// ```
22/// use s3_fifo::{S3FIFO, S3FIFOKey};
23///
24/// // The cached value must be Clone.
25/// // Hash is optional and allows using the S3FIFOKey struct
26/// #[derive(Clone, Hash)]
27/// struct Foobar { a: i32 }
28///
29/// // Create a cache with a capacity of 128 (small: 12, main: 115, ghost: 115)
30/// let mut cache = S3FIFO::new(128);
31/// let value = Foobar { a: 1 };
32/// let key = S3FIFOKey::new(&value);
33///
34/// // Check if the item is in the cache before inserting
35/// if let None = cache.get(&key) {
36///     cache.put(key.clone(), value);
37///     assert!(cache.get(&key).is_some());
38/// }
39/// ````
40pub struct S3FIFO<K, V> {
41    small: VecDeque<Item<K, V>>,
42    main: VecDeque<Item<K, V>>,
43    ghost: VecDeque<Key<K>>,
44}
45
46impl<K: PartialEq + Clone, V> S3FIFO<K, V> {
47    ///
48    /// Create a new S3FIFO cache with 10% of the capacity for
49    /// the small cache and 90% of the capacity for the main cache.
50    ///
51    /// The ghost cache is also 90% of the capacity but only holds
52    /// keys and not values.
53    ///
54    pub fn new(capacity: usize) -> Self {
55        let small_capacity = capacity / 10;
56        let main_capacity = capacity * 9 / 10;
57        S3FIFO {
58            small: VecDeque::with_capacity(small_capacity),
59            main: VecDeque::with_capacity(main_capacity),
60            ghost: VecDeque::with_capacity(main_capacity),
61        }
62    }
63
64    /// Read an item from the cache.
65    /// If the item is present, then its frequency is incremented and a reference is returned.
66    pub fn get(&self, key: &K) -> Option<&V> {
67        // Check item in small
68        if let Some(item) = self.small.iter().find(|item| item.key == *key) {
69            let _ = item
70                .freq
71                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |x| {
72                    if x > 2 {
73                        Some(3)
74                    } else {
75                        Some(x + 1)
76                    }
77                });
78            return Some(&item.value);
79        }
80
81        // Check item in main
82        if let Some(item) = self.main.iter().find(|item| item.key == *key) {
83            let _ = item
84                .freq
85                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |x| {
86                    if x > 2 {
87                        Some(3)
88                    } else {
89                        Some(x + 1)
90                    }
91                });
92            return Some(&item.value);
93        }
94
95        None
96    }
97
98    /// Read an item from the cache.
99    /// If the item is present, then its frequency is incremented and a mutable reference is returned.
100    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
101        // Check item in small
102        if let Some(item) = self.small.iter_mut().find(|item| item.key == *key) {
103            let _ = item
104                .freq
105                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |x| {
106                    if x > 2 {
107                        Some(3)
108                    } else {
109                        Some(x + 1)
110                    }
111                });
112            return Some(&mut item.value);
113        }
114
115        // Check item in main
116        if let Some(item) = self.main.iter_mut().find(|item| item.key == *key) {
117            let _ = item
118                .freq
119                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |x| {
120                    if x > 2 {
121                        Some(3)
122                    } else {
123                        Some(x + 1)
124                    }
125                });
126            return Some(&mut item.value);
127        }
128
129        None
130    }
131
132    /// Write an item to the cache.
133    /// This may evict an item from the cache.
134    /// The returnted tuple is a mutable reference to the value in the cache and any evicted value.
135    pub fn put(&mut self, key: K, value: V) -> (&mut V, Option<V>) {
136        // Check if the item is in the cache to maintain consistency
137        if let Some(item) = self.get_mut(&key) {
138            // Borrow checker would say that this item borrows self mutably for '1 lifetime
139            // That would mean all of the immutable borrows below would be invalid even though
140            // they are not and we are just returning here.
141            // V lives safely in this container and this referenc is now bound to the lifetime of the container in this scope.
142            let item = item as *mut V;
143            return (unsafe { &mut *item }, None);
144        }
145
146        // Check if item is in ghost to decide where to insert
147        let mut evicted = None;
148        if let Some(key) = self.ghost.iter().find(|k| k.key == key) {
149            let item = Item {
150                key: key.key.clone(),
151                value,
152                freq: key.freq.load(Ordering::Relaxed).into(),
153            };
154            if self.main.capacity() == self.main.len() {
155                evicted = self.evict_main();
156            }
157            self.main.push_front(item);
158            return (&mut self.main.front_mut().unwrap().value, evicted);
159        } else {
160            let item = Item {
161                key,
162                value,
163                freq: 0.into(),
164            };
165            if self.small.capacity() == self.small.len() {
166                evicted = self.evict_small();
167            }
168            self.small.push_front(item);
169            return (&mut self.small.front_mut().unwrap().value, evicted);
170        }
171    }
172
173    /// Remove an item from the cache.
174    pub fn pop(&mut self) -> Option<V> {
175        // Popping from small may move an item to main
176        while !self.small.is_empty() {
177            if let Some(value) = self.evict_small() {
178                return Some(value);
179            }
180        }
181
182        // Try evicting from main to keep ghost up to date
183        self.evict_main()
184    }
185
186    /// Remove all items from the cache, leaving it empty and with the same capacity.
187    pub fn drain(&mut self) -> Vec<V> {
188        self.ghost.clear();
189        let mut values = Vec::with_capacity(self.small.len() + self.main.len());
190        values.extend(self.small.drain(..).map(|item| item.value));
191        values.extend(self.main.drain(..).map(|item| item.value));
192        values
193    }
194
195    fn evict_small(&mut self) -> Option<V> {
196        if self.small.is_empty() {
197            return None;
198        }
199        let item = self.small.pop_back().unwrap();
200        let freq = item.freq.load(Ordering::Relaxed);
201        if freq > 1 {
202            let mut value = None;
203            if self.main.capacity() == self.main.len() {
204                value = self.evict_main();
205            }
206            self.main.push_front(item);
207            value
208        } else {
209            let Item { key, value, freq } = item;
210            if self.ghost.capacity() == self.ghost.len() {
211                self.ghost.pop_back();
212            }
213            self.ghost.push_front(Key { key, freq });
214            Some(value)
215        }
216    }
217
218    fn evict_main(&mut self) -> Option<V> {
219        // The maximum freq is 3, so if the main cache is full and all items have freq 3,
220        // then the maximum number of iterations is 3 * main.len() + 1
221        let mut iters = (3 * self.main.len() + 1) as isize;
222        while iters > 0 {
223            let Some(item) = self.main.pop_back() else {
224                return None;
225            };
226            iters -= 1;
227            let freq = item.freq.load(Ordering::Relaxed);
228            if freq > 0 {
229                item.freq.fetch_sub(1, Ordering::Relaxed);
230                self.main.push_front(item);
231            } else {
232                return Some(item.value);
233            }
234        }
235        None
236    }
237}
238
239struct Item<K, V> {
240    key: K,
241    value: V,
242    freq: AtomicI8, // not thread-safe
243}
244
245struct Key<K> {
246    key: K,
247    freq: AtomicI8, // not thread-safe
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use std::{
254        collections::hash_map::DefaultHasher,
255        hash::{Hash, Hasher},
256    };
257
258    #[derive(Hash, Clone)]
259    struct Abc {
260        a: u8,
261        b: u16,
262        c: u32,
263    }
264
265    #[test]
266    fn can_use_cache() {
267        let test_value = Abc { a: 1, b: 2, c: 3 };
268        let mut hasher = DefaultHasher::new();
269        test_value.hash(&mut hasher);
270        let test_key = hasher.finish();
271
272        // Create a cache with a capacity of 10
273        let mut cache = S3FIFO::new(10);
274        cache.put(test_key, test_value);
275        assert!(cache.get(&test_key).is_some());
276    }
277
278    #[test]
279    fn can_fill_cache() {
280        // Create a cache with a capacity of 10
281        let mut cache = S3FIFO::new(10);
282        for i in 0..10 {
283            let test_value = Abc {
284                a: i as u8,
285                b: i as u16,
286                c: i as u32,
287            };
288            let test_key = S3FIFOKey::new(&test_value);
289
290            assert!(cache.get(&test_key).is_none());
291            cache.put(test_key.clone(), test_value);
292            assert!(cache.get(&test_key).is_some());
293        }
294
295        assert_eq!(cache.small.len(), 1);
296        assert_eq!(cache.ghost.len(), 9);
297
298        // Promote to main
299        let repeat_value = Abc { a: 0, b: 0, c: 0 };
300        let repeat_key = S3FIFOKey::new(&repeat_value);
301        assert!(cache.get(&repeat_key).is_none());
302        cache.put(repeat_key, repeat_value);
303
304        assert_eq!(cache.small.len(), 1);
305        assert_eq!(cache.ghost.len(), 9);
306        assert_eq!(cache.main.len(), 1);
307
308        // Increment main
309        let repeat_value = Abc { a: 0, b: 0, c: 0 };
310        let repeat_key = S3FIFOKey::new(&repeat_value);
311        assert!(cache.get(&repeat_key).is_some());
312        // Do not insert again or else duplicate keys will be in the cache
313        // cache.put(repeat_key, repeat_value);
314
315        assert_eq!(cache.small.len(), 1);
316        assert_eq!(cache.ghost.len(), 9);
317        assert_eq!(cache.main.len(), 1);
318    }
319}
s3_fifo/lib.rs

s3_fifo/
lib.rs
