lance_core/cache/
moka.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::pin::Pin;
5use std::sync::Arc;
6use std::sync::atomic::{AtomicBool, Ordering};
7
8use async_trait::async_trait;
9use futures::Future;
10
11use crate::Result;
12use crate::error::CloneableError;
13
14use super::CacheCodec;
15use super::backend::{CacheBackend, CacheEntry, CacheKeyIterator, InternalCacheKey};
16
17/// Internal record stored in the moka cache.
18#[derive(Clone, Debug)]
19struct MokaCacheEntry {
20    entry: CacheEntry,
21    size_bytes: usize,
22}
23
24/// Per-entry key cost for eviction: the struct plus the unique `key` bytes.
25/// Excludes the shared `prefix` `Arc<str>`, which isn't freed per eviction.
26fn key_footprint(key: &InternalCacheKey) -> usize {
27    std::mem::size_of::<InternalCacheKey>() + key.key().len()
28}
29
30/// Default [`CacheBackend`] backed by a [moka](https://crates.io/crates/moka) cache.
31///
32/// Provides weighted-capacity eviction and concurrent-load deduplication
33/// via moka's built-in `optionally_get_with`.
34pub struct MokaCacheBackend {
35    cache: moka::future::Cache<InternalCacheKey, MokaCacheEntry>,
36}
37
38impl std::fmt::Debug for MokaCacheBackend {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        f.debug_struct("MokaCacheBackend")
41            .field("entry_count", &self.cache.entry_count())
42            .finish()
43    }
44}
45
46impl MokaCacheBackend {
47    pub fn with_capacity(capacity: usize) -> Self {
48        let cache = moka::future::Cache::builder()
49            .max_capacity(capacity as u64)
50            .weigher(|key: &InternalCacheKey, entry: &MokaCacheEntry| {
51                key_footprint(key)
52                    .saturating_add(entry.size_bytes)
53                    .try_into()
54                    .unwrap_or(u32::MAX)
55            })
56            .support_invalidation_closures()
57            .build();
58        Self { cache }
59    }
60
61    pub fn no_cache() -> Self {
62        Self {
63            cache: moka::future::Cache::new(0),
64        }
65    }
66}
67
68#[async_trait]
69impl CacheBackend for MokaCacheBackend {
70    async fn get(&self, key: &InternalCacheKey, _codec: Option<CacheCodec>) -> Option<CacheEntry> {
71        self.cache.get(key).await.map(|r| r.entry)
72    }
73
74    async fn insert(
75        &self,
76        key: &InternalCacheKey,
77        entry: CacheEntry,
78        size_bytes: usize,
79        _codec: Option<CacheCodec>,
80    ) {
81        self.cache
82            .insert(key.clone(), MokaCacheEntry { entry, size_bytes })
83            .await;
84    }
85
86    async fn get_or_insert<'a>(
87        &self,
88        key: &InternalCacheKey,
89        loader: Pin<Box<dyn Future<Output = Result<(CacheEntry, usize)>> + Send + 'a>>,
90        _codec: Option<CacheCodec>,
91    ) -> Result<(CacheEntry, bool)> {
92        // Track whether the loader actually ran (= cache miss).
93        let was_miss = Arc::new(AtomicBool::new(false));
94        let was_miss_clone = was_miss.clone();
95
96        let init = async move {
97            was_miss_clone.store(true, Ordering::Relaxed);
98            loader
99                .await
100                .map(|(entry, size_bytes)| MokaCacheEntry { entry, size_bytes })
101                .map_err(CloneableError)
102        };
103
104        let owned_key = key.clone();
105        match self.cache.try_get_with(owned_key, init).await {
106            Ok(record) => {
107                let was_cached = !was_miss.load(Ordering::Relaxed);
108                Ok((record.entry, was_cached))
109            }
110            Err(error) => Err(Arc::unwrap_or_clone(error).0),
111        }
112    }
113
114    async fn invalidate_prefix(&self, prefix: &str) {
115        let prefix = prefix.to_owned();
116        self.cache
117            .invalidate_entries_if(move |key, _value| key.starts_with(&prefix))
118            .expect("Cache configured correctly");
119    }
120
121    async fn clear(&self) {
122        self.cache.invalidate_all();
123        self.cache.run_pending_tasks().await;
124    }
125
126    async fn keys(&self) -> Option<CacheKeyIterator<'_>> {
127        self.cache.run_pending_tasks().await;
128        Some(Box::new(
129            self.cache.iter().map(|(key, _)| key.as_ref().clone()),
130        ))
131    }
132
133    async fn num_entries(&self) -> usize {
134        self.cache.run_pending_tasks().await;
135        self.cache.entry_count() as usize
136    }
137
138    async fn size_bytes(&self) -> usize {
139        self.cache.run_pending_tasks().await;
140        self.cache.weighted_size() as usize
141    }
142
143    fn approx_num_entries(&self) -> usize {
144        self.cache.entry_count() as usize
145    }
146
147    fn approx_size_bytes(&self) -> usize {
148        // Iterate rather than using `weighted_size()` because moka's
149        // weighted_size can be stale without `run_pending_tasks()`, which
150        // is async and can't be called from this synchronous context.
151        self.cache
152            .iter()
153            .map(|(key, entry)| key_footprint(key.as_ref()) + entry.size_bytes)
154            .sum()
155    }
156}
lance_core/cache/moka.rs

lance_core/cache/
moka.rs
