tinyufo/
lib.rs

1// Copyright 2025 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//! A In-memory cache implementation with TinyLFU as the admission policy and [S3-FIFO](https://s3fifo.com/) as the eviction policy.
16//!
17//! TinyUFO improves cache hit ratio noticeably compared to LRU.
18//!
19//! TinyUFO is lock-free. It is very fast in the systems with a lot concurrent reads and/or writes
20
21use ahash::RandomState;
22use crossbeam_queue::SegQueue;
23use std::marker::PhantomData;
24use std::sync::atomic::AtomicUsize;
25use std::sync::atomic::{
26    AtomicBool, AtomicU8,
27    Ordering::{Acquire, Relaxed, SeqCst},
28};
29mod buckets;
30mod estimation;
31
32use buckets::Buckets;
33use estimation::TinyLfu;
34use std::hash::Hash;
35
36const SMALL: bool = false;
37const MAIN: bool = true;
38
39// Indicate which queue an item is located
40#[derive(Debug, Default)]
41struct Location(AtomicBool);
42
43impl Location {
44    fn new_small() -> Self {
45        Self(AtomicBool::new(SMALL))
46    }
47
48    fn value(&self) -> bool {
49        self.0.load(Relaxed)
50    }
51
52    fn is_main(&self) -> bool {
53        self.value()
54    }
55
56    fn move_to_main(&self) {
57        self.0.store(true, Relaxed);
58    }
59}
60
61// We have 8 bits to spare but we still cap at 3. This is to make sure that the main queue
62// in the worst case can find something to evict quickly
63const USES_CAP: u8 = 3;
64
65#[derive(Debug, Default)]
66struct Uses(AtomicU8);
67
68impl Uses {
69    pub fn inc_uses(&self) -> u8 {
70        loop {
71            let uses = self.uses();
72            if uses >= USES_CAP {
73                return uses;
74            }
75            if let Err(new) = self.0.compare_exchange(uses, uses + 1, Acquire, Relaxed) {
76                // someone else beat us to it
77                if new >= USES_CAP {
78                    // already above cap
79                    return new;
80                } // else, try again
81            } else {
82                return uses + 1;
83            }
84        }
85    }
86
87    // decrease uses, return the previous value
88    pub fn decr_uses(&self) -> u8 {
89        loop {
90            let uses = self.uses();
91            if uses == 0 {
92                return 0;
93            }
94            if let Err(new) = self.0.compare_exchange(uses, uses - 1, Acquire, Relaxed) {
95                // someone else beat us to it
96                if new == 0 {
97                    return 0;
98                } // else, try again
99            } else {
100                return uses;
101            }
102        }
103    }
104
105    pub fn uses(&self) -> u8 {
106        self.0.load(Relaxed)
107    }
108}
109
110type Key = u64;
111type Weight = u16;
112
113/// The key-value pair returned from cache eviction
114#[derive(Clone)]
115pub struct KV<T> {
116    /// NOTE: that we currently don't store the Actual key in the cache. This returned value
117    /// is just the hash of it.
118    pub key: Key,
119    pub data: T,
120    pub weight: Weight,
121}
122
123// the data and its metadata
124pub struct Bucket<T> {
125    uses: Uses,
126    queue: Location,
127    weight: Weight,
128    data: T,
129}
130
131const SMALL_QUEUE_PERCENTAGE: f32 = 0.1;
132
133struct FiFoQueues<T> {
134    total_weight_limit: usize,
135
136    small: SegQueue<Key>,
137    small_weight: AtomicUsize,
138
139    main: SegQueue<Key>,
140    main_weight: AtomicUsize,
141
142    // this replaces the ghost queue of S3-FIFO with similar goal: track the evicted assets
143    estimator: TinyLfu,
144
145    _t: PhantomData<T>,
146}
147
148impl<T: Clone + Send + Sync + 'static> FiFoQueues<T> {
149    fn admit(
150        &self,
151        key: Key,
152        data: T,
153        weight: u16,
154        ignore_lfu: bool,
155        buckets: &Buckets<T>,
156    ) -> Vec<KV<T>> {
157        // Note that we only use TinyLFU during cache admission but not cache read.
158        // So effectively we mostly sketch the popularity of less popular assets.
159        // In this way the sketch is a bit more accurate on these assets.
160        // Also we don't need another separated window cache to address the sparse burst issue as
161        // this sketch doesn't favor very popular assets much.
162        let new_freq = self.estimator.incr(key);
163
164        assert!(weight > 0);
165        let new_bucket = {
166            let Some((uses, queue, weight)) = buckets.get_map(&key, |bucket| {
167                // the item exists, in case weight changes
168                let old_weight = bucket.weight;
169                let uses = bucket.uses.inc_uses();
170
171                fn update_atomic(weight: &AtomicUsize, old: u16, new: u16) {
172                    if old == new {
173                        return;
174                    }
175                    if old > new {
176                        weight.fetch_sub((old - new) as usize, SeqCst);
177                    } else {
178                        weight.fetch_add((new - old) as usize, SeqCst);
179                    }
180                }
181                let queue = bucket.queue.is_main();
182                if queue == MAIN {
183                    update_atomic(&self.main_weight, old_weight, weight);
184                } else {
185                    update_atomic(&self.small_weight, old_weight, weight);
186                }
187                (uses, queue, weight)
188            }) else {
189                let mut evicted = self.evict_to_limit(weight, buckets);
190                // TODO: figure out the right way to compare frequencies of different weights across
191                // many evicted assets. For now TinyLFU is only used when only evicting 1 item.
192                let (key, data, weight) = if !ignore_lfu && evicted.len() == 1 {
193                    // Apply the admission algorithm of TinyLFU: compare the incoming new item
194                    // and the evicted one. The more popular one is admitted to cache
195                    let evicted_first = &evicted[0];
196                    let evicted_freq = self.estimator.get(evicted_first.key);
197                    if evicted_freq > new_freq {
198                        // put it back
199                        let first = evicted.pop().expect("just check non-empty");
200                        // return the put value
201                        evicted.push(KV { key, data, weight });
202                        (first.key, first.data, first.weight)
203                    } else {
204                        (key, data, weight)
205                    }
206                } else {
207                    (key, data, weight)
208                };
209
210                let bucket = Bucket {
211                    queue: Location::new_small(),
212                    weight,
213                    uses: Default::default(), // 0
214                    data,
215                };
216                let old = buckets.insert(key, bucket);
217                if old.is_none() {
218                    // Always push key first before updating weight
219                    // If doing the other order, another concurrent thread might not
220                    // find things to evict
221                    self.small.push(key);
222                    self.small_weight.fetch_add(weight as usize, SeqCst);
223                } // else: two threads are racing adding the item
224                  // TODO: compare old.weight and update accordingly
225                return evicted;
226            };
227            Bucket {
228                queue: Location(queue.into()),
229                weight,
230                uses: Uses(uses.into()),
231                data,
232            }
233        };
234
235        // replace the existing one
236        buckets.insert(key, new_bucket);
237
238        // NOTE: there is a chance that the item itself is evicted if it happens to be the one selected
239        // by the algorithm. We could avoid this by checking if the item is in the returned evicted items,
240        // and then add it back. But to keep the code simple we just allow it to happen.
241        self.evict_to_limit(0, buckets)
242    }
243
244    // the `extra_weight` is to essentially tell the cache to reserve that amount of weight for
245    // admission. It is used when calling `evict_to_limit` before admitting the asset itself.
246    fn evict_to_limit(&self, extra_weight: Weight, buckets: &Buckets<T>) -> Vec<KV<T>> {
247        let mut evicted = if self.total_weight_limit
248            < self.small_weight.load(SeqCst) + self.main_weight.load(SeqCst) + extra_weight as usize
249        {
250            Vec::with_capacity(1)
251        } else {
252            vec![]
253        };
254        while self.total_weight_limit
255            < self.small_weight.load(SeqCst) + self.main_weight.load(SeqCst) + extra_weight as usize
256        {
257            if let Some(evicted_item) = self.evict_one(buckets) {
258                evicted.push(evicted_item);
259            } else {
260                break;
261            }
262        }
263
264        evicted
265    }
266
267    fn evict_one(&self, buckets: &Buckets<T>) -> Option<KV<T>> {
268        let evict_small = self.small_weight_limit() <= self.small_weight.load(SeqCst);
269
270        if evict_small {
271            let evicted = self.evict_one_from_small(buckets);
272            // evict_one_from_small could just promote everything to main without evicting any
273            // so need to evict_one_from_main if nothing evicted
274            if evicted.is_some() {
275                return evicted;
276            }
277        }
278        self.evict_one_from_main(buckets)
279    }
280
281    fn small_weight_limit(&self) -> usize {
282        (self.total_weight_limit as f32 * SMALL_QUEUE_PERCENTAGE).floor() as usize + 1
283    }
284
285    fn evict_one_from_small(&self, buckets: &Buckets<T>) -> Option<KV<T>> {
286        loop {
287            let Some(to_evict) = self.small.pop() else {
288                // empty queue, this is caught between another pop() and fetch_sub()
289                return None;
290            };
291
292            let v = buckets
293                .get_map(&to_evict, |bucket| {
294                    let weight = bucket.weight;
295                    self.small_weight.fetch_sub(weight as usize, SeqCst);
296
297                    if bucket.uses.uses() > 1 {
298                        // move to main
299                        bucket.queue.move_to_main();
300                        self.main.push(to_evict);
301                        self.main_weight.fetch_add(weight as usize, SeqCst);
302                        // continue until find one to evict
303                        None
304                    } else {
305                        let data = bucket.data.clone();
306                        let weight = bucket.weight;
307                        buckets.remove(&to_evict);
308                        Some(KV {
309                            key: to_evict,
310                            data,
311                            weight,
312                        })
313                    }
314                })
315                .flatten();
316            if v.is_some() {
317                // found the one to evict, break
318                return v;
319            }
320        }
321    }
322
323    fn evict_one_from_main(&self, buckets: &Buckets<T>) -> Option<KV<T>> {
324        loop {
325            let to_evict = self.main.pop()?;
326
327            if let Some(v) = buckets
328                .get_map(&to_evict, |bucket| {
329                    if bucket.uses.decr_uses() > 0 {
330                        // put it back
331                        self.main.push(to_evict);
332                        // continue the loop
333                        None
334                    } else {
335                        // evict
336                        let weight = bucket.weight;
337                        self.main_weight.fetch_sub(weight as usize, SeqCst);
338                        let data = bucket.data.clone();
339                        buckets.remove(&to_evict);
340                        Some(KV {
341                            key: to_evict,
342                            data,
343                            weight,
344                        })
345                    }
346                })
347                .flatten()
348            {
349                // found the one to evict, break
350                return Some(v);
351            }
352        }
353    }
354}
355
356/// [TinyUfo] cache
357pub struct TinyUfo<K, T> {
358    queues: FiFoQueues<T>,
359    buckets: Buckets<T>,
360    random_status: RandomState,
361    _k: PhantomData<K>,
362}
363impl<K: Hash, T: Clone + Send + Sync + 'static> TinyUfo<K, T> {
364    /// Create a new TinyUfo cache with the given weight limit and the given
365    /// size limit of the ghost queue.
366    pub fn new(total_weight_limit: usize, estimated_size: usize) -> Self {
367        let queues = FiFoQueues {
368            small: SegQueue::new(),
369            small_weight: 0.into(),
370            main: SegQueue::new(),
371            main_weight: 0.into(),
372            total_weight_limit,
373            estimator: TinyLfu::new(estimated_size),
374            _t: PhantomData,
375        };
376        TinyUfo {
377            queues,
378            buckets: Buckets::new_fast(estimated_size),
379            random_status: RandomState::new(),
380            _k: PhantomData,
381        }
382    }
383
384    /// Create a new TinyUfo cache but with more memory efficient data structures.
385    /// The trade-off is that the the get() is slower by a constant factor.
386    /// The cache hit ratio could be higher as this type of TinyUFO allows to store
387    /// more assets with the same memory.
388    pub fn new_compact(total_weight_limit: usize, estimated_size: usize) -> Self {
389        let queues = FiFoQueues {
390            small: SegQueue::new(),
391            small_weight: 0.into(),
392            main: SegQueue::new(),
393            main_weight: 0.into(),
394            total_weight_limit,
395            estimator: TinyLfu::new_compact(estimated_size),
396            _t: PhantomData,
397        };
398        TinyUfo {
399            queues,
400            buckets: Buckets::new_compact(estimated_size, 32),
401            random_status: RandomState::new(),
402            _k: PhantomData,
403        }
404    }
405
406    // TODO: with_capacity()
407
408    /// Read the given key
409    ///
410    /// Return Some(T) if the key exists
411    pub fn get(&self, key: &K) -> Option<T> {
412        let key = self.random_status.hash_one(key);
413        self.buckets.get_map(&key, |p| {
414            p.uses.inc_uses();
415            p.data.clone()
416        })
417    }
418
419    /// Put the key value to the [TinyUfo]
420    ///
421    /// Return a list of [KV] of key and `T` that are evicted
422    pub fn put(&self, key: K, data: T, weight: Weight) -> Vec<KV<T>> {
423        let key = self.random_status.hash_one(&key);
424        self.queues.admit(key, data, weight, false, &self.buckets)
425    }
426
427    /// Remove the given key from the cache if it exists
428    ///
429    /// Returns Some(T) if the key was found and removed, None otherwise
430    pub fn remove(&self, key: &K) -> Option<T> {
431        let key = self.random_status.hash_one(key);
432
433        // Get data and update weights
434        let result = self.buckets.get_map(&key, |bucket| {
435            let data = bucket.data.clone();
436            let weight = bucket.weight;
437
438            // Update weight based on queue location
439            if bucket.queue.is_main() {
440                self.queues.main_weight.fetch_sub(weight as usize, SeqCst);
441            } else {
442                self.queues.small_weight.fetch_sub(weight as usize, SeqCst);
443            }
444
445            data
446        });
447
448        // If we found and processed the item, remove it from buckets
449        if result.is_some() {
450            self.buckets.remove(&key);
451        }
452
453        result
454    }
455    /// Always put the key value to the [TinyUfo]
456    ///
457    /// Return a list of [KV] of key and `T` that are evicted
458    ///
459    /// Similar to [Self::put] but guarantee the assertion of the asset.
460    /// In [Self::put], the TinyLFU check may reject putting the current asset if it is less
461    /// popular than the once being evicted.
462    ///
463    /// In some real world use cases, a few reads to the same asset may be pending for the put action
464    /// to be finished so that they can read the asset from cache. Neither the above behaviors are ideal
465    /// for this use case.
466    ///
467    /// Compared to [Self::put], the hit ratio when using this function is reduced by about 0.5pp or less in
468    /// under zipf workloads.
469    pub fn force_put(&self, key: K, data: T, weight: Weight) -> Vec<KV<T>> {
470        let key = self.random_status.hash_one(&key);
471        self.queues.admit(key, data, weight, true, &self.buckets)
472    }
473
474    #[cfg(test)]
475    fn peek_queue(&self, key: K) -> Option<bool> {
476        let key = self.random_status.hash_one(&key);
477        self.buckets.get_queue(&key)
478    }
479}
480
481#[cfg(test)]
482mod tests {
483    use super::*;
484
485    #[test]
486    fn test_uses() {
487        let uses: Uses = Default::default();
488        assert_eq!(uses.uses(), 0);
489        uses.inc_uses();
490        assert_eq!(uses.uses(), 1);
491        for _ in 0..USES_CAP {
492            uses.inc_uses();
493        }
494        assert_eq!(uses.uses(), USES_CAP);
495
496        for _ in 0..USES_CAP + 2 {
497            uses.decr_uses();
498        }
499        assert_eq!(uses.uses(), 0);
500    }
501
502    #[test]
503    fn test_evict_from_small() {
504        let mut cache = TinyUfo::new(5, 5);
505        cache.random_status = RandomState::with_seeds(2, 3, 4, 5);
506        cache.queues.estimator = TinyLfu::new_seeded(5);
507
508        cache.put(1, 1, 1);
509        cache.put(2, 2, 2);
510        cache.put(3, 3, 2);
511        // cache full now
512
513        assert_eq!(cache.peek_queue(1), Some(SMALL));
514        assert_eq!(cache.peek_queue(2), Some(SMALL));
515        assert_eq!(cache.peek_queue(3), Some(SMALL));
516
517        let evicted = cache.put(4, 4, 3);
518        assert_eq!(evicted.len(), 2);
519        assert_eq!(evicted[0].data, 1);
520        assert_eq!(evicted[1].data, 2);
521
522        assert_eq!(cache.peek_queue(1), None);
523        assert_eq!(cache.peek_queue(2), None);
524        assert_eq!(cache.peek_queue(3), Some(SMALL));
525    }
526
527    #[test]
528    fn test_evict_from_small_to_main() {
529        let mut cache = TinyUfo::new(5, 5);
530        cache.random_status = RandomState::with_seeds(2, 3, 4, 5);
531        cache.queues.estimator = TinyLfu::new_seeded(5);
532
533        cache.put(1, 1, 1);
534        cache.put(2, 2, 2);
535        cache.put(3, 3, 2);
536        // cache full now
537
538        cache.get(&1);
539        cache.get(&1); // 1 will be moved to main during next eviction
540
541        assert_eq!(cache.peek_queue(1), Some(SMALL));
542        assert_eq!(cache.peek_queue(2), Some(SMALL));
543        assert_eq!(cache.peek_queue(3), Some(SMALL));
544
545        let evicted = cache.put(4, 4, 2);
546        assert_eq!(evicted.len(), 1);
547        assert_eq!(evicted[0].weight, 2);
548
549        assert_eq!(cache.peek_queue(1), Some(MAIN));
550        // either 2, 3, or 4 was evicted. Check evicted for which.
551        let mut remaining = vec![2, 3, 4];
552        remaining.remove(
553            remaining
554                .iter()
555                .position(|x| *x == evicted[0].data)
556                .unwrap(),
557        );
558        assert_eq!(cache.peek_queue(evicted[0].key), None);
559        for k in remaining {
560            assert_eq!(cache.peek_queue(k), Some(SMALL));
561        }
562    }
563
564    #[test]
565    fn test_evict_reentry() {
566        let mut cache = TinyUfo::new(5, 5);
567        cache.random_status = RandomState::with_seeds(2, 3, 4, 5);
568        cache.queues.estimator = TinyLfu::new_seeded(5);
569
570        cache.put(1, 1, 1);
571        cache.put(2, 2, 2);
572        cache.put(3, 3, 2);
573        // cache full now
574
575        assert_eq!(cache.peek_queue(1), Some(SMALL));
576        assert_eq!(cache.peek_queue(2), Some(SMALL));
577        assert_eq!(cache.peek_queue(3), Some(SMALL));
578
579        let evicted = cache.put(4, 4, 1);
580        assert_eq!(evicted.len(), 1);
581        assert_eq!(evicted[0].data, 1);
582
583        assert_eq!(cache.peek_queue(1), None);
584        assert_eq!(cache.peek_queue(2), Some(SMALL));
585        assert_eq!(cache.peek_queue(3), Some(SMALL));
586        assert_eq!(cache.peek_queue(4), Some(SMALL));
587
588        let evicted = cache.put(1, 1, 1);
589        assert_eq!(evicted.len(), 1);
590        assert_eq!(evicted[0].data, 2);
591
592        assert_eq!(cache.peek_queue(1), Some(SMALL));
593        assert_eq!(cache.peek_queue(2), None);
594        assert_eq!(cache.peek_queue(3), Some(SMALL));
595        assert_eq!(cache.peek_queue(4), Some(SMALL));
596    }
597
598    #[test]
599    fn test_evict_entry_denied() {
600        let mut cache = TinyUfo::new(5, 5);
601        cache.random_status = RandomState::with_seeds(2, 3, 4, 5);
602        cache.queues.estimator = TinyLfu::new_seeded(5);
603
604        cache.put(1, 1, 1);
605        cache.put(2, 2, 2);
606        cache.put(3, 3, 2);
607        // cache full now
608
609        assert_eq!(cache.peek_queue(1), Some(SMALL));
610        assert_eq!(cache.peek_queue(2), Some(SMALL));
611        assert_eq!(cache.peek_queue(3), Some(SMALL));
612
613        // trick: put a few times to bump their frequencies
614        cache.put(1, 1, 1);
615        cache.put(2, 2, 2);
616        cache.put(3, 3, 2);
617
618        let evicted = cache.put(4, 4, 1);
619        assert_eq!(evicted.len(), 1);
620        assert_eq!(evicted[0].data, 4); // 4 is returned
621
622        assert_eq!(cache.peek_queue(1), Some(SMALL));
623        assert_eq!(cache.peek_queue(2), Some(SMALL));
624        assert_eq!(cache.peek_queue(3), Some(SMALL));
625        assert_eq!(cache.peek_queue(4), None);
626    }
627
628    #[test]
629    fn test_force_put() {
630        let mut cache = TinyUfo::new(5, 5);
631        cache.random_status = RandomState::with_seeds(2, 3, 4, 5);
632        cache.queues.estimator = TinyLfu::new_seeded(5);
633
634        cache.put(1, 1, 1);
635        cache.put(2, 2, 2);
636        cache.put(3, 3, 2);
637        // cache full now
638
639        assert_eq!(cache.peek_queue(1), Some(SMALL));
640        assert_eq!(cache.peek_queue(2), Some(SMALL));
641        assert_eq!(cache.peek_queue(3), Some(SMALL));
642
643        // trick: put a few times to bump their frequencies
644        cache.put(1, 1, 1);
645        cache.put(2, 2, 2);
646        cache.put(3, 3, 2);
647
648        // force put will replace 1 with 4 even through 1 is more popular
649        let evicted = cache.force_put(4, 4, 1);
650        assert_eq!(evicted.len(), 1);
651        assert_eq!(evicted[0].data, 1); // 1 is returned
652
653        assert_eq!(cache.peek_queue(1), None);
654        assert_eq!(cache.peek_queue(2), Some(SMALL));
655        assert_eq!(cache.peek_queue(3), Some(SMALL));
656        assert_eq!(cache.peek_queue(4), Some(SMALL));
657    }
658
659    #[test]
660    fn test_evict_from_main() {
661        let mut cache = TinyUfo::new(5, 5);
662        cache.random_status = RandomState::with_seeds(2, 3, 4, 5);
663        cache.queues.estimator = TinyLfu::new_seeded(5);
664
665        cache.put(1, 1, 1);
666        cache.put(2, 2, 2);
667        cache.put(3, 3, 2);
668        // cache full now
669
670        // all 3 will qualify to main
671        cache.get(&1);
672        cache.get(&1);
673        cache.get(&2);
674        cache.get(&2);
675        cache.get(&3);
676        cache.get(&3);
677
678        let evicted = cache.put(4, 4, 1);
679        assert_eq!(evicted.len(), 1);
680        assert_eq!(evicted[0].data, 1);
681
682        // 1 kicked from main
683        assert_eq!(cache.peek_queue(1), None);
684        assert_eq!(cache.peek_queue(2), Some(MAIN));
685        assert_eq!(cache.peek_queue(3), Some(MAIN));
686        assert_eq!(cache.peek_queue(4), Some(SMALL));
687
688        let evicted = cache.put(1, 1, 1);
689        assert_eq!(evicted.len(), 1);
690        assert_eq!(evicted[0].data, 4);
691
692        assert_eq!(cache.peek_queue(1), Some(SMALL));
693        assert_eq!(cache.peek_queue(2), Some(MAIN));
694        assert_eq!(cache.peek_queue(3), Some(MAIN));
695        assert_eq!(cache.peek_queue(4), None);
696    }
697
698    #[test]
699    fn test_evict_from_small_compact() {
700        let mut cache = TinyUfo::new(5, 5);
701        cache.random_status = RandomState::with_seeds(2, 3, 4, 5);
702        cache.queues.estimator = TinyLfu::new_compact_seeded(5);
703
704        cache.put(1, 1, 1);
705        cache.put(2, 2, 2);
706        cache.put(3, 3, 2);
707        // cache full now
708
709        assert_eq!(cache.peek_queue(1), Some(SMALL));
710        assert_eq!(cache.peek_queue(2), Some(SMALL));
711        assert_eq!(cache.peek_queue(3), Some(SMALL));
712
713        let evicted = cache.put(4, 4, 3);
714        assert_eq!(evicted.len(), 2);
715        assert_eq!(evicted[0].data, 1);
716        assert_eq!(evicted[1].data, 2);
717
718        assert_eq!(cache.peek_queue(1), None);
719        assert_eq!(cache.peek_queue(2), None);
720        assert_eq!(cache.peek_queue(3), Some(SMALL));
721    }
722
723    #[test]
724    fn test_evict_from_small_to_main_compact() {
725        let mut cache = TinyUfo::new(5, 5);
726        cache.random_status = RandomState::with_seeds(2, 3, 4, 5);
727        cache.queues.estimator = TinyLfu::new_compact_seeded(5);
728
729        cache.put(1, 1, 1);
730        cache.put(2, 2, 2);
731        cache.put(3, 3, 2);
732        // cache full now
733
734        cache.get(&1);
735        cache.get(&1); // 1 will be moved to main during next eviction
736
737        assert_eq!(cache.peek_queue(1), Some(SMALL));
738        assert_eq!(cache.peek_queue(2), Some(SMALL));
739        assert_eq!(cache.peek_queue(3), Some(SMALL));
740
741        let evicted = cache.put(4, 4, 2);
742        assert_eq!(evicted.len(), 1);
743        assert_eq!(evicted[0].weight, 2);
744
745        assert_eq!(cache.peek_queue(1), Some(MAIN));
746        // either 2, 3, or 4 was evicted. Check evicted for which.
747        let mut remaining = vec![2, 3, 4];
748        remaining.remove(
749            remaining
750                .iter()
751                .position(|x| *x == evicted[0].data)
752                .unwrap(),
753        );
754        assert_eq!(cache.peek_queue(evicted[0].key), None);
755        for k in remaining {
756            assert_eq!(cache.peek_queue(k), Some(SMALL));
757        }
758    }
759    #[test]
760    fn test_remove() {
761        let mut cache = TinyUfo::new(5, 5);
762        cache.random_status = RandomState::with_seeds(2, 3, 4, 5);
763
764        cache.put(1, 1, 1);
765        cache.put(2, 2, 2);
766        cache.put(3, 3, 2);
767
768        assert_eq!(cache.remove(&1), Some(1));
769        assert_eq!(cache.remove(&3), Some(3));
770        assert_eq!(cache.get(&1), None);
771        assert_eq!(cache.get(&3), None);
772
773        // Verify empty keys get evicted when cache fills up
774        // Fill cache to trigger eviction
775        cache.put(5, 5, 2);
776        cache.put(6, 6, 2);
777        cache.put(7, 7, 2);
778
779        // The removed items (1, 3) should be naturally evicted now
780        // and new items should be in cache
781        assert_eq!(cache.get(&1), None);
782        assert_eq!(cache.get(&3), None);
783        assert!(cache.get(&5).is_some() || cache.get(&6).is_some() || cache.get(&7).is_some());
784
785        // Test weights after eviction cycles
786        let total_weight =
787            cache.queues.small_weight.load(SeqCst) + cache.queues.main_weight.load(SeqCst);
788        assert!(total_weight <= 5); // Should not exceed limit
789    }
790}
tinyufo/lib.rs

tinyufo/
lib.rs
