saluki_core/cache/
mod.rs

1//! A concurrent cache with optional expiration and telemetry.
2
3use std::{marker::PhantomData, num::NonZeroUsize, sync::Arc, time::Duration};
4
5use saluki_common::{hash::FastBuildHasher, task::spawn_traced};
6use saluki_error::GenericError;
7use saluki_metrics::{static_metrics, Counter, Gauge, Histogram};
8use tokio::time::sleep;
9use tokio_util::sync::{CancellationToken, DropGuard};
10use tracing::debug;
11
12mod expiry;
13use self::expiry::{Expiration, ExpirationBuilder, ExpiryCapableLifecycle};
14
15pub mod weight;
16use self::weight::{ItemCountWeighter, Weighter, WrappedWeighter};
17
18type RawCache<K, V, W, H> = quick_cache::sync::Cache<K, V, WrappedWeighter<W>, H, ExpiryCapableLifecycle<K>>;
19
20#[static_metrics(prefix = cache, labels(cache_id))]
21#[derive(Clone)]
22struct Telemetry {
23    current_items: Gauge,
24    current_weight: Gauge,
25    weight_limit: Gauge,
26    hits_total: Counter,
27    misses_total: Counter,
28    items_evicted_total: Counter,
29    #[metric(level = debug)]
30    items_inserted_total: Counter,
31    #[metric(level = debug)]
32    items_removed_total: Counter,
33    #[metric(level = debug)]
34    items_expired_total: Counter,
35    #[metric(level = trace)]
36    items_expired_batch_size: Histogram,
37}
38
39struct InnerCache<K, V, W, H> {
40    cache: Arc<RawCache<K, V, W, H>>,
41    _task_shutdown_guard: DropGuard,
42}
43
44/// Builder for creating a [`Cache`].
45pub struct CacheBuilder<K, V, W = ItemCountWeighter, H = FastBuildHasher> {
46    identifier: String,
47    capacity: NonZeroUsize,
48    weighter: W,
49    idle_period: Option<Duration>,
50    expiration_interval: Option<Duration>,
51    telemetry_enabled: bool,
52    _key: PhantomData<K>,
53    _value: PhantomData<V>,
54    _hasher: PhantomData<H>,
55}
56
57impl<K, V> CacheBuilder<K, V> {
58    /// Creates a new `CacheBuilder` with the given cache identifier.
59    ///
60    /// The cache identifier _should_ be unique, but it isn't required to be. Metrics for the cache will be emitted
61    /// using the given identifier, so in cases where the identifier isn't unique, those metrics will be aggregated
62    /// together and it won't be possible to distinguish between the different caches.
63    ///
64    /// # Errors
65    ///
66    /// If the given cache identifier is empty, an error is returned.
67    pub fn from_identifier<N: Into<String>>(identifier: N) -> Result<CacheBuilder<K, V>, GenericError> {
68        let identifier = identifier.into();
69        if identifier.is_empty() {
70            return Err(GenericError::msg("cache identifier must not be empty"));
71        }
72
73        Ok(CacheBuilder {
74            identifier,
75            capacity: NonZeroUsize::MAX,
76            weighter: ItemCountWeighter,
77            idle_period: None,
78            expiration_interval: None,
79            telemetry_enabled: true,
80            _key: PhantomData,
81            _value: PhantomData,
82            _hasher: PhantomData,
83        })
84    }
85
86    /// Configures a [`CacheBuilder`] that's suitable for tests.
87    ///
88    /// This configures the builder with the following defaults:
89    ///
90    /// - cache identifier of "noop"
91    /// - unlimited cache size
92    /// - telemetry disabled
93    ///
94    /// This is generally only useful for testing purposes, and is exposed publicly in order to be used in cross-crate
95    /// testing scenarios.
96    pub fn for_tests() -> CacheBuilder<K, V> {
97        CacheBuilder::from_identifier("noop")
98            .expect("cache identifier not empty")
99            .with_telemetry(false)
100    }
101}
102
103impl<K, V, W, H> CacheBuilder<K, V, W, H> {
104    /// Sets the capacity for the cache.
105    ///
106    /// The capacity is used, in conjunction with the item weighter, to determine how many items can be held in the
107    /// cache and when items should be evicted to make room for new items.
108    ///
109    /// See [`with_item_weighter`][Self::with_item_weighter] for more information on how the item weighter is used.
110    ///
111    /// Defaults to unlimited capacity.
112    pub fn with_capacity(mut self, capacity: NonZeroUsize) -> Self {
113        self.capacity = capacity;
114        self
115    }
116
117    /// Enables expiration of cached items based on how long since they were last accessed.
118    ///
119    /// Items which haven't been accessed within the configured duration will be marked for expiration, and be removed
120    /// from the cache shortly thereafter. For the purposes of expiration, "accessed" is either when the item was first
121    /// inserted or when it was last read.
122    ///
123    /// If the given value is `None`, expiration is disabled.
124    ///
125    /// Defaults to no expiration.
126    pub fn with_time_to_idle(mut self, idle_period: Option<Duration>) -> Self {
127        self.idle_period = idle_period;
128
129        // Make sure we have an expiration interval set if expiration is enabled.
130        if self.idle_period.is_some() {
131            self.expiration_interval = self.expiration_interval.or(Some(Duration::from_secs(1)));
132        }
133
134        self
135    }
136
137    /// Sets the interval at which the expiration process will run.
138    ///
139    /// This controls how often the expiration process will run to check for expired items. While items become
140    /// _eligible_ for expiration after the configured duration, they're not _guaranteed_ to be
141    /// removed immediately: the expiration process must still run to actually find the expired items and remove them.
142    ///
143    /// This means that the rough upper bound for how long an item may be kept alive is the sum of
144    /// both the configured expiration duration and the expiration interval.
145    ///
146    /// This value is only relevant if expiration is enabled.
147    ///
148    /// Defaults to 1 second.
149    pub fn with_expiration_interval(mut self, expiration_interval: Duration) -> Self {
150        self.expiration_interval = Some(expiration_interval);
151        self
152    }
153
154    /// Sets the item weighter for the cache.
155    ///
156    /// The item weighter is used to determine the "weight" of each item in the cache, which is used during
157    /// insertion/eviction to determine if an item can be held in the cache without first having to evict other items to
158    /// stay within the configured capacity.
159    ///
160    /// For example, if the configured capacity is set to 10,000, and the "item count" weighter is used, then the cache
161    /// will operate in a way that aims to simply ensure that no more than 10,000 items are held in the cache at any given
162    /// time. This allows defining custom weighters that can be used to track other aspects of the items in the cache,
163    /// such as their size in bytes, or some other metric that's relevant to the intended caching behavior.
164    ///
165    /// Defaults to "item count" weighter.
166    pub fn with_item_weighter<W2>(self, weighter: W2) -> CacheBuilder<K, V, W2, H> {
167        CacheBuilder {
168            identifier: self.identifier,
169            capacity: self.capacity,
170            weighter,
171            idle_period: self.idle_period,
172            expiration_interval: self.expiration_interval,
173            telemetry_enabled: self.telemetry_enabled,
174            _key: PhantomData,
175            _value: PhantomData,
176            _hasher: PhantomData,
177        }
178    }
179
180    /// Sets the item hasher for the cache.
181    ///
182    /// As cache keys are hashed before performing any reads or writes, the chosen hasher can potentially impact the
183    /// performance of those operations. In some scenarios, it may be desirable to use a different hasher than the
184    /// default one in order to optimize for specific key types or access patterns.
185    ///
186    /// Defaults to a fast, non-cryptographic hasher: [`FastBuildHasher`].
187    pub fn with_hasher<H2>(self) -> CacheBuilder<K, V, W, H2> {
188        CacheBuilder {
189            identifier: self.identifier,
190            capacity: self.capacity,
191            weighter: self.weighter,
192            idle_period: self.idle_period,
193            expiration_interval: self.expiration_interval,
194            telemetry_enabled: self.telemetry_enabled,
195            _key: PhantomData,
196            _value: PhantomData,
197            _hasher: PhantomData,
198        }
199    }
200
201    /// Sets whether or not to enable telemetry for this cache.
202    ///
203    /// Reporting the telemetry of the cache requires running an asynchronous task to override adding additional
204    /// overhead in the hot path of reading or writing to the cache. In some cases, it may be cumbersome to always
205    /// create the cache in an asynchronous context so that the telemetry task can be spawned. This method allows
206    /// disabling telemetry reporting in those cases.
207    ///
208    /// Defaults to telemetry enabled.
209    pub fn with_telemetry(mut self, enabled: bool) -> Self {
210        self.telemetry_enabled = enabled;
211        self
212    }
213}
214
215impl<K, V, W, H> CacheBuilder<K, V, W, H>
216where
217    K: Eq + std::hash::Hash + Clone + Send + Sync + 'static,
218    V: Clone + Send + Sync + 'static,
219    W: Weighter<K, V> + Clone + Send + Sync + 'static,
220    H: std::hash::BuildHasher + Clone + Default + Send + Sync + 'static,
221{
222    /// Builds a [`Cache`] from the current configuration.
223    pub fn build(self) -> Cache<K, V, W, H> {
224        let capacity = self.capacity.get();
225
226        let telemetry = Telemetry::new(self.identifier);
227        telemetry.weight_limit().set(capacity as f64);
228
229        // Configure expiration if enabled.
230        let eviction_counter = if self.telemetry_enabled {
231            telemetry.items_evicted_total().clone()
232        } else {
233            Counter::noop()
234        };
235        let mut expiration_builder = ExpirationBuilder::new(eviction_counter);
236        if let Some(time_to_idle) = self.idle_period {
237            expiration_builder = expiration_builder.with_time_to_idle(time_to_idle);
238        }
239        let (expiration, expiry_lifecycle) = expiration_builder.build();
240
241        // Create the underlying cache and shutdown signal.
242        let shutdown_token = CancellationToken::new();
243        let raw_cache = Arc::new(RawCache::with(
244            capacity,
245            capacity as u64,
246            WrappedWeighter::from(self.weighter),
247            H::default(),
248            expiry_lifecycle,
249        ));
250
251        let cache = Cache {
252            inner: Arc::new(InnerCache {
253                cache: Arc::clone(&raw_cache),
254                _task_shutdown_guard: shutdown_token.clone().drop_guard(),
255            }),
256            expiration: expiration.clone(),
257            telemetry: telemetry.clone(),
258        };
259
260        // If expiration is enabled, spawn a background task to actually drive expiration.
261        if let Some(expiration_interval) = self.expiration_interval {
262            let expiration = expiration.clone();
263
264            spawn_traced(drive_expiration(
265                Arc::clone(&raw_cache),
266                telemetry.clone(),
267                expiration,
268                expiration_interval,
269                shutdown_token.clone(),
270            ));
271        }
272
273        // If telemetry is enabled, spawn a background task to drive telemetry reporting.
274        if self.telemetry_enabled {
275            spawn_traced(drive_telemetry(Arc::clone(&raw_cache), telemetry, shutdown_token));
276        }
277
278        cache
279    }
280}
281
282/// A simple concurrent cache.
283#[derive(Clone)]
284pub struct Cache<K, V, W = ItemCountWeighter, H = FastBuildHasher> {
285    inner: Arc<InnerCache<K, V, W, H>>,
286    expiration: Expiration<K>,
287    telemetry: Telemetry,
288}
289
290impl<K, V, W, H> Cache<K, V, W, H>
291where
292    K: Eq + std::hash::Hash + Clone,
293    V: Clone,
294    W: Weighter<K, V> + Clone,
295    H: std::hash::BuildHasher + Clone,
296{
297    /// Returns `true` if the cache is empty.
298    pub fn is_empty(&self) -> bool {
299        self.inner.cache.is_empty()
300    }
301
302    /// Returns the number of items currently in the cache.
303    pub fn len(&self) -> usize {
304        self.inner.cache.len()
305    }
306
307    /// Returns the total weight of all items in the cache.
308    pub fn weight(&self) -> u64 {
309        self.inner.cache.weight()
310    }
311
312    /// Inserts an item into the cache with the given key and value.
313    ///
314    /// If an item with the same key already exists, it will be overwritten and the old value will be dropped. If the
315    /// cache is full, one or more items will be evicted to make room for the new item, based on the configured item
316    /// weighter and the weight of the new item.
317    pub fn insert(&self, key: K, value: V) {
318        self.inner.cache.insert(key.clone(), value);
319        self.expiration.mark_entry_accessed(key);
320        self.telemetry.items_inserted_total().increment(1);
321    }
322
323    /// Gets an item from the cache by its key.
324    ///
325    /// If the item is found, it's cloned and `Some(value)` is returned. Otherwise, `None` is returned.
326    pub fn get(&self, key: &K) -> Option<V> {
327        let value = self.inner.cache.get(key);
328        if value.is_some() {
329            self.expiration.mark_entry_accessed(key.clone());
330            self.telemetry.hits_total().increment(1);
331        } else {
332            self.telemetry.misses_total().increment(1);
333        }
334        value
335    }
336
337    /// Removes an item from the cache by its key.
338    pub fn remove(&self, key: &K) {
339        self.inner.cache.remove(key);
340        self.expiration.mark_entry_removed(key.clone());
341        self.telemetry.items_removed_total().increment(1);
342    }
343}
344
345async fn drive_expiration<K, V, W, H>(
346    cache: Arc<RawCache<K, V, W, H>>, telemetry: Telemetry, expiration: Expiration<K>, expiration_interval: Duration,
347    shutdown: CancellationToken,
348) where
349    K: Eq + std::hash::Hash + Clone,
350    V: Clone,
351    W: Weighter<K, V> + Clone,
352    H: std::hash::BuildHasher + Clone,
353{
354    let mut expired_item_keys = Vec::new();
355
356    loop {
357        tokio::select! {
358            _ = shutdown.cancelled() => break,
359            _ = sleep(expiration_interval) => {}
360        }
361
362        // Drain all expired items that have been queued up for the cache.
363        expiration.drain_expired_items(&mut expired_item_keys);
364
365        let num_expired_items = expired_item_keys.len();
366        if num_expired_items != 0 {
367            telemetry.items_expired_total().increment(num_expired_items as u64);
368            telemetry.items_expired_batch_size().record(num_expired_items as f64);
369        }
370
371        debug!(num_expired_items, "Found expired items.");
372
373        for item_key in expired_item_keys.drain(..) {
374            cache.remove(&item_key);
375            telemetry.items_removed_total().increment(1);
376            expiration.mark_entry_removed(item_key);
377        }
378
379        debug!(num_expired_items, "Removed expired items.");
380    }
381}
382
383async fn drive_telemetry<K, V, W, H>(
384    cache: Arc<RawCache<K, V, W, H>>, telemetry: Telemetry, shutdown: CancellationToken,
385) where
386    K: Eq + std::hash::Hash + Clone,
387    V: Clone,
388    W: Weighter<K, V> + Clone,
389    H: std::hash::BuildHasher + Clone,
390{
391    loop {
392        tokio::select! {
393            _ = shutdown.cancelled() => break,
394            _ = sleep(Duration::from_secs(1)) => {}
395        }
396
397        telemetry.current_items().set(cache.len() as f64);
398        telemetry.current_weight().set(cache.weight() as f64);
399    }
400}
401
402#[cfg(test)]
403mod tests {
404    use super::*;
405
406    #[derive(Clone)]
407    pub struct ItemValueWeighter;
408
409    impl<K> Weighter<K, usize> for ItemValueWeighter {
410        fn item_weight(&self, _key: &K, value: &usize) -> u64 {
411            *value as u64
412        }
413    }
414
415    #[test]
416    fn empty_cache_identifier() {
417        let result = CacheBuilder::<u64, u64>::from_identifier("");
418        assert!(result.is_err(), "expected error for empty cache identifier");
419    }
420
421    #[test]
422    fn basic() {
423        const CACHE_KEY: usize = 42;
424        const CACHE_VALUE: &str = "value1";
425
426        let cache = CacheBuilder::for_tests().build();
427
428        assert_eq!(cache.len(), 0);
429        assert_eq!(cache.weight(), 0);
430
431        cache.insert(CACHE_KEY, CACHE_VALUE);
432        assert_eq!(cache.len(), 1);
433        assert_eq!(cache.weight(), 1);
434
435        assert_eq!(cache.get(&CACHE_KEY), Some(CACHE_VALUE));
436
437        cache.remove(&CACHE_KEY);
438        assert_eq!(cache.len(), 0);
439        assert_eq!(cache.weight(), 0);
440    }
441
442    #[test]
443    fn evict_at_capacity() {
444        const CAPACITY: usize = 3;
445
446        let cache = CacheBuilder::for_tests()
447            .with_capacity(NonZeroUsize::new(CAPACITY).unwrap())
448            .build();
449
450        // Insert items up to the capacity.
451        for i in 0..CAPACITY {
452            cache.insert(i, "value");
453        }
454
455        assert_eq!(cache.len(), CAPACITY);
456        assert_eq!(cache.weight(), CAPACITY as u64);
457
458        // Inserting another item should evict something else to make room, leaving it such that the cache still has the
459        // same number of items.
460        cache.insert(CAPACITY, "new_value");
461        assert_eq!(cache.len(), CAPACITY);
462        assert_eq!(cache.weight(), CAPACITY as u64);
463
464        let mut evicted = false;
465        for i in 0..CAPACITY {
466            if cache.get(&i).is_none() {
467                evicted = true;
468                break;
469            }
470        }
471        assert!(evicted, "expected at least one original item to be evicted");
472    }
473
474    #[test]
475    fn overweight_item() {
476        const CAPACITY: usize = 10;
477
478        // Create our cache using an "item value" weighter, which uses the item value itself as the weight.
479        let cache = CacheBuilder::for_tests()
480            .with_capacity(NonZeroUsize::new(CAPACITY).unwrap())
481            .with_item_weighter(ItemValueWeighter)
482            .build();
483
484        // We should fail to insert an item that is too heavy for the cache by itself.
485        assert_eq!(cache.len(), 0);
486        assert_eq!(cache.weight(), 0);
487
488        cache.insert(1, CAPACITY + 1);
489        assert_eq!(cache.len(), 0);
490        assert_eq!(cache.weight(), 0);
491        assert_eq!(cache.get(&1), None);
492    }
493
494    #[test]
495    fn evict_on_insert_by_weight() {
496        const CAPACITY: usize = 10;
497
498        // Create our cache using an "item value" weighter, which uses the item value itself as the weight.
499        let cache = CacheBuilder::for_tests()
500            .with_capacity(NonZeroUsize::new(CAPACITY).unwrap())
501            .with_item_weighter(ItemValueWeighter)
502            .build();
503
504        // Insert three items which together have a weight equal to the cache capacity.
505        assert_eq!(cache.len(), 0);
506        assert_eq!(cache.weight(), 0);
507
508        cache.insert(1, 3);
509        cache.insert(2, 4);
510        cache.insert(3, 3);
511        assert_eq!(cache.len(), 3);
512        assert_eq!(cache.weight(), CAPACITY as u64);
513
514        // Now try to insert an item that has a weight that is smaller than the cache capacity, but larger than all
515        // prior items combined, which should evict all prior items to make room for the new item.
516        cache.insert(4, CAPACITY - 1);
517        assert_eq!(cache.len(), 1);
518        assert_eq!(cache.weight(), (CAPACITY - 1) as u64);
519
520        assert_eq!(cache.get(&1), None);
521        assert_eq!(cache.get(&2), None);
522        assert_eq!(cache.get(&3), None);
523        assert_eq!(cache.get(&4), Some(CAPACITY - 1));
524    }
525
526    #[tokio::test]
527    async fn tasks_stop_when_cache_dropped() {
528        let cache = CacheBuilder::<u64, u64>::from_identifier("test-drop")
529            .expect("valid identifier")
530            .with_time_to_idle(Some(Duration::from_secs(60)))
531            .with_expiration_interval(Duration::from_millis(50))
532            .build();
533
534        // Grab a weak reference to the raw cache data held by the background tasks.
535        let weak_cache = Arc::downgrade(&cache.inner.cache);
536
537        drop(cache);
538
539        // When `InnerCache` is dropped, the cancellation token's drop guard is also dropped, which triggers
540        // cancellation, so both tasks should wake up immediately and exit, releasing their Arc<RawCache> references.
541        //
542        // TODO: There's no good way to assert the tasks have shutdown besides sleeping and checking the weak cache is
543        // gone. It would be nice if there was a way to asynchronously _and_ fallibly shutdown the runtime with a
544        // timeout, such that we could detect if they shutdown cleanly... but alas.
545        sleep(Duration::from_millis(100)).await;
546
547        assert!(
548            weak_cache.upgrade().is_none(),
549            "raw cache should be released after background tasks exit"
550        );
551    }
552}
saluki_core/cache/mod.rs

saluki_core/cache/
mod.rs
