diff --git a/rust/lance-core/src/cache/backend.rs b/rust/lance-core/src/cache/backend.rs index cae4cc119ec..d73f9f72b69 100644 --- a/rust/lance-core/src/cache/backend.rs +++ b/rust/lance-core/src/cache/backend.rs @@ -20,6 +20,16 @@ //! Prefix invalidation and key inventory are intentionally not part of this //! interface: one-way digests cannot support either operation without //! retaining the logical strings that fixed-size keys are designed to remove. +//! +//! # Pinned entries +//! +//! A backend may keep entries in RAM while callers lease them (see the +//! [`pin`](super::pin) module): [`insert_pinned`](CacheBackend::insert_pinned), +//! [`get_or_insert_pinned`](CacheBackend::get_or_insert_pinned) and +//! [`get_leased`](CacheBackend::get_leased). The defaults store such entries +//! like any other, charged and evictable, so existing backends keep working; +//! a pinning backend records each admission with [`CachePin::record`] in a +//! [`PinBudget`](super::PinBudget) and skips pinned entries on eviction. //! Existing callers should migrate removed symbols as follows: //! - replace `with_backend_and_prefix(backend, prefix)` with //! [`LanceCache::with_backend`](super::LanceCache::with_backend) followed by @@ -40,11 +50,17 @@ use futures::Future; use crate::Result; use crate::deepsize::Context; +use super::pin::{CachePin, PinnedStats}; use super::{CacheCodec, InternalCacheKey}; /// A type-erased cache entry. pub type CacheEntry = Arc; +/// The loader [`CacheBackend::get_or_insert_pinned`] receives: the entry, +/// its size in bytes and the pin its leases hold. +pub type PinnedEntryLoader<'a> = + Pin)>> + Send + 'a>>; + /// The tier that would serve a read of an entry, as /// [`CacheBackend::peek_tier`] predicts it. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -162,6 +178,52 @@ pub trait CacheBackend: Send + Sync + std::fmt::Debug { codec: Option, ) -> Result<(CacheEntry, bool)>; + /// Look up RAM only, as an access that refreshes the entry's recency, so + /// that a caller can lease an entry it finds (see + /// [`CachePin::lease`]): never reads a persistent tier and counts no miss. + /// The default uses [`get_resident`](Self::get_resident). + async fn get_leased(&self, key: &InternalCacheKey) -> Option { + self.get_resident(key).await + } + + /// Store a RAM-only entry that stays resident while `pin` is leased and + /// the backend's pin budget has room for it. The default stores it like + /// any entry without a codec: charged, evictable, never pinned. + async fn insert_pinned( + &self, + key: &InternalCacheKey, + entry: CacheEntry, + size_bytes: usize, + _pin: &Arc, + ) { + self.insert(key, entry, size_bytes, None).await + } + + /// Get a RAM-only entry or compute it from `loader`, admitting a loaded + /// entry as [`insert_pinned`](Self::insert_pinned) does. Loads are + /// deduplicated as in [`get_or_insert`](Self::get_or_insert), which the + /// default uses, storing the entry without a pin. + async fn get_or_insert_pinned<'a>( + &self, + key: &InternalCacheKey, + loader: PinnedEntryLoader<'a>, + ) -> Result<(CacheEntry, bool)> { + let loader = Box::pin(async move { loader.await.map(|(entry, size, _)| (entry, size)) }); + self.get_or_insert(key, loader, None).await + } + + /// Size in bytes of the largest entry the backend admits to RAM, or + /// `None` when it sets no limit below its capacity. A caller can keep a + /// value that would take most of it out of the cache. + fn max_entry_bytes(&self) -> Option { + None + } + + /// What the backend's pin budget holds. The default pins nothing. + fn pinned_stats(&self) -> PinnedStats { + PinnedStats::default() + } + /// Remove all entries. async fn clear(&self); diff --git a/rust/lance-core/src/cache/mod.rs b/rust/lance-core/src/cache/mod.rs index b6b22aea548..700f741c049 100644 --- a/rust/lance-core/src/cache/mod.rs +++ b/rust/lance-core/src/cache/mod.rs @@ -51,12 +51,13 @@ pub mod codec; mod entry_io; mod key; mod moka; +pub mod pin; mod priority; mod quick; -pub use priority::PriorityEntries; +pub use priority::{PINNED_PRIORITY, PriorityEntries}; mod registry; -pub use backend::{CacheBackend, CacheEntry, CacheTier}; +pub use backend::{CacheBackend, CacheEntry, CacheTier, PinnedEntryLoader}; pub use backend_uri::{build_from_uri, parse_backend_uri}; pub use codec::{ CacheCodec, CacheCodecImpl, CacheDecode, CacheMissReason, CacheRangeReader, MAGIC, @@ -65,6 +66,10 @@ pub use codec::{ pub use entry_io::{CacheEntryReader, CacheEntryWriter}; pub use key::{CACHE_KEY_FORMAT, CacheKeySchema, CacheNamespace, InternalCacheKey, KeyBuilder}; pub use moka::MokaCacheBackend; +pub use pin::{ + CacheLease, CachePin, PINNED_CAP_FRACTION, PinBudget, PinRecord, PinnedStats, PinnedValue, + pinned_partition_cap, +}; pub use quick::{QuickCacheBackend, QuickCacheShardPolicy, recommended_cache_shards}; pub use registry::{BackendBuildFn, BackendConfig, build_from_config, register_backend}; @@ -72,7 +77,7 @@ use std::any::TypeId; use std::borrow::Cow; use std::collections::HashMap; use std::sync::{ - Arc, RwLock, Weak, + Arc, Mutex, RwLock, Weak, atomic::{AtomicU64, Ordering}, }; @@ -553,6 +558,161 @@ impl LanceCache { Ok((entry, was_cached)) } + // -- Pinned entries ------------------------------------------------------- + + /// Get a RAM-only entry or load it, with a lease that keeps it resident + /// while held, on a backend that pins entries (see [`pin`]). A loaded + /// value is leased before it is admitted, so it enters the cache pinned; + /// a cached value, or one a concurrent caller loaded, is leased as found, + /// and admitted again if it left RAM before the lease. The boolean is + /// `true` when this call did not run the loader, as in + /// [`get_or_insert_with_key_hit`](Self::get_or_insert_with_key_hit). + pub async fn get_or_insert_leased_with_key( + &self, + cache_key: K, + loader: F, + ) -> Result<(Arc, CacheLease, bool)> + where + K: CacheKey, + K::ValueType: PinnedValue + DeepSizeOf + Send + Sync + 'static, + F: FnOnce() -> Fut + Send, + Fut: Future> + Send, + { + let key = self.sized_key(&cache_key); + let state = self.state.clone(); + let preleased: Arc>> = Arc::default(); + let loaded_lease = preleased.clone(); + let typed_loader = Box::pin(async move { + let value = Arc::new(loader().await?); + let pin = value.cache_pin().clone(); + *loaded_lease.lock().unwrap_or_else(|e| e.into_inner()) = Some(CachePin::lease(&pin)); + let size = state.entry_size(value.as_ref()); + Ok((value as CacheEntry, size, pin)) + }); + let (entry, was_cached) = self + .state + .backend + .get_or_insert_pinned(&key, typed_loader) + .await?; + let value = entry.downcast::().map_err(|_| { + self.state.misses.fetch_add(1, Ordering::Relaxed); + Error::io(format!( + "cache backend returned a value with the wrong concrete type for key type {:?}", + K::stable_type_id() + )) + })?; + if was_cached { + self.state.hits.fetch_add(1, Ordering::Relaxed); + } else { + self.state.misses.fetch_add(1, Ordering::Relaxed); + } + let preleased = preleased.lock().unwrap_or_else(|e| e.into_inner()).take(); + let lease = match preleased { + Some(lease) if Arc::ptr_eq(lease.pin(), value.cache_pin()) => lease, + _ => CachePin::lease(value.cache_pin()), + }; + self.ensure_pinned(&key, || value.clone()).await; + Ok((value, lease, was_cached)) + } + + /// Lease an entry RAM holds, found as an access that refreshes its + /// recency: never reads a persistent tier and counts no hit or miss. + pub async fn get_resident_leased_with_key( + &self, + cache_key: &K, + ) -> Option<(Arc, CacheLease)> + where + K: CacheKey, + K::ValueType: PinnedValue + DeepSizeOf + Send + Sync + 'static, + { + let key = self.sized_key(cache_key); + let value = self + .state + .backend + .get_leased(&key) + .await? + .downcast::() + .ok()?; + let lease = CachePin::lease(value.cache_pin()); + Some((value, lease)) + } + + /// Admit a RAM-only value that stays resident while its pin is leased; + /// see [`CacheBackend::insert_pinned`]. + pub async fn insert_pinned_with_key(&self, cache_key: &K, value: Arc) + where + K: CacheKey, + K::ValueType: PinnedValue + DeepSizeOf + Send + Sync + 'static, + { + let key = self.sized_key(cache_key); + let size = self.state.entry_size(value.as_ref()); + let pin = value.cache_pin().clone(); + self.state + .backend + .insert_pinned(&key, value, size, &pin) + .await; + } + + /// Admit the value `make` gives unless RAM holds an entry under the key, + /// as after an eviction, a clear or a refused admission. The check is an + /// access that refreshes the entry's recency, so a caller that leases + /// the entry keeps it through the next eviction pass. Returns whether it + /// admitted the value. + /// + /// The check and the admission are one + /// [`CacheBackend::get_or_insert_pinned`], so an entry a concurrent load + /// is admitting counts as held. A single-flight backend can hand the + /// loaded value to the load's waiters before the entry is in RAM, as + /// quick_cache's guard does; a RAM lookup there misses, and admitting the + /// value again would replace the load's entry with a second admission. + pub async fn ensure_pinned_with_key( + &self, + cache_key: &K, + make: impl FnOnce() -> Arc + Send, + ) -> bool + where + K: CacheKey, + K::ValueType: PinnedValue + DeepSizeOf + Send + Sync + 'static, + { + let key = self.sized_key(cache_key); + self.ensure_pinned(&key, make).await + } + + async fn ensure_pinned( + &self, + key: &InternalCacheKey, + make: impl FnOnce() -> Arc + Send, + ) -> bool + where + V: PinnedValue + DeepSizeOf + Send + Sync + 'static, + { + let state = self.state.clone(); + let admit = Box::pin(async move { + let value = make(); + let size = state.entry_size(value.as_ref()); + let pin = value.cache_pin().clone(); + Ok((value as CacheEntry, size, pin)) + }); + match self.state.backend.get_or_insert_pinned(key, admit).await { + Ok((_, was_cached)) => !was_cached, + Err(error) => { + log::warn!("LanceCache: admitting a pinned entry failed: {error}"); + false + } + } + } + + /// Size of the largest entry the backend admits to RAM; see + /// [`CacheBackend::max_entry_bytes`]. + pub fn max_entry_bytes(&self) -> Option { + self.state.backend.max_entry_bytes() + } + + /// What the backend's pin budget holds; see [`CacheBackend::pinned_stats`]. + pub fn pinned_stats(&self) -> PinnedStats { + self.state.backend.pinned_stats() + } + pub async fn insert_unsized_with_key(&self, cache_key: &K, metadata: Arc) where K: UnsizedCacheKey, @@ -819,6 +979,63 @@ impl WeakLanceCache { cache.insert_unsized_with_key(cache_key, value).await; } + /// [`LanceCache::get_or_insert_leased_with_key`], or `None` without + /// running the loader once the cache is dropped. + pub async fn get_or_insert_leased_with_key( + &self, + cache_key: K, + loader: F, + ) -> Result, CacheLease, bool)>> + where + K: CacheKey, + K::ValueType: PinnedValue + DeepSizeOf + Send + Sync + 'static, + F: FnOnce() -> Fut + Send, + Fut: Future> + Send, + { + let Some(cache) = self.upgrade() else { + return Ok(None); + }; + cache + .get_or_insert_leased_with_key(cache_key, loader) + .await + .map(Some) + } + + /// [`LanceCache::get_resident_leased_with_key`]; a dropped cache holds nothing. + pub async fn get_resident_leased_with_key( + &self, + cache_key: &K, + ) -> Option<(Arc, CacheLease)> + where + K: CacheKey, + K::ValueType: PinnedValue + DeepSizeOf + Send + Sync + 'static, + { + self.upgrade()? + .get_resident_leased_with_key(cache_key) + .await + } + + /// [`LanceCache::ensure_pinned_with_key`]; a dropped cache admits nothing. + pub async fn ensure_pinned_with_key( + &self, + cache_key: &K, + make: impl FnOnce() -> Arc + Send, + ) -> bool + where + K: CacheKey, + K::ValueType: PinnedValue + DeepSizeOf + Send + Sync + 'static, + { + match self.upgrade() { + Some(cache) => cache.ensure_pinned_with_key(cache_key, make).await, + None => false, + } + } + + /// [`LanceCache::max_entry_bytes`]; a dropped cache sets no limit. + pub fn max_entry_bytes(&self) -> Option { + self.upgrade()?.max_entry_bytes() + } + fn upgrade(&self) -> Option { Some(LanceCache { state: self.state.upgrade()?, @@ -1233,6 +1450,210 @@ mod tests { assert!(weak.plane_admission_gated()); } + struct PinnedTestValue(Arc); + + impl DeepSizeOf for PinnedTestValue { + fn deep_size_of_children(&self, _context: &mut Context) -> usize { + 0 + } + } + + impl PinnedValue for PinnedTestValue { + fn cache_pin(&self) -> &Arc { + &self.0 + } + } + + struct PinnedTestKey(u64); + + impl CacheKey for PinnedTestKey { + type ValueType = PinnedTestValue; + + fn key(&self) -> Cow<'_, str> { + self.0.to_string().into() + } + + fn type_name() -> &'static str { + "test.PinnedValue" + } + } + + /// A backend without pinning stores leased values like any entry: + /// charged, never pinned, and with no pin budget or entry limit. One + /// without a RAM lookup has nothing to lease without loading. + #[tokio::test] + async fn default_backend_ignores_pin() { + let hash_map = LanceCache::with_backend(Arc::new(HashMapBackend::default())); + let moka = TestBackendKind::Moka.cache(4096); + for (label, cache) in [("hash map", hash_map), ("moka", moka)] { + let (value, lease, hit) = cache + .get_or_insert_leased_with_key(PinnedTestKey(1), || async { + Ok(PinnedTestValue(CachePin::new())) + }) + .await + .unwrap(); + assert!(!hit && !lease.is_pinned(), "{label}"); + assert_eq!(value.0.holders(), 1, "{label}"); + assert!( + cache.get_with_key(&PinnedTestKey(1)).await.is_some(), + "{label}" + ); + assert_eq!(cache.pinned_stats(), PinnedStats::default(), "{label}"); + assert_eq!(cache.max_entry_bytes(), None, "{label}"); + let resident = cache + .get_resident_leased_with_key(&PinnedTestKey(1)) + .await + .is_some(); + assert_eq!(resident, label == "moka"); + } + } + + /// A RAM backend whose single-flight load hands the loaded value to the + /// callers that join it before admitting the entry, as quick_cache's + /// guard notifies its waiters before it takes the shard lock. The load + /// waits on `admit` in between, and counts its admissions and the + /// entries they replaced. + #[derive(Debug, Default)] + struct PublishFirstBackend { + resident: std::sync::Mutex>, + /// Loaded values handed out but not admitted yet. + published: std::sync::Mutex>, + loaded: tokio::sync::Notify, + admit: tokio::sync::Notify, + admissions: AtomicUsize, + replaced: AtomicUsize, + } + + impl PublishFirstBackend { + fn admit_entry(&self, key: &InternalCacheKey, entry: CacheEntry) { + self.admissions.fetch_add(1, Ordering::SeqCst); + if self.resident.lock().unwrap().insert(*key, entry).is_some() { + self.replaced.fetch_add(1, Ordering::SeqCst); + } + } + } + + #[async_trait::async_trait] + impl CacheBackend for PublishFirstBackend { + async fn get_resident(&self, key: &InternalCacheKey) -> Option { + self.resident.lock().unwrap().get(key).cloned() + } + + async fn get( + &self, + key: &InternalCacheKey, + _codec: Option, + ) -> Option { + self.get_resident(key).await + } + + async fn insert( + &self, + key: &InternalCacheKey, + entry: CacheEntry, + _size_bytes: usize, + _codec: Option, + ) { + self.admit_entry(key, entry); + } + + async fn get_or_insert<'a>( + &self, + key: &InternalCacheKey, + loader: Pin> + Send + 'a>>, + _codec: Option, + ) -> Result<(CacheEntry, bool)> { + if let Some(entry) = self.get_resident(key).await { + return Ok((entry, true)); + } + let (entry, _) = loader.await?; + self.admit_entry(key, entry.clone()); + Ok((entry, false)) + } + + async fn insert_pinned( + &self, + key: &InternalCacheKey, + entry: CacheEntry, + _size_bytes: usize, + _pin: &Arc, + ) { + self.admit_entry(key, entry); + } + + async fn get_or_insert_pinned<'a>( + &self, + key: &InternalCacheKey, + loader: PinnedEntryLoader<'a>, + ) -> Result<(CacheEntry, bool)> { + let held = self.get_resident(key).await; + if let Some(entry) = held.or_else(|| self.published.lock().unwrap().get(key).cloned()) { + return Ok((entry, true)); + } + let (entry, _, _) = loader.await?; + self.published.lock().unwrap().insert(*key, entry.clone()); + self.loaded.notify_one(); + self.admit.notified().await; + self.admit_entry(key, entry.clone()); + self.published.lock().unwrap().remove(key); + Ok((entry, false)) + } + + async fn clear(&self) { + self.resident.lock().unwrap().clear(); + } + + async fn num_entries(&self) -> usize { + self.resident.lock().unwrap().len() + } + + async fn size_bytes(&self) -> usize { + 0 + } + } + + /// A caller that joins a load whose value the backend handed out before + /// admitting it leases the value without admitting it again: the load's + /// admission is the only one and replaces nothing. + #[tokio::test] + async fn leased_join_of_an_admitting_load_admits_once() { + let backend = Arc::new(PublishFirstBackend::default()); + let cache = LanceCache::with_backend(backend.clone()); + let load = tokio::spawn({ + let cache = cache.clone(); + async move { + cache + .get_or_insert_leased_with_key(PinnedTestKey(1), || async { + Ok(PinnedTestValue(CachePin::new())) + }) + .await + .unwrap() + } + }); + backend.loaded.notified().await; + let (joined, joined_lease, joined_hit) = cache + .get_or_insert_leased_with_key(PinnedTestKey(1), || async { + Ok(PinnedTestValue(CachePin::new())) + }) + .await + .unwrap(); + backend.admit.notify_one(); + let (loaded, lease, loaded_hit) = load.await.unwrap(); + + assert!(joined_hit && !loaded_hit); + assert!(Arc::ptr_eq(&joined, &loaded)); + assert_eq!(backend.admissions.load(Ordering::SeqCst), 1); + assert_eq!(backend.replaced.load(Ordering::SeqCst), 0); + assert_eq!(loaded.0.holders(), 2); + drop((joined_lease, lease)); + assert!( + !cache + .ensure_pinned_with_key(&PinnedTestKey(1), || panic!("admitted twice")) + .await + ); + assert_eq!(backend.admissions.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn residency_peek_passes_through() { let cache = LanceCache::with_capacity(4096); diff --git a/rust/lance-core/src/cache/moka.rs b/rust/lance-core/src/cache/moka.rs index 92b32a5534d..d8b68a7b5ef 100644 --- a/rust/lance-core/src/cache/moka.rs +++ b/rust/lance-core/src/cache/moka.rs @@ -50,6 +50,10 @@ fn entry_weight(key: &InternalCacheKey, size_bytes: usize, weight_unit: usize) - /// /// Nonzero capacities provide weighted eviction and concurrent-load deduplication. /// A zero capacity invokes each loader independently without caching. +/// +/// Pinned entries keep the [`CacheBackend`] defaults: moka offers no way to +/// skip an entry on eviction, so a leased entry is charged and evictable +/// like any other, and its users re-admit it when they find it gone. pub struct MokaCacheBackend { cache: moka::future::Cache, capacity: usize, diff --git a/rust/lance-core/src/cache/pin.rs b/rust/lance-core/src/cache/pin.rs new file mode 100644 index 00000000000..65bf3896288 --- /dev/null +++ b/rust/lance-core/src/cache/pin.rs @@ -0,0 +1,648 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Pinned cache entries: values that stay in RAM while their users hold them. +//! +//! A [`CachePin`] belongs to one cached value ([`PinnedValue`]), and every +//! [`CacheLease`] on it is a holder that uses the value. A backend that pins +//! entries records each admission of the value in its [`PinBudget`] +//! ([`PinRecord`]), and does not evict the entry while it is leased and the +//! budget holds its bytes. An entry nobody leases is evictable like any +//! other. +//! +//! The budget caps what pins hold at [`PINNED_CAP_FRACTION`] of each of its +//! partitions (a backend's shards), so eviction always finds unpinned weight +//! to free and pins never push admissions over capacity. A lease that finds +//! no room leaves its entry evictable and counts an overflow; the next lease +//! of an unleased value, its next admission, or the end of another pinned +//! admission of that value in the same budget partition tries again. + +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; + +/// Share of each budget partition that pinned entries may hold. The rest +/// always holds evictable weight, so an admission never finds only pinned +/// entries to evict. +pub const PINNED_CAP_FRACTION: f64 = 0.5; + +/// Bytes the pinned entries of one budget partition may hold, for a +/// partition whose share of the capacity is `share_bytes`: the cap a lease +/// compares an entry's charged bytes with. A backend whose largest +/// admissible entry is its partition's share +/// ([`CacheBackend::max_entry_bytes`](super::CacheBackend::max_entry_bytes)) +/// pins an entry only when it charges at most this much of it. +pub fn pinned_partition_cap(share_bytes: u64) -> u64 { + (share_bytes as f64 * PINNED_CAP_FRACTION) as u64 +} + +/// What a [`PinBudget`] holds, as [`PinBudget::stats`] reports it. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct PinnedStats { + /// Bytes pinned entries may hold at most, over every partition. + pub cap_bytes: u64, + /// Entries the budget holds pinned: admitted, leased and within the cap. + pub pinned_entries: u64, + /// Bytes of those entries. + pub pinned_bytes: u64, + /// Admitted entries with at least one lease, pinned or overflowed. + pub leased_entries: u64, + /// Bytes of those entries. + pub leased_bytes: u64, + /// Leases and admissions of leased entries that found no room under the + /// cap and left their entry evictable. + pub overflow: u64, +} + +#[derive(Debug, Default)] +struct BudgetState { + /// Pinned bytes per partition. + pinned: Vec, + pinned_entries: u64, + leased_entries: u64, + leased_bytes: u64, +} + +/// Bytes a cache backend lets leased entries pin, split into partitions +/// that each cap their pins at [`PINNED_CAP_FRACTION`] of their share of the +/// capacity. A backend whose shards evict independently uses one partition +/// per shard, so that no shard can fill with pinned entries. +#[derive(Debug)] +pub struct PinBudget { + partition_cap: u64, + state: Mutex, + overflow: AtomicU64, +} + +impl PinBudget { + /// A budget over `capacity` bytes split into `partitions` equal shares + /// (at least one). + pub fn new(capacity: u64, partitions: usize) -> Self { + let partitions = partitions.max(1); + let share = capacity / partitions as u64; + Self { + partition_cap: pinned_partition_cap(share), + state: Mutex::new(BudgetState { + pinned: vec![0; partitions], + ..Default::default() + }), + overflow: AtomicU64::new(0), + } + } + + /// What the budget holds now, and the overflows it counted. + pub fn stats(&self) -> PinnedStats { + let state = self.lock(); + PinnedStats { + cap_bytes: self.partition_cap * state.pinned.len() as u64, + pinned_entries: state.pinned_entries, + pinned_bytes: state.pinned.iter().sum(), + leased_entries: state.leased_entries, + leased_bytes: state.leased_bytes, + overflow: self.overflow.load(Ordering::Relaxed), + } + } + + fn lock(&self) -> MutexGuard<'_, BudgetState> { + self.state.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// Reserve `bytes` of `partition`, or count an overflow when they do not + /// fit under its cap. + fn try_reserve(&self, partition: usize, bytes: u64) -> bool { + let mut state = self.lock(); + let partition_cap = self.partition_cap; + let reserved = match state.pinned.get_mut(partition) { + Some(pinned) if pinned.saturating_add(bytes) <= partition_cap => { + *pinned += bytes; + true + } + _ => false, + }; + if reserved { + state.pinned_entries += 1; + } else { + self.overflow.fetch_add(1, Ordering::Relaxed); + } + reserved + } + + fn release(&self, partition: usize, bytes: u64) { + let mut state = self.lock(); + if let Some(pinned) = state.pinned.get_mut(partition) { + *pinned = pinned.saturating_sub(bytes); + } + state.pinned_entries = state.pinned_entries.saturating_sub(1); + } + + fn count_leased(&self, bytes: u64, leased: bool) { + let mut state = self.lock(); + if leased { + state.leased_entries += 1; + state.leased_bytes += bytes; + } else { + state.leased_entries = state.leased_entries.saturating_sub(1); + state.leased_bytes = state.leased_bytes.saturating_sub(bytes); + } + } +} + +/// One admission's charge and eviction protection. The owning pin's state +/// lock serializes updates; its record reads the flag without that lock. +#[derive(Debug)] +struct PinAdmission { + budget: Arc, + partition: usize, + bytes: u64, + pinned: Arc, + counted_leased: bool, +} + +impl PinAdmission { + fn sync(&mut self, is_leased: bool) { + if !is_leased { + self.release(); + return; + } + if !self.counted_leased { + self.budget.count_leased(self.bytes, true); + self.counted_leased = true; + } + if !self.pinned.load(Ordering::Relaxed) { + let reserved = self.budget.try_reserve(self.partition, self.bytes); + self.pinned.store(reserved, Ordering::Release); + } + } + + fn release(&mut self) { + // Make the entry evictable before another admission can reuse its + // budget, so eviction never protects uncharged bytes. + if self.pinned.swap(false, Ordering::Release) { + self.budget.release(self.partition, self.bytes); + } + if self.counted_leased { + self.budget.count_leased(self.bytes, false); + self.counted_leased = false; + } + } +} + +#[derive(Debug, Default)] +struct PinState { + /// Live leases. + holders: usize, + /// Every live cache admission has its own reservation, including aliases + /// of the same value under different keys or in different backends. + admissions: Vec, +} + +/// The pin of one cached value; see the [module docs](self). Leases and +/// admissions share lease ownership, while each admission independently +/// reserves its own cache charge. A value evicted while leased (an overflow) +/// can be pinned again on its next admission if that budget has room. +#[derive(Debug, Default)] +pub struct CachePin { + state: Mutex, + /// Whether any admission is pinned. Eviction uses the admission-specific + /// flag in its `PinRecord`, never this value-wide status. + pinned: AtomicBool, +} + +impl CachePin { + /// A new pin, with neither leases nor admissions. + pub fn new() -> Arc { + Arc::new(Self::default()) + } + + /// Take a lease: the value's entry stays in RAM, on a backend that pins, + /// until the lease and its clones drop. The first lease of an admitted + /// value reserves each admission's bytes independently, or counts an + /// overflow in that admission's budget. + pub fn lease(pin: &Arc) -> CacheLease { + let mut state = pin.lock(); + state.holders += 1; + if state.holders == 1 { + pin.sync(&mut state, true); + } + CacheLease { pin: pin.clone() } + } + + /// Record an admission of the value in `budget`, in `partition`, charged + /// `bytes`. The entry is pinned while the returned record and a lease + /// live, if the budget has room; a leased entry that finds none counts an + /// overflow and stays evictable. Each call creates a separate reservation; + /// admitting the value again never transfers an existing admission's charge. + pub fn record( + pin: &Arc, + budget: &Arc, + partition: usize, + bytes: u64, + ) -> PinRecord { + let pinned = Arc::new(AtomicBool::new(false)); + let mut admission = PinAdmission { + budget: budget.clone(), + partition, + bytes, + pinned: pinned.clone(), + counted_leased: false, + }; + let mut state = pin.lock(); + admission.sync(state.holders > 0); + state.admissions.push(admission); + pin.update_pinned(&state); + PinRecord { + pin: pin.clone(), + pinned, + } + } + + /// Whether at least one admission is leased and within its budget's cap. + /// This does not imply that every admission is protected. Backends must + /// use [`PinRecord::is_pinned`] when deciding whether to evict an entry. + pub fn is_pinned(&self) -> bool { + self.pinned.load(Ordering::Acquire) + } + + /// Live leases on the value. + pub fn holders(&self) -> usize { + self.lock().holders + } + + /// Live admissions of the value: the [`PinRecord`]s its backends hold, + /// one per admission however many copies of the entry share it. A + /// backend dropping an entry whose record is the only one left drops + /// the value from the cache; one replaced by another admission of the + /// same value does not. + pub fn admissions(&self) -> usize { + self.lock().admissions.len() + } + + /// Whether any leased admission found no room under its budget's cap + /// and remains evictable. Other admissions may still be pinned. + pub fn is_overflowed(&self) -> bool { + let state = self.lock(); + state.holders > 0 + && state + .admissions + .iter() + .any(|admission| !admission.pinned.load(Ordering::Relaxed)) + } + + fn lock(&self) -> MutexGuard<'_, PinState> { + self.state.lock().unwrap_or_else(|e| e.into_inner()) + } + + fn release_holder(&self) { + let mut state = self.lock(); + state.holders = state.holders.saturating_sub(1); + if state.holders == 0 { + self.sync(&mut state, false); + } + } + + fn drop_record(&self, pinned: &Arc) { + let mut state = self.lock(); + let Some(index) = state + .admissions + .iter() + .position(|admission| Arc::ptr_eq(&admission.pinned, pinned)) + else { + debug_assert!(false, "a live pin record must have an admission"); + return; + }; + let mut removed = state.admissions.swap_remove(index); + let is_reserved = removed.pinned.load(Ordering::Relaxed); + removed.release(); + if is_reserved && state.holders > 0 { + // A replacement can arrive before the old record drops. Once + // its charge is released, retry the value's remaining admissions + // in that partition without disturbing other budgets. + for admission in &mut state.admissions { + if Arc::ptr_eq(&admission.budget, &removed.budget) + && admission.partition == removed.partition + && !admission.pinned.load(Ordering::Relaxed) + { + admission.sync(true); + } + } + } + self.update_pinned(&state); + } + + fn sync(&self, state: &mut PinState, is_leased: bool) { + for admission in &mut state.admissions { + admission.sync(is_leased); + } + self.update_pinned(state); + } + + fn update_pinned(&self, state: &PinState) { + self.pinned.store( + state + .admissions + .iter() + .any(|admission| admission.pinned.load(Ordering::Relaxed)), + Ordering::Release, + ); + } +} + +/// One admission of a pinned value in a backend; see [`CachePin::record`]. +/// The backend keeps it with the entry, and every copy of the entry shares +/// it, so the admission ends when the last copy leaves the backend: evicted, +/// removed, replaced or cleared. +#[derive(Debug)] +pub struct PinRecord { + pin: Arc, + pinned: Arc, +} + +impl PinRecord { + /// Whether this admission is leased and holds its own budget reservation. + /// Another admission of the same value cannot provide eviction protection + /// for this record. + pub fn is_pinned(&self) -> bool { + self.pinned.load(Ordering::Acquire) + } + + /// The pin of the recorded value. + pub fn pin(&self) -> &Arc { + &self.pin + } +} + +impl Drop for PinRecord { + fn drop(&mut self) { + self.pin.drop_record(&self.pinned); + } +} + +/// A holder of a pinned value; see [`CachePin::lease`]. Clones are holders +/// too; the value's entry is evictable again once every lease dropped. +#[derive(Debug)] +pub struct CacheLease { + pin: Arc, +} + +impl CacheLease { + /// The pin this lease holds. + pub fn pin(&self) -> &Arc { + &self.pin + } + + /// Whether any admission of the leased value is pinned; see + /// [`CachePin::is_pinned`]. + pub fn is_pinned(&self) -> bool { + self.pin.is_pinned() + } +} + +impl Clone for CacheLease { + fn clone(&self) -> Self { + CachePin::lease(&self.pin) + } +} + +impl Drop for CacheLease { + fn drop(&mut self) { + self.pin.release_holder(); + } +} + +/// A cache value leased through a pin of its own, so that every admission +/// of the value is pinned by the same leases; see +/// [`LanceCache::get_or_insert_leased_with_key`](super::LanceCache::get_or_insert_leased_with_key). +pub trait PinnedValue { + /// The value's pin. + fn cache_pin(&self) -> &Arc; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn leased_admission_pins_within_the_cap() { + let budget = Arc::new(PinBudget::new(1000, 1)); + let pin = CachePin::new(); + // A lease before admission pins nothing yet. + let lease = CachePin::lease(&pin); + assert!(!pin.is_pinned()); + let record = CachePin::record(&pin, &budget, 0, 400); + assert!(pin.is_pinned() && lease.is_pinned()); + let stats = budget.stats(); + assert_eq!( + (stats.cap_bytes, stats.pinned_entries, stats.pinned_bytes), + (500, 1, 400) + ); + assert_eq!((stats.leased_entries, stats.leased_bytes), (1, 400)); + // A second holder adds nothing; the last one unpins. + let clone = lease.clone(); + assert_eq!(pin.holders(), 2); + drop(lease); + assert!(pin.is_pinned()); + drop(clone); + assert!(!pin.is_pinned()); + assert_eq!(budget.stats().pinned_bytes, 0); + assert_eq!(budget.stats().leased_entries, 0); + // An admitted entry is pinned by its next first lease. + let lease = CachePin::lease(&pin); + assert!(pin.is_pinned()); + drop(record); + assert!(!lease.is_pinned()); + assert_eq!( + budget.stats(), + PinnedStats { + cap_bytes: 500, + ..Default::default() + } + ); + } + + #[test] + fn overflow_leaves_the_entry_unpinned_and_retries() { + let budget = Arc::new(PinBudget::new(1000, 1)); + let (first, second) = (CachePin::new(), CachePin::new()); + let _first_record = CachePin::record(&first, &budget, 0, 300); + let _second_record = CachePin::record(&second, &budget, 0, 300); + let first_lease = CachePin::lease(&first); + let second_lease = CachePin::lease(&second); + assert!(first.is_pinned()); + assert!(!second.is_pinned() && second.is_overflowed()); + let stats = budget.stats(); + assert_eq!((stats.pinned_bytes, stats.overflow), (300, 1)); + assert_eq!((stats.leased_entries, stats.leased_bytes), (2, 600)); + // Another holder does not retry; the next first lease does. + let extra = second_lease.clone(); + assert_eq!(budget.stats().overflow, 1); + drop(first_lease); + drop((second_lease, extra)); + let _second_lease = CachePin::lease(&second); + assert!(second.is_pinned()); + } + + #[test] + fn budget_partitions_cap_independently() { + let budget = Arc::new(PinBudget::new(1000, 2)); + assert_eq!(budget.stats().cap_bytes, 500); + let pins = [CachePin::new(), CachePin::new(), CachePin::new()]; + let records: Vec<_> = pins + .iter() + .zip([0, 0, 1]) + .map(|(pin, partition)| CachePin::record(pin, &budget, partition, 200)) + .collect(); + let leases: Vec<_> = pins.iter().map(CachePin::lease).collect(); + // 250 bytes per partition: the second entry of partition 0 overflows. + let pinned: Vec<_> = pins.iter().map(|pin| pin.is_pinned()).collect(); + assert_eq!(pinned, [true, false, true]); + drop((records, leases)); + assert_eq!(budget.stats().pinned_bytes, 0); + } + + #[rstest::rstest] + #[case::same_budget_same_charge(true, 100)] + #[case::same_budget_different_charge(true, 200)] + #[case::different_budgets(false, 200)] + fn admissions_keep_their_own_reservations( + #[case] is_same_budget: bool, + #[case] second_bytes: u64, + #[values(false, true)] is_reverse_drop: bool, + ) { + let first = Arc::new(PinBudget::new(1000, 1)); + let second = if is_same_budget { + first.clone() + } else { + Arc::new(PinBudget::new(1000, 1)) + }; + let pin = CachePin::new(); + let lease = CachePin::lease(&pin); + let first_record = CachePin::record(&pin, &first, 0, 100); + let second_record = CachePin::record(&pin, &second, 0, second_bytes); + assert!(first_record.is_pinned() && second_record.is_pinned()); + assert_eq!(pin.admissions(), 2); + assert_eq!( + first.stats().pinned_bytes, + 100 + if is_same_budget { second_bytes } else { 0 } + ); + assert_eq!( + second.stats().pinned_bytes, + second_bytes + if is_same_budget { 100 } else { 0 } + ); + + // All admissions share lease ownership, but release their own charges. + drop(lease); + assert!(!first_record.is_pinned() && !second_record.is_pinned()); + assert_eq!(first.stats().pinned_bytes, 0); + assert_eq!(second.stats().leased_bytes, 0); + let _lease = CachePin::lease(&pin); + assert!(first_record.is_pinned() && second_record.is_pinned()); + + let (removed, kept, removed_budget, kept_budget, kept_bytes) = if is_reverse_drop { + (second_record, first_record, &second, &first, 100) + } else { + (first_record, second_record, &first, &second, second_bytes) + }; + drop(removed); + assert_eq!(pin.admissions(), 1); + assert!(kept.is_pinned() && pin.is_pinned()); + assert_eq!(kept_budget.stats().pinned_bytes, kept_bytes); + assert_eq!(kept_budget.stats().leased_bytes, kept_bytes); + assert_eq!(kept_budget.stats().pinned_entries, 1); + if !is_same_budget { + assert_eq!(removed_budget.stats().pinned_bytes, 0); + assert_eq!(removed_budget.stats().leased_entries, 0); + } + drop(kept); + assert_eq!(pin.admissions(), 0); + assert!(!pin.is_pinned()); + assert_eq!(first.stats().pinned_bytes, 0); + assert_eq!(second.stats().leased_entries, 0); + } + + #[test] + fn admission_overflow_and_retry_stay_in_their_partition() { + let budget = Arc::new(PinBudget::new(1000, 2)); + let pin = CachePin::new(); + let _lease = CachePin::lease(&pin); + let first = CachePin::record(&pin, &budget, 0, 200); + let overflow = CachePin::record(&pin, &budget, 0, 200); + let other_partition = CachePin::record(&pin, &budget, 1, 200); + assert!(first.is_pinned() && other_partition.is_pinned()); + assert!(!overflow.is_pinned()); + assert!(pin.is_pinned() && pin.is_overflowed()); + assert_eq!(budget.stats().pinned_bytes, 400); + assert_eq!(budget.stats().leased_bytes, 600); + + drop(other_partition); + assert!(first.is_pinned()); + assert!( + !overflow.is_pinned(), + "partitions cannot borrow pin capacity" + ); + drop(first); + assert!(overflow.is_pinned()); + assert!(!pin.is_overflowed()); + assert_eq!(budget.stats().pinned_entries, 1); + drop(overflow); + assert_eq!( + budget.stats(), + PinnedStats { + cap_bytes: 500, + overflow: 1, + ..Default::default() + } + ); + } + + /// Admissions count records, not the copies of an entry that share one. + #[test] + fn admissions_count_records() { + let budget = Arc::new(PinBudget::new(1000, 1)); + let pin = CachePin::new(); + let _lease = CachePin::lease(&pin); + assert_eq!(pin.admissions(), 0); + let record = Arc::new(CachePin::record(&pin, &budget, 0, 100)); + let copy = record.clone(); + assert_eq!(pin.admissions(), 1); + assert_eq!(budget.stats().pinned_bytes, 100); + let replacement = CachePin::record(&pin, &budget, 0, 100); + assert_eq!(pin.admissions(), 2); + assert_eq!(budget.stats().pinned_bytes, 200); + drop((record, copy)); + assert_eq!(pin.admissions(), 1); + assert_eq!(budget.stats().pinned_bytes, 100); + drop(replacement); + assert_eq!(pin.admissions(), 0); + assert_eq!(budget.stats().pinned_bytes, 0); + } + + #[test] + fn concurrent_admissions_reserve_and_release_independently() { + let budget = Arc::new(PinBudget::new(1000, 2)); + let pin = CachePin::new(); + let barrier = std::sync::Barrier::new(4); + std::thread::scope(|scope| { + for _ in 0..4 { + scope.spawn(|| { + let _lease = CachePin::lease(&pin); + let record = CachePin::record(&pin, &budget, 0, 100); + barrier.wait(); + assert!(budget.stats().pinned_bytes <= 250); + drop(record); + for _ in 0..32 { + let record = CachePin::record(&pin, &budget, 0, 100); + if record.is_pinned() { + assert!(budget.stats().pinned_bytes >= 100); + } + assert!(budget.stats().pinned_bytes <= 250); + std::thread::yield_now(); + } + }); + } + }); + assert_eq!(pin.admissions(), 0); + assert_eq!(pin.holders(), 0); + assert_eq!(budget.stats().pinned_bytes, 0); + assert_eq!(budget.stats().leased_entries, 0); + assert!(budget.stats().overflow > 0); + } +} diff --git a/rust/lance-core/src/cache/priority.rs b/rust/lance-core/src/cache/priority.rs index 2f45ee51f52..9b964b3b180 100644 --- a/rust/lance-core/src/cache/priority.rs +++ b/rust/lance-core/src/cache/priority.rs @@ -6,6 +6,12 @@ use super::InternalCacheKey; use std::collections::{BTreeMap, HashMap}; +/// Priority of a pinned-kind entry (a value admitted with a +/// [`CachePin`](super::CachePin)): above sign planes (3), so an idle pinned +/// entry is evicted after every plane. While leased within its budget it is +/// never evicted at all. +pub const PINNED_PRIORITY: u8 = 4; + pub struct PriorityEntries { entries: HashMap, order: BTreeMap<(u8, u64), InternalCacheKey>, @@ -65,12 +71,28 @@ impl PriorityEntries { bytes: usize, priority: u8, capacity: usize, + ) -> Vec { + self.insert_with(key, value, bytes, priority, capacity, |_| false) + } + + /// [`insert`](Self::insert) that never evicts an entry `is_pinned` + /// reports pinned, the inserted one included: pinned entries stay even + /// when only they are left over capacity. The pin budget keeps them to a + /// share of the capacity. + pub fn insert_with( + &mut self, + key: InternalCacheKey, + value: T, + bytes: usize, + priority: u8, + capacity: usize, + is_pinned: impl Fn(&T) -> bool, ) -> Vec { let mut dropped = Vec::new(); if let Some(old) = self.remove(&key) { dropped.push(old); } - if bytes > capacity { + if bytes > capacity && !is_pinned(&value) { dropped.push(value); return dropped; } @@ -80,9 +102,19 @@ impl PriorityEntries { self.order.insert((priority, self.sequence), key); self.bytes += bytes as u128; while self.bytes > capacity as u128 { - if let Some(key) = self.order.first_key_value().map(|(_, key)| *key) - && let Some(value) = self.remove(&key) - { + let victim = self + .order + .values() + .find(|key| { + self.entries + .get(key) + .is_some_and(|(_, _, _, value)| !is_pinned(value)) + }) + .copied(); + let Some(victim) = victim else { + break; + }; + if let Some(value) = self.remove(&victim) { dropped.push(value); } } @@ -181,4 +213,43 @@ mod tests { assert_eq!(touched.insert(c, 3, 40, 1, 100), vec![2]); assert!(touched.contains(&a)); } + + /// Eviction passes over pinned entries, however old or low their + /// priority, and a pinned entry heavier than the budget is admitted; only + /// pinned entries can hold the budget over capacity. + #[test] + fn eviction_skips_pinned_entries() { + let key = |id| InternalCacheKey::from_bytes([id; 16]); + let is_pinned = |value: &(u32, bool)| value.1; + let mut cache = PriorityEntries::default(); + cache.insert_with(key(1), (1, true), 40, 1, 100, is_pinned); + cache.insert_with(key(2), (2, false), 40, 3, 100, is_pinned); + // The oldest, lowest entry is pinned: the sign entry goes instead. + let dropped = cache.insert_with(key(3), (3, false), 40, 3, 100, is_pinned); + assert_eq!(dropped, vec![(2, false)]); + assert!(cache.contains(&key(1)) && cache.contains(&key(3))); + // Nothing unpinned is left to evict but the new entry itself. + let dropped = cache.insert_with(key(4), (4, true), 120, 1, 100, is_pinned); + assert_eq!(dropped, vec![(3, false)]); + assert_eq!(cache.bytes(), 160); + assert!(cache.contains(&key(1)) && cache.contains(&key(4))); + // Once unpinned, an entry heavier than the budget is refused. + let dropped = cache.insert_with(key(5), (5, false), 120, 4, 100, is_pinned); + assert_eq!(dropped, vec![(5, false)]); + } + + /// An idle pinned-kind entry outranks sign planes, so a sign plane is + /// evicted first; among pinned-kind entries the oldest goes first. + #[test] + fn pinned_priority_outranks_sign() { + let key = |id| InternalCacheKey::from_bytes([id; 16]); + let mut cache = PriorityEntries::default(); + cache.insert(key(1), 1, 40, PINNED_PRIORITY, 100); + cache.insert(key(2), 2, 40, 3, 100); + assert_eq!(cache.insert(key(3), 3, 40, 3, 100), vec![2]); + assert_eq!(cache.insert(key(4), 4, 40, PINNED_PRIORITY, 100), vec![3]); + assert_eq!(cache.insert(key(5), 5, 40, PINNED_PRIORITY, 100), vec![1]); + // Pinned-kind entries are not plane statistics. + assert_eq!(cache.stats(), [(0, 0); 3]); + } } diff --git a/rust/lance-core/src/cache/quick.rs b/rust/lance-core/src/cache/quick.rs index 20ce7b7766a..9fefee9479a 100644 --- a/rust/lance-core/src/cache/quick.rs +++ b/rust/lance-core/src/cache/quick.rs @@ -4,17 +4,23 @@ //! [`CacheBackend`] backed by [quick_cache](https://crates.io/crates/quick_cache). //! A hit takes a shard read lock, clones the cached value, and marks it as //! accessed. There is no read-operation channel or inline eviction work. +//! +//! Entries admitted with a [`CachePin`] stay while leased: quick_cache skips +//! them on eviction ([`PinLifecycle`]), and the strict priority tier does +//! too. A [`PinBudget`] with one partition per shard caps what pins hold, so +//! a shard always has unpinned weight to evict. -use super::PriorityEntries; +use super::{PINNED_PRIORITY, PriorityEntries}; use std::pin::Pin; -use std::sync::Mutex; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; use async_trait::async_trait; use futures::Future; -use super::backend::{CacheBackend, CacheEntry}; +use super::backend::{CacheBackend, CacheEntry, PinnedEntryLoader}; use super::moka::key_footprint; +use super::pin::{CachePin, PinBudget, PinRecord, PinnedStats}; use super::{CacheCodec, InternalCacheKey}; use crate::Result; use crate::deepsize::Context; @@ -23,6 +29,31 @@ use crate::deepsize::Context; struct QuickEntry { entry: CacheEntry, size_bytes: usize, + /// The admission of a pinned-kind entry in the backend's pin budget, + /// shared by every copy of the entry, so that it lasts while any tier + /// or in-flight load still holds the entry. + pin: Option>, +} + +impl QuickEntry { + fn is_pinned(&self) -> bool { + self.pin.as_ref().is_some_and(|record| record.is_pinned()) + } +} + +/// Keeps pinned entries through quick_cache's eviction passes, which check +/// every candidate as they visit it. +#[derive(Clone, Copy, Debug, Default)] +struct PinLifecycle; + +impl quick_cache::Lifecycle for PinLifecycle { + type RequestState = (); + + fn is_pinned(&self, _key: &InternalCacheKey, value: &QuickEntry) -> bool { + value.is_pinned() + } + + fn begin_request(&self) -> Self::RequestState {} } #[derive(Clone)] @@ -35,12 +66,22 @@ impl quick_cache::Weighter for EntryWeighter { } } +type QuickCache = quick_cache::sync::Cache< + InternalCacheKey, + QuickEntry, + EntryWeighter, + quick_cache::DefaultHashBuilder, + PinLifecycle, +>; + pub struct QuickCacheBackend { capacity: usize, generation: AtomicU64, priority_active: AtomicBool, priority: Mutex>, - cache: quick_cache::sync::Cache, + cache: QuickCache, + /// Caps the bytes leased entries pin, per shard. + pins: Arc, } /// Controls how a [`QuickCacheBackend`] divides its weight budget. @@ -140,18 +181,42 @@ impl QuickCacheBackend { .build() // Only errors when weight/item capacity is missing; both are set. .expect("quick_cache options"); - let cache = quick_cache::sync::Cache::with_options( - options, - EntryWeighter, - Default::default(), - Default::default(), - ); + let cache = + QuickCache::with_options(options, EntryWeighter, Default::default(), PinLifecycle); + let pins = Arc::new(PinBudget::new(capacity as u64, cache.num_shards())); Self { cache, capacity, generation: AtomicU64::new(0), priority_active: AtomicBool::new(false), priority: Mutex::new(PriorityEntries::default()), + pins, + } + } + + /// An admission of a pinned-kind entry, charged what the weigher charges. + fn pin_record( + &self, + key: &InternalCacheKey, + size_bytes: usize, + pin: &Arc, + ) -> Arc { + let bytes = key_footprint(key).saturating_add(size_bytes) as u64; + Arc::new(CachePin::record( + pin, + &self.pins, + self.cache.shard_index(key), + bytes, + )) + } + + /// Priority of `item` in the strict tier: pinned-kind entries outrank + /// every plane, and entries without a plane priority rank with signs. + fn strict_priority(item: &QuickEntry, priority: u8) -> u8 { + match (item.pin.is_some(), priority) { + (true, _) => PINNED_PRIORITY, + (false, 0) => 3, + (false, priority) => priority, } } fn admit_priority( @@ -174,19 +239,38 @@ impl QuickCacheBackend { let existing: Vec<_> = self.cache.iter().collect(); for (key, value) in existing { let size = key_footprint(&key).saturating_add(value.size_bytes); - dropped.extend(entries.insert(key, value, size, 3, self.capacity)); + if value.pin.is_some() { + // Eviction below cannot move a pinned entry, so move + // it here rather than hold it in both tiers. + dropped.extend( + self.cache + .remove_if(&key, |held| held.pin.is_some()) + .map(|(_, held)| held), + ); + } + let priority = Self::strict_priority(&value, 0); + dropped.extend(entries.insert_with( + key, + value, + size, + priority, + self.capacity, + QuickEntry::is_pinned, + )); } // Evict resident values without invalidating single-flight // placeholders: their loaders still belong to this generation. self.cache.set_capacity(0); } let size = key_footprint(&key).saturating_add(item.size_bytes); - dropped.extend(entries.insert( + let priority = Self::strict_priority(&item, priority); + dropped.extend(entries.insert_with( key, item, size, - if priority == 0 { 3 } else { priority }, + priority, self.capacity, + QuickEntry::is_pinned, )); dropped }; @@ -247,7 +331,11 @@ impl CacheBackend for QuickCacheBackend { codec: Option, ) { let priority = codec.map(|c| c.memory_priority()).unwrap_or(0); - let item = QuickEntry { entry, size_bytes }; + let item = QuickEntry { + entry, + size_bytes, + pin: None, + }; if priority > 0 || self.priority_active.load(Ordering::Acquire) { self.admit_priority( *key, @@ -284,6 +372,7 @@ impl CacheBackend for QuickCacheBackend { let item = QuickEntry { entry: entry.clone(), size_bytes, + pin: None, }; if guard.insert(item.clone()).is_ok() && (priority > 0 || self.priority_active.load(Ordering::Acquire)) @@ -295,6 +384,81 @@ impl CacheBackend for QuickCacheBackend { } } + /// A RAM hit refreshes the entry's recency like [`get`](Self::get): a + /// quick_cache hit sets its reference bit, which an eviction pass would + /// otherwise have cleared while the entry was pinned. + async fn get_leased(&self, key: &InternalCacheKey) -> Option { + self.get_resident(key).await + } + + async fn insert_pinned( + &self, + key: &InternalCacheKey, + entry: CacheEntry, + size_bytes: usize, + pin: &Arc, + ) { + let item = QuickEntry { + entry, + size_bytes, + pin: Some(self.pin_record(key, size_bytes, pin)), + }; + if self.priority_active.load(Ordering::Acquire) { + self.admit_priority( + *key, + item, + PINNED_PRIORITY, + self.generation.load(Ordering::Acquire), + ); + } else { + self.cache.insert(*key, item); + } + } + + async fn get_or_insert_pinned<'a>( + &self, + key: &InternalCacheKey, + loader: PinnedEntryLoader<'a>, + ) -> Result<(CacheEntry, bool)> { + if self.priority_active.load(Ordering::Acquire) + && let Some(value) = self + .priority + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(key) + { + return Ok((value.entry, true)); + } + let generation = self.generation.load(Ordering::Acquire); + match self.cache.get_value_or_guard_async(key).await { + Ok(value) => Ok((value.entry, true)), + Err(guard) => { + let (entry, size_bytes, pin) = loader.await?; + let item = QuickEntry { + entry: entry.clone(), + size_bytes, + pin: Some(self.pin_record(key, size_bytes, &pin)), + }; + if guard.insert(item.clone()).is_ok() + && self.priority_active.load(Ordering::Acquire) + { + self.admit_priority(*key, item, PINNED_PRIORITY, generation); + } + Ok((entry, false)) + } + } + } + + /// A shard's budget: quick_cache refuses an unpinned entry heavier than + /// about its shard's weight. + fn max_entry_bytes(&self) -> Option { + Some((self.capacity / self.cache.num_shards()) as u64) + } + + fn pinned_stats(&self) -> PinnedStats { + self.pins.stats() + } + async fn clear(&self) { self.generation.fetch_add(1, Ordering::AcqRel); let dropped = { @@ -385,7 +549,7 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use super::*; - use crate::cache::{CacheKey, CacheTier, LanceCache}; + use crate::cache::{CacheKey, CacheLease, CacheTier, LanceCache}; const TEST_CAPACITY: usize = 1_000; @@ -579,12 +743,473 @@ mod tests { } } + /// A value leased through its own pin, charged `bytes` beyond its struct. + struct Pinned { + pin: Arc, + bytes: usize, + } + + impl crate::deepsize::DeepSizeOf for Pinned { + fn deep_size_of_children(&self, _context: &mut Context) -> usize { + self.bytes + } + } + + impl crate::cache::PinnedValue for Pinned { + fn cache_pin(&self) -> &Arc { + &self.pin + } + } + + /// Capacity of the caches the pinning tests churn: one shard, as the + /// benchmark's caches below 4 GiB have. + const PIN_TEST_CAPACITY: usize = 1 << 20; + /// Bytes of each churned entry, as a small plane. + const CHURN_ENTRY_BYTES: usize = 4096; + + fn single_shard_cache() -> (Arc, LanceCache) { + let backend = Arc::new(QuickCacheBackend::with_shard_policy( + PIN_TEST_CAPACITY, + QuickCacheShardPolicy::Single, + )); + assert_eq!(backend.cache.num_shards(), 1); + assert_eq!(backend.capacity_bytes(), Some(PIN_TEST_CAPACITY)); + (backend.clone(), LanceCache::with_backend(backend)) + } + + /// Admit `bytes` of entries under `prefix`, reading each back so that it + /// is promoted to the hot ring and pushes older hot entries out, as the + /// planes queries read do. + async fn churn(cache: &LanceCache, prefix: &str, bytes: usize) { + for i in 0..bytes / CHURN_ENTRY_BYTES { + let key = TestKey::>::new(&format!("{prefix}-{i}")); + cache + .insert_with_key(&key, Arc::new(vec![0u8; CHURN_ENTRY_BYTES])) + .await; + cache.get_with_key(&key).await; + } + } + + /// Load `name`, a pinned value of `bytes`, and lease it. + async fn leased(cache: &LanceCache, name: &str, bytes: usize) -> (Arc, CacheLease) { + let (value, lease, _) = cache + .get_or_insert_leased_with_key(TestKey::::new(name), || async move { + Ok(Pinned { + pin: CachePin::new(), + bytes, + }) + }) + .await + .unwrap(); + (value, lease) + } + + async fn is_resident(cache: &LanceCache, name: &str) -> bool { + cache + .peek_resident_with_key(&TestKey::::new(name)) + .await + } + + /// A store of 30% of a single-shard cache stays through ten times the + /// capacity of plane churn while leased, with planes and the store + /// within the capacity; unleased, the churn evicts it. + #[tokio::test] + async fn pinned_entry_survives_pressure_while_leased() { + let (_, cache) = single_shard_cache(); + let store_bytes = PIN_TEST_CAPACITY * 3 / 10; + let (_store, lease) = leased(&cache, "store", store_bytes).await; + assert!(lease.is_pinned()); + for round in 0..10 { + churn(&cache, &format!("planes-{round}"), PIN_TEST_CAPACITY).await; + assert!(is_resident(&cache, "store").await, "round {round}"); + assert!( + cache.size_bytes().await <= PIN_TEST_CAPACITY, + "round {round}" + ); + } + let stats = cache.pinned_stats(); + assert_eq!((stats.pinned_entries, stats.overflow), (1, 0), "{stats:?}"); + assert!(stats.pinned_bytes >= store_bytes as u64, "{stats:?}"); + + drop(lease); + churn(&cache, "idle", 10 * PIN_TEST_CAPACITY).await; + assert!(!is_resident(&cache, "store").await); + assert_eq!(cache.pinned_stats().pinned_bytes, 0); + } + + /// Entries admitted once the hot ring holds its target enter the cold + /// ring, where churn evicts an idle one first; a leased one stays. + #[tokio::test] + async fn pinned_entry_admitted_leased_above_hot_target() { + let (backend, cache) = single_shard_cache(); + for i in 0..PIN_TEST_CAPACITY / CHURN_ENTRY_BYTES { + let key = TestKey::>::new(&format!("fill-{i}")); + cache + .insert_with_key(&key, Arc::new(vec![0u8; CHURN_ENTRY_BYTES])) + .await; + } + let (_store, lease) = leased(&cache, "store", PIN_TEST_CAPACITY / 5).await; + assert!(lease.is_pinned()); + let idle = Arc::new(Pinned { + pin: CachePin::new(), + bytes: PIN_TEST_CAPACITY / 5, + }); + cache + .insert_pinned_with_key(&TestKey::::new("idle"), idle) + .await; + assert!(is_resident(&cache, "idle").await); + churn(&cache, "planes", 10 * PIN_TEST_CAPACITY).await; + assert!(is_resident(&cache, "store").await); + assert!(!is_resident(&cache, "idle").await); + assert!(backend.size_bytes().await <= PIN_TEST_CAPACITY); + } + + /// A leased lookup is an access, as a get is: an idle entry found + /// through it keeps the reference bit a peek leaves clear, so it + /// survives the eviction pass that takes a peeked entry. + #[tokio::test] + async fn leased_get_refreshes_referenced() { + const ENTRY_BYTES: usize = 100; + const HOT_ENTRIES: u8 = 9; + let key = |id| InternalCacheKey::from_bytes([id; 16]); + for lease_with_get in [false, true] { + let entry_weight = key_footprint(&key(0)) + ENTRY_BYTES; + let cache = + QuickCacheBackend::with_capacity(entry_weight * (usize::from(HOT_ENTRIES) + 1)); + for id in 0..HOT_ENTRIES { + cache + .insert(&key(id), Arc::new(id), ENTRY_BYTES, None) + .await; + } + // The only cold entry, admitted leased and then left idle. + let victim = key(HOT_ENTRIES); + let pin = CachePin::new(); + let lease = CachePin::lease(&pin); + cache + .insert_pinned(&victim, Arc::new(0u8), ENTRY_BYTES, &pin) + .await; + assert!(lease.is_pinned()); + drop(lease); + if lease_with_get { + assert!(cache.get_leased(&victim).await.is_some()); + drop(CachePin::lease(&pin)); + } else { + assert!(cache.peek_resident(&victim).await); + } + cache + .insert(&key(HOT_ENTRIES + 1), Arc::new(0u8), ENTRY_BYTES, None) + .await; + assert_eq!( + cache.peek_resident(&victim).await, + lease_with_get, + "get={lease_with_get}" + ); + } + } + + /// Two leased stores of 30% each under the 50% cap: the second counts + /// an overflow and churn evicts it, and it is pinned again when admitted + /// once the first is released. + #[tokio::test] + async fn pinned_cap_leaves_overflow_unpinned() { + let (_, cache) = single_shard_cache(); + let store_bytes = PIN_TEST_CAPACITY * 3 / 10; + let (_first, first) = leased(&cache, "first", store_bytes).await; + let (second_store, second) = leased(&cache, "second", store_bytes).await; + assert!(first.is_pinned()); + assert!(!second.is_pinned() && second.pin().is_overflowed()); + let stats = cache.pinned_stats(); + assert_eq!(stats.cap_bytes, PIN_TEST_CAPACITY as u64 / 2); + assert_eq!( + (stats.pinned_entries, stats.leased_entries, stats.overflow), + (1, 2, 1) + ); + churn(&cache, "planes", 10 * PIN_TEST_CAPACITY).await; + assert!(is_resident(&cache, "first").await); + assert!(!is_resident(&cache, "second").await); + + drop(first); + let key = TestKey::::new("second"); + assert!( + cache + .ensure_pinned_with_key(&key, || second_store.clone()) + .await + ); + assert!(second.is_pinned()); + } + + /// However many stores are leased, pins hold at most the cap, and + /// churn stays within the capacity plus one entry. + #[tokio::test] + async fn pins_never_push_inserts_over_capacity() { + let (_, cache) = single_shard_cache(); + let store_bytes = PIN_TEST_CAPACITY / 5; + let mut stores = Vec::new(); + for store in 0..4 { + stores.push(leased(&cache, &format!("store-{store}"), store_bytes).await); + } + let pinned = stores.iter().filter(|(_, lease)| lease.is_pinned()).count(); + assert_eq!(pinned, 2); + let slack = CHURN_ENTRY_BYTES + 1024; + for round in 0..10 { + churn(&cache, &format!("planes-{round}"), PIN_TEST_CAPACITY).await; + let bytes = cache.size_bytes().await; + assert!(bytes <= PIN_TEST_CAPACITY + slack, "round {round}: {bytes}"); + let stats = cache.pinned_stats(); + assert!(stats.pinned_bytes <= stats.cap_bytes, "{stats:?}"); + } + for (store, (_, lease)) in stores.iter().enumerate() { + let name = format!("store-{store}"); + assert_eq!( + is_resident(&cache, &name).await, + lease.is_pinned(), + "{name}" + ); + } + } + + /// Concurrent first callers load a pinned entry once and all lease it. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn get_or_insert_pinned_is_single_flight() { + const CALLERS: usize = 32; + let (_, cache) = single_shard_cache(); + let loads = Arc::new(AtomicUsize::new(0)); + let barrier = Arc::new(tokio::sync::Barrier::new(CALLERS)); + let callers: Vec<_> = (0..CALLERS) + .map(|_| { + let (cache, loads, barrier) = (cache.clone(), loads.clone(), barrier.clone()); + tokio::spawn(async move { + barrier.wait().await; + cache + .get_or_insert_leased_with_key( + TestKey::::new("store"), + || async move { + loads.fetch_add(1, Ordering::SeqCst); + tokio::task::yield_now().await; + Ok(Pinned { + pin: CachePin::new(), + bytes: 1024, + }) + }, + ) + .await + .unwrap() + }) + }) + .collect(); + let mut results = Vec::with_capacity(CALLERS); + for caller in callers { + results.push(caller.await.unwrap()); + } + assert_eq!(loads.load(Ordering::SeqCst), 1); + let store = results[0].0.clone(); + assert!( + results + .iter() + .all(|(value, lease, _)| Arc::ptr_eq(value, &store) && lease.is_pinned()) + ); + assert_eq!(results.iter().filter(|(_, _, hit)| !hit).count(), 1); + assert_eq!(store.pin.holders(), CALLERS); + drop(results); + assert_eq!(store.pin.holders(), 0); + assert!(!store.pin.is_pinned()); + } + + /// The largest admissible entry is a shard's weight budget, and the pin + /// cap is half of every shard's. + #[test] + fn max_entry_bytes_is_shard_budget() { + for capacity in [0, PIN_TEST_CAPACITY, 16 << 30] { + let backend = QuickCacheBackend::with_capacity(capacity); + let shards = backend.cache.num_shards(); + let shard_bytes = (capacity / shards) as u64; + assert_eq!(backend.max_entry_bytes(), Some(shard_bytes), "{capacity}"); + let cap = (shard_bytes as f64 * crate::cache::PINNED_CAP_FRACTION) as u64; + assert_eq!( + backend.pinned_stats().cap_bytes, + cap * shards as u64, + "{capacity}" + ); + } + } + + /// A clear drops a leased entry and its pin; admitting it again pins it + /// again through the leases it still has. + #[tokio::test] + async fn recharge_after_clear() { + let (_, cache) = single_shard_cache(); + let (store, lease) = leased(&cache, "store", 1 << 16).await; + assert!(lease.is_pinned()); + cache.clear().await; + assert!(!is_resident(&cache, "store").await); + assert!(!lease.is_pinned()); + assert_eq!(cache.pinned_stats().leased_entries, 0); + let key = TestKey::::new("store"); + assert!(cache.ensure_pinned_with_key(&key, || store.clone()).await); + assert!(lease.is_pinned()); + assert!( + !cache + .ensure_pinned_with_key(&key, || panic!("admitted twice")) + .await + ); + let stats = cache.pinned_stats(); + assert_eq!((stats.pinned_entries, stats.leased_entries), (1, 1)); + } + + /// Once sign-priority planes activate the strict tier, a leased store + /// stays through their churn, moved there with the pinned priority, and + /// an overflowed store outranks the planes too. + #[tokio::test] + async fn strict_tier_keeps_leased_pinned_entries() { + let (backend, cache) = single_shard_cache(); + let store_bytes = PIN_TEST_CAPACITY * 3 / 10; + let (_store, lease) = leased(&cache, "store", store_bytes).await; + let (_other, other) = leased(&cache, "other", store_bytes).await; + assert!(lease.is_pinned() && !other.is_pinned()); + let codec = CacheCodec::new("test.sign", 1, |_, _| Ok(()), |_| Ok(Arc::new(()))) + .with_memory_priority(3); + for id in 0..(10 * PIN_TEST_CAPACITY / CHURN_ENTRY_BYTES) as u64 { + let mut bytes = [0xAB; 16]; + bytes[..8].copy_from_slice(&id.to_le_bytes()); + let key = InternalCacheKey::from_bytes(bytes); + backend + .insert(&key, Arc::new(id), CHURN_ENTRY_BYTES, Some(codec)) + .await; + } + assert!(backend.priority_active.load(Ordering::Acquire)); + assert!(is_resident(&cache, "store").await); + assert!(is_resident(&cache, "other").await); + assert!(backend.size_bytes().await <= PIN_TEST_CAPACITY); + assert_eq!(cache.pinned_stats().pinned_entries, 1); + } + + async fn pin_budget_cache(is_strict: bool) -> QuickCacheBackend { + let backend = + QuickCacheBackend::with_shard_policy(TEST_CAPACITY, QuickCacheShardPolicy::Single); + if is_strict { + let codec = CacheCodec::new("test.sign", 1, |_, _| Ok(()), |_| Ok(Arc::new(()))) + .with_memory_priority(3); + backend + .insert( + &InternalCacheKey::from_bytes([u8::MAX; 16]), + Arc::new(()), + 1, + Some(codec), + ) + .await; + } + assert_eq!(backend.priority_active.load(Ordering::Acquire), is_strict); + backend + } + + #[rstest::rstest] + #[case::quick(false)] + #[case::strict(true)] + #[tokio::test] + async fn shared_pin_aliases_obey_capacity(#[case] is_strict: bool) { + let backend = pin_budget_cache(is_strict).await; + let pin = CachePin::new(); + let _lease = CachePin::lease(&pin); + let entry_bytes = TEST_CAPACITY * 3 / 10; + let value = Arc::new(vec![0u8; entry_bytes]); + for id in 0..4 { + backend + .insert_pinned( + &InternalCacheKey::from_bytes([id; 16]), + value.clone(), + entry_bytes, + &pin, + ) + .await; + } + assert!( + backend.size_bytes().await <= TEST_CAPACITY, + "resident bytes={}, pin stats={:?}", + backend.size_bytes().await, + backend.pinned_stats(), + ); + let stats = backend.pinned_stats(); + assert_eq!(stats.pinned_entries, 1); + assert_eq!( + stats.pinned_bytes, + (entry_bytes + key_footprint(&InternalCacheKey::from_bytes([0; 16]))) as u64, + ); + } + + #[rstest::rstest] + #[case::quick(false)] + #[case::strict(true)] + #[tokio::test] + async fn shared_pin_backends_obey_capacity(#[case] is_strict: bool) { + let first = pin_budget_cache(is_strict).await; + let second = + QuickCacheBackend::with_shard_policy(10 * TEST_CAPACITY, QuickCacheShardPolicy::Single); + let entry_bytes = TEST_CAPACITY * 3 / 10; + let mut leases = Vec::with_capacity(4); + for id in 0..4 { + let pin = CachePin::new(); + leases.push(CachePin::lease(&pin)); + let value = Arc::new(vec![0u8; entry_bytes]); + let key = InternalCacheKey::from_bytes([id; 16]); + first + .insert_pinned(&key, value.clone(), entry_bytes, &pin) + .await; + second.insert_pinned(&key, value, entry_bytes, &pin).await; + } + assert!(leases.iter().all(CacheLease::is_pinned)); + assert!( + first.size_bytes().await <= TEST_CAPACITY, + "first bytes={}, first stats={:?}, second stats={:?}", + first.size_bytes().await, + first.pinned_stats(), + second.pinned_stats(), + ); + let first_stats = first.pinned_stats(); + assert_eq!(first_stats.pinned_entries, 1); + assert_eq!(second.pinned_stats().pinned_entries, 4); + first.clear().await; + assert_eq!(first.pinned_stats().pinned_bytes, 0); + assert_eq!(second.pinned_stats().pinned_entries, 4); + } + + #[rstest::rstest] + #[case::quick(false)] + #[case::strict(true)] + #[tokio::test] + async fn shared_pin_replacement_reacquires_its_reservation(#[case] is_strict: bool) { + let backend = pin_budget_cache(is_strict).await; + let key = InternalCacheKey::from_bytes([0; 16]); + let entry_bytes = TEST_CAPACITY * 3 / 10; + let value = Arc::new(vec![0u8; entry_bytes]); + let pin = CachePin::new(); + let lease = CachePin::lease(&pin); + for _ in 0..4 { + backend + .insert_pinned(&key, value.clone(), entry_bytes, &pin) + .await; + assert!(backend.peek_resident(&key).await); + assert!(lease.is_pinned()); + assert_eq!(pin.admissions(), 1); + let stats = backend.pinned_stats(); + assert_eq!((stats.pinned_entries, stats.leased_entries), (1, 1)); + assert_eq!( + stats.pinned_bytes, + (entry_bytes + key_footprint(&key)) as u64 + ); + } + backend.clear().await; + assert_eq!(pin.admissions(), 0); + assert!(!lease.is_pinned()); + assert_eq!(backend.pinned_stats().leased_bytes, 0); + } + #[test] fn entry_weight_includes_fixed_key() { let key = InternalCacheKey::from_bytes([0; 16]); let entry = QuickEntry { entry: Arc::new(()), size_bytes: 7, + pin: None, }; assert_eq!( quick_cache::Weighter::weight(&EntryWeighter, &key, &entry),