threadsafe_lru/lib.rs
1#![deny(unsafe_code)]
2
3/// # threadsafe-lru
4///
5/// This is a thread-safe implementation of an LRU (Least Recently Used) cache in Rust.
6/// The `LruCache` struct uses sharding to improve concurrency by splitting the cache into
7/// multiple smaller segments, each protected by a mutex.
8///
9/// ## Example Usage
10///
11/// ```rust
12/// use threadsafe_lru::LruCache;
13///
14/// fn main() {
15/// // Create a new LRU cache with 4 shards and capacity of 2 per shard
16/// let cache = LruCache::new(4, 2);
17///
18/// // Insert items into the cache
19/// let five = 5;
20/// let six = 6;
21/// assert_eq!(cache.insert(five, 10), None);
22/// assert_eq!(cache.insert(six, 20), None);
23///
24/// // Retrieve an item from the cache
25/// assert_eq!(cache.get(&five), Some(10));
26///
27/// // Promote an item to make it more recently used
28/// cache.promote(&five);
29///
30/// // Remove an item from the cache
31/// assert_eq!(cache.remove(&five), Some(10));
32/// }
33/// ```
34///
35/// In this example, a new `LruCache` is created with 4 shards and a capacity of 2 entries per shard.
36/// Items are inserted using the `insert` method.
37/// The `get` method retrieves an item by key, promoting it to the most recently used position.
38/// Finally, the `remove` method deletes an item from the cache.
39///
40/// This implementation ensures that operations on different keys can be performed concurrently
41/// without causing race conditions due to shared state.
42///
43use std::{
44 borrow::Borrow,
45 hash::{DefaultHasher, Hash, Hasher},
46 sync::{
47 atomic::{AtomicUsize, Ordering},
48 Mutex,
49 },
50};
51
52use hashbrown::HashMap;
53use indexlist::{Index, IndexList};
54
55// A thread-safe LRU
56pub struct LruCache<K, V> {
57 shards: Vec<Mutex<Shard<K, V>>>,
58 count: AtomicUsize,
59 shards_count: usize,
60 cap_per_shard: usize,
61}
62
63struct Shard<K, V> {
64 entries: HashMap<K, (Index<K>, V)>,
65 order: IndexList<K>,
66 count: usize,
67}
68
69impl<K, V> LruCache<K, V>
70where
71 K: Clone + Eq + Hash,
72{
73 /// Creates a new instance of `LruCache`.
74 ///
75 /// The cache is divided into multiple shards to improve concurrency by distributing the entries
76 /// across different locks.
77 ///
78 /// # Arguments
79 ///
80 /// * `shards_count` - The number of shards in the cache. Each shard acts as an independent LRU
81 /// with its own capacity and order list.
82 /// * `cap_per_shard` - The capacity for each shard, representing the maximum number of items it
83 /// can hold before evicting the least recently used item(s).
84 ///
85 /// # Returns
86 ///
87 /// A new instance of `LruCache`.
88 ///
89 /// # Examples
90 ///
91 /// ```rust
92 /// use threadsafe_lru::LruCache;
93 ///
94 /// let cache: LruCache<i32, i32> = LruCache::new(4, 2); // Creates a cache with 4 shards and capacity of 2 entries per shard.
95 /// ```
96 pub fn new(shards_count: usize, cap_per_shard: usize) -> LruCache<K, V> {
97 let mut shards = Vec::default();
98 for _ in 0..shards_count {
99 shards.push(Mutex::new(Shard {
100 entries: HashMap::with_capacity(cap_per_shard),
101 order: IndexList::with_capacity(cap_per_shard),
102 count: 0,
103 }));
104 }
105 LruCache {
106 shards,
107 count: AtomicUsize::default(),
108 shards_count,
109 cap_per_shard,
110 }
111 }
112
113 /// Inserts a new key-value pair into the cache.
114 ///
115 /// If the key already exists in the cache, its value is updated and it is promoted to the most
116 /// recently used position.
117 /// If inserting the new item causes the cache to exceed its capacity, the least recently used
118 /// item will be evicted from its shard.
119 ///
120 /// # Arguments
121 ///
122 /// * `k` - The key to insert or update.
123 /// * `v` - The value associated with the key.
124 ///
125 /// # Returns
126 ///
127 /// If the key already existed in the cache and was updated, the previous value is returned.
128 /// Otherwise, `None` is returned.
129 ///
130 /// # Examples
131 ///
132 /// ```rust
133 /// use threadsafe_lru::LruCache;
134 ///
135 /// let cache = LruCache::new(4, 2);
136 ///
137 /// let five = 5;
138 /// let six = 6;
139 ///
140 /// assert_eq!(cache.insert(five, 10), None); // Inserts a new key-value pair
141 /// assert_eq!(cache.insert(six, 20), None); // Inserts another new key-value pair
142 /// assert_eq!(cache.insert(five, 30), Some(10)); // Updates an existing key with a new value and returns the old value
143 /// ```
144 pub fn insert(&self, k: K, v: V) -> Option<V> {
145 let mut shard = self.shards[self.shard(&k)].lock().unwrap();
146 let index = shard.entries.get(&k).map(|v| v.0);
147 if shard.count == self.cap_per_shard && index.is_none() {
148 if let Some(index) = shard.order.head_index() {
149 shard.entries.remove(&k);
150 shard.order.remove(index);
151 self.count.fetch_sub(1, Ordering::Relaxed);
152 shard.count -= 1;
153 }
154 }
155
156 match index {
157 Some(index) => {
158 shard.order.remove(index);
159 let index = shard.order.push_back(k.clone());
160 shard.entries.insert(k, (index, v)).map(|v| v.1)
161 }
162 None => {
163 let index = shard.order.push_back(k.clone());
164 shard.entries.insert(k, (index, v));
165 self.count.fetch_add(1, Ordering::Relaxed);
166 shard.count += 1;
167 None
168 }
169 }
170 }
171
172 /// Retrieves a value from the cache by key.
173 ///
174 /// When a key is accessed using this method, it is promoted to the most recently used position
175 /// in its shard. If the key does not exist in the cache, `None` is returned.
176 ///
177 /// # Arguments
178 ///
179 /// * `k` - The key of the item to retrieve.
180 ///
181 /// # Returns
182 ///
183 /// The value associated with the key if it exists; otherwise, `None`.
184 ///
185 /// # Examples
186 ///
187 /// ```rust
188 /// use threadsafe_lru::LruCache;
189 ///
190 /// let cache = LruCache::new(4, 2);
191 ///
192 /// let five = 5;
193 /// let six = 6;
194 ///
195 /// assert_eq!(cache.insert(five, 10), None);
196 /// assert_eq!(cache.get(&five), Some(10));
197 /// ```
198 ///
199 /// This method ensures that frequently accessed items remain more accessible while older or less
200 /// frequently accessed items are evicted when necessary.
201
202 pub fn get<Q>(&self, k: &Q) -> Option<V>
203 where
204 K: Borrow<Q>,
205 Q: Hash + Eq + ?Sized + ToOwned<Owned = K>,
206 V: Clone,
207 {
208 let mut shard = self.shards[self.shard(k)].lock().unwrap();
209 let index = shard.entries.get(k).map(|e| e.0);
210 match index {
211 Some(index) => {
212 shard.order.remove(index);
213 shard.order.push_back(k.to_owned());
214 shard.entries.get(k).map(|e| e.1.clone())
215 }
216 None => None,
217 }
218 }
219
220 /// Retrieves and mutates a value from the cache by key.
221 ///
222 /// When a key is accessed using this method, it is promoted to the most recently used position
223 /// in its shard. If the key does not exist in the cache, `None` is returned.
224 ///
225 /// This function provides mutable access to the value associated with a given key, allowing you
226 /// to modify its contents directly without needing to retrieve and re-insert it into the cache.
227 ///
228 /// # Arguments
229 ///
230 /// * `k` - The key of the item to retrieve and mutate.
231 /// * `func` - A closure that takes a mutable reference to the value (if present) and allows
232 /// for in-place modifications.
233 ///
234 /// # Examples
235 ///
236 /// ```rust
237 /// use threadsafe_lru::LruCache;
238 ///
239 /// let cache = LruCache::new(4, 2);
240 ///
241 /// let five = 5;
242 /// let six = 6;
243 ///
244 /// assert_eq!(cache.insert(five, 10), None);
245 /// cache.get_mut(&five, |v| {
246 /// if let Some(v) = v {
247 /// *v += 1
248 /// }
249 /// });
250 /// ```
251 ///
252 /// This method ensures that frequently accessed items remain more accessible while older or less
253 /// frequently accessed items are evicted when necessary.
254 pub fn get_mut<Q, F>(&self, k: &Q, mut func: F)
255 where
256 K: Borrow<Q>,
257 Q: Hash + Eq + ?Sized + ToOwned<Owned = K>,
258 F: FnMut(Option<&mut V>),
259 {
260 let mut shard = self.shards[self.shard(k)].lock().unwrap();
261 let index = shard.entries.get(k).map(|e| e.0);
262 if let Some(index) = index {
263 shard.order.remove(index);
264 shard.order.push_back(k.to_owned());
265 func(shard.entries.get_mut(k).map(|e| &mut e.1));
266 }
267 }
268
269 /// Removes a key-value pair from the cache by key.
270 ///
271 /// When an item is removed from the cache using this method, its corresponding entry is deleted
272 /// along with any references to it in the order list. If the key does not exist in the cache,
273 /// `None` is returned.
274 ///
275 /// This operation ensures that the cache maintains its integrity and correctly tracks the number of items stored.
276 ///
277 /// # Arguments
278 ///
279 /// * `k` - The key of the item to remove.
280 ///
281 /// # Returns
282 ///
283 /// The value associated with the key if it existed; otherwise, `None`.
284 ///
285 /// # Examples
286 ///
287 /// ```rust
288 /// use threadsafe_lru::LruCache;
289 ///
290 /// let cache = LruCache::new(4, 2);
291 ///
292 /// let five = 5;
293 /// let six = 6;
294 ///
295 /// assert_eq!(cache.insert(five, 10), None); // Inserts a new key-value pair
296 /// assert_eq!(cache.insert(six, 20), None); // Inserts another new key-value pair
297 ///
298 /// assert_eq!(cache.remove(&five), Some(10)); // Removes the item with key five and returns its value
299 /// assert_eq!(cache.get(&five), None); // Key five no longer exists in the cache
300 /// ```
301 ///
302 pub fn remove<Q>(&self, k: &Q) -> Option<V>
303 where
304 K: Borrow<Q>,
305 Q: Hash + Eq + ?Sized,
306 {
307 let mut shard = self.shards[self.shard(k)].lock().unwrap();
308 let entry = shard.entries.remove(k);
309 match entry {
310 Some((index, value)) => {
311 shard.order.remove(index);
312 self.count.fetch_sub(1, Ordering::Relaxed);
313 shard.count -= 1;
314 Some(value)
315 }
316 None => None,
317 }
318 }
319
320 /// Promotes a key-value pair in the cache to the most recently used position.
321 ///
322 /// When an item is accessed using this method, it is promoted to the most recently used position
323 /// in its shard. If the key does not exist in the cache, no action is taken.
324 ///
325 /// # Arguments
326 ///
327 /// * `k` - The key of the item to promote.
328 ///
329 ///
330 /// # Examples
331 ///
332 /// ```rust
333 /// use threadsafe_lru::LruCache;
334 ///
335 /// let cache = LruCache::new(4, 2);
336 ///
337 /// let five = 5;
338 /// let six = 6;
339 /// let seven = 7;
340 ///
341 /// assert_eq!(cache.insert(five, 10), None); // Inserts a new key-value pair
342 /// assert_eq!(cache.insert(six, 20), None); // Inserts another new key-value pair
343 ///
344 /// cache.promote(&five);
345 ///
346 /// assert_eq!(cache.insert(seven, 30), None); // Inserts another new key-value pair
347 /// assert_eq!(cache.get(&five), Some(10)); // Retrieving the promoted item
348 /// ```
349 ///
350 pub fn promote<Q>(&self, k: &Q)
351 where
352 K: Borrow<Q>,
353 Q: Hash + Eq + ?Sized + ToOwned<Owned = K>,
354 {
355 let mut shard = self.shards[self.shard(k)].lock().unwrap();
356 let index = shard.entries.get(k).map(|v| v.0);
357 if let Some(index) = index {
358 shard.order.remove(index);
359 shard.order.push_back(k.to_owned());
360 }
361 }
362
363 /// Returns the total number of key-value pairs currently stored in the cache.
364 ///
365 /// This method provides a quick way to check how many items are present in the cache without
366 /// iterating over its contents.
367 ///
368 /// # Examples
369 ///
370 /// ```rust
371 /// use threadsafe_lru::LruCache;
372 ///
373 /// let cache = LruCache::new(4, 2);
374 ///
375 /// let five = 5;
376 /// let six = 6;
377 ///
378 /// assert_eq!(cache.insert(five, 10), None); // Inserts a new key-value pair
379 /// assert_eq!(cache.len(), 1); // Cache now has one item
380 ///
381 /// cache.insert(six, 20);
382 /// assert_eq!(cache.len(), 2); // Cache now has two items
383 ///
384 /// cache.remove(&five);
385 /// assert_eq!(cache.len(), 1); // Cache now has only one item
386 ///
387 /// let new_cache: LruCache<i32, i32> = LruCache::new(4, 2);
388 /// assert_eq!(new_cache.len(), 0); // New cache is empty
389 /// ```
390 ///
391 pub fn len(&self) -> usize {
392 self.count.load(Ordering::Relaxed)
393 }
394
395 /// Checks if the cache is empty.
396 ///
397 /// This method provides a quick way to determine whether the cache contains any key-value pairs
398 /// without iterating over its contents.
399 ///
400 /// # Examples
401 ///
402 /// ```rust
403 /// use threadsafe_lru::LruCache;
404 ///
405 /// let cache = LruCache::new(4, 2);
406 ///
407 /// assert!(cache.is_empty()); // Cache is empty upon creation
408 ///
409 /// let five = 5;
410 /// let six = 6;
411 ///
412 /// cache.insert(five, 10); // Inserting a key-value pair
413 /// assert!(!cache.is_empty()); // Cache is no longer empty
414 ///
415 /// cache.remove(&five); // Removing the key-value pair
416 /// assert!(cache.is_empty()); // Cache is empty again after removal
417 /// ```
418 ///
419 pub fn is_empty(&self) -> bool {
420 self.len() == 0
421 }
422
423 fn shard<Q>(&self, k: &Q) -> usize
424 where
425 K: Borrow<Q>,
426 Q: Hash + Eq + ?Sized,
427 {
428 let mut hasher = DefaultHasher::new();
429 k.hash(&mut hasher);
430 hasher.finish() as usize % self.shards_count
431 }
432}
433
434#[cfg(test)]
435mod tests {
436 use super::*;
437
438 #[test]
439 fn test_new() {
440 let shards_count = 4;
441 let cap_per_shard = 10;
442
443 let cache: LruCache<u8, u8> = LruCache::new(shards_count, cap_per_shard);
444 assert_eq!(cache.shards.len(), shards_count);
445
446 for shard in &cache.shards {
447 let lock = shard.lock().unwrap();
448 assert!(lock.entries.capacity() >= cap_per_shard);
449 assert_eq!(lock.count, 0);
450 }
451
452 assert_eq!(cache.shards_count, shards_count);
453 assert_eq!(cache.cap_per_shard, cap_per_shard);
454 }
455
456 #[test]
457 fn test_insert() {
458 let shards_count = 4;
459 let cap_per_shard = 2;
460
461 let cache = LruCache::new(shards_count, cap_per_shard);
462
463 let five = 5;
464 let six = 6;
465 let nine = 9;
466 assert_eq!(cache.shard(&five), cache.shard(&six));
467 assert_eq!(cache.shard(&five), cache.shard(&nine));
468
469 assert_eq!(cache.insert(five, 10), None);
470 assert_eq!(cache.insert(five, 10), Some(10));
471 assert_eq!(cache.count.load(Ordering::Relaxed), 1);
472 assert_eq!(cache.insert(six, 20), None);
473 assert_eq!(cache.count.load(Ordering::Relaxed), 2);
474 assert_eq!(cache.insert(nine, 30), None);
475 assert_eq!(cache.count.load(Ordering::Relaxed), 2);
476 assert_eq!(cache.insert(six, 20), Some(20));
477 assert_eq!(cache.count.load(Ordering::Relaxed), 2);
478 }
479
480 #[test]
481 fn test_get() {
482 let shards_count = 4;
483 let cap_per_shard = 2;
484
485 let cache = LruCache::new(shards_count, cap_per_shard);
486
487 let five = 5;
488 let six = 6;
489 assert_eq!(cache.shard(&five), cache.shard(&six));
490 assert_eq!(cache.insert(five, 10), None);
491 assert_eq!(cache.insert(six, 20), None);
492
493 assert_eq!(cache.get(&five), Some(10));
494 assert_eq!(cache.get(&six), Some(20));
495
496 let shard = cache.shards[cache.shard(&five)].lock().unwrap();
497 assert_eq!(shard.order.head(), Some(&five));
498 assert_eq!(shard.order.tail(), Some(&six));
499
500 assert_eq!(cache.get(&3), None);
501
502 assert_eq!(cache.count.load(Ordering::Relaxed), 2);
503 }
504
505 #[test]
506 fn test_get_mut() {
507 let shards_count = 4;
508 let cap_per_shard = 2;
509
510 let cache = LruCache::new(shards_count, cap_per_shard);
511
512 let five = 5;
513 let six = 6;
514 assert_eq!(cache.shard(&five), cache.shard(&six));
515 assert_eq!(cache.insert(five, 10), None);
516 assert_eq!(cache.insert(six, 20), None);
517
518 cache.get_mut(&five, |v| {
519 if let Some(v) = v {
520 *v = 30
521 }
522 });
523 assert_eq!(cache.get(&five), Some(30));
524 assert_eq!(cache.count.load(Ordering::Relaxed), 2);
525
526 let shard = cache.shards[cache.shard(&five)].lock().unwrap();
527 assert_eq!(shard.order.tail(), Some(&five));
528 assert_eq!(shard.order.head(), Some(&six));
529
530 cache.get_mut(&3, |v| {
531 if let Some(v) = v {
532 *v = 10
533 }
534 });
535 assert_eq!(cache.count.load(Ordering::Relaxed), 2);
536 }
537
538 #[test]
539 fn test_remove() {
540 let shards_count = 4;
541 let cap_per_shard = 2;
542
543 let cache = LruCache::new(shards_count, cap_per_shard);
544
545 let five = 5;
546 let six = 6;
547 assert_eq!(cache.shard(&five), cache.shard(&six));
548 assert_eq!(cache.insert(five, 10), None);
549 assert_eq!(cache.insert(six, 20), None);
550
551 assert_eq!(cache.remove(&five), Some(10));
552 assert_eq!(cache.count.load(Ordering::Relaxed), 1);
553 let shard = cache.shards[cache.shard(&six)].lock().unwrap();
554 assert!(!shard.entries.contains_key(&five));
555 assert_eq!(shard.order.head(), Some(&six));
556 drop(shard);
557
558 assert_eq!(cache.remove(&5), None);
559 assert_eq!(cache.count.load(Ordering::Relaxed), 1);
560
561 let new_cache: LruCache<i32, i32> = LruCache::new(shards_count, cap_per_shard);
562 assert_eq!(new_cache.remove(&five), None);
563 assert_eq!(new_cache.count.load(Ordering::Relaxed), 0);
564 }
565
566 #[test]
567 fn test_promote() {
568 let shards_count = 4;
569 let cap_per_shard = 2;
570
571 let cache = LruCache::new(shards_count, cap_per_shard);
572
573 let five = 5;
574 let six = 6;
575 assert_eq!(cache.shard(&five), cache.shard(&six));
576 assert_eq!(cache.insert(five, 10), None);
577 assert_eq!(cache.insert(six, 20), None);
578
579 let shard = cache.shards[cache.shard(&five)].lock().unwrap();
580 assert_eq!(shard.order.head(), Some(&five));
581 assert_eq!(shard.order.tail(), Some(&six));
582 drop(shard);
583
584 cache.promote(&five);
585 let shard = cache.shards[cache.shard(&five)].lock().unwrap();
586 assert_eq!(shard.order.head(), Some(&six));
587 assert_eq!(shard.order.tail(), Some(&five));
588 drop(shard);
589
590 assert_eq!(cache.get(&five), Some(10));
591 let shard = cache.shards[cache.shard(&five)].lock().unwrap();
592 assert_eq!(shard.order.head(), Some(&six));
593 assert_eq!(shard.order.tail(), Some(&five));
594 drop(shard);
595 }
596
597 #[test]
598 fn test_is_empty() {
599 let shards_count = 4;
600 let cap_per_shard = 2;
601
602 let cache = LruCache::new(shards_count, cap_per_shard);
603 assert!(cache.is_empty());
604 assert_eq!(cache.len(), 0);
605
606 let five = 5;
607 assert_eq!(cache.insert(five, 10), None);
608 assert!(!cache.is_empty());
609 assert_eq!(cache.len(), 1);
610
611 cache.remove(&five);
612 assert!(cache.is_empty());
613 assert_eq!(cache.len(), 0);
614 }
615
616 #[test]
617 fn test_len() {
618 let shards_count = 4;
619 let cap_per_shard = 2;
620
621 let cache = LruCache::new(shards_count, cap_per_shard);
622
623 let five = 5;
624 let six = 6;
625 assert_eq!(cache.insert(five, 10), None);
626 assert_eq!(cache.len(), 1);
627
628 cache.insert(six, 20);
629 assert_eq!(cache.len(), 2);
630
631 cache.remove(&five);
632 assert_eq!(cache.len(), 1);
633
634 let new_cache: LruCache<i32, i32> = LruCache::new(shards_count, cap_per_shard);
635 assert_eq!(new_cache.len(), 0);
636 }
637
638 #[test]
639 fn test_concurrent_access() {
640 use std::sync::Arc;
641 use std::thread;
642
643 let shards_count = 4;
644 let cap_per_shard = 2;
645
646 let cache: Arc<LruCache<i32, i32>> = Arc::new(LruCache::new(shards_count, cap_per_shard));
647
648 const THREAD_COUNT: usize = 10;
649 const OPERATIONS_PER_THREAD: usize = 100;
650
651 let mut handles = vec![];
652
653 for _ in 0..THREAD_COUNT {
654 let cache = Arc::clone(&cache);
655
656 let handle = thread::spawn(move || {
657 for _ in 0..OPERATIONS_PER_THREAD {
658 let key = rand::random::<i32>();
659 let value = rand::random::<i32>();
660
661 let op_type = rand::random::<u8>() % 3;
662 match op_type {
663 0 => {
664 cache.insert(key, value);
665 }
666 1 => {
667 cache.get(&key);
668 }
669 2 => {
670 cache.remove(&key);
671 }
672 _ => unreachable!(),
673 }
674 }
675 });
676
677 handles.push(handle);
678 }
679
680 for handle in handles {
681 handle.join().unwrap();
682 }
683 }
684}