lance_core/cache/
quick.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! [`CacheBackend`] backed by [quick_cache](https://crates.io/crates/quick_cache),
5//! whose hit path is one atomic bit — no read-op channel or inline
6//! housekeeping. Used for the session index and metadata caches; the index
7//! cache sees thousands of cache reads per query.
8
9use std::pin::Pin;
10
11use async_trait::async_trait;
12use futures::Future;
13
14use super::backend::{CacheBackend, CacheEntry};
15use super::moka::key_footprint;
16use super::{CacheCodec, InternalCacheKey};
17use crate::Result;
18
19#[derive(Clone)]
20struct QuickEntry {
21    entry: CacheEntry,
22    size_bytes: usize,
23}
24
25#[derive(Clone)]
26struct EntryWeighter;
27
28impl quick_cache::Weighter<InternalCacheKey, QuickEntry> for EntryWeighter {
29    fn weight(&self, key: &InternalCacheKey, value: &QuickEntry) -> u64 {
30        // Same accounting as the moka backend.
31        key_footprint(key).saturating_add(value.size_bytes).max(1) as u64
32    }
33}
34
35pub struct QuickCacheBackend {
36    cache: quick_cache::sync::Cache<InternalCacheKey, QuickEntry, EntryWeighter>,
37}
38
39impl std::fmt::Debug for QuickCacheBackend {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        f.debug_struct("QuickCacheBackend")
42            .field("entry_count", &self.cache.len())
43            .finish()
44    }
45}
46
47/// Minimum weight budget (4 GiB) per shard: shards don't borrow capacity, and
48/// an entry heavier than ~its shard's budget is silently refused admission.
49const MIN_SHARD_SHARE: usize = 4 << 30;
50
51/// Recommended shard count: `min(cpus / 2, capacity / 4 GiB)`, power of two
52/// in `[1, 1024]`. The cpu term bounds lock contention; the capacity term
53/// keeps each shard's budget >= 4 GiB so large entries stay admissible.
54/// Rounded down because quick_cache rounds requests up.
55pub fn recommended_cache_shards(capacity: usize) -> usize {
56    let by_cpu = std::thread::available_parallelism()
57        .map(|n| n.get())
58        .unwrap_or(2)
59        / 2;
60    let shards = (capacity / MIN_SHARD_SHARE).min(by_cpu).max(1);
61    let shards = if shards.is_power_of_two() {
62        shards
63    } else {
64        shards.next_power_of_two() / 2
65    };
66    shards.clamp(1, 1024)
67}
68
69/// Assumed average entry size for pre-allocation sizing.
70const ESTIMATED_AVG_ENTRY_BYTES: usize = 64 << 10;
71
72impl QuickCacheBackend {
73    /// Create a backend holding up to `capacity` bytes of weighted entries
74    /// (weight = key footprint + declared size), sharded per
75    /// [`recommended_cache_shards`].
76    pub fn with_capacity(capacity: usize) -> Self {
77        let shards = recommended_cache_shards(capacity);
78        // Floor protects the shard count from quick_cache's items-per-shard
79        // heuristic; ceiling bounds pre-allocation.
80        let estimated_items = (capacity / ESTIMATED_AVG_ENTRY_BYTES).clamp(shards * 32, 1_000_000);
81        let options = quick_cache::OptionsBuilder::new()
82            .estimated_items_capacity(estimated_items)
83            .weight_capacity(capacity as u64)
84            .shards(shards)
85            .build()
86            // Only errors when weight/item capacity is missing; both are set.
87            .expect("quick_cache options");
88        let cache = quick_cache::sync::Cache::with_options(
89            options,
90            EntryWeighter,
91            Default::default(),
92            Default::default(),
93        );
94        Self { cache }
95    }
96}
97
98#[async_trait]
99impl CacheBackend for QuickCacheBackend {
100    async fn get(&self, key: &InternalCacheKey, _codec: Option<CacheCodec>) -> Option<CacheEntry> {
101        self.cache.get(key).map(|v| v.entry)
102    }
103
104    async fn insert(
105        &self,
106        key: &InternalCacheKey,
107        entry: CacheEntry,
108        size_bytes: usize,
109        _codec: Option<CacheCodec>,
110    ) {
111        self.cache.insert(*key, QuickEntry { entry, size_bytes });
112    }
113
114    async fn get_or_insert<'a>(
115        &self,
116        key: &InternalCacheKey,
117        loader: Pin<Box<dyn Future<Output = Result<(CacheEntry, usize)>> + Send + 'a>>,
118        _codec: Option<CacheCodec>,
119    ) -> Result<(CacheEntry, bool)> {
120        match self.cache.get_value_or_guard_async(key).await {
121            Ok(value) => Ok((value.entry, true)),
122            Err(guard) => {
123                let (entry, size_bytes) = loader.await?;
124                let _ = guard.insert(QuickEntry {
125                    entry: entry.clone(),
126                    size_bytes,
127                });
128                Ok((entry, false))
129            }
130        }
131    }
132
133    async fn clear(&self) {
134        self.cache.clear();
135    }
136
137    async fn num_entries(&self) -> usize {
138        self.cache.len()
139    }
140
141    async fn size_bytes(&self) -> usize {
142        self.cache.weight() as usize
143    }
144
145    fn approx_num_entries(&self) -> usize {
146        self.cache.len()
147    }
148
149    fn approx_size_bytes(&self) -> usize {
150        self.cache.weight() as usize
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use std::marker::PhantomData;
157    use std::sync::Arc;
158    use std::sync::atomic::{AtomicUsize, Ordering};
159
160    use super::*;
161    use crate::cache::{CacheKey, LanceCache};
162
163    struct TestKey<T: 'static> {
164        key: String,
165        _phantom: PhantomData<T>,
166    }
167
168    impl<T: 'static> TestKey<T> {
169        fn new(key: &str) -> Self {
170            Self {
171                key: key.to_string(),
172                _phantom: PhantomData,
173            }
174        }
175    }
176
177    impl<T: 'static> CacheKey for TestKey<T> {
178        type ValueType = T;
179        fn key(&self) -> std::borrow::Cow<'_, str> {
180            std::borrow::Cow::Borrowed(&self.key)
181        }
182        fn type_name() -> &'static str {
183            std::any::type_name::<T>()
184        }
185    }
186
187    #[test]
188    fn entry_weight_includes_fixed_key() {
189        let key = InternalCacheKey::from_bytes([0; 16]);
190        let entry = QuickEntry {
191            entry: Arc::new(()),
192            size_bytes: 7,
193        };
194        assert_eq!(
195            quick_cache::Weighter::weight(&EntryWeighter, &key, &entry),
196            23
197        );
198    }
199
200    #[tokio::test]
201    async fn test_quick_backend_roundtrip_singleflight_and_eviction() {
202        // Capacity must be large relative to one entry: quick_cache shards
203        // its weight budget, and an entry heavier than its shard's share is
204        // not admitted at all.
205        const CAPACITY: usize = 1 << 20;
206        let item = Arc::new(vec![1u8, 2, 3]);
207        let cache = LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(CAPACITY)));
208
209        // insert + get roundtrip and weighted accounting
210        cache
211            .insert_with_key(&TestKey::<Vec<u8>>::new("a"), item.clone())
212            .await;
213        assert_eq!(
214            cache
215                .get_with_key(&TestKey::<Vec<u8>>::new("a"))
216                .await
217                .as_deref(),
218            Some(&vec![1u8, 2, 3])
219        );
220        assert_eq!(cache.approx_size(), 1);
221        assert!(cache.size_bytes().await > 0);
222
223        // get_or_insert runs the loader only on a miss
224        let loads = Arc::new(AtomicUsize::new(0));
225        for _ in 0..2 {
226            let loads = loads.clone();
227            let value = cache
228                .get_or_insert_with_key(TestKey::<Vec<u8>>::new("b"), || async move {
229                    loads.fetch_add(1, Ordering::SeqCst);
230                    Ok(vec![7u8])
231                })
232                .await
233                .unwrap();
234            assert_eq!(value.as_ref(), &vec![7u8]);
235        }
236        assert_eq!(loads.load(Ordering::SeqCst), 1);
237
238        // capacity is enforced: overfill with 4x capacity of 16KiB entries
239        // and confirm eviction kept the weighted size within budget
240        for i in 0..256 {
241            cache
242                .insert_with_key(
243                    &TestKey::<Vec<u8>>::new(&format!("fill-{i}")),
244                    Arc::new(vec![0u8; 16 << 10]),
245                )
246                .await;
247        }
248        assert!(cache.size_bytes().await <= CAPACITY);
249        assert!(cache.size().await < 258);
250
251        cache.clear().await;
252        assert_eq!(cache.size().await, 0);
253    }
254
255    #[tokio::test]
256    async fn test_quick_backend_tiny_capacity() {
257        // A tiny cache must not over-provision item metadata and must still
258        // admit and evict correctly within its weight budget.
259        const CAPACITY: usize = 64 << 10;
260        let cache = LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(CAPACITY)));
261        for i in 0..64 {
262            cache
263                .insert_with_key(
264                    &TestKey::<Vec<u8>>::new(&format!("k-{i}")),
265                    Arc::new(vec![0u8; 4 << 10]),
266                )
267                .await;
268        }
269        assert!(cache.size_bytes().await <= CAPACITY);
270        assert!(cache.size().await >= 1);
271        let hit = cache
272            .get_with_key(&TestKey::<Vec<u8>>::new("k-63"))
273            .await
274            .is_some()
275            || cache
276                .get_with_key(&TestKey::<Vec<u8>>::new("k-62"))
277                .await
278                .is_some();
279        assert!(hit, "recently inserted entries should be resident");
280    }
281}
lance_core/cache/quick.rs

lance_core/cache/
quick.rs
