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

common_cache/
cache.rs
