alice.pds.demo.boltless.dev
This repository has no description
18 kB
634 lines
did:plc:3fwecdnvtcscjnrx2p4n7a…
1mod expiring;
2
3pub use expiring::{Admitted, Expiring, GroupQuota, Quotas, Rejected, TotalQuota};
4
5use std::collections::{BTreeMap, HashMap};
6use std::future::Future;
7use std::hash::Hash;
8use std::sync::{Arc, Mutex, OnceLock, Weak};
9use std::time::Duration;
10
11use knot_runtime::{Clock, UnixMicros};
12
13knot_types::scalar_newtype! {
14 pub struct EntryCount(u64);
15 pub struct Weight(u64);
16}
17
18pub trait Cache<K, V>: Send + Sync {
19 fn get(&self, key: &K) -> Option<V>;
20 fn insert(&self, key: K, value: V);
21 fn invalidate(&self, key: &K);
22 fn invalidate_all(&self);
23 fn entry_count(&self) -> EntryCount;
24 fn weighted_size(&self) -> Weight;
25}
26
27pub struct Untimed;
28
29impl Clock for Untimed {
30 fn now_unix_micros(&self) -> UnixMicros {
31 UnixMicros::new(0)
32 }
33}
34
35pub struct Moka<K, V> {
36 inner: moka::sync::Cache<K, V>,
37}
38
39impl<K, V> Moka<K, V>
40where
41 K: Hash + Eq + Send + Sync + 'static,
42 V: Clone + Send + Sync + 'static,
43{
44 pub fn by_count(max_entries: EntryCount) -> Self {
45 Self {
46 inner: moka::sync::Cache::builder()
47 .max_capacity(max_entries.get())
48 .build(),
49 }
50 }
51
52 pub fn by_weight<F>(max_weight: Weight, weigh: F) -> Self
53 where
54 F: Fn(&V) -> Weight + Send + Sync + 'static,
55 {
56 Self {
57 inner: moka::sync::Cache::builder()
58 .max_capacity(max_weight.get())
59 .weigher(move |_key: &K, value: &V| weigh(value).get().min(u32::MAX as u64) as u32)
60 .build(),
61 }
62 }
63
64 pub fn get_or_try_insert_with<E, F>(&self, key: K, init: F) -> Result<V, Arc<E>>
65 where
66 F: FnOnce() -> Result<V, E>,
67 E: Send + Sync + 'static,
68 {
69 self.inner.try_get_with(key, init)
70 }
71}
72
73impl<K, V> Cache<K, V> for Moka<K, V>
74where
75 K: Hash + Eq + Send + Sync + 'static,
76 V: Clone + Send + Sync + 'static,
77{
78 fn get(&self, key: &K) -> Option<V> {
79 self.inner.get(key)
80 }
81
82 fn insert(&self, key: K, value: V) {
83 self.inner.insert(key, value);
84 }
85
86 fn invalidate(&self, key: &K) {
87 self.inner.invalidate(key);
88 }
89
90 fn invalidate_all(&self) {
91 self.inner.invalidate_all();
92 }
93
94 fn entry_count(&self) -> EntryCount {
95 EntryCount::new(self.inner.entry_count())
96 }
97
98 fn weighted_size(&self) -> Weight {
99 Weight::new(self.inner.weighted_size())
100 }
101}
102
103#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
104struct Tick(u64);
105
106impl Tick {
107 fn issue(&mut self) -> Tick {
108 let issued = *self;
109 self.0 = self.0.saturating_add(1);
110 issued
111 }
112}
113
114struct Node<V> {
115 value: V,
116 tick: Tick,
117 weight: u64,
118 expires_at: Option<UnixMicros>,
119}
120
121struct LruInner<K, V> {
122 by_key: HashMap<K, Node<V>>,
123 order: BTreeMap<Tick, K>,
124 next: Tick,
125 total_weight: u64,
126}
127
128type Weigh<V> = Arc<dyn Fn(&V) -> Weight + Send + Sync>;
129
130pub struct Lru<K, V, C: Clock = Untimed> {
131 inner: Mutex<LruInner<K, V>>,
132 max_entries: Option<EntryCount>,
133 max_weight: Option<Weight>,
134 weigh: Option<Weigh<V>>,
135 ttl: Option<Duration>,
136 clock: C,
137}
138
139impl<K, V> Lru<K, V, Untimed>
140where
141 K: Hash + Eq + Clone,
142 V: Clone,
143{
144 pub fn by_count(max_entries: EntryCount) -> Self {
145 Self::build(Some(max_entries), None, None, None, Untimed)
146 }
147
148 pub fn by_weight<F>(max_weight: Weight, weigh: F) -> Self
149 where
150 F: Fn(&V) -> Weight + Send + Sync + 'static,
151 {
152 Self::build(None, Some(max_weight), Some(Arc::new(weigh)), None, Untimed)
153 }
154}
155
156impl<K, V, C> Lru<K, V, C>
157where
158 K: Hash + Eq + Clone,
159 V: Clone,
160 C: Clock,
161{
162 pub fn by_count_with_ttl(max_entries: EntryCount, ttl: Duration, clock: C) -> Self {
163 Self::build(Some(max_entries), None, None, Some(ttl), clock)
164 }
165
166 pub fn by_weight_with_ttl<F>(max_weight: Weight, ttl: Duration, clock: C, weigh: F) -> Self
167 where
168 F: Fn(&V) -> Weight + Send + Sync + 'static,
169 {
170 Self::build(
171 None,
172 Some(max_weight),
173 Some(Arc::new(weigh)),
174 Some(ttl),
175 clock,
176 )
177 }
178
179 pub fn with_entry_cap(mut self, max_entries: EntryCount) -> Self {
180 self.max_entries = Some(max_entries);
181 self
182 }
183
184 fn build(
185 max_entries: Option<EntryCount>,
186 max_weight: Option<Weight>,
187 weigh: Option<Weigh<V>>,
188 ttl: Option<Duration>,
189 clock: C,
190 ) -> Self {
191 Self {
192 inner: Mutex::new(LruInner {
193 by_key: HashMap::new(),
194 order: BTreeMap::new(),
195 next: Tick(0),
196 total_weight: 0,
197 }),
198 max_entries,
199 max_weight,
200 weigh,
201 ttl,
202 clock,
203 }
204 }
205
206 fn weight_of(&self, value: &V) -> u64 {
207 self.weigh.as_ref().map_or(1, |weigh| weigh(value).get())
208 }
209
210 fn lock(&self) -> std::sync::MutexGuard<'_, LruInner<K, V>> {
211 self.inner
212 .lock()
213 .unwrap_or_else(|poisoned| poisoned.into_inner())
214 }
215}
216
217impl<K, V, C> Cache<K, V> for Lru<K, V, C>
218where
219 K: Hash + Eq + Clone + Send + Sync + 'static,
220 V: Clone + Send + Sync + 'static,
221 C: Clock,
222{
223 fn get(&self, key: &K) -> Option<V> {
224 let mut guard = self.lock();
225 let inner = &mut *guard;
226 let (stale, expires_at) = inner
227 .by_key
228 .get(key)
229 .map(|node| (node.tick, node.expires_at))?;
230 if expires_at.is_some_and(|at| at.get() <= self.clock.now_unix_micros().get()) {
231 if let Some(node) = inner.by_key.remove(key) {
232 inner.total_weight = inner.total_weight.saturating_sub(node.weight);
233 }
234 inner.order.remove(&stale);
235 return None;
236 }
237 let fresh = inner.next.issue();
238 inner.order.remove(&stale);
239 inner.order.insert(fresh, key.clone());
240 let node = inner.by_key.get_mut(key).expect("hit is still present");
241 node.tick = fresh;
242 Some(node.value.clone())
243 }
244
245 fn insert(&self, key: K, value: V) {
246 let expires_at = self.ttl.map(|ttl| {
247 UnixMicros::new(
248 self.clock
249 .now_unix_micros()
250 .get()
251 .saturating_add(ttl.as_micros() as u64),
252 )
253 });
254 let weight = self.weight_of(&value);
255 let mut guard = self.lock();
256 let inner = &mut *guard;
257 if let Some(previous) = inner.by_key.remove(&key) {
258 inner.order.remove(&previous.tick);
259 inner.total_weight = inner.total_weight.saturating_sub(previous.weight);
260 }
261 let fresh = inner.next.issue();
262 inner.order.insert(fresh, key.clone());
263 inner.total_weight = inner.total_weight.saturating_add(weight);
264 inner.by_key.insert(
265 key,
266 Node {
267 value,
268 tick: fresh,
269 weight,
270 expires_at,
271 },
272 );
273 while self
274 .max_entries
275 .is_some_and(|max| inner.by_key.len() as u64 > max.get())
276 || self
277 .max_weight
278 .is_some_and(|max| inner.total_weight > max.get())
279 {
280 let Some((_, evicted)) = inner.order.pop_first() else {
281 break;
282 };
283 if let Some(node) = inner.by_key.remove(&evicted) {
284 inner.total_weight = inner.total_weight.saturating_sub(node.weight);
285 }
286 }
287 }
288
289 fn invalidate(&self, key: &K) {
290 let mut guard = self.lock();
291 if let Some(node) = guard.by_key.remove(key) {
292 guard.order.remove(&node.tick);
293 guard.total_weight = guard.total_weight.saturating_sub(node.weight);
294 }
295 }
296
297 fn invalidate_all(&self) {
298 let mut guard = self.lock();
299 guard.by_key.clear();
300 guard.order.clear();
301 guard.total_weight = 0;
302 }
303
304 fn entry_count(&self) -> EntryCount {
305 EntryCount::new(self.lock().by_key.len() as u64)
306 }
307
308 fn weighted_size(&self) -> Weight {
309 Weight::new(self.lock().total_weight)
310 }
311}
312
313pub struct Noop;
314
315impl<K, V> Cache<K, V> for Noop
316where
317 K: Send + Sync + 'static,
318 V: Send + Sync + 'static,
319{
320 fn get(&self, _key: &K) -> Option<V> {
321 None
322 }
323
324 fn insert(&self, _key: K, _value: V) {}
325
326 fn invalidate(&self, _key: &K) {}
327
328 fn invalidate_all(&self) {}
329
330 fn entry_count(&self) -> EntryCount {
331 EntryCount::new(0)
332 }
333
334 fn weighted_size(&self) -> Weight {
335 Weight::new(0)
336 }
337}
338
339pub struct Filled<V> {
340 pub value: V,
341 pub fresh: bool,
342}
343
344type DeterministicHasher = std::hash::BuildHasherDefault<std::collections::hash_map::DefaultHasher>;
345
346pub trait AsyncCache<K, V>: Send + Sync {
347 fn get(&self, key: &K) -> impl Future<Output = Option<V>> + Send;
348 fn entry_count(&self) -> EntryCount;
349}
350
351pub struct MokaFuture<K, V> {
352 inner: moka::future::Cache<K, V, DeterministicHasher>,
353}
354
355impl<K, V> MokaFuture<K, V>
356where
357 K: Hash + Eq + Send + Sync + 'static,
358 V: Clone + Send + Sync + 'static,
359{
360 pub fn by_count(max_entries: EntryCount) -> Self {
361 Self {
362 inner: moka::future::Cache::builder()
363 .max_capacity(max_entries.get())
364 .build_with_hasher(DeterministicHasher::default()),
365 }
366 }
367
368 pub async fn get_or_fill_if<Fut, P>(&self, key: K, refill_if: P, fill: Fut) -> Filled<V>
369 where
370 Fut: Future<Output = V> + Send,
371 P: FnMut(&V) -> bool + Send,
372 {
373 let entry = self
374 .inner
375 .entry(key)
376 .or_insert_with_if(fill, refill_if)
377 .await;
378 Filled {
379 fresh: entry.is_fresh(),
380 value: entry.into_value(),
381 }
382 }
383
384 pub async fn run_pending_tasks(&self) {
385 self.inner.run_pending_tasks().await;
386 }
387}
388
389impl<K, V> AsyncCache<K, V> for MokaFuture<K, V>
390where
391 K: Hash + Eq + Send + Sync + 'static,
392 V: Clone + Send + Sync + 'static,
393{
394 fn get(&self, key: &K) -> impl Future<Output = Option<V>> + Send {
395 self.inner.get(key)
396 }
397
398 fn entry_count(&self) -> EntryCount {
399 EntryCount::new(self.inner.entry_count())
400 }
401}
402
403pub trait Reclaimable: Send + Sync {
404 fn footprint(&self) -> Weight;
405 fn reclaim(&self);
406}
407
408impl<K, V> Reclaimable for Moka<K, V>
409where
410 K: Hash + Eq + Send + Sync + 'static,
411 V: Clone + Send + Sync + 'static,
412{
413 fn footprint(&self) -> Weight {
414 self.weighted_size()
415 }
416
417 fn reclaim(&self) {
418 self.invalidate_all();
419 }
420}
421
422impl<K, V, C> Reclaimable for Lru<K, V, C>
423where
424 K: Hash + Eq + Clone + Send + Sync + 'static,
425 V: Clone + Send + Sync + 'static,
426 C: Clock,
427{
428 fn footprint(&self) -> Weight {
429 self.weighted_size()
430 }
431
432 fn reclaim(&self) {
433 self.invalidate_all();
434 }
435}
436
437#[derive(Default)]
438struct Registry {
439 caches: Mutex<Vec<Weak<dyn Reclaimable>>>,
440}
441
442impl Registry {
443 fn lock(&self) -> std::sync::MutexGuard<'_, Vec<Weak<dyn Reclaimable>>> {
444 self.caches
445 .lock()
446 .unwrap_or_else(|poisoned| poisoned.into_inner())
447 }
448
449 fn register<T: Reclaimable + 'static>(&self, cache: &Arc<T>) {
450 let erased: Arc<dyn Reclaimable> = cache.clone();
451 let weak = Arc::downgrade(&erased);
452 let mut caches = self.lock();
453 caches.retain(|entry| entry.strong_count() > 0);
454 caches.push(weak);
455 }
456
457 fn reclaim_largest(&self) -> Option<Weight> {
458 let largest = self
459 .lock()
460 .iter()
461 .filter_map(Weak::upgrade)
462 .max_by_key(|cache| cache.footprint().get())?;
463 let freed = largest.footprint();
464 (freed.get() > 0).then(|| {
465 largest.reclaim();
466 freed
467 })
468 }
469}
470
471fn global_registry() -> &'static Registry {
472 static REGISTRY: OnceLock<Registry> = OnceLock::new();
473 REGISTRY.get_or_init(Registry::default)
474}
475
476pub fn register<T: Reclaimable + 'static>(cache: &Arc<T>) {
477 global_registry().register(cache);
478}
479
480pub fn reclaim_largest() -> Option<Weight> {
481 global_registry().reclaim_largest()
482}
483
484#[cfg(test)]
485mod tests {
486 use std::sync::Arc;
487
488 use knot_runtime::ManualClock;
489
490 use super::*;
491
492 #[test]
493 fn a_no_ttl_cache_retains_until_invalidated() {
494 let cache: Moka<u32, u32> = Moka::by_count(EntryCount::new(16));
495 cache.insert(1, 9);
496 cache.insert(2, 10);
497 assert_eq!(cache.get(&1), Some(9));
498 cache.invalidate(&1);
499 assert_eq!(cache.get(&1), None);
500 assert_eq!(cache.get(&2), Some(10));
501 }
502
503 #[test]
504 fn the_same_clock_sequence_yields_the_same_hits_and_misses() {
505 let run = || {
506 let clock = Arc::new(ManualClock::new(UnixMicros::new(0)));
507 let cache: Lru<u32, u32, _> = Lru::by_count_with_ttl(
508 EntryCount::new(16),
509 Duration::from_secs(5),
510 Arc::clone(&clock),
511 );
512 cache.insert(1, 100);
513 let before = cache.get(&1);
514 clock.advance(Duration::from_secs(6));
515 let after = cache.get(&1);
516 (before, after)
517 };
518 assert_eq!(run(), run());
519 }
520
521 #[test]
522 fn get_or_try_insert_with_serves_the_first_value_without_recomputing() {
523 let cache: Moka<u32, u32> = Moka::by_count(EntryCount::new(16));
524 let first = cache.get_or_try_insert_with(1, || Ok::<u32, ()>(7));
525 assert_eq!(first.unwrap(), 7);
526 let second = cache.get_or_try_insert_with(1, || Ok::<u32, ()>(99));
527 assert_eq!(second.unwrap(), 7);
528 }
529
530 #[test]
531 fn get_or_try_insert_with_propagates_the_error_and_stores_nothing() {
532 let cache: Moka<u32, u32> = Moka::by_count(EntryCount::new(16));
533 let failed = cache.get_or_try_insert_with(1, || Err::<u32, u32>(9));
534 assert_eq!(*failed.unwrap_err(), 9);
535 assert_eq!(cache.get(&1), None);
536 }
537
538 #[test]
539 fn the_lru_evicts_the_least_recently_used_key() {
540 let cache: Lru<u32, u32> = Lru::by_count(EntryCount::new(3));
541 cache.insert(0, 0);
542 cache.insert(1, 1);
543 cache.insert(2, 2);
544 assert_eq!(cache.get(&0), Some(0));
545 cache.insert(3, 3);
546 assert_eq!(
547 cache.get(&1),
548 None,
549 "the least recently used key is evicted"
550 );
551 assert_eq!(cache.get(&0), Some(0), "the touched key survives");
552 assert_eq!(cache.get(&3), Some(3));
553 assert_eq!(cache.entry_count().get(), 3);
554 }
555
556 #[test]
557 fn the_weighted_lru_evicts_oldest_until_it_fits_the_byte_budget() {
558 let cache: Lru<u32, Vec<u8>> = Lru::by_weight(Weight::new(8), |value: &Vec<u8>| {
559 Weight::new(value.len() as u64)
560 });
561 cache.insert(0, vec![0u8; 5]);
562 cache.insert(1, vec![0u8; 5]);
563 assert_eq!(cache.get(&0), None, "the oldest entry is evicted to fit");
564 assert_eq!(cache.get(&1), Some(vec![0u8; 5]));
565 assert_eq!(cache.weighted_size().get(), 5);
566 }
567
568 #[test]
569 fn the_entry_cap_bounds_a_flood_of_zero_weight_entries() {
570 let cache: Lru<u32, Vec<u8>> = Lru::by_weight(Weight::new(1_000_000), |value: &Vec<u8>| {
571 Weight::new(value.len() as u64)
572 })
573 .with_entry_cap(EntryCount::new(4));
574 (0..64).for_each(|nonce| cache.insert(nonce, Vec::new()));
575 assert!(cache.entry_count().get() <= 4);
576 }
577
578 #[test]
579 fn the_lru_expires_entries_on_the_injected_clock() {
580 let clock = Arc::new(ManualClock::new(UnixMicros::new(0)));
581 let cache: Lru<u32, u32, _> = Lru::by_count_with_ttl(
582 EntryCount::new(8),
583 Duration::from_secs(1),
584 Arc::clone(&clock),
585 );
586 cache.insert(1, 5);
587 assert_eq!(cache.get(&1), Some(5));
588 clock.advance(Duration::from_secs(2));
589 assert_eq!(cache.get(&1), None);
590 assert_eq!(cache.entry_count().get(), 0);
591 }
592
593 #[test]
594 fn the_governor_sheds_the_largest_registered_cache() {
595 let weigh = |value: &Vec<u8>| Weight::new(value.len() as u64);
596 let small: Arc<Lru<u32, Vec<u8>>> = Arc::new(Lru::by_weight(Weight::new(1_000_000), weigh));
597 let big: Arc<Lru<u32, Vec<u8>>> = Arc::new(Lru::by_weight(Weight::new(1_000_000), weigh));
598 small.insert(0, vec![0u8; 10]);
599 big.insert(0, vec![0u8; 100]);
600 let registry = Registry::default();
601 registry.register(&small);
602 registry.register(&big);
603 let freed = registry
604 .reclaim_largest()
605 .expect("a registered cache is shed");
606 assert_eq!(freed.get(), 100, "the largest footprint is reclaimed");
607 assert_eq!(big.entry_count().get(), 0, "the largest cache is emptied");
608 assert_eq!(
609 small.entry_count().get(),
610 1,
611 "the smaller cache is untouched"
612 );
613 }
614
615 #[test]
616 fn a_registered_cache_with_nothing_to_reclaim_is_not_shed() {
617 let registry = Registry::default();
618 let cache: Arc<Lru<u32, Vec<u8>>> = Arc::new(Lru::by_count(EntryCount::new(8)));
619 registry.register(&cache);
620 assert_eq!(
621 registry.reclaim_largest(),
622 None,
623 "an empty registered cache reports nothing to shed"
624 );
625 }
626
627 #[test]
628 fn the_noop_cache_never_retains() {
629 let cache: &dyn Cache<u32, u32> = &Noop;
630 cache.insert(1, 2);
631 assert_eq!(cache.get(&1), None);
632 assert_eq!(cache.entry_count().get(), 0);
633 }
634}