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

lance_core/cache/
mod.rs
