common_cache/
cache.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::borrow::Borrow;
16use std::hash::Hash;
17use std::sync::Arc;
18use std::time::Duration;
19
20use crate::metrics::{CacheMetrics, CacheSizeMetric, CacheSizeRegistration};
21
22/// Reason for releasing an entry from the cache.
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum RemovalCause {
25    /// The entry reached its expiration time.
26    Expired,
27    /// The entry was explicitly invalidated.
28    Explicit,
29    /// A new value replaced the entry.
30    Replaced,
31    /// The entry was removed to satisfy the capacity policy.
32    Size,
33}
34
35impl RemovalCause {
36    /// Returns the established eviction metric label.
37    pub fn as_str(self) -> &'static str {
38        match self {
39            Self::Expired => "expired",
40            Self::Explicit => "explicit",
41            Self::Replaced => "replaced",
42            Self::Size => "size",
43        }
44    }
45}
46
47impl From<moka::notification::RemovalCause> for RemovalCause {
48    fn from(value: moka::notification::RemovalCause) -> Self {
49        match value {
50            moka::notification::RemovalCause::Expired => Self::Expired,
51            moka::notification::RemovalCause::Explicit => Self::Explicit,
52            moka::notification::RemovalCause::Replaced => Self::Replaced,
53            moka::notification::RemovalCause::Size => Self::Size,
54        }
55    }
56}
57
58/// A weighted memory cache with backend-provided asynchronous load deduplication.
59/// Values should be inexpensive to clone, such as `Arc<T>` or `Bytes`.
60/// Capacity is maintained asynchronously; insertion does not guarantee retention.
61#[derive(Clone)]
62pub struct Cache<K, V> {
63    inner: moka::future::Cache<K, V>,
64    capacity: u64,
65    size_registration: Option<CacheSizeRegistration>,
66    metrics: Option<Arc<CacheMetrics>>,
67}
68
69impl<K, V> Cache<K, V>
70where
71    K: Eq + Hash + Send + Sync + 'static,
72    V: Clone + Send + Sync + 'static,
73{
74    /// Creates a cache with the caller's weight and eviction accounting.
75    pub fn new<W, L>(name: &str, capacity: u64, weigher: W, listener: L) -> Self
76    where
77        W: Fn(&K, &V) -> u32 + Send + Sync + 'static,
78        L: Fn(Arc<K>, V, RemovalCause) + Send + Sync + 'static,
79    {
80        Self::new_with_ttl(name, capacity, weigher, listener, None)
81    }
82
83    /// Creates a cache with an optional write-based expiration time.
84    pub fn new_with_ttl<W, L>(
85        name: &str,
86        capacity: u64,
87        weigher: W,
88        listener: L,
89        ttl: Option<Duration>,
90    ) -> Self
91    where
92        W: Fn(&K, &V) -> u32 + Send + Sync + 'static,
93        L: Fn(Arc<K>, V, RemovalCause) + Send + Sync + 'static,
94    {
95        let mut builder = moka::future::Cache::builder()
96            .name(name)
97            .max_capacity(capacity)
98            .support_invalidation_closures()
99            .weigher(weigher)
100            .eviction_listener(move |key, value, cause| listener(key, value, cause.into()));
101        if let Some(ttl) = ttl {
102            builder = builder.time_to_live(ttl);
103        }
104        Self {
105            inner: builder.build(),
106            capacity,
107            size_registration: None,
108            metrics: None,
109        }
110    }
111
112    /// Attaches caller-bound metrics without changing existing hit/miss accounting.
113    pub fn with_metrics(mut self, metrics: Arc<CacheMetrics>) -> Self {
114        self.metrics = Some(metrics);
115        self
116    }
117
118    /// Samples this cache's retained weight when the caller collects byte metrics.
119    pub fn with_size_metric(mut self, metric: CacheSizeMetric) -> Self {
120        let cache = self.inner.clone();
121        self.size_registration = Some(metric.register(move || cache.weighted_size()));
122        self
123    }
124
125    /// Gets a completed value without starting a load.
126    pub fn get<'a, Q>(&'a self, key: &'a Q) -> impl Future<Output = Option<V>> + Send + 'a
127    where
128        K: Borrow<Q>,
129        Q: Hash + Eq + Sync + ?Sized,
130    {
131        self.inner.get(key)
132    }
133
134    /// Shares a successful load among concurrent callers with the same key and error type.
135    /// Errors are shared with current waiters but are not retained as cached values.
136    /// Cancelling the active loader lets a remaining caller execute its own loader.
137    pub async fn try_get_with<Q, F, E>(&self, key: &Q, load: F) -> Result<V, Arc<E>>
138    where
139        K: Borrow<Q>,
140        Q: ToOwned<Owned = K> + Hash + Eq + ?Sized,
141        F: Future<Output = Result<V, E>>,
142        E: Send + Sync + 'static,
143    {
144        let request = self
145            .metrics
146            .as_ref()
147            .map(|metrics| metrics.requests.start());
148        let result = self
149            .inner
150            .try_get_with_by_ref(key, async {
151                // Only the backend-selected loader contributes physical load metrics.
152                let operation = self.metrics.as_ref().map(|metrics| metrics.loads.start());
153                let result = load.await;
154                if let Some(operation) = operation {
155                    operation.finish(&result);
156                }
157                result
158            })
159            .await;
160        if let Some(request) = request {
161            request.finish(&result);
162        }
163        result
164    }
165
166    /// Submits a completed value to the cache's admission and eviction policy.
167    pub async fn insert(&self, key: K, value: V) {
168        self.inner.insert(key, value).await;
169    }
170
171    /// Invalidates a completed entry. This does not cancel a concurrent loader.
172    pub async fn invalidate<Q>(&self, key: &Q)
173    where
174        K: Borrow<Q>,
175        Q: Hash + Eq + ?Sized,
176    {
177        self.inner.invalidate(key).await;
178    }
179
180    /// Invalidates completed entries matching a predicate. Pending loaders are not fenced.
181    /// Physical removal and eviction accounting run asynchronously.
182    pub fn invalidate_if<F>(&self, predicate: F) -> std::io::Result<()>
183    where
184        F: Fn(&K, &V) -> bool + Send + Sync + 'static,
185    {
186        self.inner
187            .invalidate_entries_if(predicate)
188            .map(|_| ())
189            .map_err(std::io::Error::other)
190    }
191
192    /// Returns the configured weight capacity.
193    pub fn capacity(&self) -> u64 {
194        self.capacity
195    }
196
197    /// Returns the approximate number of retained entries.
198    pub fn entry_count(&self) -> u64 {
199        self.inner.entry_count()
200    }
201
202    /// Returns the approximate total retained weight.
203    pub fn weighted_size(&self) -> u64 {
204        self.inner.weighted_size()
205    }
206
207    /// Applies pending admission, eviction, and accounting work.
208    pub async fn run_pending_tasks(&self) {
209        self.inner.run_pending_tasks().await;
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use std::sync::atomic::{AtomicUsize, Ordering};
216
217    use tokio::sync::Notify;
218
219    use super::*;
220
221    // This is also a compile-time contract for generic callers crossing async-trait boundaries.
222    fn send_lookup<'a, K, V, Q>(
223        cache: &'a Cache<K, V>,
224        key: &'a Q,
225    ) -> impl Future<Output = Option<V>> + Send + 'a
226    where
227        K: Borrow<Q> + Eq + Hash + Send + Sync + 'static,
228        V: Clone + Send + Sync + 'static,
229        Q: Hash + Eq + Sync + ?Sized,
230    {
231        cache.get(key)
232    }
233
234    #[tokio::test]
235    async fn test_borrowed_lookup_has_a_send_future() {
236        let cache = cache(1024);
237        cache
238            .insert("key".to_string(), Arc::new("value".to_string()))
239            .await;
240        let future: futures::future::BoxFuture<'_, Option<Arc<String>>> =
241            Box::pin(send_lookup(&cache, "key"));
242        assert_eq!(future.await.as_deref().map(String::as_str), Some("value"));
243    }
244
245    #[tokio::test]
246    async fn test_predicate_invalidation_preserves_other_entries() {
247        let cache = cache(1024);
248        cache
249            .insert("old".to_string(), Arc::new("a".to_string()))
250            .await;
251        cache
252            .insert("keep".to_string(), Arc::new("b".to_string()))
253            .await;
254        cache.run_pending_tasks().await;
255        cache.invalidate_if(|key, _| key == "old").unwrap();
256        assert!(cache.get("old").await.is_none());
257        assert_eq!(
258            cache.get("keep").await.as_deref().map(String::as_str),
259            Some("b")
260        );
261    }
262
263    fn cache(capacity: u64) -> Cache<String, Arc<String>> {
264        Cache::new(
265            "test",
266            capacity,
267            |_, value: &Arc<String>| value.len() as u32,
268            |_, _, _| {},
269        )
270    }
271
272    fn operation_metrics(name: &str) -> crate::metrics::OperationMetrics {
273        use prometheus::{Histogram, HistogramOpts, IntCounter, IntGauge};
274
275        crate::metrics::OperationMetrics {
276            started: IntCounter::new(format!("{name}_started"), "test").unwrap(),
277            succeeded: IntCounter::new(format!("{name}_succeeded"), "test").unwrap(),
278            failed: IntCounter::new(format!("{name}_failed"), "test").unwrap(),
279            cancelled: IntCounter::new(format!("{name}_cancelled"), "test").unwrap(),
280            active: IntGauge::new(format!("{name}_active"), "test").unwrap(),
281            elapsed: Histogram::with_opts(HistogramOpts::new(format!("{name}_seconds"), "test"))
282                .unwrap(),
283        }
284    }
285
286    fn instrumented_cache() -> (Cache<String, Arc<String>>, Arc<CacheMetrics>) {
287        let metrics = Arc::new(CacheMetrics {
288            requests: operation_metrics("requests"),
289            loads: operation_metrics("loads"),
290        });
291        (cache(1024).with_metrics(metrics.clone()), metrics)
292    }
293
294    #[tokio::test]
295    async fn test_metrics_distinguish_requests_from_executed_loads() {
296        let (cache, metrics) = instrumented_cache();
297        // Merely constructing an async call does not start an operation.
298        drop(cache.try_get_with("unused", async { Ok::<_, ()>(Arc::new(String::new())) }));
299        assert_eq!(0, metrics.requests.started.get());
300        let gate = Notify::new();
301        let mut first = Box::pin(cache.try_get_with("key", async {
302            gate.notified().await;
303            Ok::<_, ()>(Arc::new("value".to_string()))
304        }));
305        let mut second = Box::pin(cache.try_get_with("key", async { Err::<Arc<String>, _>(()) }));
306        assert!(futures::poll!(&mut first).is_pending());
307        assert!(futures::poll!(&mut second).is_pending());
308        assert_eq!(2, metrics.requests.active.get());
309        assert_eq!(1, metrics.loads.active.get());
310        gate.notify_waiters();
311        let (first, second) = tokio::join!(first, second);
312        assert!(Arc::ptr_eq(&first.unwrap(), &second.unwrap()));
313        cache
314            .try_get_with("key", async { Err::<Arc<String>, _>(()) })
315            .await
316            .unwrap();
317        assert_eq!(3, metrics.requests.started.get());
318        assert_eq!(3, metrics.requests.succeeded.get());
319        assert_eq!(1, metrics.loads.started.get());
320        assert_eq!(1, metrics.loads.succeeded.get());
321        assert_eq!(0, metrics.requests.active.get());
322        assert_eq!(0, metrics.loads.active.get());
323        assert_eq!(3, metrics.requests.elapsed.get_sample_count());
324        assert_eq!(1, metrics.loads.elapsed.get_sample_count());
325        assert_eq!(0, metrics.requests.failed.get());
326        assert_eq!(0, metrics.requests.cancelled.get());
327    }
328
329    #[tokio::test]
330    async fn test_metrics_release_active_work_on_cancellation_and_error() {
331        let (cache, metrics) = instrumented_cache();
332        let mut producer =
333            Box::pin(cache.try_get_with("key", std::future::pending::<Result<Arc<String>, ()>>()));
334        let mut waiter =
335            Box::pin(cache.try_get_with("key", std::future::pending::<Result<Arc<String>, ()>>()));
336        assert!(futures::poll!(&mut producer).is_pending());
337        assert!(futures::poll!(&mut waiter).is_pending());
338        drop(waiter);
339        assert_eq!(1, metrics.requests.cancelled.get());
340        assert_eq!(0, metrics.loads.cancelled.get());
341        drop(producer);
342        assert_eq!(2, metrics.requests.cancelled.get());
343        assert_eq!(1, metrics.loads.cancelled.get());
344        let result = cache
345            .try_get_with("key", async { Err::<Arc<String>, _>(()) })
346            .await;
347        assert!(result.is_err());
348        assert_eq!(1, metrics.requests.failed.get());
349        assert_eq!(1, metrics.loads.failed.get());
350        assert_eq!(0, metrics.requests.active.get());
351        assert_eq!(0, metrics.loads.active.get());
352        assert_eq!(3, metrics.requests.elapsed.get_sample_count());
353        assert_eq!(2, metrics.loads.elapsed.get_sample_count());
354    }
355
356    #[tokio::test]
357    async fn test_instrumented_loading_future_is_send() {
358        fn assert_send<T: Send>(value: T) -> T {
359            value
360        }
361
362        let (cache, _) = instrumented_cache();
363        assert_send(
364            cache.try_get_with("key", async { Ok::<_, ()>(Arc::new("value".to_string())) }),
365        )
366        .await
367        .unwrap();
368    }
369
370    #[tokio::test]
371    async fn test_shares_load_and_bypasses_loader_on_hit() {
372        let cache = cache(1024);
373        let calls = AtomicUsize::new(0);
374        let gate = Notify::new();
375        let load = || async {
376            calls.fetch_add(1, Ordering::Relaxed);
377            gate.notified().await;
378            Ok::<_, ()>(Arc::new("value".to_string()))
379        };
380        let mut first = Box::pin(cache.try_get_with("key", load()));
381        let mut second = Box::pin(cache.try_get_with("key", load()));
382        assert!(futures::poll!(&mut first).is_pending());
383        assert!(futures::poll!(&mut second).is_pending());
384        assert_eq!(1, calls.load(Ordering::Relaxed));
385        gate.notify_waiters();
386        let (first, second) = tokio::join!(first, second);
387        let first = first.unwrap();
388        assert!(Arc::ptr_eq(&first, &second.unwrap()));
389        let hit = cache
390            .try_get_with("key", async { Err::<Arc<String>, _>(()) })
391            .await
392            .unwrap();
393        assert!(Arc::ptr_eq(&first, &hit));
394    }
395
396    #[tokio::test]
397    async fn test_remaining_waiter_loads_after_producer_cancellation() {
398        let (cache, metrics) = instrumented_cache();
399        let mut first =
400            Box::pin(cache.try_get_with("key", std::future::pending::<Result<Arc<String>, ()>>()));
401        assert!(futures::poll!(&mut first).is_pending());
402        let mut second = Box::pin(
403            cache.try_get_with("key", async { Ok::<_, ()>(Arc::new("retry".to_string())) }),
404        );
405        assert!(futures::poll!(&mut second).is_pending());
406        drop(first);
407        assert_eq!("retry", second.await.unwrap().as_str());
408        assert_eq!(2, metrics.requests.started.get());
409        assert_eq!(1, metrics.requests.cancelled.get());
410        assert_eq!(1, metrics.requests.succeeded.get());
411        assert_eq!(2, metrics.loads.started.get());
412        assert_eq!(1, metrics.loads.cancelled.get());
413        assert_eq!(1, metrics.loads.succeeded.get());
414        assert_eq!(0, metrics.requests.active.get());
415        assert_eq!(0, metrics.loads.active.get());
416    }
417
418    #[tokio::test]
419    async fn test_errors_are_shared_but_not_cached() {
420        let (cache, metrics) = instrumented_cache();
421        let calls = AtomicUsize::new(0);
422        let gate = Notify::new();
423        let load = || async {
424            calls.fetch_add(1, Ordering::Relaxed);
425            gate.notified().await;
426            Err::<Arc<String>, _>("read failed")
427        };
428        let mut first = Box::pin(cache.try_get_with("key", load()));
429        let mut second = Box::pin(cache.try_get_with("key", load()));
430        assert!(futures::poll!(&mut first).is_pending());
431        assert!(futures::poll!(&mut second).is_pending());
432        gate.notify_waiters();
433        let (first, second) = tokio::join!(first, second);
434        assert!(Arc::ptr_eq(&first.unwrap_err(), &second.unwrap_err()));
435        assert_eq!(2, metrics.requests.failed.get());
436        assert_eq!(1, metrics.loads.failed.get());
437        assert_eq!(0, metrics.requests.active.get());
438        assert_eq!(0, metrics.loads.active.get());
439        assert_eq!(1, calls.load(Ordering::Relaxed));
440        assert!(cache.get("key").await.is_none());
441        assert_eq!(
442            "retry",
443            cache
444                .try_get_with("key", async {
445                    Ok::<_, &'static str>(Arc::new("retry".to_string()))
446                })
447                .await
448                .unwrap()
449                .as_str()
450        );
451    }
452
453    #[tokio::test]
454    async fn test_sampled_size_sums_instances_and_preserves_manual_bytes() {
455        use crate::metrics::CacheSizeMetrics;
456        use prometheus::Opts;
457        use prometheus::core::Collector;
458
459        let metrics =
460            CacheSizeMetrics::new(Opts::new("test_cache_bytes", "test"), &["type"]).unwrap();
461        let manual = metrics.with_label_values(&["shared"]);
462        manual.set(3);
463        let make_cache = || {
464            Cache::new("sampled", 100, |_: &u64, _: &Arc<String>| 7, |_, _, _| {})
465                .with_size_metric(metrics.bind(&["shared"]))
466        };
467        let a = make_cache();
468        let b = make_cache();
469        let sample = || {
470            metrics.collect()[0].get_metric()[0]
471                .get_gauge()
472                .as_ref()
473                .map_or(0.0, |gauge| gauge.value())
474        };
475        a.insert(1, Arc::new("one".into())).await;
476        b.insert(2, Arc::new("two".into())).await;
477        a.run_pending_tasks().await;
478        b.run_pending_tasks().await;
479        assert_eq!(17.0, sample());
480        assert_eq!(3, manual.get());
481        let consumer = a.get(&1).await.unwrap();
482        a.invalidate(&1).await;
483        a.run_pending_tasks().await;
484        assert_eq!(10.0, sample());
485        drop(consumer);
486        let cloned = b.clone();
487        drop(b);
488        assert_eq!(10.0, sample());
489        drop(cloned);
490        assert_eq!(3.0, sample());
491    }
492}
common_cache/cache.rs

common_cache/
cache.rs
