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, QuickCacheShardPolicy, 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    /// Weighted capacity in bytes, if the backend reports one.
339    pub fn capacity_bytes(&self) -> Option<usize> {
340        self.state.backend.capacity_bytes()
341    }
342
343    // -- Stats / clear --------------------------------------------------------
344
345    pub async fn stats(&self) -> CacheStats {
346        CacheStats {
347            hits: self.state.hits.load(Ordering::Relaxed),
348            misses: self.state.misses.load(Ordering::Relaxed),
349            num_entries: self.state.backend.num_entries().await,
350            size_bytes: self.state.backend.size_bytes().await,
351        }
352    }
353
354    pub async fn clear(&self) {
355        self.state.backend.clear().await;
356        self.state.hits.store(0, Ordering::Relaxed);
357        self.state.misses.store(0, Ordering::Relaxed);
358    }
359
360    // -- CacheKey-based methods -----------------------------------------------
361
362    pub async fn insert_with_key<K>(&self, cache_key: &K, metadata: Arc<K::ValueType>)
363    where
364        K: CacheKey,
365        K::ValueType: DeepSizeOf + Send + Sync + 'static,
366    {
367        let size = self.state.entry_size(metadata.as_ref());
368        let key = self.sized_key(cache_key);
369        self.state
370            .backend
371            .insert(&key, metadata, size, K::codec())
372            .await;
373    }
374
375    pub async fn get_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
376    where
377        K: CacheKey,
378        K::ValueType: DeepSizeOf + Send + Sync + 'static,
379    {
380        let key = self.sized_key(cache_key);
381        let Some(entry) = self.state.backend.get(&key, K::codec()).await else {
382            self.state.misses.fetch_add(1, Ordering::Relaxed);
383            return None;
384        };
385        match entry.downcast::<K::ValueType>() {
386            Ok(value) => {
387                self.state.hits.fetch_add(1, Ordering::Relaxed);
388                Some(value)
389            }
390            Err(_) => {
391                // Type mismatch: the backend returned a different concrete
392                // type than expected (e.g. a disk cache may store
393                // intermediate state). Treat as a miss.
394                log::warn!(
395                    "cache backend returned a value with the wrong concrete type for key type {:?}",
396                    K::stable_type_id()
397                );
398                self.state.misses.fetch_add(1, Ordering::Relaxed);
399                None
400            }
401        }
402    }
403
404    pub async fn get_or_insert_with_key<K, F, Fut>(
405        &self,
406        cache_key: K,
407        loader: F,
408    ) -> Result<Arc<K::ValueType>>
409    where
410        K: CacheKey,
411        K::ValueType: DeepSizeOf + Send + Sync + 'static,
412        F: FnOnce() -> Fut + Send,
413        Fut: Future<Output = Result<K::ValueType>> + Send,
414    {
415        self.get_or_insert_with_key_hit(cache_key, loader)
416            .await
417            .map(|(value, _)| value)
418    }
419
420    /// Same as [`get_or_insert_with_key`](Self::get_or_insert_with_key), but
421    /// also returns a boolean indicating whether the loader was skipped for
422    /// this call.
423    ///
424    /// - `true` means this call did **not** execute the loader. That covers
425    ///   both a true cache hit on an already-populated entry and a coalesced
426    ///   concurrent load where an in-flight loader started by a different
427    ///   caller produced the value.
428    /// - `false` means the loader ran on this call (a real cache miss).
429    ///
430    /// Callers that want strict "served from cache" semantics should treat
431    /// coalesced loads as misses; the current backend does not distinguish the
432    /// two cases. Prefer this over rolling a caller-side `Arc<AtomicBool>`
433    /// when the caller needs per-query hit/miss counters — the backend already
434    /// tracks this bit internally and this method just exposes it.
435    pub async fn get_or_insert_with_key_hit<K, F, Fut>(
436        &self,
437        cache_key: K,
438        loader: F,
439    ) -> Result<(Arc<K::ValueType>, bool)>
440    where
441        K: CacheKey,
442        K::ValueType: DeepSizeOf + Send + Sync + 'static,
443        F: FnOnce() -> Fut + Send,
444        Fut: Future<Output = Result<K::ValueType>> + Send,
445    {
446        let key = self.sized_key(&cache_key);
447        let state = self.state.clone();
448        let typed_loader = Box::pin(async move {
449            let value = Arc::new(loader().await?);
450            let size = state.entry_size(value.as_ref());
451            Ok((value as CacheEntry, size))
452        });
453
454        let (entry, was_cached) = self
455            .state
456            .backend
457            .get_or_insert(&key, typed_loader, K::codec())
458            .await?;
459        let entry = entry.downcast::<K::ValueType>().map_err(|_| {
460            self.state.misses.fetch_add(1, Ordering::Relaxed);
461            Error::io(format!(
462                "cache backend returned a value with the wrong concrete type for key type {:?}",
463                K::stable_type_id()
464            ))
465        })?;
466        if was_cached {
467            self.state.hits.fetch_add(1, Ordering::Relaxed);
468        } else {
469            self.state.misses.fetch_add(1, Ordering::Relaxed);
470        }
471        Ok((entry, was_cached))
472    }
473
474    pub async fn insert_unsized_with_key<K>(&self, cache_key: &K, metadata: Arc<K::ValueType>)
475    where
476        K: UnsizedCacheKey,
477        K::ValueType: DeepSizeOf + Send + Sync + 'static,
478    {
479        let metadata = Arc::new(metadata);
480        let size = self.state.entry_size(metadata.as_ref());
481        let key = self.unsized_key(cache_key);
482        self.state.backend.insert(&key, metadata, size, None).await;
483    }
484
485    pub async fn get_or_insert_unsized_with_key<K, F, Fut>(
486        &self,
487        cache_key: K,
488        loader: F,
489    ) -> Result<Arc<K::ValueType>>
490    where
491        K: UnsizedCacheKey,
492        K::ValueType: DeepSizeOf + Send + Sync + 'static,
493        F: FnOnce() -> Fut + Send,
494        Fut: Future<Output = Result<Arc<K::ValueType>>> + Send,
495    {
496        let key = self.unsized_key(&cache_key);
497        let state = self.state.clone();
498        let typed_loader = Box::pin(async move {
499            let value = loader().await?;
500            let size = state.entry_size(&value);
501            Ok((Arc::new(value) as CacheEntry, size))
502        });
503
504        let (entry, was_cached) = self
505            .state
506            .backend
507            .get_or_insert(&key, typed_loader, None)
508            .await?;
509        let entry = entry.downcast::<Arc<K::ValueType>>().map_err(|_| {
510            self.state.misses.fetch_add(1, Ordering::Relaxed);
511            Error::io(format!(
512                "cache backend returned a value with the wrong concrete type for unsized key type {:?}",
513                K::stable_type_id()
514            ))
515        })?;
516        if was_cached {
517            self.state.hits.fetch_add(1, Ordering::Relaxed);
518        } else {
519            self.state.misses.fetch_add(1, Ordering::Relaxed);
520        }
521        Ok(entry.as_ref().clone())
522    }
523
524    pub async fn get_unsized_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
525    where
526        K: UnsizedCacheKey,
527        K::ValueType: DeepSizeOf + Send + Sync + 'static,
528    {
529        let key = self.unsized_key(cache_key);
530        let Some(entry) = self.state.backend.get(&key, None).await else {
531            self.state.misses.fetch_add(1, Ordering::Relaxed);
532            return None;
533        };
534        match entry.downcast::<Arc<K::ValueType>>() {
535            Ok(value) => {
536                self.state.hits.fetch_add(1, Ordering::Relaxed);
537                Some(value.as_ref().clone())
538            }
539            Err(_) => {
540                // Type mismatch: the backend returned a different concrete
541                // type than expected (e.g. a disk cache may store
542                // intermediate state). Treat as a miss.
543                log::warn!(
544                    "cache backend returned a value with the wrong concrete type for unsized key type {:?}",
545                    K::stable_type_id()
546                );
547                self.state.misses.fetch_add(1, Ordering::Relaxed);
548                None
549            }
550        }
551    }
552
553    fn sized_key<K: CacheKey>(&self, cache_key: &K) -> InternalCacheKey {
554        let mut builder = KeyBuilder::new(self.namespace, K::stable_type_id(), K::schema());
555        cache_key.write_key(&mut builder);
556        builder.finish()
557    }
558
559    fn unsized_key<K: UnsizedCacheKey>(&self, cache_key: &K) -> InternalCacheKey {
560        let mut builder = KeyBuilder::new(self.namespace, K::stable_type_id(), K::schema());
561        cache_key.write_key(&mut builder);
562        builder.finish()
563    }
564}
565
566// ---------------------------------------------------------------------------
567// WeakLanceCache
568// ---------------------------------------------------------------------------
569
570/// A weak reference to a LanceCache, used by indices to avoid circular references.
571/// When the original cache is dropped, operations on this will gracefully no-op.
572#[derive(Clone, Debug)]
573pub struct WeakLanceCache {
574    state: Weak<CacheState>,
575    namespace: key::CacheNamespace,
576}
577
578impl WeakLanceCache {
579    pub fn from(cache: &LanceCache) -> Self {
580        Self {
581            state: Arc::downgrade(&cache.state),
582            namespace: cache.namespace,
583        }
584    }
585
586    pub fn with_key_prefix(&self, prefix: &str) -> Self {
587        Self {
588            state: self.state.clone(),
589            namespace: self.namespace.child(prefix),
590        }
591    }
592
593    /// Weighted capacity in bytes, if the cache is alive and its backend
594    /// reports one.
595    pub fn capacity_bytes(&self) -> Option<usize> {
596        self.upgrade()?.capacity_bytes()
597    }
598
599    pub async fn get_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
600    where
601        K: CacheKey,
602        K::ValueType: DeepSizeOf + Send + Sync + 'static,
603    {
604        self.upgrade()?.get_with_key(cache_key).await
605    }
606
607    pub async fn insert_with_key<K>(&self, cache_key: &K, value: Arc<K::ValueType>) -> bool
608    where
609        K: CacheKey,
610        K::ValueType: DeepSizeOf + Send + Sync + 'static,
611    {
612        let Some(cache) = self.upgrade() else {
613            log::warn!("WeakLanceCache: cache no longer available, unable to insert item");
614            return false;
615        };
616        cache.insert_with_key(cache_key, value).await;
617        true
618    }
619
620    /// Get or insert an item, computing it if necessary.
621    ///
622    /// Deduplication of concurrent loads is handled by the backend.
623    pub async fn get_or_insert_with_key<K, F, Fut>(
624        &self,
625        cache_key: K,
626        loader: F,
627    ) -> Result<Arc<K::ValueType>>
628    where
629        K: CacheKey,
630        K::ValueType: DeepSizeOf + Send + Sync + 'static,
631        F: FnOnce() -> Fut + Send,
632        Fut: Future<Output = Result<K::ValueType>> + Send,
633    {
634        self.get_or_insert_with_key_hit(cache_key, loader)
635            .await
636            .map(|(value, _)| value)
637    }
638
639    /// Same as [`get_or_insert_with_key`](Self::get_or_insert_with_key), but
640    /// also returns a boolean indicating whether the loader was skipped for
641    /// this call. See [`LanceCache::get_or_insert_with_key_hit`] for the
642    /// coalesced-load caveat.
643    pub async fn get_or_insert_with_key_hit<K, F, Fut>(
644        &self,
645        cache_key: K,
646        loader: F,
647    ) -> Result<(Arc<K::ValueType>, bool)>
648    where
649        K: CacheKey,
650        K::ValueType: DeepSizeOf + Send + Sync + 'static,
651        F: FnOnce() -> Fut + Send,
652        Fut: Future<Output = Result<K::ValueType>> + Send,
653    {
654        let Some(cache) = self.upgrade() else {
655            log::warn!("WeakLanceCache: cache no longer available, computing without caching");
656            return loader().await.map(|value| (Arc::new(value), false));
657        };
658        cache.get_or_insert_with_key_hit(cache_key, loader).await
659    }
660
661    pub async fn get_unsized_with_key<K>(&self, cache_key: &K) -> Option<Arc<K::ValueType>>
662    where
663        K: UnsizedCacheKey,
664        K::ValueType: DeepSizeOf + Send + Sync + 'static,
665    {
666        self.upgrade()?.get_unsized_with_key(cache_key).await
667    }
668
669    pub async fn insert_unsized_with_key<K>(&self, cache_key: &K, value: Arc<K::ValueType>)
670    where
671        K: UnsizedCacheKey,
672        K::ValueType: DeepSizeOf + Send + Sync + 'static,
673    {
674        let Some(cache) = self.upgrade() else {
675            log::warn!("WeakLanceCache: cache no longer available, unable to insert unsized item");
676            return;
677        };
678        cache.insert_unsized_with_key(cache_key, value).await;
679    }
680
681    fn upgrade(&self) -> Option<LanceCache> {
682        Some(LanceCache {
683            state: self.state.upgrade()?,
684            namespace: self.namespace,
685        })
686    }
687}
688
689// ---------------------------------------------------------------------------
690// CacheStats
691// ---------------------------------------------------------------------------
692
693#[derive(Debug, Clone)]
694pub struct CacheStats {
695    /// Number of times `get`, `get_unsized`, or `get_or_insert` found an item in the cache.
696    pub hits: u64,
697    /// Number of times `get`, `get_unsized`, or `get_or_insert` did not find an item in the cache.
698    pub misses: u64,
699    /// Number of entries currently in the cache.
700    pub num_entries: usize,
701    /// Total size in bytes of all entries in the cache.
702    pub size_bytes: usize,
703}
704
705impl CacheStats {
706    pub fn hit_ratio(&self) -> f32 {
707        if self.hits + self.misses == 0 {
708            0.0
709        } else {
710            self.hits as f32 / (self.hits + self.misses) as f32
711        }
712    }
713
714    pub fn miss_ratio(&self) -> f32 {
715        if self.hits + self.misses == 0 {
716            0.0
717        } else {
718            self.misses as f32 / (self.hits + self.misses) as f32
719        }
720    }
721}
722
723#[cfg(test)]
724mod tests {
725    use std::collections::HashMap;
726    use std::pin::Pin;
727    use std::sync::{
728        atomic::{AtomicUsize, Ordering},
729        mpsc,
730    };
731    use std::task::Poll;
732    use std::thread;
733    use std::time::Duration;
734
735    use super::*;
736
737    async fn report_first_pending<F>(
738        future: F,
739        parked: tokio::sync::oneshot::Sender<()>,
740    ) -> F::Output
741    where
742        F: Future,
743    {
744        tokio::pin!(future);
745        let mut parked = Some(parked);
746        futures::future::poll_fn(|cx| match future.as_mut().poll(cx) {
747            Poll::Pending => {
748                if let Some(parked) = parked.take() {
749                    let _ = parked.send(());
750                }
751                Poll::Pending
752            }
753            Poll::Ready(output) => Poll::Ready(output),
754        })
755        .await
756    }
757
758    #[derive(Clone)]
759    struct VersionedTestKey<const SCHEMA_VERSION: u32> {
760        id: u64,
761    }
762
763    type TestKey = VersionedTestKey<1>;
764    type TestKeyV2 = VersionedTestKey<2>;
765
766    impl<const SCHEMA_VERSION: u32> VersionedTestKey<SCHEMA_VERSION> {
767        fn new(id: u64) -> Self {
768            Self { id }
769        }
770    }
771
772    impl<const SCHEMA_VERSION: u32> CacheKey for VersionedTestKey<SCHEMA_VERSION> {
773        type ValueType = Vec<u32>;
774
775        fn key(&self) -> Cow<'_, str> {
776            self.id.to_string().into()
777        }
778
779        fn type_name() -> &'static str {
780            "test.VecU32"
781        }
782
783        fn schema() -> CacheKeySchema {
784            CacheKeySchema::new("test.vec-u32-key", SCHEMA_VERSION)
785        }
786
787        fn write_key(&self, builder: &mut KeyBuilder) {
788            builder.write_u64(self.id);
789        }
790    }
791
792    struct SharedTestValue {
793        data: Arc<Vec<u8>>,
794    }
795
796    impl DeepSizeOf for SharedTestValue {
797        fn deep_size_of_children(&self, context: &mut Context) -> usize {
798            self.data.deep_size_of_children(context)
799        }
800    }
801
802    struct SharedTestKey(u64);
803
804    impl CacheKey for SharedTestKey {
805        type ValueType = SharedTestValue;
806
807        fn key(&self) -> Cow<'_, str> {
808            self.0.to_string().into()
809        }
810
811        fn type_name() -> &'static str {
812            "test.SharedValue"
813        }
814
815        fn schema() -> CacheKeySchema {
816            CacheKeySchema::new("test.shared-value-key", 1)
817        }
818
819        fn write_key(&self, builder: &mut KeyBuilder) {
820            builder.write_u64(self.0);
821        }
822    }
823
824    struct ReentrantValue(LanceCache);
825
826    impl DeepSizeOf for ReentrantValue {
827        fn deep_size_of_children(&self, context: &mut Context) -> usize {
828            self.0.deep_size_of_children(context)
829        }
830    }
831
832    struct ReentrantKey;
833
834    impl CacheKey for ReentrantKey {
835        type ValueType = ReentrantValue;
836
837        fn key(&self) -> Cow<'_, str> {
838            Cow::Borrowed("reentrant")
839        }
840
841        fn type_name() -> &'static str {
842            "test.ReentrantValue"
843        }
844    }
845
846    #[derive(Clone, Copy, Debug)]
847    enum TestBackendKind {
848        Moka,
849        Quick,
850    }
851
852    impl TestBackendKind {
853        fn cache(self, capacity: usize) -> LanceCache {
854            match self {
855                Self::Moka => LanceCache::with_capacity(capacity),
856                Self::Quick => {
857                    LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(capacity)))
858                }
859            }
860        }
861    }
862
863    struct LegacyBridgeKey(&'static str);
864
865    impl CacheKey for LegacyBridgeKey {
866        type ValueType = Vec<u32>;
867
868        fn key(&self) -> Cow<'_, str> {
869            Cow::Borrowed(self.0)
870        }
871
872        fn type_name() -> &'static str {
873            "test.LegacyBridge"
874        }
875    }
876
877    struct ExplicitBridgeKey(&'static str);
878
879    impl CacheKey for ExplicitBridgeKey {
880        type ValueType = Vec<u32>;
881
882        fn key(&self) -> Cow<'_, str> {
883            Cow::Borrowed(self.0)
884        }
885
886        fn type_name() -> &'static str {
887            "test.LegacyBridge"
888        }
889
890        fn write_key(&self, builder: &mut KeyBuilder) {
891            builder.write_str(self.0);
892        }
893    }
894
895    trait TestDynValue: DeepSizeOf + Send + Sync {
896        fn values(&self) -> &[u32];
897    }
898
899    impl TestDynValue for Vec<u32> {
900        fn values(&self) -> &[u32] {
901            self
902        }
903    }
904
905    struct LegacyUnsizedBridgeKey(&'static str);
906
907    impl UnsizedCacheKey for LegacyUnsizedBridgeKey {
908        type ValueType = dyn TestDynValue;
909
910        fn key(&self) -> Cow<'_, str> {
911            Cow::Borrowed(self.0)
912        }
913
914        fn type_name() -> &'static str {
915            "test.LegacyUnsizedBridge"
916        }
917    }
918
919    struct ExplicitUnsizedBridgeKey(&'static str);
920
921    impl UnsizedCacheKey for ExplicitUnsizedBridgeKey {
922        type ValueType = dyn TestDynValue;
923
924        fn key(&self) -> Cow<'_, str> {
925            Cow::Borrowed(self.0)
926        }
927
928        fn type_name() -> &'static str {
929            "test.LegacyUnsizedBridge"
930        }
931
932        fn write_key(&self, builder: &mut KeyBuilder) {
933            builder.write_str(self.0);
934        }
935    }
936
937    #[derive(Debug, Default)]
938    struct HashMapBackend {
939        entries: tokio::sync::Mutex<HashMap<InternalCacheKey, (CacheEntry, usize)>>,
940    }
941
942    #[async_trait::async_trait]
943    impl CacheBackend for HashMapBackend {
944        async fn get(
945            &self,
946            key: &InternalCacheKey,
947            _codec: Option<CacheCodec>,
948        ) -> Option<CacheEntry> {
949            self.entries
950                .lock()
951                .await
952                .get(key)
953                .map(|(entry, _)| entry.clone())
954        }
955
956        async fn insert(
957            &self,
958            key: &InternalCacheKey,
959            entry: CacheEntry,
960            size_bytes: usize,
961            _codec: Option<CacheCodec>,
962        ) {
963            self.entries.lock().await.insert(*key, (entry, size_bytes));
964        }
965
966        async fn get_or_insert<'a>(
967            &self,
968            key: &InternalCacheKey,
969            loader: Pin<Box<dyn Future<Output = Result<(CacheEntry, usize)>> + Send + 'a>>,
970            codec: Option<CacheCodec>,
971        ) -> Result<(CacheEntry, bool)> {
972            if let Some(entry) = self.get(key, codec).await {
973                return Ok((entry, true));
974            }
975            let (entry, size_bytes) = loader.await?;
976            self.insert(key, entry.clone(), size_bytes, codec).await;
977            Ok((entry, false))
978        }
979
980        async fn clear(&self) {
981            self.entries.lock().await.clear();
982        }
983
984        async fn num_entries(&self) -> usize {
985            self.entries.lock().await.len()
986        }
987
988        async fn size_bytes(&self) -> usize {
989            self.entries
990                .lock()
991                .await
992                .values()
993                .map(|(_, size_bytes)| size_bytes)
994                .sum()
995        }
996    }
997
998    #[derive(Debug)]
999    struct WrongTypeBackend;
1000
1001    #[async_trait::async_trait]
1002    impl CacheBackend for WrongTypeBackend {
1003        async fn get(
1004            &self,
1005            _key: &InternalCacheKey,
1006            _codec: Option<CacheCodec>,
1007        ) -> Option<CacheEntry> {
1008            Some(Arc::new(String::from("wrong type")))
1009        }
1010
1011        async fn insert(
1012            &self,
1013            _key: &InternalCacheKey,
1014            _entry: CacheEntry,
1015            _size_bytes: usize,
1016            _codec: Option<CacheCodec>,
1017        ) {
1018        }
1019
1020        async fn get_or_insert<'a>(
1021            &self,
1022            _key: &InternalCacheKey,
1023            _loader: Pin<Box<dyn Future<Output = Result<(CacheEntry, usize)>> + Send + 'a>>,
1024            _codec: Option<CacheCodec>,
1025        ) -> Result<(CacheEntry, bool)> {
1026            Ok((Arc::new(String::from("wrong type")), true))
1027        }
1028
1029        async fn clear(&self) {}
1030
1031        async fn num_entries(&self) -> usize {
1032            0
1033        }
1034
1035        async fn size_bytes(&self) -> usize {
1036            0
1037        }
1038    }
1039
1040    #[tokio::test]
1041    async fn typed_roundtrip_stats_clear_and_namespace_isolation() {
1042        let cache = LanceCache::with_capacity(4096);
1043        let left = cache.with_key_prefix("left");
1044        let right = cache.with_key_prefix("right");
1045        left.insert_with_key(&TestKey::new(7), Arc::new(vec![1, 2, 3]))
1046            .await;
1047
1048        assert_eq!(
1049            left.get_with_key(&TestKey::new(7)).await.as_deref(),
1050            Some(&vec![1, 2, 3])
1051        );
1052        assert!(right.get_with_key(&TestKey::new(7)).await.is_none());
1053        let stats = cache.stats().await;
1054        assert_eq!((stats.hits, stats.misses, stats.num_entries), (1, 1, 1));
1055
1056        cache.clear().await;
1057        let stats = left.stats().await;
1058        assert_eq!((stats.hits, stats.misses, stats.num_entries), (0, 0, 0));
1059    }
1060
1061    #[tokio::test]
1062    async fn strong_and_weak_handles_share_state_and_namespace() {
1063        let cache = LanceCache::with_capacity(4096);
1064        let child = cache.with_key_prefix("child");
1065        let weak = WeakLanceCache::from(&child);
1066
1067        assert!(
1068            weak.insert_with_key(&TestKey::new(1), Arc::new(vec![1]))
1069                .await
1070        );
1071        assert_eq!(
1072            child.get_with_key(&TestKey::new(1)).await.as_deref(),
1073            Some(&vec![1])
1074        );
1075        child
1076            .insert_with_key(&TestKey::new(2), Arc::new(vec![2]))
1077            .await;
1078        assert_eq!(
1079            weak.get_with_key(&TestKey::new(2)).await.as_deref(),
1080            Some(&vec![2])
1081        );
1082        assert_eq!((cache.stats().await.hits, cache.size().await), (2, 2));
1083    }
1084
1085    #[tokio::test]
1086    async fn nested_namespace_segments_do_not_alias_combined_segments() {
1087        let cache = LanceCache::with_capacity(4096);
1088        let nested = cache.with_key_prefix("a").with_key_prefix("b");
1089        let combined = cache.with_key_prefix("a/b");
1090        nested
1091            .insert_with_key(&TestKey::new(1), Arc::new(vec![10]))
1092            .await;
1093        assert!(combined.get_with_key(&TestKey::new(1)).await.is_none());
1094    }
1095
1096    #[tokio::test]
1097    async fn schema_change_produces_a_cold_miss() {
1098        let cache = LanceCache::with_capacity(4096);
1099        cache
1100            .insert_with_key(&TestKey::new(1), Arc::new(vec![10]))
1101            .await;
1102        assert!(cache.get_with_key(&TestKeyV2::new(1)).await.is_none());
1103    }
1104
1105    #[tokio::test]
1106    async fn get_or_insert_with_key_hit_reports_loader_execution() {
1107        let cache = LanceCache::with_capacity(4096);
1108
1109        // Cold: loader runs, was_cached = false.
1110        let (value, was_cached) = cache
1111            .get_or_insert_with_key_hit(TestKey::new(1), || async { Ok(vec![1, 2, 3]) })
1112            .await
1113            .unwrap();
1114        assert_eq!(*value, vec![1, 2, 3]);
1115        assert!(!was_cached);
1116
1117        // Warm: loader must not run and was_cached = true.
1118        let (value, was_cached) = cache
1119            .get_or_insert_with_key_hit(TestKey::new(1), || async {
1120                panic!("should not be called")
1121            })
1122            .await
1123            .unwrap();
1124        assert_eq!(*value, vec![1, 2, 3]);
1125        assert!(was_cached);
1126    }
1127
1128    #[tokio::test]
1129    async fn default_string_bridge_matches_explicit_legacy_encoding() {
1130        let cache = LanceCache::with_capacity(4096);
1131        cache
1132            .insert_with_key(&LegacyBridgeKey("same"), Arc::new(vec![10]))
1133            .await;
1134        assert_eq!(
1135            cache
1136                .get_with_key(&ExplicitBridgeKey("same"))
1137                .await
1138                .as_deref(),
1139            Some(&vec![10])
1140        );
1141    }
1142
1143    #[tokio::test]
1144    async fn unsized_default_string_bridge_matches_explicit_legacy_encoding() {
1145        let cache = LanceCache::with_capacity(4096);
1146        let value: Arc<dyn TestDynValue> = Arc::new(vec![10, 20]);
1147        cache
1148            .insert_unsized_with_key(&LegacyUnsizedBridgeKey("same"), value)
1149            .await;
1150
1151        let cached = cache
1152            .get_unsized_with_key(&ExplicitUnsizedBridgeKey("same"))
1153            .await
1154            .unwrap();
1155        assert_eq!(cached.values(), &[10, 20]);
1156    }
1157
1158    #[tokio::test]
1159    async fn custom_backend_receives_opaque_keys_and_shared_clear() {
1160        let backend = Arc::new(HashMapBackend::default());
1161        let cache = LanceCache::with_backend(backend.clone());
1162        let child = cache.with_key_prefix("child");
1163        let value = Arc::new(vec![1, 2, 3]);
1164        let value_size = cache_entry_size(value.as_ref());
1165
1166        child.insert_with_key(&TestKey::new(7), value).await;
1167        assert_eq!(
1168            child.get_with_key(&TestKey::new(7)).await.as_deref(),
1169            Some(&vec![1, 2, 3])
1170        );
1171        assert_eq!(backend.entries.lock().await.len(), 1);
1172        assert_eq!(cache.size_bytes().await, value_size);
1173
1174        cache.clear().await;
1175        assert!(backend.entries.lock().await.is_empty());
1176        assert_eq!(child.stats().await.hits, 0);
1177    }
1178
1179    #[tokio::test]
1180    async fn backend_type_collisions_are_contextual_misses_or_errors() {
1181        let cache = LanceCache::with_backend(Arc::new(WrongTypeBackend));
1182
1183        assert!(cache.get_with_key(&TestKey::new(1)).await.is_none());
1184        let error = cache
1185            .get_or_insert_with_key(TestKey::new(2), || async { Ok(vec![2]) })
1186            .await
1187            .unwrap_err();
1188        assert!(error.to_string().contains("test.VecU32"));
1189        let stats = cache.stats().await;
1190        assert_eq!((stats.hits, stats.misses), (0, 2));
1191    }
1192
1193    #[tokio::test]
1194    async fn moka_weight_includes_the_fixed_physical_key() {
1195        let value = Arc::new(vec![0_u32; 3]);
1196        let expected = cache_entry_size(value.as_ref())
1197            .checked_add(std::mem::size_of::<InternalCacheKey>())
1198            .unwrap();
1199        let cache = LanceCache::with_capacity(expected * 2);
1200        cache.insert_with_key(&TestKey::new(1), value).await;
1201        assert_eq!(cache.size_bytes().await, expected);
1202    }
1203
1204    #[rstest::rstest]
1205    #[case::moka(TestBackendKind::Moka)]
1206    #[case::quick(TestBackendKind::Quick)]
1207    #[tokio::test]
1208    async fn deep_size_deduplicates_shared_entry_allocations(
1209        #[case] backend_kind: TestBackendKind,
1210    ) {
1211        let cache = backend_kind.cache(1 << 20);
1212        let shared_data = Arc::new(vec![0_u8; 1024]);
1213
1214        for id in 0..2 {
1215            let data = shared_data.clone();
1216            cache
1217                .get_or_insert_with_key(SharedTestKey(id), || async move {
1218                    Ok(SharedTestValue { data })
1219                })
1220                .await
1221                .unwrap();
1222        }
1223
1224        let arc_overhead = std::mem::size_of::<AtomicUsize>() * 2;
1225        let shared_allocation = std::mem::size_of::<Vec<u8>>() + shared_data.capacity();
1226        let expected_entries = 2 * std::mem::size_of::<InternalCacheKey>()
1227            + 2 * (std::mem::size_of::<SharedTestValue>() + arc_overhead)
1228            + shared_allocation;
1229
1230        let weighted_size = cache.size_bytes().await;
1231        assert_eq!(weighted_size, expected_entries + shared_allocation);
1232        assert_eq!(
1233            cache.deep_size_of(),
1234            std::mem::size_of::<LanceCache>() + expected_entries
1235        );
1236
1237        let mut context = Context::new();
1238        assert_eq!(cache.deep_size_of_children(&mut context), expected_entries);
1239        assert_eq!(
1240            cache
1241                .with_key_prefix("another-handle")
1242                .deep_size_of_children(&mut context),
1243            0
1244        );
1245    }
1246
1247    #[test]
1248    fn sizing_can_reenter_the_same_cache() {
1249        let (done_tx, done_rx) = mpsc::channel();
1250        let worker = thread::spawn(move || {
1251            let runtime = tokio::runtime::Builder::new_current_thread()
1252                .enable_all()
1253                .build()
1254                .unwrap();
1255            runtime.block_on(async move {
1256                let cache = LanceCache::with_capacity(4096);
1257                cache
1258                    .insert_with_key(&ReentrantKey, Arc::new(ReentrantValue(cache.clone())))
1259                    .await;
1260                done_tx.send(()).unwrap();
1261            });
1262        });
1263
1264        done_rx
1265            .recv_timeout(Duration::from_secs(1))
1266            .expect("cache insertion deadlocked during sizing");
1267        worker.join().unwrap();
1268    }
1269
1270    #[tokio::test]
1271    async fn no_cache_computes_each_time() {
1272        let cache = LanceCache::no_cache();
1273        let loads = Arc::new(AtomicUsize::new(0));
1274        for _ in 0..2 {
1275            let loads = loads.clone();
1276            let value = cache
1277                .get_or_insert_with_key(TestKey::new(1), move || async move {
1278                    loads.fetch_add(1, Ordering::SeqCst);
1279                    Ok(vec![42])
1280                })
1281                .await
1282                .unwrap();
1283            assert_eq!(value.as_slice(), &[42]);
1284        }
1285        assert_eq!(loads.load(Ordering::SeqCst), 2);
1286        assert_eq!(cache.size().await, 0);
1287    }
1288
1289    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1290    async fn single_flight_coalesces_success_after_contenders_are_parked() {
1291        const CONTENDERS: usize = 4;
1292
1293        let cache = Arc::new(LanceCache::with_capacity(4096));
1294        let loader_calls = Arc::new(AtomicUsize::new(0));
1295        let release = Arc::new(tokio::sync::Notify::new());
1296        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1297
1298        let owner = {
1299            let cache = cache.clone();
1300            let loader_calls = loader_calls.clone();
1301            let release = release.clone();
1302            tokio::spawn(async move {
1303                cache
1304                    .get_or_insert_with_key(TestKey::new(10), move || async move {
1305                        loader_calls.fetch_add(1, Ordering::SeqCst);
1306                        let _ = started_tx.send(());
1307                        release.notified().await;
1308                        Ok(vec![10])
1309                    })
1310                    .await
1311            })
1312        };
1313        started_rx.await.unwrap();
1314
1315        let mut contenders = Vec::new();
1316        let mut parked = Vec::new();
1317        for _ in 0..CONTENDERS {
1318            let cache = cache.clone();
1319            let loader_calls = loader_calls.clone();
1320            let (parked_tx, parked_rx) = tokio::sync::oneshot::channel();
1321            parked.push(parked_rx);
1322            contenders.push(tokio::spawn(async move {
1323                report_first_pending(
1324                    cache.get_or_insert_with_key(TestKey::new(10), move || async move {
1325                        loader_calls.fetch_add(1, Ordering::SeqCst);
1326                        Ok(vec![99])
1327                    }),
1328                    parked_tx,
1329                )
1330                .await
1331            }));
1332        }
1333        for parked in parked {
1334            parked
1335                .await
1336                .expect("contender completed instead of parking behind owner");
1337        }
1338        assert_eq!(loader_calls.load(Ordering::SeqCst), 1);
1339        assert!(contenders.iter().all(|handle| !handle.is_finished()));
1340
1341        release.notify_one();
1342        assert_eq!(owner.await.unwrap().unwrap().as_slice(), &[10]);
1343        for contender in contenders {
1344            assert_eq!(contender.await.unwrap().unwrap().as_slice(), &[10]);
1345        }
1346        let stats = cache.stats().await;
1347        assert_eq!((stats.hits, stats.misses), (CONTENDERS as u64, 1),);
1348    }
1349
1350    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1351    async fn single_flight_coalesces_errors_after_contenders_are_parked() {
1352        const CONTENDERS: usize = 4;
1353
1354        let cache = Arc::new(LanceCache::with_capacity(4096));
1355        let loader_calls = Arc::new(AtomicUsize::new(0));
1356        let release = Arc::new(tokio::sync::Notify::new());
1357        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1358
1359        let owner = {
1360            let cache = cache.clone();
1361            let loader_calls = loader_calls.clone();
1362            let release = release.clone();
1363            tokio::spawn(async move {
1364                cache
1365                    .get_or_insert_with_key(TestKey::new(20), move || async move {
1366                        loader_calls.fetch_add(1, Ordering::SeqCst);
1367                        let _ = started_tx.send(());
1368                        release.notified().await;
1369                        Err(Error::timeout("owner loader timed out"))
1370                    })
1371                    .await
1372            })
1373        };
1374        started_rx.await.unwrap();
1375
1376        let mut contenders = Vec::new();
1377        let mut parked = Vec::new();
1378        for _ in 0..CONTENDERS {
1379            let cache = cache.clone();
1380            let loader_calls = loader_calls.clone();
1381            let (parked_tx, parked_rx) = tokio::sync::oneshot::channel();
1382            parked.push(parked_rx);
1383            contenders.push(tokio::spawn(async move {
1384                report_first_pending(
1385                    cache.get_or_insert_with_key(TestKey::new(20), move || async move {
1386                        loader_calls.fetch_add(1, Ordering::SeqCst);
1387                        Err(Error::timeout("contender loader timed out"))
1388                    }),
1389                    parked_tx,
1390                )
1391                .await
1392            }));
1393        }
1394        for parked in parked {
1395            parked
1396                .await
1397                .expect("contender completed instead of parking behind owner");
1398        }
1399        assert_eq!(loader_calls.load(Ordering::SeqCst), 1);
1400        assert!(contenders.iter().all(|handle| !handle.is_finished()));
1401
1402        release.notify_one();
1403        assert!(matches!(owner.await.unwrap(), Err(Error::Timeout { .. })));
1404        for contender in contenders {
1405            assert!(matches!(
1406                contender.await.unwrap(),
1407                Err(Error::Timeout { .. })
1408            ));
1409        }
1410        assert_eq!(loader_calls.load(Ordering::SeqCst), 1);
1411    }
1412
1413    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1414    async fn single_flight_retries_after_the_owner_is_cancelled() {
1415        let cache = Arc::new(LanceCache::with_capacity(4096));
1416        let loader_calls = Arc::new(AtomicUsize::new(0));
1417        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1418
1419        let owner = {
1420            let cache = cache.clone();
1421            let loader_calls = loader_calls.clone();
1422            tokio::spawn(async move {
1423                cache
1424                    .get_or_insert_with_key(TestKey::new(30), move || async move {
1425                        loader_calls.fetch_add(1, Ordering::SeqCst);
1426                        let _ = started_tx.send(());
1427                        std::future::pending::<()>().await;
1428                        Ok(vec![30])
1429                    })
1430                    .await
1431            })
1432        };
1433        started_rx.await.unwrap();
1434
1435        let (parked_tx, parked_rx) = tokio::sync::oneshot::channel();
1436        let contender = {
1437            let cache = cache.clone();
1438            let loader_calls = loader_calls.clone();
1439            tokio::spawn(async move {
1440                report_first_pending(
1441                    cache.get_or_insert_with_key(TestKey::new(30), move || async move {
1442                        loader_calls.fetch_add(1, Ordering::SeqCst);
1443                        Ok(vec![31])
1444                    }),
1445                    parked_tx,
1446                )
1447                .await
1448            })
1449        };
1450        parked_rx
1451            .await
1452            .expect("contender completed instead of parking behind owner");
1453        assert_eq!(loader_calls.load(Ordering::SeqCst), 1);
1454        assert!(!contender.is_finished());
1455
1456        owner.abort();
1457        assert!(owner.await.unwrap_err().is_cancelled());
1458        let value = tokio::time::timeout(std::time::Duration::from_secs(5), contender)
1459            .await
1460            .expect("contender remained parked after owner cancellation")
1461            .unwrap()
1462            .unwrap();
1463        assert_eq!(value.as_slice(), &[31]);
1464        assert_eq!(loader_calls.load(Ordering::SeqCst), 2);
1465    }
1466
1467    #[tokio::test]
1468    async fn expired_weak_cache_degrades_without_retaining_state() {
1469        let cache = LanceCache::with_capacity(4096);
1470        let weak = WeakLanceCache::from(&cache);
1471        drop(cache);
1472
1473        assert!(weak.get_with_key(&TestKey::new(1)).await.is_none());
1474        assert!(
1475            !weak
1476                .insert_with_key(&TestKey::new(1), Arc::new(vec![1]))
1477                .await
1478        );
1479        let value = weak
1480            .get_or_insert_with_key(TestKey::new(1), || async { Ok(vec![7]) })
1481            .await
1482            .unwrap();
1483        assert_eq!(value.as_slice(), &[7]);
1484    }
1485}
lance_core/cache/mod.rs

lance_core/cache/
mod.rs
