Skip to main content

pingora_lru/
lib.rs

1// Copyright 2026 Cloudflare, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! An implementation of an LRU that focuses on memory efficiency, concurrency and persistence
16//!
17//! Features
18//! - keys can have different sizes
19//! - LRUs are sharded to avoid global locks.
20//! - Memory layout and usage are optimized: small and no memory fragmentation
21
22pub mod linked_list;
23
24use linked_list::{LinkedList, LinkedListIter};
25
26use hashbrown::HashMap;
27use parking_lot::RwLock;
28use std::sync::atomic::{AtomicUsize, Ordering};
29
30/// The LRU with `N` shards
31pub struct Lru<T, const N: usize> {
32    units: [RwLock<LruUnit<T>>; N],
33    weight: AtomicUsize,
34    weight_limit: usize,
35    len_watermark: Option<usize>,
36    len: AtomicUsize,
37    evicted_weight: AtomicUsize,
38    evicted_len: AtomicUsize,
39}
40
41impl<T, const N: usize> Lru<T, N> {
42    /// Create an [Lru] with the given weight limit and predicted capacity.
43    ///
44    /// The capacity is per shard (for simplicity). So the total capacity = capacity * N
45    pub fn with_capacity(weight_limit: usize, capacity: usize) -> Self {
46        Self::with_capacity_and_watermark(weight_limit, capacity, None)
47    }
48
49    /// Create an [Lru] with the given weight limit, predicted capacity and optional watermark
50    ///
51    /// The capacity is per shard (for simplicity). So the total capacity = capacity * N
52    ///
53    /// The watermark indicates at what count we should begin evicting and acts as a limit
54    /// on the total number of allowed items.
55    pub fn with_capacity_and_watermark(
56        weight_limit: usize,
57        capacity: usize,
58        len_watermark: Option<usize>,
59    ) -> Self {
60        // use the unsafe code from ArrayVec just to init the array
61        let mut units = arrayvec::ArrayVec::<_, N>::new();
62        for _ in 0..N {
63            units.push(RwLock::new(LruUnit::with_capacity(capacity)));
64        }
65        Lru {
66            units: units.into_inner().map_err(|_| "").unwrap(),
67            weight: AtomicUsize::new(0),
68            weight_limit,
69            len_watermark,
70            len: AtomicUsize::new(0),
71            evicted_weight: AtomicUsize::new(0),
72            evicted_len: AtomicUsize::new(0),
73        }
74    }
75
76    /// Admit the key value to the [Lru]
77    ///
78    /// Return the shard index which the asset is added to
79    pub fn admit(&self, key: u64, data: T, weight: usize) -> usize {
80        let shard = get_shard(key, N);
81        let unit = &mut self.units[shard].write();
82
83        // Make sure weight is positive otherwise eviction won't work
84        // TODO: Probably should use NonZeroUsize instead
85        let weight = weight.max(1);
86
87        let old_weight = unit.admit(key, data, weight);
88        if old_weight != weight {
89            self.weight.fetch_add(weight, Ordering::Relaxed);
90            if old_weight > 0 {
91                self.weight.fetch_sub(old_weight, Ordering::Relaxed);
92            } else {
93                // Assume old_weight == 0 means a new item is admitted
94                self.len.fetch_add(1, Ordering::Relaxed);
95            }
96        }
97        shard
98    }
99
100    /// Increment the weight associated with a given key, up to an optional max weight.
101    /// If a `max_weight` is provided, the weight cannot exceed this max weight. If the current
102    /// weight is higher than the max, it will be capped to the max.
103    ///
104    /// Return the total new weight. 0 indicates the key did not exist.
105    pub fn increment_weight(&self, key: u64, delta: usize, max_weight: Option<usize>) -> usize {
106        let shard = get_shard(key, N);
107        let unit = &mut self.units[shard].write();
108        if let Some((old_weight, new_weight)) = unit.increment_weight(key, delta, max_weight) {
109            if new_weight >= old_weight {
110                self.weight
111                    .fetch_add(new_weight - old_weight, Ordering::Relaxed);
112            } else {
113                self.weight
114                    .fetch_sub(old_weight - new_weight, Ordering::Relaxed);
115            }
116            new_weight
117        } else {
118            0
119        }
120    }
121
122    /// Promote the key to the head of the LRU
123    ///
124    /// Return `true` if the key exists.
125    pub fn promote(&self, key: u64) -> bool {
126        self.units[get_shard(key, N)].write().access(key)
127    }
128
129    /// Promote to the top n of the LRU
130    ///
131    /// This function is a bit more efficient in terms of reducing lock contention because it
132    /// will acquire a write lock only if the key is outside top n but only acquires a read lock
133    /// when the key is already in the top n.
134    ///
135    /// Return false if the item doesn't exist
136    pub fn promote_top_n(&self, key: u64, top: usize) -> bool {
137        let unit = &self.units[get_shard(key, N)];
138        if !unit.read().need_promote(key, top) {
139            return true;
140        }
141        unit.write().access(key)
142    }
143
144    /// Evict at most one item from the given shard
145    ///
146    /// Return the evicted asset and its size if there is anything to evict
147    pub fn evict_shard(&self, shard: u64) -> Option<(T, usize)> {
148        let evicted = self.units[get_shard(shard, N)].write().evict();
149        if let Some((_, weight)) = evicted.as_ref() {
150            self.weight.fetch_sub(*weight, Ordering::Relaxed);
151            self.len.fetch_sub(1, Ordering::Relaxed);
152            self.evicted_weight.fetch_add(*weight, Ordering::Relaxed);
153            self.evicted_len.fetch_add(1, Ordering::Relaxed);
154        }
155        evicted
156    }
157
158    /// Evict the [Lru] until the overall weight is below the limit (or the configured watermark).
159    ///
160    /// Return a list of evicted items.
161    ///
162    /// The evicted items are randomly selected from all the shards.
163    pub fn evict_to_limit(&self) -> Vec<(T, usize)> {
164        let mut evicted = vec![];
165        let mut initial_weight = self.weight();
166        let mut initial_len = self.len();
167        let mut shard_seed = rand::random(); // start from a random shard
168        let mut empty_shard = 0;
169
170        // Entries can be admitted or removed from the LRU by others during the loop below
171        // Track initial size not to over evict due to entries admitted after the loop starts
172        // self.weight() / self.len() is also used not to over evict
173        // due to entries already removed by others
174        while ((initial_weight > self.weight_limit && self.weight() > self.weight_limit)
175            || self
176                .len_watermark
177                .is_some_and(|w| initial_len > w && self.len() > w))
178            && empty_shard < N
179        {
180            if let Some(i) = self.evict_shard(shard_seed) {
181                initial_weight -= i.1;
182                initial_len = initial_len.saturating_sub(1);
183                evicted.push(i)
184            } else {
185                empty_shard += 1;
186            }
187            // move on to the next shard
188            shard_seed += 1;
189        }
190        evicted
191    }
192
193    /// Remove the given asset.
194    pub fn remove(&self, key: u64) -> Option<(T, usize)> {
195        let removed = self.units[get_shard(key, N)].write().remove(key);
196        if let Some((_, weight)) = removed.as_ref() {
197            self.weight.fetch_sub(*weight, Ordering::Relaxed);
198            self.len.fetch_sub(1, Ordering::Relaxed);
199        }
200        removed
201    }
202
203    /// Insert the item to the tail of this LRU.
204    ///
205    /// Useful to recreate an LRU in most-to-least order
206    pub fn insert_tail(&self, key: u64, data: T, weight: usize) -> bool {
207        if self.units[get_shard(key, N)]
208            .write()
209            .insert_tail(key, data, weight)
210        {
211            self.weight.fetch_add(weight, Ordering::Relaxed);
212            self.len.fetch_add(1, Ordering::Relaxed);
213            true
214        } else {
215            false
216        }
217    }
218
219    /// Check existence of a key without changing the order in LRU.
220    pub fn peek(&self, key: u64) -> bool {
221        self.units[get_shard(key, N)].read().peek(key).is_some()
222    }
223
224    /// Check the weight of a key without changing the order in LRU.
225    pub fn peek_weight(&self, key: u64) -> Option<usize> {
226        self.units[get_shard(key, N)].read().peek_weight(key)
227    }
228
229    /// Return the current total weight.
230    pub fn weight(&self) -> usize {
231        self.weight.load(Ordering::Relaxed)
232    }
233
234    /// Return the total weight of items evicted from this [Lru].
235    pub fn evicted_weight(&self) -> usize {
236        self.evicted_weight.load(Ordering::Relaxed)
237    }
238
239    /// Return the total count of items evicted from this [Lru].
240    pub fn evicted_len(&self) -> usize {
241        self.evicted_len.load(Ordering::Relaxed)
242    }
243
244    /// The number of items inside this [Lru].
245    #[allow(clippy::len_without_is_empty)]
246    pub fn len(&self) -> usize {
247        self.len.load(Ordering::Relaxed)
248    }
249
250    /// Scan a shard with the given function F
251    pub fn iter_for_each<F>(&self, shard: usize, f: F)
252    where
253        F: FnMut((&T, usize)),
254    {
255        assert!(shard < N);
256        self.units[shard].read().iter().for_each(f);
257    }
258
259    /// Get the total number of shards
260    pub const fn shards(&self) -> usize {
261        N
262    }
263
264    /// Get the number of items inside a shard
265    pub fn shard_len(&self, shard: usize) -> usize {
266        self.units[shard].read().len()
267    }
268
269    /// Get the weight (total size) inside a shard
270    pub fn shard_weight(&self, shard: usize) -> usize {
271        self.units[shard].read().used_weight
272    }
273}
274
275#[inline]
276fn get_shard(key: u64, n_shards: usize) -> usize {
277    (key % n_shards as u64) as usize
278}
279
280struct LruNode<T> {
281    data: T,
282    list_index: usize,
283    weight: usize,
284}
285
286struct LruUnit<T> {
287    lookup_table: HashMap<u64, Box<LruNode<T>>>,
288    order: LinkedList,
289    used_weight: usize,
290}
291
292impl<T> LruUnit<T> {
293    fn with_capacity(capacity: usize) -> Self {
294        LruUnit {
295            lookup_table: HashMap::with_capacity(capacity),
296            order: LinkedList::with_capacity(capacity),
297            used_weight: 0,
298        }
299    }
300
301    /// Peek data associated with key, if it exists.
302    pub fn peek(&self, key: u64) -> Option<&T> {
303        self.lookup_table.get(&key).map(|n| &n.data)
304    }
305
306    /// Peek weight associated with key, if it exists.
307    pub fn peek_weight(&self, key: u64) -> Option<usize> {
308        self.lookup_table.get(&key).map(|n| n.weight)
309    }
310
311    /// Admit into LRU, return old weight if there was any.
312    pub fn admit(&mut self, key: u64, data: T, weight: usize) -> usize {
313        if let Some(node) = self.lookup_table.get_mut(&key) {
314            let old_weight = Self::adjust_weight(node, &mut self.used_weight, weight);
315            node.data = data;
316            self.order.promote(node.list_index);
317            return old_weight;
318        }
319        self.used_weight += weight;
320        let list_index = self.order.push_head(key);
321        let node = Box::new(LruNode {
322            data,
323            list_index,
324            weight,
325        });
326        self.lookup_table.insert(key, node);
327        0
328    }
329
330    /// Increase the weight of an existing key. Returns the new weight or 0 if the key did not
331    /// exist, along with the new weight (or 0).
332    ///
333    /// If a `max_weight` is provided, the weight cannot exceed this max weight. If the current
334    /// weight is higher than the max, it will be capped to the max.
335    pub fn increment_weight(
336        &mut self,
337        key: u64,
338        delta: usize,
339        max_weight: Option<usize>,
340    ) -> Option<(usize, usize)> {
341        if let Some(node) = self.lookup_table.get_mut(&key) {
342            let new_weight =
343                max_weight.map_or(node.weight + delta, |m| (node.weight + delta).min(m));
344            let old_weight = Self::adjust_weight(node, &mut self.used_weight, new_weight);
345            self.order.promote(node.list_index);
346            return Some((old_weight, new_weight));
347        }
348        None
349    }
350
351    pub fn access(&mut self, key: u64) -> bool {
352        if let Some(node) = self.lookup_table.get(&key) {
353            self.order.promote(node.list_index);
354            true
355        } else {
356            false
357        }
358    }
359
360    // Check if a key is already in the top n most recently used nodes.
361    // this is a heuristic to reduce write, which requires exclusive locks, for promotion,
362    // especially on very populate nodes
363    // NOTE: O(n) search here so limit needs to be small
364    pub fn need_promote(&self, key: u64, limit: usize) -> bool {
365        !self.order.exist_near_head(key, limit)
366    }
367
368    // try to evict 1 node
369    pub fn evict(&mut self) -> Option<(T, usize)> {
370        self.order.pop_tail().map(|key| {
371            // unwrap is safe because we always insert in both the hashtable and the list
372            let node = self.lookup_table.remove(&key).unwrap();
373            self.used_weight -= node.weight;
374            (node.data, node.weight)
375        })
376    }
377    // TODO: scan the tail up to K elements to decide which ones to evict
378
379    pub fn remove(&mut self, key: u64) -> Option<(T, usize)> {
380        self.lookup_table.remove(&key).map(|node| {
381            let list_key = self.order.remove(node.list_index);
382            assert_eq!(key, list_key);
383            self.used_weight -= node.weight;
384            (node.data, node.weight)
385        })
386    }
387
388    pub fn insert_tail(&mut self, key: u64, data: T, weight: usize) -> bool {
389        if self.lookup_table.contains_key(&key) {
390            return false;
391        }
392        let list_index = self.order.push_tail(key);
393        let node = Box::new(LruNode {
394            data,
395            list_index,
396            weight,
397        });
398        self.lookup_table.insert(key, node);
399        self.used_weight += weight;
400        true
401    }
402
403    pub fn len(&self) -> usize {
404        assert_eq!(self.lookup_table.len(), self.order.len());
405        self.lookup_table.len()
406    }
407
408    #[cfg(test)]
409    pub fn used_weight(&self) -> usize {
410        self.used_weight
411    }
412
413    pub fn iter(&self) -> LruUnitIter<'_, T> {
414        LruUnitIter {
415            unit: self,
416            iter: self.order.iter(),
417        }
418    }
419
420    // Adjusts node weight to the new given weight.
421    // Returns old weight.
422    #[inline]
423    fn adjust_weight(node: &mut LruNode<T>, used_weight: &mut usize, weight: usize) -> usize {
424        let old_weight = node.weight;
425        if weight != old_weight {
426            *used_weight += weight;
427            *used_weight -= old_weight;
428            node.weight = weight;
429        }
430        old_weight
431    }
432}
433
434struct LruUnitIter<'a, T> {
435    unit: &'a LruUnit<T>,
436    iter: LinkedListIter<'a>,
437}
438
439impl<'a, T> Iterator for LruUnitIter<'a, T> {
440    type Item = (&'a T, usize);
441
442    fn next(&mut self) -> Option<Self::Item> {
443        self.iter.next().map(|key| {
444            // safe because we always items in table and list are always 1:1
445            let node = self.unit.lookup_table.get(key).unwrap();
446            (&node.data, node.weight)
447        })
448    }
449
450    fn size_hint(&self) -> (usize, Option<usize>) {
451        self.iter.size_hint()
452    }
453}
454
455impl<T> DoubleEndedIterator for LruUnitIter<'_, T> {
456    fn next_back(&mut self) -> Option<Self::Item> {
457        self.iter.next_back().map(|key| {
458            // safe because we always items in table and list are always 1:1
459            let node = self.unit.lookup_table.get(key).unwrap();
460            (&node.data, node.weight)
461        })
462    }
463}
464
465#[cfg(test)]
466mod test_lru {
467    use super::*;
468
469    fn assert_lru<T: Copy + PartialEq + std::fmt::Debug, const N: usize>(
470        lru: &Lru<T, N>,
471        values: &[T],
472        shard: usize,
473    ) {
474        let mut list_values = vec![];
475        lru.iter_for_each(shard, |(v, _)| list_values.push(*v));
476        assert_eq!(values, &list_values)
477    }
478
479    #[test]
480    fn test_admit() {
481        let lru = Lru::<_, 2>::with_capacity(30, 10);
482        assert_eq!(lru.len(), 0);
483
484        lru.admit(2, 2, 3);
485        assert_eq!(lru.len(), 1);
486        assert_eq!(lru.weight(), 3);
487
488        lru.admit(2, 2, 1);
489        assert_eq!(lru.len(), 1);
490        assert_eq!(lru.weight(), 1);
491
492        lru.admit(2, 2, 2); // admit again with different weight
493        assert_eq!(lru.len(), 1);
494        assert_eq!(lru.weight(), 2);
495
496        lru.admit(3, 3, 3);
497        lru.admit(4, 4, 4);
498
499        assert_eq!(lru.weight(), 2 + 3 + 4);
500        assert_eq!(lru.len(), 3);
501    }
502
503    #[test]
504    fn test_promote() {
505        let lru = Lru::<_, 2>::with_capacity(30, 10);
506
507        lru.admit(2, 2, 2);
508        lru.admit(3, 3, 3);
509        lru.admit(4, 4, 4);
510        lru.admit(5, 5, 5);
511        lru.admit(6, 6, 6);
512        assert_lru(&lru, &[6, 4, 2], 0);
513        assert_lru(&lru, &[5, 3], 1);
514
515        assert!(lru.promote(3));
516        assert_lru(&lru, &[3, 5], 1);
517        assert!(lru.promote(3));
518        assert_lru(&lru, &[3, 5], 1);
519
520        assert!(lru.promote(2));
521        assert_lru(&lru, &[2, 6, 4], 0);
522
523        assert!(!lru.promote(7)); // 7 doesn't exist
524        assert_lru(&lru, &[2, 6, 4], 0);
525        assert_lru(&lru, &[3, 5], 1);
526
527        // promote 2 to top 1, already there
528        assert!(lru.promote_top_n(2, 1));
529        assert_lru(&lru, &[2, 6, 4], 0);
530
531        // promote 4 to top 3, already there
532        assert!(lru.promote_top_n(4, 3));
533        assert_lru(&lru, &[2, 6, 4], 0);
534
535        // promote 4 to top 2
536        assert!(lru.promote_top_n(4, 2));
537        assert_lru(&lru, &[4, 2, 6], 0);
538
539        // promote 2 to top 1
540        assert!(lru.promote_top_n(2, 1));
541        assert_lru(&lru, &[2, 4, 6], 0);
542
543        assert!(!lru.promote_top_n(7, 1)); // 7 doesn't exist
544    }
545
546    #[test]
547    fn test_evict() {
548        let lru = Lru::<_, 2>::with_capacity(14, 10);
549
550        // same weight to make the random eviction less random
551        lru.admit(2, 2, 2);
552        lru.admit(3, 3, 2);
553        lru.admit(4, 4, 4);
554        lru.admit(5, 5, 4);
555        lru.admit(6, 6, 2);
556        lru.admit(7, 7, 2);
557
558        assert_lru(&lru, &[6, 4, 2], 0);
559        assert_lru(&lru, &[7, 5, 3], 1);
560
561        assert_eq!(lru.weight(), 16);
562        assert_eq!(lru.len(), 6);
563
564        let evicted = lru.evict_to_limit();
565        assert_eq!(lru.weight(), 14);
566        assert_eq!(lru.len(), 5);
567        assert_eq!(lru.evicted_weight(), 2);
568        assert_eq!(lru.evicted_len(), 1);
569        assert_eq!(evicted.len(), 1);
570        assert_eq!(evicted[0].1, 2); //weight
571        assert!(evicted[0].0 == 2 || evicted[0].0 == 3); //either 2 or 3 are evicted
572
573        let lru = Lru::<_, 2>::with_capacity(6, 10);
574
575        // same weight random eviction less random
576        lru.admit(2, 2, 2);
577        lru.admit(3, 3, 2);
578        lru.admit(4, 4, 2);
579        lru.admit(5, 5, 2);
580        lru.admit(6, 6, 2);
581        lru.admit(7, 7, 2);
582        assert_eq!(lru.weight(), 12);
583        assert_eq!(lru.len(), 6);
584
585        let evicted = lru.evict_to_limit();
586        // NOTE: there is a low chance this test would fail see the TODO in evict_to_limit
587        assert_eq!(lru.weight(), 6);
588        assert_eq!(lru.len(), 3);
589        assert_eq!(lru.evicted_weight(), 6);
590        assert_eq!(lru.evicted_len(), 3);
591        assert_eq!(evicted.len(), 3);
592    }
593
594    #[test]
595    fn test_increment_weight() {
596        let lru = Lru::<_, 2>::with_capacity(6, 10);
597        lru.admit(1, 1, 1);
598        lru.increment_weight(1, 1, None);
599        assert_eq!(lru.weight(), 1 + 1);
600
601        lru.increment_weight(0, 1000, None);
602        assert_eq!(lru.weight(), 1 + 1);
603
604        lru.admit(2, 2, 2);
605        lru.increment_weight(2, 2, None);
606        assert_eq!(lru.weight(), 1 + 1 + 2 + 2);
607
608        lru.increment_weight(2, 2, Some(3));
609        assert_eq!(lru.weight(), 1 + 1 + 3);
610    }
611
612    #[test]
613    fn test_remove() {
614        let lru = Lru::<_, 2>::with_capacity(30, 10);
615        lru.admit(2, 2, 2);
616        lru.admit(3, 3, 3);
617        lru.admit(4, 4, 4);
618        lru.admit(5, 5, 5);
619        lru.admit(6, 6, 6);
620
621        assert_eq!(lru.weight(), 2 + 3 + 4 + 5 + 6);
622        assert_eq!(lru.len(), 5);
623        assert_lru(&lru, &[6, 4, 2], 0);
624        assert_lru(&lru, &[5, 3], 1);
625
626        let node = lru.remove(6).unwrap();
627        assert_eq!(node.0, 6); // data
628        assert_eq!(node.1, 6); // weight
629        assert_eq!(lru.weight(), 2 + 3 + 4 + 5);
630        assert_eq!(lru.len(), 4);
631        assert_lru(&lru, &[4, 2], 0);
632
633        let node = lru.remove(3).unwrap();
634        assert_eq!(node.0, 3); // data
635        assert_eq!(node.1, 3); // weight
636        assert_eq!(lru.weight(), 2 + 4 + 5);
637        assert_eq!(lru.len(), 3);
638        assert_lru(&lru, &[5], 1);
639
640        assert!(lru.remove(7).is_none());
641    }
642
643    #[test]
644    fn test_peek() {
645        let lru = Lru::<_, 2>::with_capacity(30, 10);
646        lru.admit(2, 2, 2);
647        lru.admit(3, 3, 3);
648        lru.admit(4, 4, 4);
649
650        assert!(lru.peek(4));
651        assert!(lru.peek(3));
652        assert!(lru.peek(2));
653
654        assert_lru(&lru, &[4, 2], 0);
655        assert_lru(&lru, &[3], 1);
656    }
657
658    #[test]
659    fn test_insert_tail() {
660        let lru = Lru::<_, 2>::with_capacity(30, 10);
661        lru.admit(2, 2, 2);
662        lru.admit(3, 3, 3);
663        lru.admit(4, 4, 4);
664        lru.admit(5, 5, 5);
665        lru.admit(6, 6, 6);
666
667        assert_eq!(lru.weight(), 2 + 3 + 4 + 5 + 6);
668        assert_eq!(lru.len(), 5);
669        assert_lru(&lru, &[6, 4, 2], 0);
670        assert_lru(&lru, &[5, 3], 1);
671
672        assert!(lru.insert_tail(7, 7, 7));
673        assert_eq!(lru.weight(), 2 + 3 + 4 + 5 + 6 + 7);
674        assert_eq!(lru.len(), 6);
675        assert_lru(&lru, &[5, 3, 7], 1);
676
677        // ignore existing ones
678        assert!(!lru.insert_tail(6, 6, 7));
679    }
680
681    #[test]
682    fn test_watermark_eviction() {
683        const WEIGHT_LIMIT: usize = usize::MAX / 2;
684        let lru = Lru::<u64, 2>::with_capacity_and_watermark(WEIGHT_LIMIT, 10, Some(4));
685
686        // admit 6 items, each weight 1
687        for k in [2u64, 3, 4, 5, 6, 7] {
688            lru.admit(k, k, 1);
689        }
690
691        assert!(lru.weight() < WEIGHT_LIMIT);
692        assert_eq!(lru.len(), 6);
693
694        let evicted = lru.evict_to_limit();
695        assert_eq!(lru.len(), 4);
696        assert_eq!(evicted.len(), 2);
697        assert_eq!(lru.evicted_len(), 2);
698    }
699}
700
701#[cfg(test)]
702mod test_lru_unit {
703    use super::*;
704
705    fn assert_lru<T: Copy + PartialEq + std::fmt::Debug>(lru: &LruUnit<T>, values: &[T]) {
706        let list_values: Vec<_> = lru.iter().map(|(v, _)| *v).collect();
707        assert_eq!(values, &list_values)
708    }
709
710    #[test]
711    fn test_admit() {
712        let mut lru = LruUnit::with_capacity(10);
713        assert_eq!(lru.len(), 0);
714        assert!(lru.peek(0).is_none());
715
716        lru.admit(2, 2, 1);
717        assert_eq!(lru.len(), 1);
718        assert_eq!(lru.peek(2).unwrap(), &2);
719        assert_eq!(lru.used_weight(), 1);
720
721        lru.admit(2, 2, 2); // admit again with different weight
722        assert_eq!(lru.used_weight(), 2);
723
724        lru.admit(3, 3, 3);
725        lru.admit(4, 4, 4);
726
727        assert_eq!(lru.used_weight(), 2 + 3 + 4);
728        assert_lru(&lru, &[4, 3, 2]);
729    }
730
731    #[test]
732    fn test_access() {
733        let mut lru = LruUnit::with_capacity(10);
734
735        lru.admit(2, 2, 2);
736        lru.admit(3, 3, 3);
737        lru.admit(4, 4, 4);
738        assert_lru(&lru, &[4, 3, 2]);
739
740        assert!(lru.access(3));
741        assert_lru(&lru, &[3, 4, 2]);
742        assert!(lru.access(3));
743        assert_lru(&lru, &[3, 4, 2]);
744        assert!(lru.access(2));
745        assert_lru(&lru, &[2, 3, 4]);
746
747        assert!(!lru.access(5)); // 5 doesn't exist
748        assert_lru(&lru, &[2, 3, 4]);
749
750        assert!(!lru.need_promote(2, 1));
751        assert!(lru.need_promote(3, 1));
752        assert!(!lru.need_promote(4, 9999));
753    }
754
755    #[test]
756    fn test_evict() {
757        let mut lru = LruUnit::with_capacity(10);
758
759        lru.admit(2, 2, 2);
760        lru.admit(3, 3, 3);
761        lru.admit(4, 4, 4);
762        assert_lru(&lru, &[4, 3, 2]);
763
764        assert!(lru.access(3));
765        assert!(lru.access(3));
766        assert!(lru.access(2));
767        assert_lru(&lru, &[2, 3, 4]);
768
769        assert_eq!(lru.used_weight(), 2 + 3 + 4);
770        assert_eq!(lru.evict(), Some((4, 4)));
771        assert_eq!(lru.used_weight(), 2 + 3);
772        assert_lru(&lru, &[2, 3]);
773
774        assert_eq!(lru.evict(), Some((3, 3)));
775        assert_eq!(lru.used_weight(), 2);
776        assert_lru(&lru, &[2]);
777
778        assert_eq!(lru.evict(), Some((2, 2)));
779        assert_eq!(lru.used_weight(), 0);
780        assert_lru(&lru, &[]);
781
782        assert_eq!(lru.evict(), None);
783        assert_eq!(lru.used_weight(), 0);
784        assert_lru(&lru, &[]);
785    }
786
787    #[test]
788    fn test_increment_weight() {
789        let mut lru = LruUnit::with_capacity(10);
790        lru.admit(1, 1, 1);
791        lru.increment_weight(1, 1, None);
792        assert_eq!(lru.used_weight(), 1 + 1);
793
794        lru.increment_weight(0, 1000, None);
795        assert_eq!(lru.used_weight(), 1 + 1);
796
797        lru.admit(2, 2, 2);
798        lru.increment_weight(2, 2, None);
799        assert_eq!(lru.used_weight(), 1 + 1 + 2 + 2);
800
801        lru.admit(3, 3, 3);
802        lru.increment_weight(3, 3, Some(5));
803        assert_eq!(lru.used_weight(), 1 + 1 + 2 + 2 + 3 + 2);
804
805        lru.increment_weight(3, 3, Some(3));
806        assert_eq!(lru.used_weight(), 1 + 1 + 2 + 2 + 3);
807    }
808
809    #[test]
810    fn test_remove() {
811        let mut lru = LruUnit::with_capacity(10);
812
813        lru.admit(2, 2, 2);
814        lru.admit(3, 3, 3);
815        lru.admit(4, 4, 4);
816        lru.admit(5, 5, 5);
817        assert_lru(&lru, &[5, 4, 3, 2]);
818
819        assert!(lru.access(4));
820        assert!(lru.access(3));
821        assert!(lru.access(3));
822        assert!(lru.access(2));
823        assert_lru(&lru, &[2, 3, 4, 5]);
824
825        assert_eq!(lru.used_weight(), 2 + 3 + 4 + 5);
826        assert_eq!(lru.remove(2), Some((2, 2)));
827        assert_eq!(lru.used_weight(), 3 + 4 + 5);
828        assert_lru(&lru, &[3, 4, 5]);
829
830        assert_eq!(lru.remove(4), Some((4, 4)));
831        assert_eq!(lru.used_weight(), 3 + 5);
832        assert_lru(&lru, &[3, 5]);
833
834        assert_eq!(lru.remove(5), Some((5, 5)));
835        assert_eq!(lru.used_weight(), 3);
836        assert_lru(&lru, &[3]);
837
838        assert_eq!(lru.remove(1), None);
839        assert_eq!(lru.used_weight(), 3);
840        assert_lru(&lru, &[3]);
841
842        assert_eq!(lru.remove(3), Some((3, 3)));
843        assert_eq!(lru.used_weight(), 0);
844        assert_lru(&lru, &[]);
845    }
846
847    #[test]
848    fn test_insert_tail() {
849        let mut lru = LruUnit::with_capacity(10);
850        assert_eq!(lru.len(), 0);
851        assert!(lru.peek(0).is_none());
852
853        assert!(lru.insert_tail(2, 2, 1));
854        assert_eq!(lru.len(), 1);
855        assert_eq!(lru.peek(2).unwrap(), &2);
856        assert_eq!(lru.used_weight(), 1);
857
858        assert!(!lru.insert_tail(2, 2, 2));
859        assert!(lru.insert_tail(3, 3, 3));
860        assert_eq!(lru.used_weight(), 1 + 3);
861        assert_lru(&lru, &[2, 3]);
862
863        assert!(lru.insert_tail(4, 4, 4));
864        assert!(lru.insert_tail(5, 5, 5));
865        assert_eq!(lru.used_weight(), 1 + 3 + 4 + 5);
866        assert_lru(&lru, &[2, 3, 4, 5]);
867    }
868}