lance_core/cache/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Lance cache system.
5//!
6//! ## For cache users
7//!
8//! Use [`LanceCache`] (or [`WeakLanceCache`]) to store and retrieve typed
9//! values. Define a [`CacheKey`] (or [`UnsizedCacheKey`] for trait objects) to
10//! describe what you're caching and its type.
11//!
12//! To make a value type serializable (so persistent backends can store it),
13//! implement [`CacheCodecImpl`] on the type, then override [`CacheKey::codec`]:
14//!
15//! ```ignore
16//! impl CacheCodecImpl for MyData {
17//!     fn serialize(&self, w: &mut dyn Write) -> Result<()> { /* ... */ }
18//!     fn deserialize(data: &Bytes) -> Result<Self> { /* ... */ }
19//! }
20//!
21//! impl CacheKey for MyDataKey {
22//!     type ValueType = MyData;
23//!     fn key(&self) -> Cow<'_, str> { /* ... */ }
24//!     fn type_name() -> &'static str { "MyData" }
25//!     fn codec() -> Option<CacheCodec> {
26//!         Some(CacheCodec::from_impl::<MyData>())
27//!     }
28//! }
29//! ```
30//!
31//! ## For backend implementors
32//!
33//! Implement [`CacheBackend`] to provide a custom storage layer (disk, Redis,
34//! etc.). Backends receive [`InternalCacheKey`] keys and type-erased
35//! [`CacheEntry`] values — the typed wrapping is handled by [`LanceCache`].
36//! See the [`backend`] module for details.
37//!
38//! ## Serialization flow
39//!
40//! When a [`CacheKey`] provides a codec via [`CacheKey::codec`]:
41//!
42//! 1. [`LanceCache`] wraps the [`CacheCodec`] and passes it to the backend
43//!    alongside the entry on `insert` and `get` calls.
44//! 2. In-memory backends (like [`MokaCacheBackend`]) ignore the codec.
45//! 3. Persistent backends use `codec.serialize(entry, writer)` on insert and
46//!    `codec.deserialize(reader)` on get to persist entries across restarts.
47
48pub mod backend;
49pub mod codec;
50mod entry_io;
51mod moka;
52
53pub use backend::{CacheBackend, CacheEntry, CacheKeyIterator, InternalCacheKey};
54pub use codec::{
55    CacheCodec, CacheCodecImpl, CacheDecode, CacheMissReason, MAGIC, has_cache_envelope,
56};
57pub use entry_io::{CacheEntryReader, CacheEntryWriter};
58pub use moka::MokaCacheBackend;
59
60use std::borrow::Cow;
61use std::sync::{
62    Arc,
63    atomic::{AtomicU64, Ordering},
64};
65
66use futures::{Future, FutureExt};
67
68use crate::Result;
69
70pub use crate::deepsize::{Context, DeepSizeOf};
71
72// ---------------------------------------------------------------------------
73// CacheKey / UnsizedCacheKey — typed key traits for cache users
74// ---------------------------------------------------------------------------
75
76/// Typed cache key for sized value types.
77///
78/// Implement this trait to define a new type of cached entry. [`LanceCache`]
79/// uses the key string and type name to construct an [`InternalCacheKey`]
80/// for the backend.
81///
82/// # Example
83///
84/// ```ignore
85/// struct MyKey { id: u64 }
86///
87/// impl CacheKey for MyKey {
88///     type ValueType = MyData;
89///     fn key(&self) -> Cow<'_, str> { self.id.to_string().into() }
90///     fn type_name() -> &'static str { "MyData" }
91/// }
92/// ```
93pub trait CacheKey {
94    type ValueType: 'static;
95
96    fn key(&self) -> Cow<'_, str>;
97
98    /// Short, stable string identifying this value type.
99    ///
100    /// Two `CacheKey` impls that store different `ValueType`s **must** return
101    /// different type names; if they collide, gets will silently return `None`
102    /// due to failed downcasts.
103    ///
104    /// Use a short literal (e.g. `"Vec<IndexMetadata>"`), not
105    /// `std::any::type_name` — the latter is not guaranteed stable across
106    /// compiler versions or build configurations.
107    fn type_name() -> &'static str;
108
109    /// Optional codec for serializing/deserializing this key's value type.
110    ///
111    /// Returns `None` by default. Cache backends that support persistence
112    /// (e.g. disk-backed caches) use this to serialize entries on insert and
113    /// deserialize on get. Types without a codec will only be stored in-memory.
114    ///
115    /// [`CacheCodec`] is `Copy` (two plain function pointers), so returning it
116    /// by value is cheap — no allocation needed.
117    fn codec() -> Option<CacheCodec> {
118        None
119    }
120}
121
122/// Like [`CacheKey`] but for unsized value types (e.g. `dyn Trait`).
123///
124/// The cache wraps values in an extra `Arc` layer internally; callers pass
125/// and receive `Arc<T>` where `T: ?Sized`.
126///
127/// Unsized cache entries are always in-memory only (no serialization codec).
128/// For serializable entries, use a sized [`CacheKey`] instead.
129pub trait UnsizedCacheKey {
130    type ValueType: 'static + ?Sized;
131
132    fn key(&self) -> Cow<'_, str>;
133
134    /// Short, stable string identifying this value type.
135    /// See [`CacheKey::type_name`] for requirements.
136    fn type_name() -> &'static str;
137}
138
139// ---------------------------------------------------------------------------
140// Internal helpers
141// ---------------------------------------------------------------------------
142
143/// Size of a cached `Arc<T>`, accounting for the Arc overhead (two atomic counters).
144fn cache_entry_size<T: DeepSizeOf + ?Sized>(value: &T) -> usize {
145    value.deep_size_of() + std::mem::size_of::<std::sync::atomic::AtomicUsize>() * 2
146}
147
148/// Build an [`InternalCacheKey`] from a cache's prefix, a user key string,
149/// and a type name.
150fn build_key(prefix: &Arc<str>, key: &str, type_name: &'static str) -> InternalCacheKey {
151    InternalCacheKey::new(prefix.clone(), Arc::from(key), type_name)
152}
153
154// ---------------------------------------------------------------------------
155// LanceCache — typed wrapper around dyn CacheBackend
156// ---------------------------------------------------------------------------
157
158/// Typed cache wrapper that handles key construction and type safety.
159///
160/// Internally delegates to a [`CacheBackend`]. The default backend is
161/// [`MokaCacheBackend`]; pass a custom backend via [`LanceCache::with_backend`].
162#[derive(Clone)]
163pub struct LanceCache {
164    cache: Arc<dyn CacheBackend>,
165    prefix: Arc<str>,
166    hits: Arc<AtomicU64>,
167    misses: Arc<AtomicU64>,
168}
169
170impl std::fmt::Debug for LanceCache {
171    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
172        f.debug_struct("LanceCache")
173            .field("cache", &self.cache)
174            .finish()
175    }
176}
177
178impl DeepSizeOf for LanceCache {
179    fn deep_size_of_children(&self, _: &mut Context) -> usize {
180        self.cache.approx_size_bytes()
181    }
182}
183
184impl LanceCache {
185    pub fn with_capacity(capacity: usize) -> Self {
186        Self {
187            cache: Arc::new(MokaCacheBackend::with_capacity(capacity)),
188            prefix: Arc::from(""),
189            hits: Arc::new(AtomicU64::new(0)),
190            misses: Arc::new(AtomicU64::new(0)),
191        }
192    }
193
194    /// Create a cache backed by a custom [`CacheBackend`].
195    pub fn with_backend(backend: Arc<dyn CacheBackend>) -> Self {
196        Self {
197            cache: backend,
198            prefix: Arc::from(""),
199            hits: Arc::new(AtomicU64::new(0)),
200            misses: Arc::new(AtomicU64::new(0)),
201        }
202    }
203
204    pub fn no_cache() -> Self {
205        Self {
206            cache: Arc::new(MokaCacheBackend::no_cache()),
207            prefix: Arc::from(""),
208            hits: Arc::new(AtomicU64::new(0)),
209            misses: Arc::new(AtomicU64::new(0)),
210        }
211    }
212
213    /// Create a cache with the given backend and an exact prefix string.
214    /// Unlike `with_key_prefix`, this sets the prefix verbatim (no trailing slash added).
215    pub fn with_backend_and_prefix(backend: Arc<dyn CacheBackend>, prefix: String) -> Self {
216        Self {
217            cache: backend,
218            prefix: Arc::from(prefix),
219            hits: Arc::new(AtomicU64::new(0)),
220            misses: Arc::new(AtomicU64::new(0)),
221        }
222    }
223
224    /// Appends a prefix to the cache key.
225    pub fn with_key_prefix(&self, prefix: &str) -> Self {
226        Self {
227            cache: self.cache.clone(),
228            prefix: Arc::from(format!("{}{}/", self.prefix, prefix)),
229            hits: self.hits.clone(),
230            misses: self.misses.clone(),
231        }
232    }
233
234    /// Invalidate all entries whose prefix starts with the given string.
235    pub async fn invalidate_prefix(&self, prefix: &str) {
236        let full_prefix = format!("{}{}", self.prefix, prefix);
237        self.cache.invalidate_prefix(&full_prefix).await;
238    }
239
240    pub async fn size(&self) -> usize {
241        self.cache.num_entries().await
242    }
243
244    pub fn approx_size(&self) -> usize {
245        self.cache.approx_num_entries()
246    }
247
248    pub async fn size_bytes(&self) -> usize {
249        self.cache.size_bytes().await
250    }
251
252    /// Return an iterator over keys currently stored under this cache's prefix.
253    ///
254    /// Returns `None` when the backend does not support key inventory. The
255    /// iterator is intended for diagnostics and may be weakly consistent with
256    /// concurrent cache mutations.
257    ///
258    /// # Examples
259    ///
260    /// ```
261    /// # use std::{borrow::Cow, sync::Arc};
262    /// # use lance_core::cache::{CacheKey, LanceCache};
263    /// # struct MyKey;
264    /// # impl CacheKey for MyKey {
265    /// #     type ValueType = Vec<i32>;
266    /// #     fn key(&self) -> Cow<'_, str> { Cow::Borrowed("my-key") }
267    /// #     fn type_name() -> &'static str { "VecI32" }
268    /// # }
269    /// # async fn example() {
270    /// let cache = LanceCache::with_capacity(1024);
271    /// cache.insert_with_key(&MyKey, Arc::new(vec![1, 2, 3])).await;
272    ///
273    /// let mut keys = cache.keys().await.expect("Moka supports key inventory");
274    /// assert_eq!(keys.next().unwrap().key(), "my-key");
275    /// # }
276    /// ```
277    pub async fn keys(&self) -> Option<CacheKeyIterator<'_>> {
278        Some(Box::new(
279            self.cache
280                .keys()
281                .await?
282                .filter(|key| key.starts_with(&self.prefix)),
283        ))
284    }
285
286    // -- Sized insert/get (internal, shared by sized and unsized paths) --------
287
288    async fn insert_with_id<T: DeepSizeOf + Send + Sync + 'static>(
289        &self,
290        key: &str,
291        type_name: &'static str,
292        codec: Option<CacheCodec>,
293        metadata: Arc<T>,
294    ) {
295        let size = cache_entry_size(&*metadata);
296        let cache_key = build_key(&self.prefix, key, type_name);
297        self.cache.insert(&cache_key, metadata, size, codec).await;
298    }
299
300    async fn get_with_id<T: Send + Sync + 'static>(
301        &self,
302        key: &str,
303        type_name: &'static str,
304        codec: Option<CacheCodec>,
305    ) -> Option<Arc<T>> {
306        let cache_key = build_key(&self.prefix, key, type_name);
307        if let Some(entry) = self.cache.get(&cache_key, codec).await {
308            match entry.downcast::<T>() {
309                Ok(val) => {
310                    self.hits.fetch_add(1, Ordering::Relaxed);
311                    Some(val)
312                }
313                Err(_) => {
314                    // Type mismatch: the backend returned a different concrete
315                    // type than expected (e.g. a disk cache may store
316                    // intermediate state). Treat as a miss.
317                    self.misses.fetch_add(1, Ordering::Relaxed);
318                    None
319                }
320            }
321        } else {
322            self.misses.fetch_add(1, Ordering::Relaxed);
323            None
324        }
325    }
326
327    // -- Stats / clear --------------------------------------------------------
328
329    pub async fn stats(&self) -> CacheStats {
330        CacheStats {
331            hits: self.hits.load(Ordering::Relaxed),
332            misses: self.misses.load(Ordering::Relaxed),
333            num_entries: self.cache.num_entries().await,
334            size_bytes: self.cache.size_bytes().await,
335        }
336    }
337
338    pub async fn clear(&self) {
339        self.cache.clear().await;
340        self.hits.store(0, Ordering::Relaxed);
341        self.misses.store(0, Ordering::Relaxed);
342    }
343
344    // -- CacheKey-based methods -----------------------------------------------
345
346    pub async fn insert_with_key<K>(&self, cache_key: &K, metadata: Arc<K::ValueType>)
347    where
348        K: CacheKey,
349        K::ValueType: DeepSizeOf + Send + Sync + 'static,
350    {
351        self.insert_with_id(&cache_key.key(), K::type_name(), K::codec(), metadata)
352            .boxed()
353            .await
354    }
355
356    pub async fn get_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
357    where
358        K: CacheKey,
359        K::ValueType: DeepSizeOf + Send + Sync + 'static,
360    {
361        self.get_with_id::<K::ValueType>(&cache_key.key(), K::type_name(), K::codec())
362            .boxed()
363            .await
364    }
365
366    pub async fn get_or_insert_with_key<K, F, Fut>(
367        &self,
368        cache_key: K,
369        loader: F,
370    ) -> Result<Arc<K::ValueType>>
371    where
372        K: CacheKey,
373        K::ValueType: DeepSizeOf + Send + Sync + 'static,
374        F: FnOnce() -> Fut + Send,
375        Fut: Future<Output = Result<K::ValueType>> + Send,
376    {
377        let key = build_key(&self.prefix, &cache_key.key(), K::type_name());
378
379        let typed_loader = Box::pin(async move {
380            let value = loader().await?;
381            let arc = Arc::new(value);
382            let size = cache_entry_size(&*arc);
383            Ok((arc as CacheEntry, size))
384        });
385
386        let (entry, was_cached) = self
387            .cache
388            .get_or_insert(&key, typed_loader, K::codec())
389            .await?;
390
391        if was_cached {
392            self.hits.fetch_add(1, Ordering::Relaxed);
393        } else {
394            self.misses.fetch_add(1, Ordering::Relaxed);
395        }
396
397        Ok(entry.downcast::<K::ValueType>().unwrap())
398    }
399
400    pub async fn insert_unsized_with_key<K>(&self, cache_key: &K, metadata: Arc<K::ValueType>)
401    where
402        K: UnsizedCacheKey,
403        K::ValueType: DeepSizeOf + Send + Sync + 'static,
404    {
405        self.insert_with_id(&cache_key.key(), K::type_name(), None, Arc::new(metadata))
406            .boxed()
407            .await
408    }
409
410    pub async fn get_unsized_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
411    where
412        K: UnsizedCacheKey,
413        K::ValueType: DeepSizeOf + Send + Sync + 'static,
414    {
415        let outer = self
416            .get_with_id::<Arc<K::ValueType>>(&cache_key.key(), K::type_name(), None)
417            .boxed()
418            .await?;
419        Some(outer.as_ref().clone())
420    }
421}
422
423// ---------------------------------------------------------------------------
424// WeakLanceCache
425// ---------------------------------------------------------------------------
426
427/// A weak reference to a LanceCache, used by indices to avoid circular references.
428/// When the original cache is dropped, operations on this will gracefully no-op.
429#[derive(Clone, Debug)]
430pub struct WeakLanceCache {
431    inner: std::sync::Weak<dyn CacheBackend>,
432    prefix: Arc<str>,
433    hits: Arc<AtomicU64>,
434    misses: Arc<AtomicU64>,
435}
436
437impl WeakLanceCache {
438    pub fn from(cache: &LanceCache) -> Self {
439        Self {
440            inner: Arc::downgrade(&cache.cache),
441            prefix: cache.prefix.clone(),
442            hits: cache.hits.clone(),
443            misses: cache.misses.clone(),
444        }
445    }
446
447    pub fn with_key_prefix(&self, prefix: &str) -> Self {
448        Self {
449            inner: self.inner.clone(),
450            prefix: Arc::from(format!("{}{}/", self.prefix, prefix)),
451            hits: self.hits.clone(),
452            misses: self.misses.clone(),
453        }
454    }
455
456    /// The key prefix used for all entries in this cache.
457    pub fn prefix(&self) -> &str {
458        &self.prefix
459    }
460
461    pub async fn get_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
462    where
463        K: CacheKey,
464        K::ValueType: DeepSizeOf + Send + Sync + 'static,
465    {
466        let cache = self.inner.upgrade()?;
467        let key = build_key(&self.prefix, &cache_key.key(), K::type_name());
468        if let Some(entry) = cache.get(&key, K::codec()).await {
469            self.hits.fetch_add(1, Ordering::Relaxed);
470            Some(entry.downcast::<K::ValueType>().unwrap())
471        } else {
472            self.misses.fetch_add(1, Ordering::Relaxed);
473            None
474        }
475    }
476
477    pub async fn insert_with_key<K>(&self, cache_key: &K, value: Arc<K::ValueType>) -> bool
478    where
479        K: CacheKey,
480        K::ValueType: DeepSizeOf + Send + Sync + 'static,
481    {
482        if let Some(cache) = self.inner.upgrade() {
483            let size = cache_entry_size(&*value);
484            let key = build_key(&self.prefix, &cache_key.key(), K::type_name());
485            cache.insert(&key, value, size, K::codec()).await;
486            true
487        } else {
488            log::warn!("WeakLanceCache: cache no longer available, unable to insert item");
489            false
490        }
491    }
492
493    /// Get or insert an item, computing it if necessary.
494    ///
495    /// Deduplication of concurrent loads is handled by the backend.
496    pub async fn get_or_insert_with_key<K, F, Fut>(
497        &self,
498        cache_key: K,
499        loader: F,
500    ) -> Result<Arc<K::ValueType>>
501    where
502        K: CacheKey,
503        K::ValueType: DeepSizeOf + Send + Sync + 'static,
504        F: FnOnce() -> Fut + Send,
505        Fut: Future<Output = Result<K::ValueType>> + Send,
506    {
507        if let Some(cache) = self.inner.upgrade() {
508            let key = build_key(&self.prefix, &cache_key.key(), K::type_name());
509            let typed_loader = Box::pin(async move {
510                let value = loader().await?;
511                let arc = Arc::new(value);
512                let size = cache_entry_size(&*arc);
513                Ok((arc as CacheEntry, size))
514            });
515            let (entry, was_cached) = cache.get_or_insert(&key, typed_loader, K::codec()).await?;
516            if was_cached {
517                self.hits.fetch_add(1, Ordering::Relaxed);
518            } else {
519                self.misses.fetch_add(1, Ordering::Relaxed);
520            }
521            Ok(entry.downcast::<K::ValueType>().unwrap())
522        } else {
523            log::warn!("WeakLanceCache: cache no longer available, computing without caching");
524            loader().await.map(Arc::new)
525        }
526    }
527
528    pub async fn get_unsized_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
529    where
530        K: UnsizedCacheKey,
531        K::ValueType: DeepSizeOf + Send + Sync + 'static,
532    {
533        let cache = self.inner.upgrade()?;
534        let key = build_key(&self.prefix, &cache_key.key(), K::type_name());
535        if let Some(entry) = cache.get(&key, None).await {
536            entry
537                .downcast::<Arc<K::ValueType>>()
538                .ok()
539                .map(|arc| arc.as_ref().clone())
540        } else {
541            None
542        }
543    }
544
545    pub async fn insert_unsized_with_key<K>(&self, cache_key: &K, value: Arc<K::ValueType>)
546    where
547        K: UnsizedCacheKey,
548        K::ValueType: DeepSizeOf + Send + Sync + 'static,
549    {
550        if let Some(cache) = self.inner.upgrade() {
551            let wrapper = Arc::new(value);
552            let size = cache_entry_size(&*wrapper);
553            let key = build_key(&self.prefix, &cache_key.key(), K::type_name());
554            cache.insert(&key, wrapper, size, None).await;
555        } else {
556            log::warn!("WeakLanceCache: cache no longer available, unable to insert unsized item");
557        }
558    }
559}
560
561// ---------------------------------------------------------------------------
562// CacheStats
563// ---------------------------------------------------------------------------
564
565#[derive(Debug, Clone)]
566pub struct CacheStats {
567    /// Number of times `get`, `get_unsized`, or `get_or_insert` found an item in the cache.
568    pub hits: u64,
569    /// Number of times `get`, `get_unsized`, or `get_or_insert` did not find an item in the cache.
570    pub misses: u64,
571    /// Number of entries currently in the cache.
572    pub num_entries: usize,
573    /// Total size in bytes of all entries in the cache.
574    pub size_bytes: usize,
575}
576
577impl CacheStats {
578    pub fn hit_ratio(&self) -> f32 {
579        if self.hits + self.misses == 0 {
580            0.0
581        } else {
582            self.hits as f32 / (self.hits + self.misses) as f32
583        }
584    }
585
586    pub fn miss_ratio(&self) -> f32 {
587        if self.hits + self.misses == 0 {
588            0.0
589        } else {
590            self.misses as f32 / (self.hits + self.misses) as f32
591        }
592    }
593}
594
595#[cfg(test)]
596mod tests {
597    use super::*;
598    use crate::Error;
599    use std::collections::{BTreeSet, HashMap};
600    use std::marker::PhantomData;
601    use std::sync::atomic::AtomicUsize;
602
603    struct TestKey<T: 'static> {
604        key: String,
605        _phantom: PhantomData<T>,
606    }
607
608    impl<T: 'static> TestKey<T> {
609        fn new(key: &str) -> Self {
610            Self {
611                key: key.to_string(),
612                _phantom: PhantomData,
613            }
614        }
615    }
616
617    impl<T: 'static> CacheKey for TestKey<T> {
618        type ValueType = T;
619        fn key(&self) -> std::borrow::Cow<'_, str> {
620            std::borrow::Cow::Borrowed(&self.key)
621        }
622        fn type_name() -> &'static str {
623            std::any::type_name::<T>()
624        }
625    }
626
627    /// Test helper: an UnsizedCacheKey for trait object values.
628    struct TestUnsizedKey<T: 'static + ?Sized> {
629        key: String,
630        _phantom: PhantomData<T>,
631    }
632
633    impl<T: 'static + ?Sized> TestUnsizedKey<T> {
634        fn new(key: &str) -> Self {
635            Self {
636                key: key.to_string(),
637                _phantom: PhantomData,
638            }
639        }
640    }
641
642    impl<T: 'static + ?Sized> UnsizedCacheKey for TestUnsizedKey<T> {
643        type ValueType = T;
644        fn key(&self) -> std::borrow::Cow<'_, str> {
645            std::borrow::Cow::Borrowed(&self.key)
646        }
647        fn type_name() -> &'static str {
648            std::any::type_name::<T>()
649        }
650    }
651
652    fn key_fields(keys: &[InternalCacheKey]) -> BTreeSet<(String, String, &'static str)> {
653        keys.iter()
654            .map(|key| {
655                (
656                    key.prefix().to_string(),
657                    key.key().to_string(),
658                    key.type_name(),
659                )
660            })
661            .collect()
662    }
663
664    #[tokio::test]
665    async fn test_cache_bytes() {
666        let item = Arc::new(vec![1, 2, 3]);
667        let item_size = item.deep_size_of();
668        let capacity = 10 * item_size;
669        let cache = LanceCache::with_capacity(capacity);
670
671        cache
672            .insert_with_key(&TestKey::<Vec<i32>>::new("key"), item.clone())
673            .await;
674        assert_eq!(cache.size().await, 1);
675
676        let retrieved = cache
677            .get_with_key(&TestKey::<Vec<i32>>::new("key"))
678            .await
679            .unwrap();
680        assert_eq!(*retrieved, *item);
681
682        for i in 0..20 {
683            cache
684                .insert_with_key(
685                    &TestKey::<Vec<i32>>::new(&format!("key_{}", i)),
686                    Arc::new(vec![i, i, i]),
687                )
688                .await;
689        }
690        assert!(cache.size_bytes().await <= capacity);
691    }
692
693    #[tokio::test]
694    async fn test_cache_weighs_key_footprint() {
695        // Weighted size charges the key's unique bytes, not just the value.
696        let cache = LanceCache::with_capacity(usize::MAX);
697        let key = "k".repeat(10_000);
698        let value = Arc::new(vec![1_i32]);
699        let expected =
700            std::mem::size_of::<InternalCacheKey>() + key.len() + cache_entry_size(&*value);
701        cache
702            .insert_with_key(&TestKey::<Vec<i32>>::new(&key), value)
703            .await;
704        assert_eq!(cache.size_bytes().await, expected);
705    }
706
707    #[tokio::test]
708    async fn test_cache_shared_prefix_not_charged_per_entry() {
709        // The shared prefix contributes nothing per entry (it isn't freed on a
710        // single eviction); only struct + unique key + value are charged.
711        let cache = LanceCache::with_capacity(usize::MAX).with_key_prefix(&"p".repeat(10_000));
712        for i in 0..100 {
713            cache
714                .insert_with_key(
715                    &TestKey::<Vec<i32>>::new(&i.to_string()),
716                    Arc::new(vec![1_i32]),
717                )
718                .await;
719        }
720        let value_cost = cache_entry_size(&vec![1_i32]);
721        let key_bytes: usize = (0..100).map(|i| i.to_string().len()).sum();
722        let expected = 100 * (std::mem::size_of::<InternalCacheKey>() + value_cost) + key_bytes;
723        assert_eq!(cache.size_bytes().await, expected);
724    }
725
726    #[tokio::test]
727    async fn test_cache_trait_objects() {
728        #[derive(Debug, DeepSizeOf)]
729        struct MyType(i32);
730
731        trait MyTrait: DeepSizeOf + Send + Sync + std::any::Any {
732            fn as_any(&self) -> &dyn std::any::Any;
733        }
734
735        impl MyTrait for MyType {
736            fn as_any(&self) -> &dyn std::any::Any {
737                self
738            }
739        }
740
741        let item: Arc<dyn MyTrait> = Arc::new(MyType(42));
742        let cache = LanceCache::with_capacity(1000);
743        cache
744            .insert_unsized_with_key(&TestUnsizedKey::<dyn MyTrait>::new("test"), item)
745            .await;
746
747        let retrieved = cache
748            .get_unsized_with_key(&TestUnsizedKey::<dyn MyTrait>::new("test"))
749            .await
750            .unwrap();
751        assert_eq!(retrieved.as_any().downcast_ref::<MyType>().unwrap().0, 42);
752    }
753
754    #[tokio::test]
755    async fn test_cache_stats_basic() {
756        let cache = LanceCache::with_capacity(1000);
757        assert_eq!(cache.stats().await.hits, 0);
758
759        // Miss
760        assert!(
761            cache
762                .get_with_key(&TestKey::<Vec<i32>>::new("x"))
763                .await
764                .is_none()
765        );
766        assert_eq!(cache.stats().await.misses, 1);
767
768        // Insert then hit
769        cache
770            .insert_with_key(&TestKey::new("k"), Arc::new(vec![1, 2, 3]))
771            .await;
772        assert!(
773            cache
774                .get_with_key(&TestKey::<Vec<i32>>::new("k"))
775                .await
776                .is_some()
777        );
778        assert_eq!(cache.stats().await.hits, 1);
779    }
780
781    #[tokio::test]
782    async fn test_cache_stats_with_prefixes() {
783        let base = LanceCache::with_capacity(1000);
784        let prefixed = base.with_key_prefix("ns");
785
786        assert!(
787            prefixed
788                .get_with_key(&TestKey::<Vec<i32>>::new("k"))
789                .await
790                .is_none()
791        );
792        assert_eq!(base.stats().await.misses, 1);
793
794        prefixed
795            .insert_with_key(&TestKey::new("k"), Arc::new(vec![1]))
796            .await;
797        assert!(
798            prefixed
799                .get_with_key(&TestKey::<Vec<i32>>::new("k"))
800                .await
801                .is_some()
802        );
803        assert_eq!(base.stats().await.hits, 1);
804    }
805
806    #[tokio::test]
807    async fn test_cache_keys_with_prefixes() {
808        let base = LanceCache::with_capacity(1000);
809        let prefixed = base.with_key_prefix("ns");
810        let nested = prefixed.with_key_prefix("index");
811        let other = base.with_key_prefix("ns-other");
812
813        base.insert_with_key(&TestKey::new("root"), Arc::new(vec![0]))
814            .await;
815        prefixed
816            .insert_with_key(&TestKey::new("child"), Arc::new(vec![1]))
817            .await;
818        nested
819            .insert_with_key(&TestKey::new("nested"), Arc::new(vec![2]))
820            .await;
821        other
822            .insert_with_key(&TestKey::new("other"), Arc::new(vec![3]))
823            .await;
824
825        let base_keys = base.keys().await.unwrap().collect::<Vec<_>>();
826        assert_eq!(
827            key_fields(&base_keys),
828            BTreeSet::from([
829                (
830                    "".to_string(),
831                    "root".to_string(),
832                    TestKey::<Vec<i32>>::type_name()
833                ),
834                (
835                    "ns/".to_string(),
836                    "child".to_string(),
837                    TestKey::<Vec<i32>>::type_name()
838                ),
839                (
840                    "ns/index/".to_string(),
841                    "nested".to_string(),
842                    TestKey::<Vec<i32>>::type_name()
843                ),
844                (
845                    "ns-other/".to_string(),
846                    "other".to_string(),
847                    TestKey::<Vec<i32>>::type_name()
848                ),
849            ])
850        );
851
852        let prefixed_keys = prefixed.keys().await.unwrap().collect::<Vec<_>>();
853        assert_eq!(
854            key_fields(&prefixed_keys),
855            BTreeSet::from([
856                (
857                    "ns/".to_string(),
858                    "child".to_string(),
859                    TestKey::<Vec<i32>>::type_name()
860                ),
861                (
862                    "ns/index/".to_string(),
863                    "nested".to_string(),
864                    TestKey::<Vec<i32>>::type_name()
865                ),
866            ])
867        );
868    }
869
870    #[tokio::test]
871    async fn test_cache_keys_reflect_invalidation_and_clear() {
872        let base = LanceCache::with_capacity(1000);
873        let prefixed = base.with_key_prefix("ns");
874        let other = base.with_key_prefix("other");
875
876        prefixed
877            .insert_with_key(&TestKey::new("child"), Arc::new(vec![1]))
878            .await;
879        other
880            .insert_with_key(&TestKey::new("other"), Arc::new(vec![2]))
881            .await;
882        assert_eq!(base.keys().await.unwrap().count(), 2);
883
884        prefixed.invalidate_prefix("").await;
885        let keys = base.keys().await.unwrap().collect::<Vec<_>>();
886        assert_eq!(
887            key_fields(&keys),
888            BTreeSet::from([(
889                "other/".to_string(),
890                "other".to_string(),
891                TestKey::<Vec<i32>>::type_name()
892            )])
893        );
894
895        base.clear().await;
896        assert_eq!(base.keys().await.unwrap().count(), 0);
897    }
898
899    #[tokio::test]
900    async fn test_cache_get_or_insert() {
901        let cache = LanceCache::with_capacity(1000);
902
903        let v: Arc<Vec<i32>> = cache
904            .get_or_insert_with_key(TestKey::<Vec<i32>>::new("k"), || async {
905                Ok(vec![1, 2, 3])
906            })
907            .await
908            .unwrap();
909        assert_eq!(*v, vec![1, 2, 3]);
910        assert_eq!(cache.stats().await.misses, 1);
911        assert_eq!(cache.stats().await.hits, 0);
912
913        // Second call should not invoke loader and should be a hit
914        let v: Arc<Vec<i32>> = cache
915            .get_or_insert_with_key(TestKey::<Vec<i32>>::new("k"), || async {
916                panic!("should not be called")
917            })
918            .await
919            .unwrap();
920        assert_eq!(*v, vec![1, 2, 3]);
921        assert_eq!(cache.stats().await.hits, 1);
922    }
923
924    #[tokio::test]
925    async fn test_cache_coalesces_concurrent_loader_errors() {
926        let cache = LanceCache::with_capacity(1000);
927        let barrier = Arc::new(tokio::sync::Barrier::new(5));
928        let loader_calls = Arc::new(AtomicUsize::new(0));
929
930        let results = futures::future::join_all((0..5).map(|_| {
931            let cache = cache.clone();
932            let barrier = barrier.clone();
933            let loader_calls = loader_calls.clone();
934            async move {
935                barrier.wait().await;
936                cache
937                    .get_or_insert_with_key(TestKey::<Vec<i32>>::new("error"), || async move {
938                        loader_calls.fetch_add(1, Ordering::Relaxed);
939                        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
940                        Err(Error::timeout("metadata read timed out"))
941                    })
942                    .await
943            }
944        }))
945        .await;
946
947        assert_eq!(loader_calls.load(Ordering::Relaxed), 1);
948        assert!(
949            results
950                .iter()
951                .all(|result| matches!(result, Err(Error::Timeout { .. })))
952        );
953    }
954
955    #[tokio::test]
956    async fn test_custom_backend() {
957        use async_trait::async_trait;
958        use tokio::sync::Mutex;
959
960        #[derive(Debug)]
961        struct HashMapBackend {
962            map: Mutex<HashMap<InternalCacheKey, (CacheEntry, usize)>>,
963        }
964
965        impl HashMapBackend {
966            fn new() -> Self {
967                Self {
968                    map: Mutex::new(HashMap::new()),
969                }
970            }
971        }
972
973        #[async_trait]
974        impl CacheBackend for HashMapBackend {
975            async fn get(
976                &self,
977                key: &InternalCacheKey,
978                _codec: Option<CacheCodec>,
979            ) -> Option<CacheEntry> {
980                self.map.lock().await.get(key).map(|(e, _)| e.clone())
981            }
982            async fn insert(
983                &self,
984                key: &InternalCacheKey,
985                entry: CacheEntry,
986                size_bytes: usize,
987                _codec: Option<CacheCodec>,
988            ) {
989                self.map
990                    .lock()
991                    .await
992                    .insert(key.clone(), (entry, size_bytes));
993            }
994            async fn get_or_insert<'a>(
995                &self,
996                key: &InternalCacheKey,
997                loader: std::pin::Pin<
998                    Box<dyn futures::Future<Output = Result<(CacheEntry, usize)>> + Send + 'a>,
999                >,
1000                _codec: Option<CacheCodec>,
1001            ) -> Result<(CacheEntry, bool)> {
1002                if let Some((entry, _)) = self.map.lock().await.get(key) {
1003                    Ok((entry.clone(), true))
1004                } else {
1005                    let (entry, size) = loader.await?;
1006                    self.map
1007                        .lock()
1008                        .await
1009                        .insert(key.clone(), (entry.clone(), size));
1010                    Ok((entry, false))
1011                }
1012            }
1013            async fn invalidate_prefix(&self, prefix: &str) {
1014                self.map.lock().await.retain(|k, _| !k.starts_with(prefix));
1015            }
1016            async fn clear(&self) {
1017                self.map.lock().await.clear();
1018            }
1019            async fn num_entries(&self) -> usize {
1020                self.map.lock().await.len()
1021            }
1022            async fn size_bytes(&self) -> usize {
1023                self.map.lock().await.values().map(|(_, s)| *s).sum()
1024            }
1025        }
1026
1027        let cache = LanceCache::with_backend(Arc::new(HashMapBackend::new()));
1028
1029        cache
1030            .insert_with_key(&TestKey::new("k"), Arc::new(vec![1, 2, 3]))
1031            .await;
1032        assert!(
1033            cache
1034                .get_with_key(&TestKey::<Vec<i32>>::new("k"))
1035                .await
1036                .is_some()
1037        );
1038        // Different type at same key = miss
1039        assert!(
1040            cache
1041                .get_with_key(&TestKey::<Vec<u8>>::new("k"))
1042                .await
1043                .is_none()
1044        );
1045        assert!(cache.keys().await.is_none());
1046    }
1047
1048    #[tokio::test]
1049    async fn test_get_or_insert_dedup() {
1050        use std::sync::atomic::AtomicUsize;
1051
1052        let load_count = Arc::new(AtomicUsize::new(0));
1053        let cache = LanceCache::with_capacity(10000);
1054
1055        let (barrier_tx, _) = tokio::sync::broadcast::channel::<()>(1);
1056        let mut handles = Vec::new();
1057        for _ in 0..5 {
1058            let cache = cache.clone();
1059            let load_count = load_count.clone();
1060            let mut barrier_rx = barrier_tx.subscribe();
1061            handles.push(tokio::spawn(async move {
1062                barrier_rx.recv().await.ok();
1063                cache
1064                    .get_or_insert_with_key(TestKey::<Vec<i32>>::new("key"), || {
1065                        let load_count = load_count.clone();
1066                        async move {
1067                            load_count.fetch_add(1, Ordering::SeqCst);
1068                            tokio::task::yield_now().await;
1069                            Ok(vec![1, 2, 3])
1070                        }
1071                    })
1072                    .await
1073            }));
1074        }
1075        barrier_tx.send(()).unwrap();
1076        for h in handles {
1077            let result: Arc<Vec<i32>> = h.await.unwrap().unwrap();
1078            assert_eq!(*result, vec![1, 2, 3]);
1079        }
1080
1081        assert_eq!(load_count.load(Ordering::SeqCst), 1);
1082    }
1083}
lance_core/cache/mod.rs

lance_core/cache/
mod.rs
