Skip to main content

wp_knowledge/
cache.rs

1use std::collections::HashMap;
2use std::net::IpAddr;
3use std::num::NonZeroUsize;
4
5use lru::LruCache;
6use wp_model_core::model::{DataField, FValueStr, Value};
7
8/// 本地缓存去重索引表上限:触顶后整体重置,防止运行期无界增长。
9const MAX_LOCAL_IDX: usize = 100_000;
10
11#[derive(Debug, Clone)]
12pub struct FieldQueryCache {
13    str_idx: HashMap<FValueStr, usize>,
14    i64_idx: HashMap<i64, usize>,
15    ip_idx: HashMap<IpAddr, usize>,
16    cache_data: LruCache<LocalCacheKey, Vec<DataField>>,
17    idx_num: usize,
18    /// `scope -> generation`:本地缓存按 scope(provider 粒度)跟踪代际,
19    /// 避免多 provider 切换时整缓存重置。
20    generations: HashMap<u64, u64>,
21}
22
23pub type QueryLocalCache = FieldQueryCache;
24
25impl Default for FieldQueryCache {
26    fn default() -> Self {
27        Self::with_capacity(100)
28    }
29}
30
31impl FieldQueryCache {
32    pub fn with_capacity(size: usize) -> Self {
33        let size = size.max(1);
34        Self {
35            str_idx: HashMap::new(),
36            i64_idx: HashMap::new(),
37            ip_idx: HashMap::new(),
38            cache_data: LruCache::new(NonZeroUsize::new(size).expect("non-zero cache size")),
39            idx_num: 0,
40            generations: HashMap::new(),
41        }
42    }
43
44    fn get_idx(&self, param: &DataField) -> Option<usize> {
45        match param.get_value() {
46            Value::Chars(v) => self.str_idx.get(v).copied(),
47            Value::Digit(v) => self.i64_idx.get(v).copied(),
48            Value::IpAddr(v) => self.ip_idx.get(v).copied(),
49            _ => None,
50        }
51    }
52
53    fn try_up_idx(&mut self, param: &DataField) -> Option<usize> {
54        if self.idx_num >= MAX_LOCAL_IDX {
55            // 索引表触顶:整体重置,防止无界增长(代价是清空一次本地缓存)
56            self.str_idx.clear();
57            self.i64_idx.clear();
58            self.ip_idx.clear();
59            self.cache_data.clear();
60            self.generations.clear();
61            self.idx_num = 0;
62        }
63        match param.get_value() {
64            Value::Chars(v) => {
65                if let Some(idx) = self.str_idx.get(v) {
66                    Some(*idx)
67                } else {
68                    self.idx_num += 1;
69                    self.str_idx.insert(v.clone(), self.idx_num);
70                    Some(self.idx_num)
71                }
72            }
73            Value::Digit(v) => {
74                if let Some(idx) = self.i64_idx.get(v) {
75                    Some(*idx)
76                } else {
77                    self.idx_num += 1;
78                    self.i64_idx.insert(*v, self.idx_num);
79                    Some(self.idx_num)
80                }
81            }
82            Value::IpAddr(v) => {
83                if let Some(idx) = self.ip_idx.get(v) {
84                    Some(*idx)
85                } else {
86                    self.idx_num += 1;
87                    self.ip_idx.insert(*v, self.idx_num);
88                    Some(self.idx_num)
89                }
90            }
91            _ => None,
92        }
93    }
94
95    /// 仅当 scope 的 generation 变化时,清理该 scope 的缓存条目;
96    /// 其它 provider 的本地缓存保留。
97    fn prepare_generation_scoped(&mut self, scope: u64, generation: u64) {
98        if self.generations.get(&scope) == Some(&generation) {
99            return;
100        }
101        self.generations.insert(scope, generation);
102        let stale: Vec<LocalCacheKey> = self
103            .cache_data
104            .iter()
105            .filter(|(key, _)| key.scope_hash == scope)
106            .map(|(key, _)| key.clone())
107            .collect();
108        for key in stale {
109            self.cache_data.pop(&key);
110        }
111    }
112}
113
114#[derive(PartialEq, Eq, Hash, Debug, Clone)]
115pub enum EnumSizeIndex {
116    Idx1(usize),
117    Idx2(usize, usize),
118    Idx3(usize, usize, usize),
119    Idx4(usize, usize, usize, usize),
120    Idx5(usize, usize, usize, usize, usize),
121    Idx6(usize, usize, usize, usize, usize, usize),
122}
123
124#[derive(PartialEq, Eq, Hash, Debug, Clone)]
125struct LocalCacheKey {
126    scope_hash: u64,
127    idxs: EnumSizeIndex,
128}
129
130pub trait CacheAble<P, T, const N: usize> {
131    /// 通知缓存某 scope 进入新 generation(通常因 provider 重载)。实现可只清理该 scope。
132    fn prepare_generation(&mut self, _scope: u64, _generation: u64) {}
133    fn save_scoped(&mut self, _scope_hash: u64, params: &[P; N], result: T) {
134        self.save(params, result);
135    }
136    fn fetch_scoped(&self, _scope_hash: u64, params: &[P; N]) -> Option<&T> {
137        self.fetch(params)
138    }
139    fn save(&mut self, params: &[P; N], result: T);
140    fn fetch(&self, params: &[P; N]) -> Option<&T>;
141}
142
143impl CacheAble<DataField, Vec<DataField>, 1> for FieldQueryCache {
144    fn prepare_generation(&mut self, scope: u64, generation: u64) {
145        self.prepare_generation_scoped(scope, generation);
146    }
147    fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 1], result: Vec<DataField>) {
148        if let Some(i0) = self.try_up_idx(&params[0]) {
149            self.cache_data.put(
150                LocalCacheKey {
151                    scope_hash,
152                    idxs: EnumSizeIndex::Idx1(i0),
153                },
154                result,
155            );
156        }
157    }
158
159    fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 1]) -> Option<&Vec<DataField>> {
160        if let Some(i0) = self.get_idx(&params[0]) {
161            return self.cache_data.peek(&LocalCacheKey {
162                scope_hash,
163                idxs: EnumSizeIndex::Idx1(i0),
164            });
165        }
166        None
167    }
168
169    fn save(&mut self, params: &[DataField; 1], result: Vec<DataField>) {
170        self.save_scoped(0, params, result);
171    }
172
173    fn fetch(&self, params: &[DataField; 1]) -> Option<&Vec<DataField>> {
174        self.fetch_scoped(0, params)
175    }
176}
177
178impl CacheAble<DataField, Vec<DataField>, 2> for FieldQueryCache {
179    fn prepare_generation(&mut self, scope: u64, generation: u64) {
180        self.prepare_generation_scoped(scope, generation);
181    }
182
183    fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 2], result: Vec<DataField>) {
184        if let (Some(i0), Some(i1)) = (self.try_up_idx(&params[0]), self.try_up_idx(&params[1])) {
185            self.cache_data.put(
186                LocalCacheKey {
187                    scope_hash,
188                    idxs: EnumSizeIndex::Idx2(i0, i1),
189                },
190                result,
191            );
192        }
193    }
194
195    fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 2]) -> Option<&Vec<DataField>> {
196        if let (Some(i0), Some(i1)) = (self.get_idx(&params[0]), self.get_idx(&params[1])) {
197            return self.cache_data.peek(&LocalCacheKey {
198                scope_hash,
199                idxs: EnumSizeIndex::Idx2(i0, i1),
200            });
201        }
202        None
203    }
204
205    fn save(&mut self, params: &[DataField; 2], result: Vec<DataField>) {
206        self.save_scoped(0, params, result);
207    }
208
209    fn fetch(&self, params: &[DataField; 2]) -> Option<&Vec<DataField>> {
210        self.fetch_scoped(0, params)
211    }
212}
213
214impl CacheAble<DataField, Vec<DataField>, 3> for FieldQueryCache {
215    fn prepare_generation(&mut self, scope: u64, generation: u64) {
216        self.prepare_generation_scoped(scope, generation);
217    }
218
219    fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 3], result: Vec<DataField>) {
220        if let (Some(i0), Some(i1), Some(i2)) = (
221            self.try_up_idx(&params[0]),
222            self.try_up_idx(&params[1]),
223            self.try_up_idx(&params[2]),
224        ) {
225            self.cache_data.put(
226                LocalCacheKey {
227                    scope_hash,
228                    idxs: EnumSizeIndex::Idx3(i0, i1, i2),
229                },
230                result,
231            );
232        }
233    }
234
235    fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 3]) -> Option<&Vec<DataField>> {
236        if let (Some(i0), Some(i1), Some(i2)) = (
237            self.get_idx(&params[0]),
238            self.get_idx(&params[1]),
239            self.get_idx(&params[2]),
240        ) {
241            return self.cache_data.peek(&LocalCacheKey {
242                scope_hash,
243                idxs: EnumSizeIndex::Idx3(i0, i1, i2),
244            });
245        }
246        None
247    }
248
249    fn save(&mut self, params: &[DataField; 3], result: Vec<DataField>) {
250        self.save_scoped(0, params, result);
251    }
252
253    fn fetch(&self, params: &[DataField; 3]) -> Option<&Vec<DataField>> {
254        self.fetch_scoped(0, params)
255    }
256}
257
258impl CacheAble<DataField, Vec<DataField>, 4> for FieldQueryCache {
259    fn prepare_generation(&mut self, scope: u64, generation: u64) {
260        self.prepare_generation_scoped(scope, generation);
261    }
262
263    fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 4], result: Vec<DataField>) {
264        if let (Some(i0), Some(i1), Some(i2), Some(i3)) = (
265            self.try_up_idx(&params[0]),
266            self.try_up_idx(&params[1]),
267            self.try_up_idx(&params[2]),
268            self.try_up_idx(&params[3]),
269        ) {
270            self.cache_data.put(
271                LocalCacheKey {
272                    scope_hash,
273                    idxs: EnumSizeIndex::Idx4(i0, i1, i2, i3),
274                },
275                result,
276            );
277        }
278    }
279
280    fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 4]) -> Option<&Vec<DataField>> {
281        if let (Some(i0), Some(i1), Some(i2), Some(i3)) = (
282            self.get_idx(&params[0]),
283            self.get_idx(&params[1]),
284            self.get_idx(&params[2]),
285            self.get_idx(&params[3]),
286        ) {
287            return self.cache_data.peek(&LocalCacheKey {
288                scope_hash,
289                idxs: EnumSizeIndex::Idx4(i0, i1, i2, i3),
290            });
291        }
292        None
293    }
294
295    fn save(&mut self, params: &[DataField; 4], result: Vec<DataField>) {
296        self.save_scoped(0, params, result);
297    }
298
299    fn fetch(&self, params: &[DataField; 4]) -> Option<&Vec<DataField>> {
300        self.fetch_scoped(0, params)
301    }
302}
303
304impl CacheAble<DataField, Vec<DataField>, 5> for FieldQueryCache {
305    fn prepare_generation(&mut self, scope: u64, generation: u64) {
306        self.prepare_generation_scoped(scope, generation);
307    }
308
309    fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 5], result: Vec<DataField>) {
310        if let (Some(i0), Some(i1), Some(i2), Some(i3), Some(i4)) = (
311            self.try_up_idx(&params[0]),
312            self.try_up_idx(&params[1]),
313            self.try_up_idx(&params[2]),
314            self.try_up_idx(&params[3]),
315            self.try_up_idx(&params[4]),
316        ) {
317            self.cache_data.put(
318                LocalCacheKey {
319                    scope_hash,
320                    idxs: EnumSizeIndex::Idx5(i0, i1, i2, i3, i4),
321                },
322                result,
323            );
324        }
325    }
326
327    fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 5]) -> Option<&Vec<DataField>> {
328        if let (Some(i0), Some(i1), Some(i2), Some(i3), Some(i4)) = (
329            self.get_idx(&params[0]),
330            self.get_idx(&params[1]),
331            self.get_idx(&params[2]),
332            self.get_idx(&params[3]),
333            self.get_idx(&params[4]),
334        ) {
335            return self.cache_data.peek(&LocalCacheKey {
336                scope_hash,
337                idxs: EnumSizeIndex::Idx5(i0, i1, i2, i3, i4),
338            });
339        }
340        None
341    }
342
343    fn save(&mut self, params: &[DataField; 5], result: Vec<DataField>) {
344        self.save_scoped(0, params, result);
345    }
346
347    fn fetch(&self, params: &[DataField; 5]) -> Option<&Vec<DataField>> {
348        self.fetch_scoped(0, params)
349    }
350}
351
352impl CacheAble<DataField, Vec<DataField>, 6> for FieldQueryCache {
353    fn prepare_generation(&mut self, scope: u64, generation: u64) {
354        self.prepare_generation_scoped(scope, generation);
355    }
356
357    fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 6], result: Vec<DataField>) {
358        if let (Some(i0), Some(i1), Some(i2), Some(i3), Some(i4), Some(i5)) = (
359            self.try_up_idx(&params[0]),
360            self.try_up_idx(&params[1]),
361            self.try_up_idx(&params[2]),
362            self.try_up_idx(&params[3]),
363            self.try_up_idx(&params[4]),
364            self.try_up_idx(&params[5]),
365        ) {
366            self.cache_data.put(
367                LocalCacheKey {
368                    scope_hash,
369                    idxs: EnumSizeIndex::Idx6(i0, i1, i2, i3, i4, i5),
370                },
371                result,
372            );
373        }
374    }
375
376    fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 6]) -> Option<&Vec<DataField>> {
377        if let (Some(i0), Some(i1), Some(i2), Some(i3), Some(i4), Some(i5)) = (
378            self.get_idx(&params[0]),
379            self.get_idx(&params[1]),
380            self.get_idx(&params[2]),
381            self.get_idx(&params[3]),
382            self.get_idx(&params[4]),
383            self.get_idx(&params[5]),
384        ) {
385            return self.cache_data.peek(&LocalCacheKey {
386                scope_hash,
387                idxs: EnumSizeIndex::Idx6(i0, i1, i2, i3, i4, i5),
388            });
389        }
390        None
391    }
392
393    fn save(&mut self, params: &[DataField; 6], result: Vec<DataField>) {
394        self.save_scoped(0, params, result);
395    }
396
397    fn fetch(&self, params: &[DataField; 6]) -> Option<&Vec<DataField>> {
398        self.fetch_scoped(0, params)
399    }
400}