k_cache/
cache.rs

1use std::hash::Hash;
2use std::marker::PhantomData;
3use std::sync::atomic::AtomicBool;
4use std::sync::Arc;
5use std::{
6    borrow::Borrow,
7    collections::{HashMap, VecDeque},
8    hash::BuildHasher,
9};
10
11pub trait Weigher<K, V> {
12    fn weigh(_k: &K, _v: &V) -> usize {
13        1
14    }
15}
16
17#[derive(Debug, Clone)]
18pub struct One;
19impl<K, V> Weigher<K, V> for One {}
20
21pub trait Lifecycle<K, V> {
22    fn on_eviction(&self, _key: K, _value: V) {}
23}
24
25#[derive(Debug, Clone, Copy, Default)]
26pub struct DefaultLifecycle;
27impl<K, V> Lifecycle<K, V> for DefaultLifecycle {}
28
29#[derive(Debug)]
30struct SieveEntry<D> {
31    data: D,
32    visited: Arc<AtomicBool>,
33}
34
35#[derive(Debug)]
36pub struct Cache<K, V, S, W: Weigher<K, V> = One, L: Lifecycle<K, V> = DefaultLifecycle> {
37    map: HashMap<K, SieveEntry<V>, S>,
38    sieve_pool: VecDeque<SieveEntry<K>>,
39    sieve_hand: usize,
40    max_weight: usize,
41    weight: usize,
42    lifecycle: L,
43    _phantom: PhantomData<W>,
44}
45
46impl<K, V, S, W, L> Cache<K, V, S, W, L>
47where
48    K: Eq + Hash + Clone,
49    S: BuildHasher,
50    W: Weigher<K, V>,
51    L: Lifecycle<K, V> + Default,
52{
53    pub fn new(hasher: S, max_weight: usize) -> Self {
54        Self {
55            map: HashMap::with_hasher(hasher),
56            sieve_pool: VecDeque::new(),
57            sieve_hand: 0,
58            max_weight,
59            weight: 0,
60            lifecycle: Default::default(),
61            _phantom: PhantomData,
62        }
63    }
64}
65
66impl<K, V, S, W, L> Cache<K, V, S, W, L>
67where
68    K: Eq + Hash + Clone,
69    S: BuildHasher,
70    W: Weigher<K, V>,
71    L: Lifecycle<K, V>,
72{
73    pub fn new_with_lifecycle(hasher: S, max_weight: usize, lifecycle: L) -> Self {
74        Self {
75            map: HashMap::with_hasher(hasher),
76            sieve_pool: VecDeque::new(),
77            sieve_hand: 0,
78            max_weight,
79            weight: 0,
80            lifecycle,
81            _phantom: PhantomData,
82        }
83    }
84
85    pub fn put(&mut self, key: K, value: V) {
86        let new_weight = self.make_room_for(&key, &value);
87        self.weight += new_weight;
88        let visited = Arc::new(AtomicBool::new(true));
89        if let Some(replaced) = self.map.insert(
90            key.clone(),
91            SieveEntry {
92                data: value,
93                visited: visited.clone(),
94            },
95        ) {
96            let replaced_weight = W::weigh(&key, &replaced.data);
97            match self.weight.checked_sub(replaced_weight) {
98                Some(new_weight) => self.weight = new_weight,
99                None => {
100                    log::error!("weight underflow");
101                    self.weight = 0;
102                }
103            }
104        }
105        self.sieve_pool.push_back(SieveEntry {
106            data: key.clone(),
107            visited,
108        });
109    }
110
111    pub fn get<Q>(&self, key: &Q) -> Option<&V>
112    where
113        K: Borrow<Q>,
114        Q: Hash + Eq + ?Sized,
115    {
116        match self.map.get(key) {
117            Some(entry) => {
118                entry
119                    .visited
120                    .store(true, std::sync::atomic::Ordering::Relaxed);
121                Some(&entry.data)
122            }
123            None => None,
124        }
125    }
126
127    pub fn remove(&mut self, key: &K) -> Option<V> {
128        match self.map.remove(key) {
129            Some(removed) => {
130                let removed_weight = W::weigh(key, &removed.data);
131                match self.weight.checked_sub(removed_weight) {
132                    Some(new_weight) => self.weight = new_weight,
133                    None => {
134                        log::error!("weight underflow");
135                        self.weight = 0;
136                    }
137                };
138                Some(removed.data)
139            }
140            None => {
141                // This can happen when the entry was already removed
142                log::debug!("garbage collecting sieve entry as {}", self.sieve_hand);
143                None
144            }
145        }
146    }
147
148    fn make_room_for(&mut self, key: &K, value: &V) -> usize {
149        let entry_weight = W::weigh(key, value);
150        while self.max_weight < self.weight + entry_weight {
151            let sieve_entry = &mut self.sieve_pool[self.sieve_hand];
152            let visited = sieve_entry
153                .visited
154                .swap(false, std::sync::atomic::Ordering::Relaxed);
155            if visited {
156                self.sieve_hand = (self.sieve_hand + 1) % self.sieve_pool.len();
157            } else {
158                let sieve_key_entry = self
159                    .sieve_pool
160                    .swap_remove_back(self.sieve_hand)
161                    .expect("the index must be present");
162                let removed = self.remove(&sieve_key_entry.data);
163                if let Some(removed_value) = removed {
164                    self.lifecycle
165                        .on_eviction(sieve_key_entry.data, removed_value);
166                }
167
168                if self.sieve_hand == self.sieve_pool.len() {
169                    self.sieve_hand = 0;
170                }
171            }
172        }
173        entry_weight
174    }
175}
k_cache/cache.rs

k_cache/
cache.rs
