alice.pds.demo.boltless.dev
This repository has no description
18 kB
646 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 fn by_weight<F>(max_weight: Weight, weigh: F) -> Self
369 where
370 F: Fn(&V) -> Weight + Send + Sync + 'static,
371 {
372 Self {
373 inner: moka::future::Cache::builder()
374 .max_capacity(max_weight.get())
375 .weigher(move |_key: &K, value: &V| weigh(value).get().min(u32::MAX as u64) as u32)
376 .build_with_hasher(DeterministicHasher::default()),
377 }
378 }
379
380 pub async fn get_or_fill_if<Fut, P>(&self, key: K, refill_if: P, fill: Fut) -> Filled<V>
381 where
382 Fut: Future<Output = V> + Send,
383 P: FnMut(&V) -> bool + Send,
384 {
385 let entry = self
386 .inner
387 .entry(key)
388 .or_insert_with_if(fill, refill_if)
389 .await;
390 Filled {
391 fresh: entry.is_fresh(),
392 value: entry.into_value(),
393 }
394 }
395
396 pub async fn run_pending_tasks(&self) {
397 self.inner.run_pending_tasks().await;
398 }
399}
400
401impl<K, V> AsyncCache<K, V> for MokaFuture<K, V>
402where
403 K: Hash + Eq + Send + Sync + 'static,
404 V: Clone + Send + Sync + 'static,
405{
406 fn get(&self, key: &K) -> impl Future<Output = Option<V>> + Send {
407 self.inner.get(key)
408 }
409
410 fn entry_count(&self) -> EntryCount {
411 EntryCount::new(self.inner.entry_count())
412 }
413}
414
415pub trait Reclaimable: Send + Sync {
416 fn footprint(&self) -> Weight;
417 fn reclaim(&self);
418}
419
420impl<K, V> Reclaimable for Moka<K, V>
421where
422 K: Hash + Eq + Send + Sync + 'static,
423 V: Clone + Send + Sync + 'static,
424{
425 fn footprint(&self) -> Weight {
426 self.weighted_size()
427 }
428
429 fn reclaim(&self) {
430 self.invalidate_all();
431 }
432}
433
434impl<K, V, C> Reclaimable for Lru<K, V, C>
435where
436 K: Hash + Eq + Clone + Send + Sync + 'static,
437 V: Clone + Send + Sync + 'static,
438 C: Clock,
439{
440 fn footprint(&self) -> Weight {
441 self.weighted_size()
442 }
443
444 fn reclaim(&self) {
445 self.invalidate_all();
446 }
447}
448
449#[derive(Default)]
450struct Registry {
451 caches: Mutex<Vec<Weak<dyn Reclaimable>>>,
452}
453
454impl Registry {
455 fn lock(&self) -> std::sync::MutexGuard<'_, Vec<Weak<dyn Reclaimable>>> {
456 self.caches
457 .lock()
458 .unwrap_or_else(|poisoned| poisoned.into_inner())
459 }
460
461 fn register<T: Reclaimable + 'static>(&self, cache: &Arc<T>) {
462 let erased: Arc<dyn Reclaimable> = cache.clone();
463 let weak = Arc::downgrade(&erased);
464 let mut caches = self.lock();
465 caches.retain(|entry| entry.strong_count() > 0);
466 caches.push(weak);
467 }
468
469 fn reclaim_largest(&self) -> Option<Weight> {
470 let largest = self
471 .lock()
472 .iter()
473 .filter_map(Weak::upgrade)
474 .max_by_key(|cache| cache.footprint().get())?;
475 let freed = largest.footprint();
476 (freed.get() > 0).then(|| {
477 largest.reclaim();
478 freed
479 })
480 }
481}
482
483fn global_registry() -> &'static Registry {
484 static REGISTRY: OnceLock<Registry> = OnceLock::new();
485 REGISTRY.get_or_init(Registry::default)
486}
487
488pub fn register<T: Reclaimable + 'static>(cache: &Arc<T>) {
489 global_registry().register(cache);
490}
491
492pub fn reclaim_largest() -> Option<Weight> {
493 global_registry().reclaim_largest()
494}
495
496#[cfg(test)]
497mod tests {
498 use std::sync::Arc;
499
500 use knot_runtime::ManualClock;
501
502 use super::*;
503
504 #[test]
505 fn a_no_ttl_cache_retains_until_invalidated() {
506 let cache: Moka<u32, u32> = Moka::by_count(EntryCount::new(16));
507 cache.insert(1, 9);
508 cache.insert(2, 10);
509 assert_eq!(cache.get(&1), Some(9));
510 cache.invalidate(&1);
511 assert_eq!(cache.get(&1), None);
512 assert_eq!(cache.get(&2), Some(10));
513 }
514
515 #[test]
516 fn the_same_clock_sequence_yields_the_same_hits_and_misses() {
517 let run = || {
518 let clock = Arc::new(ManualClock::new(UnixMicros::new(0)));
519 let cache: Lru<u32, u32, _> = Lru::by_count_with_ttl(
520 EntryCount::new(16),
521 Duration::from_secs(5),
522 Arc::clone(&clock),
523 );
524 cache.insert(1, 100);
525 let before = cache.get(&1);
526 clock.advance(Duration::from_secs(6));
527 let after = cache.get(&1);
528 (before, after)
529 };
530 assert_eq!(run(), run());
531 }
532
533 #[test]
534 fn get_or_try_insert_with_serves_the_first_value_without_recomputing() {
535 let cache: Moka<u32, u32> = Moka::by_count(EntryCount::new(16));
536 let first = cache.get_or_try_insert_with(1, || Ok::<u32, ()>(7));
537 assert_eq!(first.unwrap(), 7);
538 let second = cache.get_or_try_insert_with(1, || Ok::<u32, ()>(99));
539 assert_eq!(second.unwrap(), 7);
540 }
541
542 #[test]
543 fn get_or_try_insert_with_propagates_the_error_and_stores_nothing() {
544 let cache: Moka<u32, u32> = Moka::by_count(EntryCount::new(16));
545 let failed = cache.get_or_try_insert_with(1, || Err::<u32, u32>(9));
546 assert_eq!(*failed.unwrap_err(), 9);
547 assert_eq!(cache.get(&1), None);
548 }
549
550 #[test]
551 fn the_lru_evicts_the_least_recently_used_key() {
552 let cache: Lru<u32, u32> = Lru::by_count(EntryCount::new(3));
553 cache.insert(0, 0);
554 cache.insert(1, 1);
555 cache.insert(2, 2);
556 assert_eq!(cache.get(&0), Some(0));
557 cache.insert(3, 3);
558 assert_eq!(
559 cache.get(&1),
560 None,
561 "the least recently used key is evicted"
562 );
563 assert_eq!(cache.get(&0), Some(0), "the touched key survives");
564 assert_eq!(cache.get(&3), Some(3));
565 assert_eq!(cache.entry_count().get(), 3);
566 }
567
568 #[test]
569 fn the_weighted_lru_evicts_oldest_until_it_fits_the_byte_budget() {
570 let cache: Lru<u32, Vec<u8>> = Lru::by_weight(Weight::new(8), |value: &Vec<u8>| {
571 Weight::new(value.len() as u64)
572 });
573 cache.insert(0, vec![0u8; 5]);
574 cache.insert(1, vec![0u8; 5]);
575 assert_eq!(cache.get(&0), None, "the oldest entry is evicted to fit");
576 assert_eq!(cache.get(&1), Some(vec![0u8; 5]));
577 assert_eq!(cache.weighted_size().get(), 5);
578 }
579
580 #[test]
581 fn the_entry_cap_bounds_a_flood_of_zero_weight_entries() {
582 let cache: Lru<u32, Vec<u8>> = Lru::by_weight(Weight::new(1_000_000), |value: &Vec<u8>| {
583 Weight::new(value.len() as u64)
584 })
585 .with_entry_cap(EntryCount::new(4));
586 (0..64).for_each(|nonce| cache.insert(nonce, Vec::new()));
587 assert!(cache.entry_count().get() <= 4);
588 }
589
590 #[test]
591 fn the_lru_expires_entries_on_the_injected_clock() {
592 let clock = Arc::new(ManualClock::new(UnixMicros::new(0)));
593 let cache: Lru<u32, u32, _> = Lru::by_count_with_ttl(
594 EntryCount::new(8),
595 Duration::from_secs(1),
596 Arc::clone(&clock),
597 );
598 cache.insert(1, 5);
599 assert_eq!(cache.get(&1), Some(5));
600 clock.advance(Duration::from_secs(2));
601 assert_eq!(cache.get(&1), None);
602 assert_eq!(cache.entry_count().get(), 0);
603 }
604
605 #[test]
606 fn the_governor_sheds_the_largest_registered_cache() {
607 let weigh = |value: &Vec<u8>| Weight::new(value.len() as u64);
608 let small: Arc<Lru<u32, Vec<u8>>> = Arc::new(Lru::by_weight(Weight::new(1_000_000), weigh));
609 let big: Arc<Lru<u32, Vec<u8>>> = Arc::new(Lru::by_weight(Weight::new(1_000_000), weigh));
610 small.insert(0, vec![0u8; 10]);
611 big.insert(0, vec![0u8; 100]);
612 let registry = Registry::default();
613 registry.register(&small);
614 registry.register(&big);
615 let freed = registry
616 .reclaim_largest()
617 .expect("a registered cache is shed");
618 assert_eq!(freed.get(), 100, "the largest footprint is reclaimed");
619 assert_eq!(big.entry_count().get(), 0, "the largest cache is emptied");
620 assert_eq!(
621 small.entry_count().get(),
622 1,
623 "the smaller cache is untouched"
624 );
625 }
626
627 #[test]
628 fn a_registered_cache_with_nothing_to_reclaim_is_not_shed() {
629 let registry = Registry::default();
630 let cache: Arc<Lru<u32, Vec<u8>>> = Arc::new(Lru::by_count(EntryCount::new(8)));
631 registry.register(&cache);
632 assert_eq!(
633 registry.reclaim_largest(),
634 None,
635 "an empty registered cache reports nothing to shed"
636 );
637 }
638
639 #[test]
640 fn the_noop_cache_never_retains() {
641 let cache: &dyn Cache<u32, u32> = &Noop;
642 cache.insert(1, 2);
643 assert_eq!(cache.get(&1), None);
644 assert_eq!(cache.entry_count().get(), 0);
645 }
646}