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 opaque, fixed-size [`InternalCacheKey`] values and
35//! type-erased [`CacheEntry`] values. The typed wrapping is handled by
36//! [`LanceCache`]. See the [`backend`] module for migration 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 key;
52mod moka;
53mod quick;
54
55pub use backend::{CacheBackend, CacheEntry};
56pub use codec::{
57    CacheCodec, CacheCodecImpl, CacheDecode, CacheMissReason, MAGIC, has_cache_envelope,
58};
59pub use entry_io::{CacheEntryReader, CacheEntryWriter};
60pub use key::{CACHE_KEY_FORMAT, CacheKeySchema, CacheNamespace, InternalCacheKey, KeyBuilder};
61pub use moka::MokaCacheBackend;
62pub use quick::{QuickCacheBackend, recommended_cache_shards};
63
64use std::borrow::Cow;
65use std::sync::{
66    Arc, Weak,
67    atomic::{AtomicU64, Ordering},
68};
69
70use futures::Future;
71
72use crate::{Error, Result};
73
74pub use crate::deepsize::{Context, DeepSizeOf};
75
76// ---------------------------------------------------------------------------
77// CacheKey / UnsizedCacheKey — typed key traits for cache users
78// ---------------------------------------------------------------------------
79
80/// Typed cache key for sized value types.
81///
82/// Existing implementations can continue returning a logical string from
83/// [`key`](Self::key). Performance-sensitive implementations should also
84/// provide a stable schema and stream typed fields through
85/// [`write_key`](Self::write_key), avoiding construction of that string.
86///
87/// # Example
88///
89/// ```ignore
90/// struct MyKey { id: u64 }
91///
92/// impl CacheKey for MyKey {
93///     type ValueType = MyData;
94///     fn key(&self) -> Cow<'_, str> { self.id.to_string().into() }
95///     fn type_name() -> &'static str { "MyData" }
96/// }
97/// ```
98pub trait CacheKey {
99    type ValueType: 'static;
100
101    fn key(&self) -> Cow<'_, str>;
102
103    /// Short, stable string identifying this value type.
104    ///
105    /// Two `CacheKey` impls that store different `ValueType`s **must** return
106    /// different type names.
107    ///
108    /// Use a short literal (e.g. `"Vec<IndexMetadata>"`), not
109    /// `std::any::type_name` — the latter is not guaranteed stable across
110    /// compiler versions or build configurations.
111    fn type_name() -> &'static str;
112
113    /// Stable identity included in the physical key.
114    ///
115    /// The compatibility default preserves existing implementations by using
116    /// their author-assigned [`type_name`](Self::type_name).
117    fn stable_type_id() -> &'static str {
118        Self::type_name()
119    }
120
121    /// Versioned schema for the logical key fields.
122    fn schema() -> CacheKeySchema {
123        CacheKeySchema::LEGACY_TEXT
124    }
125
126    /// Stream the logical key fields into the canonical key builder.
127    ///
128    /// The compatibility default hashes the existing string key. In-tree hot
129    /// paths override this with typed, allocation-free field encoding.
130    fn write_key(&self, builder: &mut KeyBuilder) {
131        builder.write_str(self.key().as_ref());
132    }
133
134    /// Optional codec for serializing/deserializing this key's value type.
135    ///
136    /// Returns `None` by default. Cache backends that support persistence
137    /// (e.g. disk-backed caches) use this to serialize entries on insert and
138    /// deserialize on get. Types without a codec will only be stored in-memory.
139    ///
140    /// [`CacheCodec`] is `Copy` (two plain function pointers), so returning it
141    /// by value is cheap — no allocation needed.
142    fn codec() -> Option<CacheCodec> {
143        None
144    }
145}
146
147/// Like [`CacheKey`] but for unsized value types (e.g. `dyn Trait`).
148///
149/// The cache wraps values in an extra `Arc` layer internally; callers pass
150/// and receive `Arc<T>` where `T: ?Sized`.
151///
152/// Unsized cache entries are always in-memory only (no serialization codec).
153/// For serializable entries, use a sized [`CacheKey`] instead.
154pub trait UnsizedCacheKey {
155    type ValueType: 'static + ?Sized;
156
157    fn key(&self) -> Cow<'_, str>;
158
159    /// Short, stable string identifying this value type.
160    /// See [`CacheKey::type_name`] for requirements.
161    fn type_name() -> &'static str;
162
163    /// Stable identity included in the physical key.
164    fn stable_type_id() -> &'static str {
165        Self::type_name()
166    }
167
168    /// Versioned schema for the logical key fields.
169    fn schema() -> CacheKeySchema {
170        CacheKeySchema::LEGACY_TEXT
171    }
172
173    /// Stream the logical key fields into the canonical key builder.
174    fn write_key(&self, builder: &mut KeyBuilder) {
175        builder.write_str(self.key().as_ref());
176    }
177}
178
179// ---------------------------------------------------------------------------
180// Internal helpers
181// ---------------------------------------------------------------------------
182
183/// Size of a cached `Arc<T>`, accounting for the Arc overhead (two atomic counters).
184fn cache_entry_size<T: DeepSizeOf + ?Sized>(value: &T) -> usize {
185    value.deep_size_of() + std::mem::size_of::<std::sync::atomic::AtomicUsize>() * 2
186}
187
188#[derive(Debug)]
189struct CacheState {
190    backend: Arc<dyn CacheBackend>,
191    hits: AtomicU64,
192    misses: AtomicU64,
193}
194
195impl CacheState {
196    fn new(backend: Arc<dyn CacheBackend>) -> Self {
197        Self {
198            backend,
199            hits: AtomicU64::new(0),
200            misses: AtomicU64::new(0),
201        }
202    }
203}
204
205// ---------------------------------------------------------------------------
206// LanceCache — typed wrapper around dyn CacheBackend
207// ---------------------------------------------------------------------------
208
209/// Typed cache wrapper that handles key construction and type safety.
210///
211/// Internally delegates to a [`CacheBackend`]. The default backend is
212/// [`MokaCacheBackend`]; pass a custom backend via [`LanceCache::with_backend`].
213#[derive(Clone)]
214pub struct LanceCache {
215    state: Arc<CacheState>,
216    namespace: key::CacheNamespace,
217}
218
219impl std::fmt::Debug for LanceCache {
220    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
221        f.debug_struct("LanceCache")
222            .field("backend", &self.state.backend)
223            .finish_non_exhaustive()
224    }
225}
226
227impl DeepSizeOf for LanceCache {
228    fn deep_size_of_children(&self, _: &mut Context) -> usize {
229        self.state.backend.approx_size_bytes()
230    }
231}
232
233impl LanceCache {
234    pub fn with_capacity(capacity: usize) -> Self {
235        Self::with_backend(Arc::new(MokaCacheBackend::with_capacity(capacity)))
236    }
237
238    /// Create a cache backed by a custom [`CacheBackend`].
239    pub fn with_backend(backend: Arc<dyn CacheBackend>) -> Self {
240        Self {
241            state: Arc::new(CacheState::new(backend)),
242            namespace: key::CacheNamespace::root(),
243        }
244    }
245
246    pub fn no_cache() -> Self {
247        Self::with_backend(Arc::new(MokaCacheBackend::no_cache()))
248    }
249
250    /// Derive a child namespace for all keys in the returned cache handle.
251    ///
252    /// Each call adds one framed hierarchy segment. Consequently,
253    /// `cache.with_key_prefix("a").with_key_prefix("b")` is deliberately
254    /// distinct from `cache.with_key_prefix("a/b")`.
255    pub fn with_key_prefix(&self, prefix: &str) -> Self {
256        Self {
257            state: self.state.clone(),
258            namespace: self.namespace.child(prefix),
259        }
260    }
261
262    pub async fn size(&self) -> usize {
263        self.state.backend.num_entries().await
264    }
265
266    pub fn approx_size(&self) -> usize {
267        self.state.backend.approx_num_entries()
268    }
269
270    pub async fn size_bytes(&self) -> usize {
271        self.state.backend.size_bytes().await
272    }
273
274    // -- Stats / clear --------------------------------------------------------
275
276    pub async fn stats(&self) -> CacheStats {
277        CacheStats {
278            hits: self.state.hits.load(Ordering::Relaxed),
279            misses: self.state.misses.load(Ordering::Relaxed),
280            num_entries: self.state.backend.num_entries().await,
281            size_bytes: self.state.backend.size_bytes().await,
282        }
283    }
284
285    pub async fn clear(&self) {
286        self.state.backend.clear().await;
287        self.state.hits.store(0, Ordering::Relaxed);
288        self.state.misses.store(0, Ordering::Relaxed);
289    }
290
291    // -- CacheKey-based methods -----------------------------------------------
292
293    pub async fn insert_with_key<K>(&self, cache_key: &K, metadata: Arc<K::ValueType>)
294    where
295        K: CacheKey,
296        K::ValueType: DeepSizeOf + Send + Sync + 'static,
297    {
298        let size = cache_entry_size(metadata.as_ref());
299        let key = self.sized_key(cache_key);
300        self.state
301            .backend
302            .insert(&key, metadata, size, K::codec())
303            .await;
304    }
305
306    pub async fn get_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
307    where
308        K: CacheKey,
309        K::ValueType: DeepSizeOf + Send + Sync + 'static,
310    {
311        let key = self.sized_key(cache_key);
312        let Some(entry) = self.state.backend.get(&key, K::codec()).await else {
313            self.state.misses.fetch_add(1, Ordering::Relaxed);
314            return None;
315        };
316        match entry.downcast::<K::ValueType>() {
317            Ok(value) => {
318                self.state.hits.fetch_add(1, Ordering::Relaxed);
319                Some(value)
320            }
321            Err(_) => {
322                // Type mismatch: the backend returned a different concrete
323                // type than expected (e.g. a disk cache may store
324                // intermediate state). Treat as a miss.
325                log::warn!(
326                    "cache backend returned a value with the wrong concrete type for key type {:?}",
327                    K::stable_type_id()
328                );
329                self.state.misses.fetch_add(1, Ordering::Relaxed);
330                None
331            }
332        }
333    }
334
335    pub async fn get_or_insert_with_key<K, F, Fut>(
336        &self,
337        cache_key: K,
338        loader: F,
339    ) -> Result<Arc<K::ValueType>>
340    where
341        K: CacheKey,
342        K::ValueType: DeepSizeOf + Send + Sync + 'static,
343        F: FnOnce() -> Fut + Send,
344        Fut: Future<Output = Result<K::ValueType>> + Send,
345    {
346        self.get_or_insert_with_key_hit(cache_key, loader)
347            .await
348            .map(|(value, _)| value)
349    }
350
351    /// Same as [`get_or_insert_with_key`](Self::get_or_insert_with_key), but
352    /// also returns a boolean indicating whether the loader was skipped for
353    /// this call.
354    ///
355    /// - `true` means this call did **not** execute the loader. That covers
356    ///   both a true cache hit on an already-populated entry and a coalesced
357    ///   concurrent load where an in-flight loader started by a different
358    ///   caller produced the value.
359    /// - `false` means the loader ran on this call (a real cache miss).
360    ///
361    /// Callers that want strict "served from cache" semantics should treat
362    /// coalesced loads as misses; the current backend does not distinguish the
363    /// two cases. Prefer this over rolling a caller-side `Arc<AtomicBool>`
364    /// when the caller needs per-query hit/miss counters — the backend already
365    /// tracks this bit internally and this method just exposes it.
366    pub async fn get_or_insert_with_key_hit<K, F, Fut>(
367        &self,
368        cache_key: K,
369        loader: F,
370    ) -> Result<(Arc<K::ValueType>, bool)>
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 = self.sized_key(&cache_key);
378        let typed_loader = Box::pin(async move {
379            let value = Arc::new(loader().await?);
380            let size = cache_entry_size(value.as_ref());
381            Ok((value as CacheEntry, size))
382        });
383
384        let (entry, was_cached) = self
385            .state
386            .backend
387            .get_or_insert(&key, typed_loader, K::codec())
388            .await?;
389        let entry = entry.downcast::<K::ValueType>().map_err(|_| {
390            self.state.misses.fetch_add(1, Ordering::Relaxed);
391            Error::io(format!(
392                "cache backend returned a value with the wrong concrete type for key type {:?}",
393                K::stable_type_id()
394            ))
395        })?;
396        if was_cached {
397            self.state.hits.fetch_add(1, Ordering::Relaxed);
398        } else {
399            self.state.misses.fetch_add(1, Ordering::Relaxed);
400        }
401        Ok((entry, was_cached))
402    }
403
404    pub async fn insert_unsized_with_key<K>(&self, cache_key: &K, metadata: Arc<K::ValueType>)
405    where
406        K: UnsizedCacheKey,
407        K::ValueType: DeepSizeOf + Send + Sync + 'static,
408    {
409        let metadata = Arc::new(metadata);
410        let size = cache_entry_size(metadata.as_ref());
411        let key = self.unsized_key(cache_key);
412        self.state.backend.insert(&key, metadata, size, None).await;
413    }
414
415    pub async fn get_unsized_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
416    where
417        K: UnsizedCacheKey,
418        K::ValueType: DeepSizeOf + Send + Sync + 'static,
419    {
420        let key = self.unsized_key(cache_key);
421        let Some(entry) = self.state.backend.get(&key, None).await else {
422            self.state.misses.fetch_add(1, Ordering::Relaxed);
423            return None;
424        };
425        match entry.downcast::<Arc<K::ValueType>>() {
426            Ok(value) => {
427                self.state.hits.fetch_add(1, Ordering::Relaxed);
428                Some(value.as_ref().clone())
429            }
430            Err(_) => {
431                // Type mismatch: the backend returned a different concrete
432                // type than expected (e.g. a disk cache may store
433                // intermediate state). Treat as a miss.
434                log::warn!(
435                    "cache backend returned a value with the wrong concrete type for unsized key type {:?}",
436                    K::stable_type_id()
437                );
438                self.state.misses.fetch_add(1, Ordering::Relaxed);
439                None
440            }
441        }
442    }
443
444    fn sized_key<K: CacheKey>(&self, cache_key: &K) -> InternalCacheKey {
445        let mut builder = KeyBuilder::new(self.namespace, K::stable_type_id(), K::schema());
446        cache_key.write_key(&mut builder);
447        builder.finish()
448    }
449
450    fn unsized_key<K: UnsizedCacheKey>(&self, cache_key: &K) -> InternalCacheKey {
451        let mut builder = KeyBuilder::new(self.namespace, K::stable_type_id(), K::schema());
452        cache_key.write_key(&mut builder);
453        builder.finish()
454    }
455}
456
457// ---------------------------------------------------------------------------
458// WeakLanceCache
459// ---------------------------------------------------------------------------
460
461/// A weak reference to a LanceCache, used by indices to avoid circular references.
462/// When the original cache is dropped, operations on this will gracefully no-op.
463#[derive(Clone, Debug)]
464pub struct WeakLanceCache {
465    state: Weak<CacheState>,
466    namespace: key::CacheNamespace,
467}
468
469impl WeakLanceCache {
470    pub fn from(cache: &LanceCache) -> Self {
471        Self {
472            state: Arc::downgrade(&cache.state),
473            namespace: cache.namespace,
474        }
475    }
476
477    pub fn with_key_prefix(&self, prefix: &str) -> Self {
478        Self {
479            state: self.state.clone(),
480            namespace: self.namespace.child(prefix),
481        }
482    }
483
484    pub async fn get_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
485    where
486        K: CacheKey,
487        K::ValueType: DeepSizeOf + Send + Sync + 'static,
488    {
489        self.upgrade()?.get_with_key(cache_key).await
490    }
491
492    pub async fn insert_with_key<K>(&self, cache_key: &K, value: Arc<K::ValueType>) -> bool
493    where
494        K: CacheKey,
495        K::ValueType: DeepSizeOf + Send + Sync + 'static,
496    {
497        let Some(cache) = self.upgrade() else {
498            log::warn!("WeakLanceCache: cache no longer available, unable to insert item");
499            return false;
500        };
501        cache.insert_with_key(cache_key, value).await;
502        true
503    }
504
505    /// Get or insert an item, computing it if necessary.
506    ///
507    /// Deduplication of concurrent loads is handled by the backend.
508    pub async fn get_or_insert_with_key<K, F, Fut>(
509        &self,
510        cache_key: K,
511        loader: F,
512    ) -> Result<Arc<K::ValueType>>
513    where
514        K: CacheKey,
515        K::ValueType: DeepSizeOf + Send + Sync + 'static,
516        F: FnOnce() -> Fut + Send,
517        Fut: Future<Output = Result<K::ValueType>> + Send,
518    {
519        self.get_or_insert_with_key_hit(cache_key, loader)
520            .await
521            .map(|(value, _)| value)
522    }
523
524    /// Same as [`get_or_insert_with_key`](Self::get_or_insert_with_key), but
525    /// also returns a boolean indicating whether the loader was skipped for
526    /// this call. See [`LanceCache::get_or_insert_with_key_hit`] for the
527    /// coalesced-load caveat.
528    pub async fn get_or_insert_with_key_hit<K, F, Fut>(
529        &self,
530        cache_key: K,
531        loader: F,
532    ) -> Result<(Arc<K::ValueType>, bool)>
533    where
534        K: CacheKey,
535        K::ValueType: DeepSizeOf + Send + Sync + 'static,
536        F: FnOnce() -> Fut + Send,
537        Fut: Future<Output = Result<K::ValueType>> + Send,
538    {
539        let Some(cache) = self.upgrade() else {
540            log::warn!("WeakLanceCache: cache no longer available, computing without caching");
541            return loader().await.map(|value| (Arc::new(value), false));
542        };
543        cache.get_or_insert_with_key_hit(cache_key, loader).await
544    }
545
546    pub async fn get_unsized_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
547    where
548        K: UnsizedCacheKey,
549        K::ValueType: DeepSizeOf + Send + Sync + 'static,
550    {
551        self.upgrade()?.get_unsized_with_key(cache_key).await
552    }
553
554    pub async fn insert_unsized_with_key<K>(&self, cache_key: &K, value: Arc<K::ValueType>)
555    where
556        K: UnsizedCacheKey,
557        K::ValueType: DeepSizeOf + Send + Sync + 'static,
558    {
559        let Some(cache) = self.upgrade() else {
560            log::warn!("WeakLanceCache: cache no longer available, unable to insert unsized item");
561            return;
562        };
563        cache.insert_unsized_with_key(cache_key, value).await;
564    }
565
566    fn upgrade(&self) -> Option<LanceCache> {
567        Some(LanceCache {
568            state: self.state.upgrade()?,
569            namespace: self.namespace,
570        })
571    }
572}
573
574// ---------------------------------------------------------------------------
575// CacheStats
576// ---------------------------------------------------------------------------
577
578#[derive(Debug, Clone)]
579pub struct CacheStats {
580    /// Number of times `get`, `get_unsized`, or `get_or_insert` found an item in the cache.
581    pub hits: u64,
582    /// Number of times `get`, `get_unsized`, or `get_or_insert` did not find an item in the cache.
583    pub misses: u64,
584    /// Number of entries currently in the cache.
585    pub num_entries: usize,
586    /// Total size in bytes of all entries in the cache.
587    pub size_bytes: usize,
588}
589
590impl CacheStats {
591    pub fn hit_ratio(&self) -> f32 {
592        if self.hits + self.misses == 0 {
593            0.0
594        } else {
595            self.hits as f32 / (self.hits + self.misses) as f32
596        }
597    }
598
599    pub fn miss_ratio(&self) -> f32 {
600        if self.hits + self.misses == 0 {
601            0.0
602        } else {
603            self.misses as f32 / (self.hits + self.misses) as f32
604        }
605    }
606}
607
608#[cfg(test)]
609mod tests {
610    use std::collections::HashMap;
611    use std::pin::Pin;
612    use std::sync::atomic::{AtomicUsize, Ordering};
613    use std::task::Poll;
614
615    use super::*;
616
617    async fn report_first_pending<F>(
618        future: F,
619        parked: tokio::sync::oneshot::Sender<()>,
620    ) -> F::Output
621    where
622        F: Future,
623    {
624        tokio::pin!(future);
625        let mut parked = Some(parked);
626        futures::future::poll_fn(|cx| match future.as_mut().poll(cx) {
627            Poll::Pending => {
628                if let Some(parked) = parked.take() {
629                    let _ = parked.send(());
630                }
631                Poll::Pending
632            }
633            Poll::Ready(output) => Poll::Ready(output),
634        })
635        .await
636    }
637
638    #[derive(Clone)]
639    struct VersionedTestKey<const SCHEMA_VERSION: u32> {
640        id: u64,
641    }
642
643    type TestKey = VersionedTestKey<1>;
644    type TestKeyV2 = VersionedTestKey<2>;
645
646    impl<const SCHEMA_VERSION: u32> VersionedTestKey<SCHEMA_VERSION> {
647        fn new(id: u64) -> Self {
648            Self { id }
649        }
650    }
651
652    impl<const SCHEMA_VERSION: u32> CacheKey for VersionedTestKey<SCHEMA_VERSION> {
653        type ValueType = Vec<u32>;
654
655        fn key(&self) -> Cow<'_, str> {
656            self.id.to_string().into()
657        }
658
659        fn type_name() -> &'static str {
660            "test.VecU32"
661        }
662
663        fn schema() -> CacheKeySchema {
664            CacheKeySchema::new("test.vec-u32-key", SCHEMA_VERSION)
665        }
666
667        fn write_key(&self, builder: &mut KeyBuilder) {
668            builder.write_u64(self.id);
669        }
670    }
671
672    struct LegacyBridgeKey(&'static str);
673
674    impl CacheKey for LegacyBridgeKey {
675        type ValueType = Vec<u32>;
676
677        fn key(&self) -> Cow<'_, str> {
678            Cow::Borrowed(self.0)
679        }
680
681        fn type_name() -> &'static str {
682            "test.LegacyBridge"
683        }
684    }
685
686    struct ExplicitBridgeKey(&'static str);
687
688    impl CacheKey for ExplicitBridgeKey {
689        type ValueType = Vec<u32>;
690
691        fn key(&self) -> Cow<'_, str> {
692            Cow::Borrowed(self.0)
693        }
694
695        fn type_name() -> &'static str {
696            "test.LegacyBridge"
697        }
698
699        fn write_key(&self, builder: &mut KeyBuilder) {
700            builder.write_str(self.0);
701        }
702    }
703
704    trait TestDynValue: DeepSizeOf + Send + Sync {
705        fn values(&self) -> &[u32];
706    }
707
708    impl TestDynValue for Vec<u32> {
709        fn values(&self) -> &[u32] {
710            self
711        }
712    }
713
714    struct LegacyUnsizedBridgeKey(&'static str);
715
716    impl UnsizedCacheKey for LegacyUnsizedBridgeKey {
717        type ValueType = dyn TestDynValue;
718
719        fn key(&self) -> Cow<'_, str> {
720            Cow::Borrowed(self.0)
721        }
722
723        fn type_name() -> &'static str {
724            "test.LegacyUnsizedBridge"
725        }
726    }
727
728    struct ExplicitUnsizedBridgeKey(&'static str);
729
730    impl UnsizedCacheKey for ExplicitUnsizedBridgeKey {
731        type ValueType = dyn TestDynValue;
732
733        fn key(&self) -> Cow<'_, str> {
734            Cow::Borrowed(self.0)
735        }
736
737        fn type_name() -> &'static str {
738            "test.LegacyUnsizedBridge"
739        }
740
741        fn write_key(&self, builder: &mut KeyBuilder) {
742            builder.write_str(self.0);
743        }
744    }
745
746    #[derive(Debug, Default)]
747    struct HashMapBackend {
748        entries: tokio::sync::Mutex<HashMap<InternalCacheKey, (CacheEntry, usize)>>,
749    }
750
751    #[async_trait::async_trait]
752    impl CacheBackend for HashMapBackend {
753        async fn get(
754            &self,
755            key: &InternalCacheKey,
756            _codec: Option<CacheCodec>,
757        ) -> Option<CacheEntry> {
758            self.entries
759                .lock()
760                .await
761                .get(key)
762                .map(|(entry, _)| entry.clone())
763        }
764
765        async fn insert(
766            &self,
767            key: &InternalCacheKey,
768            entry: CacheEntry,
769            size_bytes: usize,
770            _codec: Option<CacheCodec>,
771        ) {
772            self.entries.lock().await.insert(*key, (entry, size_bytes));
773        }
774
775        async fn get_or_insert<'a>(
776            &self,
777            key: &InternalCacheKey,
778            loader: Pin<Box<dyn Future<Output = Result<(CacheEntry, usize)>> + Send + 'a>>,
779            codec: Option<CacheCodec>,
780        ) -> Result<(CacheEntry, bool)> {
781            if let Some(entry) = self.get(key, codec).await {
782                return Ok((entry, true));
783            }
784            let (entry, size_bytes) = loader.await?;
785            self.insert(key, entry.clone(), size_bytes, codec).await;
786            Ok((entry, false))
787        }
788
789        async fn clear(&self) {
790            self.entries.lock().await.clear();
791        }
792
793        async fn num_entries(&self) -> usize {
794            self.entries.lock().await.len()
795        }
796
797        async fn size_bytes(&self) -> usize {
798            self.entries
799                .lock()
800                .await
801                .values()
802                .map(|(_, size_bytes)| size_bytes)
803                .sum()
804        }
805    }
806
807    #[derive(Debug)]
808    struct WrongTypeBackend;
809
810    #[async_trait::async_trait]
811    impl CacheBackend for WrongTypeBackend {
812        async fn get(
813            &self,
814            _key: &InternalCacheKey,
815            _codec: Option<CacheCodec>,
816        ) -> Option<CacheEntry> {
817            Some(Arc::new(String::from("wrong type")))
818        }
819
820        async fn insert(
821            &self,
822            _key: &InternalCacheKey,
823            _entry: CacheEntry,
824            _size_bytes: usize,
825            _codec: Option<CacheCodec>,
826        ) {
827        }
828
829        async fn get_or_insert<'a>(
830            &self,
831            _key: &InternalCacheKey,
832            _loader: Pin<Box<dyn Future<Output = Result<(CacheEntry, usize)>> + Send + 'a>>,
833            _codec: Option<CacheCodec>,
834        ) -> Result<(CacheEntry, bool)> {
835            Ok((Arc::new(String::from("wrong type")), true))
836        }
837
838        async fn clear(&self) {}
839
840        async fn num_entries(&self) -> usize {
841            0
842        }
843
844        async fn size_bytes(&self) -> usize {
845            0
846        }
847    }
848
849    #[tokio::test]
850    async fn typed_roundtrip_stats_clear_and_namespace_isolation() {
851        let cache = LanceCache::with_capacity(4096);
852        let left = cache.with_key_prefix("left");
853        let right = cache.with_key_prefix("right");
854        left.insert_with_key(&TestKey::new(7), Arc::new(vec![1, 2, 3]))
855            .await;
856
857        assert_eq!(
858            left.get_with_key(&TestKey::new(7)).await.as_deref(),
859            Some(&vec![1, 2, 3])
860        );
861        assert!(right.get_with_key(&TestKey::new(7)).await.is_none());
862        let stats = cache.stats().await;
863        assert_eq!((stats.hits, stats.misses, stats.num_entries), (1, 1, 1));
864
865        cache.clear().await;
866        let stats = left.stats().await;
867        assert_eq!((stats.hits, stats.misses, stats.num_entries), (0, 0, 0));
868    }
869
870    #[tokio::test]
871    async fn strong_and_weak_handles_share_state_and_namespace() {
872        let cache = LanceCache::with_capacity(4096);
873        let child = cache.with_key_prefix("child");
874        let weak = WeakLanceCache::from(&child);
875
876        assert!(
877            weak.insert_with_key(&TestKey::new(1), Arc::new(vec![1]))
878                .await
879        );
880        assert_eq!(
881            child.get_with_key(&TestKey::new(1)).await.as_deref(),
882            Some(&vec![1])
883        );
884        child
885            .insert_with_key(&TestKey::new(2), Arc::new(vec![2]))
886            .await;
887        assert_eq!(
888            weak.get_with_key(&TestKey::new(2)).await.as_deref(),
889            Some(&vec![2])
890        );
891        assert_eq!((cache.stats().await.hits, cache.size().await), (2, 2));
892    }
893
894    #[tokio::test]
895    async fn nested_namespace_segments_do_not_alias_combined_segments() {
896        let cache = LanceCache::with_capacity(4096);
897        let nested = cache.with_key_prefix("a").with_key_prefix("b");
898        let combined = cache.with_key_prefix("a/b");
899        nested
900            .insert_with_key(&TestKey::new(1), Arc::new(vec![10]))
901            .await;
902        assert!(combined.get_with_key(&TestKey::new(1)).await.is_none());
903    }
904
905    #[tokio::test]
906    async fn schema_change_produces_a_cold_miss() {
907        let cache = LanceCache::with_capacity(4096);
908        cache
909            .insert_with_key(&TestKey::new(1), Arc::new(vec![10]))
910            .await;
911        assert!(cache.get_with_key(&TestKeyV2::new(1)).await.is_none());
912    }
913
914    #[tokio::test]
915    async fn get_or_insert_with_key_hit_reports_loader_execution() {
916        let cache = LanceCache::with_capacity(4096);
917
918        // Cold: loader runs, was_cached = false.
919        let (value, was_cached) = cache
920            .get_or_insert_with_key_hit(TestKey::new(1), || async { Ok(vec![1, 2, 3]) })
921            .await
922            .unwrap();
923        assert_eq!(*value, vec![1, 2, 3]);
924        assert!(!was_cached);
925
926        // Warm: loader must not run and was_cached = true.
927        let (value, was_cached) = cache
928            .get_or_insert_with_key_hit(TestKey::new(1), || async {
929                panic!("should not be called")
930            })
931            .await
932            .unwrap();
933        assert_eq!(*value, vec![1, 2, 3]);
934        assert!(was_cached);
935    }
936
937    #[tokio::test]
938    async fn default_string_bridge_matches_explicit_legacy_encoding() {
939        let cache = LanceCache::with_capacity(4096);
940        cache
941            .insert_with_key(&LegacyBridgeKey("same"), Arc::new(vec![10]))
942            .await;
943        assert_eq!(
944            cache
945                .get_with_key(&ExplicitBridgeKey("same"))
946                .await
947                .as_deref(),
948            Some(&vec![10])
949        );
950    }
951
952    #[tokio::test]
953    async fn unsized_default_string_bridge_matches_explicit_legacy_encoding() {
954        let cache = LanceCache::with_capacity(4096);
955        let value: Arc<dyn TestDynValue> = Arc::new(vec![10, 20]);
956        cache
957            .insert_unsized_with_key(&LegacyUnsizedBridgeKey("same"), value)
958            .await;
959
960        let cached = cache
961            .get_unsized_with_key(&ExplicitUnsizedBridgeKey("same"))
962            .await
963            .unwrap();
964        assert_eq!(cached.values(), &[10, 20]);
965    }
966
967    #[tokio::test]
968    async fn custom_backend_receives_opaque_keys_and_shared_clear() {
969        let backend = Arc::new(HashMapBackend::default());
970        let cache = LanceCache::with_backend(backend.clone());
971        let child = cache.with_key_prefix("child");
972        let value = Arc::new(vec![1, 2, 3]);
973        let value_size = cache_entry_size(value.as_ref());
974
975        child.insert_with_key(&TestKey::new(7), value).await;
976        assert_eq!(
977            child.get_with_key(&TestKey::new(7)).await.as_deref(),
978            Some(&vec![1, 2, 3])
979        );
980        assert_eq!(backend.entries.lock().await.len(), 1);
981        assert_eq!(cache.size_bytes().await, value_size);
982
983        cache.clear().await;
984        assert!(backend.entries.lock().await.is_empty());
985        assert_eq!(child.stats().await.hits, 0);
986    }
987
988    #[tokio::test]
989    async fn backend_type_collisions_are_contextual_misses_or_errors() {
990        let cache = LanceCache::with_backend(Arc::new(WrongTypeBackend));
991
992        assert!(cache.get_with_key(&TestKey::new(1)).await.is_none());
993        let error = cache
994            .get_or_insert_with_key(TestKey::new(2), || async { Ok(vec![2]) })
995            .await
996            .unwrap_err();
997        assert!(error.to_string().contains("test.VecU32"));
998        let stats = cache.stats().await;
999        assert_eq!((stats.hits, stats.misses), (0, 2));
1000    }
1001
1002    #[tokio::test]
1003    async fn moka_weight_includes_the_fixed_physical_key() {
1004        let value = Arc::new(vec![0_u32; 3]);
1005        let expected = cache_entry_size(value.as_ref())
1006            .checked_add(std::mem::size_of::<InternalCacheKey>())
1007            .unwrap();
1008        let cache = LanceCache::with_capacity(expected * 2);
1009        cache.insert_with_key(&TestKey::new(1), value).await;
1010        assert_eq!(cache.size_bytes().await, expected);
1011    }
1012
1013    #[tokio::test]
1014    async fn no_cache_computes_each_time() {
1015        let cache = LanceCache::no_cache();
1016        let loads = Arc::new(AtomicUsize::new(0));
1017        for _ in 0..2 {
1018            let loads = loads.clone();
1019            let value = cache
1020                .get_or_insert_with_key(TestKey::new(1), move || async move {
1021                    loads.fetch_add(1, Ordering::SeqCst);
1022                    Ok(vec![42])
1023                })
1024                .await
1025                .unwrap();
1026            assert_eq!(value.as_slice(), &[42]);
1027        }
1028        assert_eq!(loads.load(Ordering::SeqCst), 2);
1029        assert_eq!(cache.size().await, 0);
1030    }
1031
1032    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1033    async fn single_flight_coalesces_success_after_contenders_are_parked() {
1034        const CONTENDERS: usize = 4;
1035
1036        let cache = Arc::new(LanceCache::with_capacity(4096));
1037        let loader_calls = Arc::new(AtomicUsize::new(0));
1038        let release = Arc::new(tokio::sync::Notify::new());
1039        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1040
1041        let owner = {
1042            let cache = cache.clone();
1043            let loader_calls = loader_calls.clone();
1044            let release = release.clone();
1045            tokio::spawn(async move {
1046                cache
1047                    .get_or_insert_with_key(TestKey::new(10), move || async move {
1048                        loader_calls.fetch_add(1, Ordering::SeqCst);
1049                        let _ = started_tx.send(());
1050                        release.notified().await;
1051                        Ok(vec![10])
1052                    })
1053                    .await
1054            })
1055        };
1056        started_rx.await.unwrap();
1057
1058        let mut contenders = Vec::new();
1059        let mut parked = Vec::new();
1060        for _ in 0..CONTENDERS {
1061            let cache = cache.clone();
1062            let loader_calls = loader_calls.clone();
1063            let (parked_tx, parked_rx) = tokio::sync::oneshot::channel();
1064            parked.push(parked_rx);
1065            contenders.push(tokio::spawn(async move {
1066                report_first_pending(
1067                    cache.get_or_insert_with_key(TestKey::new(10), move || async move {
1068                        loader_calls.fetch_add(1, Ordering::SeqCst);
1069                        Ok(vec![99])
1070                    }),
1071                    parked_tx,
1072                )
1073                .await
1074            }));
1075        }
1076        for parked in parked {
1077            parked
1078                .await
1079                .expect("contender completed instead of parking behind owner");
1080        }
1081        assert_eq!(loader_calls.load(Ordering::SeqCst), 1);
1082        assert!(contenders.iter().all(|handle| !handle.is_finished()));
1083
1084        release.notify_one();
1085        assert_eq!(owner.await.unwrap().unwrap().as_slice(), &[10]);
1086        for contender in contenders {
1087            assert_eq!(contender.await.unwrap().unwrap().as_slice(), &[10]);
1088        }
1089        let stats = cache.stats().await;
1090        assert_eq!((stats.hits, stats.misses), (CONTENDERS as u64, 1),);
1091    }
1092
1093    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1094    async fn single_flight_coalesces_errors_after_contenders_are_parked() {
1095        const CONTENDERS: usize = 4;
1096
1097        let cache = Arc::new(LanceCache::with_capacity(4096));
1098        let loader_calls = Arc::new(AtomicUsize::new(0));
1099        let release = Arc::new(tokio::sync::Notify::new());
1100        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1101
1102        let owner = {
1103            let cache = cache.clone();
1104            let loader_calls = loader_calls.clone();
1105            let release = release.clone();
1106            tokio::spawn(async move {
1107                cache
1108                    .get_or_insert_with_key(TestKey::new(20), move || async move {
1109                        loader_calls.fetch_add(1, Ordering::SeqCst);
1110                        let _ = started_tx.send(());
1111                        release.notified().await;
1112                        Err(Error::timeout("owner loader timed out"))
1113                    })
1114                    .await
1115            })
1116        };
1117        started_rx.await.unwrap();
1118
1119        let mut contenders = Vec::new();
1120        let mut parked = Vec::new();
1121        for _ in 0..CONTENDERS {
1122            let cache = cache.clone();
1123            let loader_calls = loader_calls.clone();
1124            let (parked_tx, parked_rx) = tokio::sync::oneshot::channel();
1125            parked.push(parked_rx);
1126            contenders.push(tokio::spawn(async move {
1127                report_first_pending(
1128                    cache.get_or_insert_with_key(TestKey::new(20), move || async move {
1129                        loader_calls.fetch_add(1, Ordering::SeqCst);
1130                        Err(Error::timeout("contender loader timed out"))
1131                    }),
1132                    parked_tx,
1133                )
1134                .await
1135            }));
1136        }
1137        for parked in parked {
1138            parked
1139                .await
1140                .expect("contender completed instead of parking behind owner");
1141        }
1142        assert_eq!(loader_calls.load(Ordering::SeqCst), 1);
1143        assert!(contenders.iter().all(|handle| !handle.is_finished()));
1144
1145        release.notify_one();
1146        assert!(matches!(owner.await.unwrap(), Err(Error::Timeout { .. })));
1147        for contender in contenders {
1148            assert!(matches!(
1149                contender.await.unwrap(),
1150                Err(Error::Timeout { .. })
1151            ));
1152        }
1153        assert_eq!(loader_calls.load(Ordering::SeqCst), 1);
1154    }
1155
1156    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1157    async fn single_flight_retries_after_the_owner_is_cancelled() {
1158        let cache = Arc::new(LanceCache::with_capacity(4096));
1159        let loader_calls = Arc::new(AtomicUsize::new(0));
1160        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1161
1162        let owner = {
1163            let cache = cache.clone();
1164            let loader_calls = loader_calls.clone();
1165            tokio::spawn(async move {
1166                cache
1167                    .get_or_insert_with_key(TestKey::new(30), move || async move {
1168                        loader_calls.fetch_add(1, Ordering::SeqCst);
1169                        let _ = started_tx.send(());
1170                        std::future::pending::<()>().await;
1171                        Ok(vec![30])
1172                    })
1173                    .await
1174            })
1175        };
1176        started_rx.await.unwrap();
1177
1178        let (parked_tx, parked_rx) = tokio::sync::oneshot::channel();
1179        let contender = {
1180            let cache = cache.clone();
1181            let loader_calls = loader_calls.clone();
1182            tokio::spawn(async move {
1183                report_first_pending(
1184                    cache.get_or_insert_with_key(TestKey::new(30), move || async move {
1185                        loader_calls.fetch_add(1, Ordering::SeqCst);
1186                        Ok(vec![31])
1187                    }),
1188                    parked_tx,
1189                )
1190                .await
1191            })
1192        };
1193        parked_rx
1194            .await
1195            .expect("contender completed instead of parking behind owner");
1196        assert_eq!(loader_calls.load(Ordering::SeqCst), 1);
1197        assert!(!contender.is_finished());
1198
1199        owner.abort();
1200        assert!(owner.await.unwrap_err().is_cancelled());
1201        let value = tokio::time::timeout(std::time::Duration::from_secs(5), contender)
1202            .await
1203            .expect("contender remained parked after owner cancellation")
1204            .unwrap()
1205            .unwrap();
1206        assert_eq!(value.as_slice(), &[31]);
1207        assert_eq!(loader_calls.load(Ordering::SeqCst), 2);
1208    }
1209
1210    #[tokio::test]
1211    async fn expired_weak_cache_degrades_without_retaining_state() {
1212        let cache = LanceCache::with_capacity(4096);
1213        let weak = WeakLanceCache::from(&cache);
1214        drop(cache);
1215
1216        assert!(weak.get_with_key(&TestKey::new(1)).await.is_none());
1217        assert!(
1218            !weak
1219                .insert_with_key(&TestKey::new(1), Arc::new(vec![1]))
1220                .await
1221        );
1222        let value = weak
1223            .get_or_insert_with_key(TestKey::new(1), || async { Ok(vec![7]) })
1224            .await
1225            .unwrap();
1226        assert_eq!(value.as_slice(), &[7]);
1227    }
1228}
lance_core/cache/mod.rs

lance_core/cache/
mod.rs
