From eb0525a54ef326c1505d9eba63d0f70a8ec802fc Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Fri, 9 Oct 2026 21:27:38 +0800 Subject: [PATCH 1/3] feat(cache): add tier probes and selective entry reads --- rust/lance-core/src/cache/backend.rs | 70 ++++++ rust/lance-core/src/cache/codec.rs | 241 +++++++++++++++++- rust/lance-core/src/cache/mod.rs | 229 +++++++++++++++++- rust/lance-core/src/cache/moka.rs | 10 + rust/lance-core/src/cache/priority.rs | 184 ++++++++++++++ rust/lance-core/src/cache/quick.rs | 336 ++++++++++++++++++++++++-- 6 files changed, 1049 insertions(+), 21 deletions(-) create mode 100644 rust/lance-core/src/cache/priority.rs diff --git a/rust/lance-core/src/cache/backend.rs b/rust/lance-core/src/cache/backend.rs index ffffa08ae25..d2d57cfd53f 100644 --- a/rust/lance-core/src/cache/backend.rs +++ b/rust/lance-core/src/cache/backend.rs @@ -45,6 +45,20 @@ use super::{CacheCodec, InternalCacheKey}; /// A type-erased cache entry. pub type CacheEntry = Arc; +/// The tier that would serve a read of an entry, as +/// [`CacheBackend::peek_tier`] predicts it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum CacheTier { + /// In RAM: a read is a memory hit. + Resident, + /// Not in RAM but held by a local tier below it, such as a persistent + /// store or a buffer of recently evicted entries: a read costs local I/O + /// at most, never the caller's loader. + Local, + /// Held by no tier: a read runs the caller's loader. + Absent, +} + /// Low-level pluggable cache backend. /// /// Implementations store entries keyed by [`InternalCacheKey`] and return @@ -53,6 +67,62 @@ pub type CacheEntry = Arc; /// backend authors only need to implement storage and eviction. #[async_trait] pub trait CacheBackend: Send + Sync + std::fmt::Debug { + /// Look up RAM only, without reading or promoting a persistent entry. + /// Backends without a resident lookup return a miss safely. + async fn get_resident(&self, _key: &InternalCacheKey) -> Option { + None + } + + /// Report whether `key` is resident in RAM without counting as an access, + /// so a residency probe does not change what the backend evicts next. + /// The default uses [`get_resident`](Self::get_resident), which backends + /// with recency state may treat as an access. + async fn peek_resident(&self, key: &InternalCacheKey) -> bool { + self.get_resident(key).await.is_some() + } + + /// Predict which tier would serve a read of `key`, without reading the + /// entry or counting as an access, so the check changes neither what the + /// backend evicts next nor its hit statistics. The default knows RAM only, + /// through [`peek_resident`](Self::peek_resident), and reports every other + /// entry as [`CacheTier::Absent`]; backends with a local tier below RAM + /// override it. + async fn peek_tier(&self, key: &InternalCacheKey) -> CacheTier { + if self.peek_resident(key).await { + CacheTier::Resident + } else { + CacheTier::Absent + } + } + + /// Whether layered index planes should gate lower-plane RAM admission on + /// their sign plane being resident. Backends that admit plane entries + /// through their ordinary policy return `false`, so lower planes are + /// loaded and admitted like any other entry. + fn plane_admission_gated(&self) -> bool { + true + } + + /// Read an entry without admitting a persistent hit into RAM. + /// Backends without this capability safely fall back to resident entries. + async fn get_without_promotion( + &self, + key: &InternalCacheKey, + _codec: Option, + ) -> Option { + self.get_resident(key).await + } + + /// Gather selected rows from persistent storage, without RAM admission. + async fn get_rows( + &self, + _key: &InternalCacheKey, + _rows: &[u32], + _codec: Option, + ) -> Option { + None + } + /// Look up an entry by its key. /// /// `codec` is provided so that persistent backends can deserialize the diff --git a/rust/lance-core/src/cache/codec.rs b/rust/lance-core/src/cache/codec.rs index eff1b5e2b31..ee2d610b1b6 100644 --- a/rust/lance-core/src/cache/codec.rs +++ b/rust/lance-core/src/cache/codec.rs @@ -188,6 +188,17 @@ impl CacheDecode { // CacheCodecImpl — trait for serializable cache entry types // --------------------------------------------------------------------------- +/// Bounded random reads from a persistent cache payload. Backends may coalesce +/// and align physical I/O while codecs request logical byte ranges. +pub trait CacheRangeReader { + fn read_range(&self, range: std::ops::Range) -> Result; +} +impl) -> Result> CacheRangeReader for F { + fn read_range(&self, range: std::ops::Range) -> Result { + self(range) + } +} + /// Serialization trait for cache entries. /// /// **Experimental**: the serialized format is not yet covered by a stability @@ -207,6 +218,40 @@ impl CacheDecode { /// The read sequence mirroring the write sequence for each `type_version` is /// the invariant the implementor owns. pub trait CacheCodecImpl: Send + Sync { + /// Whether this body supports gathering sorted, unique row offsets. + const SUPPORTS_ROW_SELECTION: bool = false; + fn deserialize_rows( + _reader: &dyn CacheRangeReader, + _body_offset: usize, + _version: u32, + _rows: &[u32], + ) -> Result + where + Self: Sized, + { + Err(Error::invalid_input( + "cache codec does not support row selection", + )) + } + + /// Payload byte ranges that [`deserialize_rows`](Self::deserialize_rows) + /// requests for `rows`, in request order, excluding the envelope and any + /// body header it reads through `reader` while planning. + /// + /// Lets a backend issue those reads concurrently before decoding. The + /// default `None` means the body cannot plan, and backends read serially. + fn plan_row_ranges( + _reader: &dyn CacheRangeReader, + _body_offset: usize, + _version: u32, + _rows: &[u32], + ) -> Result>>> + where + Self: Sized, + { + Ok(None) + } + /// Stable identity for this entry type. **Must not change once shipped.** /// This is a deliberate author-assigned string, not `std::any::type_name` /// (which is not stable across compiler versions). @@ -234,15 +279,28 @@ pub trait CacheCodecImpl: Send + Sync { pub(crate) type ArcAny = Arc; +type RowDecoder = fn(&dyn CacheRangeReader, usize, u32, &[u32]) -> Result; +type RowPlanner = + fn(&dyn CacheRangeReader, usize, u32, &[u32]) -> Result>>>; + +/// Envelope bytes before `type_id`: magic, envelope version and `type_id_len`. +const ENVELOPE_ID_OFFSET: usize = 4 + 1 + 2; +/// Envelope bytes that follow `type_id`: `type_version`. +const ENVELOPE_VERSION_BYTES: usize = 4; + /// Type-erased codec for serializing and deserializing cache entries. /// -/// `CacheCodec` carries the entry's stable `type_id`/`version` plus two plain +/// `CacheCodec` carries the entry's stable `type_id`/`version` plus plain /// function pointers — it is `Copy` and has no heap allocation. Construct one /// via [`CacheCodec::from_impl`] for types that implement [`CacheCodecImpl`], /// or [`CacheCodec::new`] for custom cases (e.g. when the orphan rule prevents /// a direct impl). #[derive(Copy, Clone)] pub struct CacheCodec { + row_decoder: Option, + row_planner: Option, + memory_priority: u8, + plane_tag: Option, type_id: &'static str, version: u32, serialize_body: fn(&ArcAny, &mut CacheEntryWriter<'_>) -> Result<()>, @@ -289,6 +347,10 @@ impl CacheCodec { deserialize_body: fn(&mut CacheEntryReader<'_>) -> Result, ) -> Self { Self { + row_decoder: None, + row_planner: None, + memory_priority: 0, + plane_tag: None, type_id, version, serialize_body, @@ -299,6 +361,20 @@ impl CacheCodec { /// Create a `CacheCodec` from a [`CacheCodecImpl`] implementation. pub fn from_impl() -> Self { Self { + row_decoder: if T::SUPPORTS_ROW_SELECTION { + Some(|reader, offset, version, rows| { + Ok(Arc::new(T::deserialize_rows(reader, offset, version, rows)?) as ArcAny) + }) + } else { + None + }, + row_planner: if T::SUPPORTS_ROW_SELECTION { + Some(T::plan_row_ranges) + } else { + None + }, + memory_priority: 0, + plane_tag: None, type_id: T::TYPE_ID, version: T::CURRENT_VERSION, serialize_body: serialize_via_impl::, @@ -306,6 +382,28 @@ impl CacheCodec { } } + /// Hint for priority-aware memory tiers. Larger values survive pressure first. + pub fn with_memory_priority(mut self, priority: u8) -> Self { + self.memory_priority = priority; + self + } + + pub fn memory_priority(&self) -> u8 { + self.memory_priority + } + + /// Tag the entry as one plane of a layered index so backends can report + /// per-plane metrics without deriving the plane from an admission priority. + pub fn with_plane_tag(mut self, plane: u8) -> Self { + self.plane_tag = Some(plane); + self + } + + /// The layered plane this entry belongs to, or `None` for other entries. + pub fn plane_tag(&self) -> Option { + self.plane_tag + } + /// Return the stable entry type identity. /// /// Persistent and tiered backends can use this metadata to route encoded @@ -321,6 +419,62 @@ impl CacheCodec { (self.serialize_body)(value, &mut entry_writer) } + pub fn supports_row_selection(&self) -> bool { + self.row_decoder.is_some() + } + + /// Read and validate the envelope of a range-read entry, returning the + /// body offset and the entry's `type_version`. + fn read_row_envelope(&self, reader: &dyn CacheRangeReader) -> Result<(usize, u32)> { + let prefix = reader.read_range(0..ENVELOPE_ID_OFFSET)?; + let id_bytes = prefix + .get(ENVELOPE_ID_OFFSET - 2..ENVELOPE_ID_OFFSET) + .ok_or_else(|| Error::invalid_input("short cache envelope"))?; + let id_len = u16::from_le_bytes(id_bytes.try_into().unwrap()) as usize; + let header = + reader.read_range(0..(ENVELOPE_ID_OFFSET + id_len + ENVELOPE_VERSION_BYTES))?; + let envelope = parse_envelope(&header) + .ok_or_else(|| Error::invalid_input("invalid cache envelope"))?; + if envelope.type_id != self.type_id || envelope.type_version > self.version { + return Err(Error::invalid_input("incompatible cache entry")); + } + Ok((envelope.body_offset, envelope.type_version)) + } + + /// Gather rows without loading the complete persistent entry. + pub fn deserialize_rows( + &self, + reader: &dyn CacheRangeReader, + rows: &[u32], + ) -> CacheDecode { + let decode = || -> Result { + let decoder = self + .row_decoder + .ok_or_else(|| Error::invalid_input("row selection unsupported"))?; + let (body_offset, version) = self.read_row_envelope(reader)?; + decoder(reader, body_offset, version, rows) + }; + match decode() { + Ok(entry) => CacheDecode::Hit(entry), + Err(_) => CacheDecode::Miss(CacheMissReason::BodyError), + } + } + + /// Payload ranges that [`deserialize_rows`](Self::deserialize_rows) would + /// request for `rows`, beyond the envelope and body header it reads from + /// `reader` here. `None` when the codec cannot plan, the entry is not + /// compatible, or the rows are invalid; callers then decode serially, + /// which reports invalid rows as a miss. + pub fn plan_rows( + &self, + reader: &dyn CacheRangeReader, + rows: &[u32], + ) -> Option>> { + let planner = self.row_planner?; + let (body_offset, version) = self.read_row_envelope(reader).ok()?; + planner(reader, body_offset, version, rows).ok().flatten() + } + /// Deserialize an entry from `data`. /// /// Never fails: any non-fatal failure to interpret the bytes becomes a @@ -536,6 +690,91 @@ mod tests { ); } + #[test] + fn plane_tag_is_metadata_only() { + let plain = CacheCodec::from_impl::(); + assert_eq!(plain.plane_tag(), None); + + let tagged = plain.with_plane_tag(2).with_memory_priority(1); + assert_eq!(tagged.plane_tag(), Some(2)); + assert_eq!(tagged.memory_priority(), 1); + let copied = tagged; + assert_eq!(copied.plane_tag(), Some(2)); + + // The tag never reaches the wire format. + let value: ArcAny = Arc::new(Widget { n: 11 }); + let (mut plain_bytes, mut tagged_bytes) = (Vec::new(), Vec::new()); + plain.serialize(&value, &mut plain_bytes).unwrap(); + tagged.serialize(&value, &mut tagged_bytes).unwrap(); + assert_eq!(plain_bytes, tagged_bytes); + let decoded = tagged + .deserialize(&Bytes::from(tagged_bytes)) + .hit() + .unwrap(); + assert_eq!(decoded.downcast_ref::(), Some(&Widget { n: 11 })); + } + + /// Row-selectable body that keeps the default (absent) range planner. + struct Rows; + + impl CacheCodecImpl for Rows { + const TYPE_ID: &'static str = "test.Rows"; + const CURRENT_VERSION: u32 = 1; + const SUPPORTS_ROW_SELECTION: bool = true; + + fn serialize(&self, _writer: &mut CacheEntryWriter<'_>) -> Result<()> { + Ok(()) + } + + fn deserialize(_reader: &mut CacheEntryReader<'_>) -> Result { + Ok(Self) + } + + fn deserialize_rows( + _reader: &dyn CacheRangeReader, + _body_offset: usize, + _version: u32, + _rows: &[u32], + ) -> Result { + Ok(Self) + } + } + + #[test] + fn plan_rows_defaults_to_serial_decode() { + let value: ArcAny = Arc::new(Widget { n: 1 }); + let widget = CacheCodec::from_impl::(); + let mut widget_bytes = Vec::new(); + widget.serialize(&value, &mut widget_bytes).unwrap(); + let widget_bytes = Bytes::from(widget_bytes); + let read = |range: std::ops::Range| -> Result { + widget_bytes + .get(range.clone()) + .map(|_| widget_bytes.slice(range)) + .ok_or_else(|| Error::invalid_input("out of range")) + }; + assert!(!widget.supports_row_selection()); + assert!(widget.plan_rows(&read, &[0]).is_none()); + + let rows = CacheCodec::from_impl::(); + let mut rows_bytes = Vec::new(); + rows.serialize(&(Arc::new(Rows) as ArcAny), &mut rows_bytes) + .unwrap(); + let rows_bytes = Bytes::from(rows_bytes); + let read = |range: std::ops::Range| -> Result { + rows_bytes + .get(range.clone()) + .map(|_| rows_bytes.slice(range)) + .ok_or_else(|| Error::invalid_input("out of range")) + }; + assert!(rows.supports_row_selection()); + assert!(rows.plan_rows(&read, &[0]).is_none()); + assert!(matches!( + rows.deserialize_rows(&read, &[0]), + CacheDecode::Hit(_) + )); + } + #[test] fn reader_exposes_envelope_version() { // type_version travels through the envelope to reader.version(). diff --git a/rust/lance-core/src/cache/mod.rs b/rust/lance-core/src/cache/mod.rs index 44b5604bd1a..0028caafe0b 100644 --- a/rust/lance-core/src/cache/mod.rs +++ b/rust/lance-core/src/cache/mod.rs @@ -51,13 +51,16 @@ pub mod codec; mod entry_io; mod key; mod moka; +mod priority; mod quick; +pub use priority::PriorityEntries; mod registry; -pub use backend::{CacheBackend, CacheEntry}; +pub use backend::{CacheBackend, CacheEntry, CacheTier}; pub use backend_uri::{build_from_uri, parse_backend_uri}; pub use codec::{ - CacheCodec, CacheCodecImpl, CacheDecode, CacheMissReason, MAGIC, has_cache_envelope, + CacheCodec, CacheCodecImpl, CacheDecode, CacheMissReason, CacheRangeReader, MAGIC, + has_cache_envelope, }; pub use entry_io::{CacheEntryReader, CacheEntryWriter}; pub use key::{CACHE_KEY_FORMAT, CacheKeySchema, CacheNamespace, InternalCacheKey, KeyBuilder}; @@ -137,6 +140,11 @@ pub trait CacheKey { builder.write_str(self.key().as_ref()); } + /// Per-key cache policy, defaulting to the value type's codec. + fn codec_for_key(&self) -> Option { + Self::codec() + } + /// Optional codec for serializing/deserializing this key's value type. /// /// Returns `None` by default. Cache backends that support persistence @@ -363,17 +371,91 @@ impl LanceCache { let key = self.sized_key(cache_key); self.state .backend - .insert(&key, metadata, size, K::codec()) + .insert(&key, metadata, size, cache_key.codec_for_key()) .await; } + /// Read a resident entry without loading its disk copy. + pub async fn get_resident_with_key(&self, cache_key: &K) -> Option> + where + K: CacheKey, + K::ValueType: DeepSizeOf + Send + Sync + 'static, + { + let key = self.sized_key(cache_key); + self.state + .backend + .get_resident(&key) + .await? + .downcast::() + .ok() + } + + /// Report whether an entry is resident without counting as an access. + pub async fn peek_resident_with_key(&self, cache_key: &K) -> bool { + let key = self.sized_key(cache_key); + self.state.backend.peek_resident(&key).await + } + + /// Predict which tier would serve a read of the entry, without reading it + /// or counting as an access; see [`CacheBackend::peek_tier`]. + pub async fn peek_tier_with_key(&self, cache_key: &K) -> CacheTier { + let key = self.sized_key(cache_key); + self.state.backend.peek_tier(&key).await + } + + /// Whether lower layered planes gate RAM admission on a resident sign plane. + pub fn plane_admission_gated(&self) -> bool { + self.state.backend.plane_admission_gated() + } + + /// Read memory or persistent storage without promoting a persistent hit. + pub async fn get_without_promotion_with_key( + &self, + cache_key: &K, + ) -> Option> + where + K: CacheKey, + K::ValueType: DeepSizeOf + Send + Sync + 'static, + { + let key = self.sized_key(cache_key); + self.state + .backend + .get_without_promotion(&key, cache_key.codec_for_key()) + .await? + .downcast::() + .ok() + } + + /// Read selected rows from a persistent entry without memory admission. + pub async fn get_rows_with_key( + &self, + cache_key: &K, + rows: &[u32], + ) -> Option> + where + K: CacheKey, + K::ValueType: DeepSizeOf + Send + Sync + 'static, + { + self.state + .backend + .get_rows(&self.sized_key(cache_key), rows, cache_key.codec_for_key()) + .await? + .downcast::() + .ok() + } + pub async fn get_with_key(&self, cache_key: &K) -> Option> where K: CacheKey, K::ValueType: DeepSizeOf + Send + Sync + 'static, { let key = self.sized_key(cache_key); - let Some(entry) = self.state.backend.get(&key, K::codec()).await else { + let Some(entry) = self + .state + .backend + .get(&key, cache_key.codec_for_key()) + .await + else { self.state.misses.fetch_add(1, Ordering::Relaxed); return None; }; @@ -449,7 +531,7 @@ impl LanceCache { let (entry, was_cached) = self .state .backend - .get_or_insert(&key, typed_loader, K::codec()) + .get_or_insert(&key, typed_loader, cache_key.codec_for_key()) .await?; let entry = entry.downcast::().map_err(|_| { self.state.misses.fetch_add(1, Ordering::Relaxed); @@ -585,6 +667,65 @@ impl WeakLanceCache { } } + /// Read RAM only, avoiding whole-entry promotion for candidate range reads. + pub async fn get_resident_with_key(&self, cache_key: &K) -> Option> + where + K: CacheKey, + K::ValueType: DeepSizeOf + Send + Sync + 'static, + { + self.upgrade()?.get_resident_with_key(cache_key).await + } + + /// Report residency without counting as an access; a dropped cache holds nothing. + pub async fn peek_resident_with_key(&self, cache_key: &K) -> bool { + match self.upgrade() { + Some(cache) => cache.peek_resident_with_key(cache_key).await, + None => false, + } + } + + /// Predict the serving tier without counting as an access; a dropped + /// cache holds nothing. + pub async fn peek_tier_with_key(&self, cache_key: &K) -> CacheTier { + match self.upgrade() { + Some(cache) => cache.peek_tier_with_key(cache_key).await, + None => CacheTier::Absent, + } + } + + /// Whether lower layered planes gate RAM admission; a dropped cache keeps the default. + pub fn plane_admission_gated(&self) -> bool { + self.upgrade() + .is_none_or(|cache| cache.plane_admission_gated()) + } + + /// Read a cached candidate plane without whole-entry RAM admission. + pub async fn get_without_promotion_with_key( + &self, + cache_key: &K, + ) -> Option> + where + K: CacheKey, + K::ValueType: DeepSizeOf + Send + Sync + 'static, + { + self.upgrade()? + .get_without_promotion_with_key(cache_key) + .await + } + + /// Read selected cached rows without promoting a whole plane. + pub async fn get_rows_with_key( + &self, + cache_key: &K, + rows: &[u32], + ) -> Option> + where + K: CacheKey, + K::ValueType: DeepSizeOf + Send + Sync + 'static, + { + self.upgrade()?.get_rows_with_key(cache_key, rows).await + } + pub async fn get_with_key(&self, cache_key: &K) -> Option> where K: CacheKey, @@ -1071,6 +1212,84 @@ mod tests { assert_eq!((cache.stats().await.hits, cache.size().await), (2, 2)); } + #[test] + fn plane_admission_gating_passes_through() { + let cache = LanceCache::with_capacity(4096); + let weak = WeakLanceCache::from(&cache); + assert!(cache.plane_admission_gated() && weak.plane_admission_gated()); + + drop(cache); + assert!(weak.plane_admission_gated()); + } + + #[tokio::test] + async fn residency_peek_passes_through() { + let cache = LanceCache::with_capacity(4096); + let weak = WeakLanceCache::from(&cache); + assert!(!cache.peek_resident_with_key(&TestKey::new(1)).await); + cache + .insert_with_key(&TestKey::new(1), Arc::new(vec![1])) + .await; + assert!(cache.peek_resident_with_key(&TestKey::new(1)).await); + assert!(weak.peek_resident_with_key(&TestKey::new(1)).await); + assert!(!weak.peek_resident_with_key(&TestKey::new(2)).await); + + drop(cache); + assert!(!weak.peek_resident_with_key(&TestKey::new(1)).await); + } + + /// Lance's in-memory backends have no tier below RAM, so the default tier + /// peek reports resident entries and nothing else as held. Peeks are not + /// accesses: they count neither hits nor misses. + #[tokio::test] + async fn tier_peek_defaults_to_residency_without_access() { + for kind in [TestBackendKind::Moka, TestBackendKind::Quick] { + let cache = kind.cache(4096); + let weak = WeakLanceCache::from(&cache); + assert_eq!( + cache.peek_tier_with_key(&TestKey::new(1)).await, + CacheTier::Absent, + "{kind:?}" + ); + cache + .insert_with_key(&TestKey::new(1), Arc::new(vec![1])) + .await; + assert_eq!( + cache.peek_tier_with_key(&TestKey::new(1)).await, + CacheTier::Resident, + "{kind:?}" + ); + assert_eq!( + weak.peek_tier_with_key(&TestKey::new(1)).await, + CacheTier::Resident, + "{kind:?}" + ); + assert_eq!( + weak.peek_tier_with_key(&TestKey::new(2)).await, + CacheTier::Absent, + "{kind:?}" + ); + let stats = cache.stats().await; + assert_eq!((stats.hits, stats.misses), (0, 0), "{kind:?}"); + + cache.clear().await; + assert_eq!( + weak.peek_tier_with_key(&TestKey::new(1)).await, + CacheTier::Absent, + "{kind:?}" + ); + cache + .insert_with_key(&TestKey::new(1), Arc::new(vec![1])) + .await; + drop(cache); + assert_eq!( + weak.peek_tier_with_key(&TestKey::new(1)).await, + CacheTier::Absent, + "{kind:?}" + ); + } + } + #[tokio::test] async fn nested_namespace_segments_do_not_alias_combined_segments() { 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 f9fba02e2eb..9c297ce38f3 100644 --- a/rust/lance-core/src/cache/moka.rs +++ b/rust/lance-core/src/cache/moka.rs @@ -105,6 +105,16 @@ impl MokaCacheBackend { #[async_trait] impl CacheBackend for MokaCacheBackend { + async fn get_resident(&self, key: &InternalCacheKey) -> Option { + self.cache.get(key).await.map(|r| r.entry) + } + + /// Moka does not count `contains_key` as a read: the entry's popularity + /// estimate and recency order stay unchanged. + async fn peek_resident(&self, key: &InternalCacheKey) -> bool { + self.cache.contains_key(key) + } + async fn get(&self, key: &InternalCacheKey, _codec: Option) -> Option { self.cache.get(key).await.map(|r| r.entry) } diff --git a/rust/lance-core/src/cache/priority.rs b/rust/lance-core/src/cache/priority.rs new file mode 100644 index 00000000000..2f45ee51f52 --- /dev/null +++ b/rust/lance-core/src/cache/priority.rs @@ -0,0 +1,184 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Strict priority admission for layered index planes, sharing the cache byte budget. + +use super::InternalCacheKey; +use std::collections::{BTreeMap, HashMap}; + +pub struct PriorityEntries { + entries: HashMap, + order: BTreeMap<(u8, u64), InternalCacheKey>, + sequence: u64, + bytes: u128, +} + +impl Default for PriorityEntries { + fn default() -> Self { + Self { + entries: HashMap::new(), + order: BTreeMap::new(), + sequence: 0, + bytes: 0, + } + } +} + +impl PriorityEntries { + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + pub fn len(&self) -> usize { + self.entries.len() + } + pub fn bytes(&self) -> usize { + self.bytes as usize + } + pub fn get(&mut self, key: &InternalCacheKey) -> Option { + let (priority, stamp, _, value) = self.entries.get_mut(key)?; + self.order.remove(&(*priority, *stamp)); + self.sequence += 1; + *stamp = self.sequence; + self.order.insert((*priority, *stamp), *key); + Some(value.clone()) + } + /// Look up an entry without refreshing its recency. + pub fn peek(&self, key: &InternalCacheKey) -> Option<&T> { + self.entries.get(key).map(|(_, _, _, value)| value) + } + + /// Membership only; unlike [`get`](Self::get), recency is unchanged. + pub fn contains(&self, key: &InternalCacheKey) -> bool { + self.entries.contains_key(key) + } + pub fn remove(&mut self, key: &InternalCacheKey) -> Option { + let (priority, stamp, bytes, value) = self.entries.remove(key)?; + self.order.remove(&(priority, stamp)); + self.bytes -= bytes as u128; + Some(value) + } + pub fn insert( + &mut self, + key: InternalCacheKey, + value: T, + bytes: usize, + priority: u8, + capacity: usize, + ) -> Vec { + let mut dropped = Vec::new(); + if let Some(old) = self.remove(&key) { + dropped.push(old); + } + if bytes > capacity { + dropped.push(value); + return dropped; + } + self.sequence += 1; + self.entries + .insert(key, (priority, self.sequence, bytes, value)); + 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) + { + dropped.push(value); + } + } + dropped + } + pub fn clear(&mut self) -> Vec { + self.order.clear(); + self.bytes = 0; + self.entries + .drain() + .map(|(_, (_, _, _, value))| value) + .collect() + } + pub fn snapshot(&self) -> Vec<(InternalCacheKey, usize, T)> { + self.entries + .iter() + .map(|(key, (_, _, bytes, value))| (*key, *bytes, value.clone())) + .collect() + } + pub fn stats(&self) -> [(usize, usize); 3] { + let mut stats = [(0, 0); 3]; + for (priority, _, bytes, _) in self.entries.values() { + if (1..=3).contains(priority) { + let s = &mut stats[(3 - priority) as usize]; + s.0 += 1; + s.1 += bytes; + } + } + stats + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn eviction_preserves_plane_priority_and_budget() { + let mut cache = PriorityEntries::default(); + let sign = InternalCacheKey::from_bytes([1; 16]); + let high = InternalCacheKey::from_bytes([2; 16]); + let low = InternalCacheKey::from_bytes([3; 16]); + cache.insert(low, 1, 30, 1, 100); + cache.insert(high, 2, 50, 2, 100); + assert_eq!(cache.insert(sign, 3, 40, 3, 100), vec![1]); + assert_eq!(cache.get(&sign), Some(3)); + assert_eq!(cache.get(&high), Some(2)); + assert_eq!(cache.get(&low), None); + assert_eq!(cache.insert(low, 4, 30, 1, 100), vec![4]); + assert_eq!(cache.bytes(), 90); + assert_eq!(cache.stats(), [(1, 40), (1, 50), (0, 0)]); + assert_eq!( + cache.insert(InternalCacheKey::from_bytes([4; 16]), 5, 60, 3, 100), + vec![2] + ); + assert_eq!(cache.bytes(), 100); + } + + #[test] + fn contains_reports_self_evicted_admissions() { + let (sign, low) = ( + InternalCacheKey::from_bytes([1; 16]), + InternalCacheKey::from_bytes([3; 16]), + ); + let mut cache = PriorityEntries::default(); + cache.insert(sign, 1, 80, 3, 100); + assert!(cache.contains(&sign)); + // The low plane is the minimum of the strict order, so its own + // admission evicts it and the sign plane stays resident. + assert_eq!(cache.insert(low, 2, 40, 1, 100), vec![2]); + assert!(!cache.contains(&low)); + assert!(cache.contains(&sign)); + } + + #[test] + fn peek_does_not_refresh_recency() { + let (a, b, c) = ( + InternalCacheKey::from_bytes([1; 16]), + InternalCacheKey::from_bytes([2; 16]), + InternalCacheKey::from_bytes([3; 16]), + ); + let mut peeked = PriorityEntries::default(); + peeked.insert(a, 1, 40, 1, 100); + peeked.insert(b, 2, 40, 1, 100); + assert_eq!(peeked.peek(&a), Some(&1)); + assert!(peeked.contains(&a)); + assert_eq!(peeked.insert(c, 3, 40, 1, 100), vec![1]); + assert!(!peeked.contains(&a) && peeked.peek(&a).is_none()); + assert!(peeked.contains(&b) && peeked.contains(&c)); + + // A real access protects the same entry, so the test observes recency. + let mut touched = PriorityEntries::default(); + touched.insert(a, 1, 40, 1, 100); + touched.insert(b, 2, 40, 1, 100); + assert_eq!(touched.get(&a), Some(1)); + assert_eq!(touched.insert(c, 3, 40, 1, 100), vec![2]); + assert!(touched.contains(&a)); + } +} diff --git a/rust/lance-core/src/cache/quick.rs b/rust/lance-core/src/cache/quick.rs index 53e0d0e8870..1dd78dd63ff 100644 --- a/rust/lance-core/src/cache/quick.rs +++ b/rust/lance-core/src/cache/quick.rs @@ -6,7 +6,10 @@ //! housekeeping. Used for the session index and metadata caches; the index //! cache sees thousands of cache reads per query. +use super::PriorityEntries; use std::pin::Pin; +use std::sync::Mutex; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use async_trait::async_trait; use futures::Future; @@ -34,6 +37,10 @@ impl quick_cache::Weighter for EntryWeighter { } pub struct QuickCacheBackend { + capacity: usize, + generation: AtomicU64, + priority_active: AtomicBool, + priority: Mutex>, cache: quick_cache::sync::Cache, } @@ -92,13 +99,86 @@ impl QuickCacheBackend { Default::default(), Default::default(), ); - Self { cache } + Self { + cache, + capacity, + generation: AtomicU64::new(0), + priority_active: AtomicBool::new(false), + priority: Mutex::new(PriorityEntries::default()), + } + } + fn admit_priority( + &self, + key: InternalCacheKey, + item: QuickEntry, + priority: u8, + generation: u64, + ) { + self.cache.remove(&key); + let dropped = { + let mut entries = self.priority.lock().unwrap_or_else(|e| e.into_inner()); + if self.generation.load(Ordering::Acquire) != generation { + return; + } + let mut dropped = Vec::new(); + if !self.priority_active.swap(true, Ordering::AcqRel) { + // Metadata and existing entries compete with signs, ahead of ex planes. + // Reserving the whole budget for planes would otherwise evict the IVF model. + 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)); + } + // 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( + key, + item, + size, + if priority == 0 { 3 } else { priority }, + self.capacity, + )); + dropped + }; + drop(dropped); } } #[async_trait] impl CacheBackend for QuickCacheBackend { - async fn get(&self, key: &InternalCacheKey, _codec: Option) -> Option { + async fn get_resident(&self, key: &InternalCacheKey) -> Option { + self.priority + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(key) + .or_else(|| self.cache.get(key)) + .map(|r| r.entry) + } + + /// Membership checks only: priority stamps and quick_cache's reference + /// bits stay unchanged, so a probed entry is evicted as if never probed. + async fn peek_resident(&self, key: &InternalCacheKey) -> bool { + self.priority + .lock() + .unwrap_or_else(|e| e.into_inner()) + .contains(key) + || self.cache.contains_key(key) + } + + async fn get(&self, key: &InternalCacheKey, codec: Option) -> Option { + if (self.priority_active.load(Ordering::Acquire) + || codec.is_some_and(|c| c.memory_priority() > 0)) + && let Some(value) = self + .priority + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(key) + { + return Some(value.entry); + } self.cache.get(key).map(|v| v.entry) } @@ -107,48 +187,104 @@ impl CacheBackend for QuickCacheBackend { key: &InternalCacheKey, entry: CacheEntry, size_bytes: usize, - _codec: Option, + codec: Option, ) { - self.cache.insert(*key, QuickEntry { entry, size_bytes }); + let priority = codec.map(|c| c.memory_priority()).unwrap_or(0); + let item = QuickEntry { entry, size_bytes }; + if priority > 0 || self.priority_active.load(Ordering::Acquire) { + self.admit_priority( + *key, + item, + priority, + self.generation.load(Ordering::Acquire), + ); + } else { + self.cache.insert(*key, item); + } } async fn get_or_insert<'a>( &self, key: &InternalCacheKey, loader: Pin> + Send + 'a>>, - _codec: Option, + codec: Option, ) -> Result<(CacheEntry, bool)> { + let priority = codec.map(|c| c.memory_priority()).unwrap_or(0); + if (priority > 0 || 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) = loader.await?; - let _ = guard.insert(QuickEntry { + let item = QuickEntry { entry: entry.clone(), size_bytes, - }); + }; + if guard.insert(item.clone()).is_ok() + && (priority > 0 || self.priority_active.load(Ordering::Acquire)) + { + self.admit_priority(*key, item, priority, generation); + } Ok((entry, false)) } } } async fn clear(&self) { + self.generation.fetch_add(1, Ordering::AcqRel); + let dropped = { + let mut entries = self.priority.lock().unwrap_or_else(|e| e.into_inner()); + let dropped = entries.clear(); + self.priority_active.store(false, Ordering::Release); + self.cache.set_capacity(self.capacity as u64); + dropped + }; + drop(dropped); self.cache.clear(); } async fn num_entries(&self) -> usize { self.cache.len() + + self + .priority + .lock() + .unwrap_or_else(|e| e.into_inner()) + .len() } async fn size_bytes(&self) -> usize { self.cache.weight() as usize + + self + .priority + .lock() + .unwrap_or_else(|e| e.into_inner()) + .bytes() } fn approx_num_entries(&self) -> usize { self.cache.len() + + self + .priority + .lock() + .unwrap_or_else(|e| e.into_inner()) + .len() } fn approx_size_bytes(&self) -> usize { self.cache.weight() as usize + + self + .priority + .lock() + .unwrap_or_else(|e| e.into_inner()) + .bytes() } fn deep_size_of_entries( @@ -156,14 +292,27 @@ impl CacheBackend for QuickCacheBackend { context: &mut Context, size_of_entry: &dyn Fn(&CacheEntry, &mut Context) -> Option, ) -> Option { + let prioritized: usize = self + .priority + .lock() + .unwrap_or_else(|e| e.into_inner()) + .snapshot() + .into_iter() + .map(|(key, _, value)| { + key_footprint(&key) + + size_of_entry(&value.entry, context).unwrap_or(value.size_bytes) + }) + .sum(); Some( - self.cache - .iter() - .map(|(key, record)| { - key_footprint(&key) - + size_of_entry(&record.entry, context).unwrap_or(record.size_bytes) - }) - .sum(), + prioritized + + self + .cache + .iter() + .map(|(key, record)| { + key_footprint(&key) + + size_of_entry(&record.entry, context).unwrap_or(record.size_bytes) + }) + .sum::(), ) } } @@ -175,7 +324,7 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use super::*; - use crate::cache::{CacheKey, LanceCache}; + use crate::cache::{CacheKey, CacheTier, LanceCache}; struct TestKey { key: String, @@ -201,6 +350,163 @@ mod tests { } } + #[tokio::test] + async fn priority_transition_preserves_inflight_loads() { + let cache = Arc::new(QuickCacheBackend::with_capacity(1024)); + let key = |id| InternalCacheKey::from_bytes([id; 16]); + let codec = CacheCodec::new("test.priority", 1, |_, _| Ok(()), |_| Ok(Arc::new(()))) + .with_memory_priority(3); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let loading_cache = cache.clone(); + let loading = tokio::spawn(async move { + loading_cache + .get_or_insert( + &key(1), + Box::pin(async move { + started_tx.send(()).unwrap(); + release_rx.await.unwrap(); + Ok((Arc::new(1u64) as CacheEntry, 32)) + }), + Some(codec), + ) + .await + .unwrap() + }); + started_rx.await.unwrap(); + cache.insert(&key(2), Arc::new(2u64), 32, Some(codec)).await; + release_tx.send(()).unwrap(); + loading.await.unwrap(); + assert!(cache.get_resident(&key(1)).await.is_some()); + assert!(cache.get_resident(&key(2)).await.is_some()); + } + + #[tokio::test] + async fn priority_budget_retains_metadata_and_singleflight_entries() { + let cache = QuickCacheBackend::with_capacity(160); + let key = |id| InternalCacheKey::from_bytes([id; 16]); + let codec = CacheCodec::new("test.priority", 1, |_, _| Ok(()), |_| Ok(Arc::new(()))); + cache.insert(&key(0), Arc::new(0u64), 32, None).await; + for (id, priority) in [(1, 3), (2, 2), (3, 1)] { + cache + .insert( + &key(id), + Arc::new(id), + 32, + Some(codec.with_memory_priority(priority)), + ) + .await; + } + assert!(cache.get_resident(&key(0)).await.is_some()); + assert!(cache.get_resident(&key(1)).await.is_some()); + assert!(cache.get_resident(&key(2)).await.is_some()); + assert!(cache.get_resident(&key(3)).await.is_none()); + let (_, hit) = cache + .get_or_insert( + &key(4), + Box::pin(async { Ok((Arc::new(4u64) as CacheEntry, 32)) }), + None, + ) + .await + .unwrap(); + assert!(!hit); + let (_, hit) = cache + .get_or_insert( + &key(4), + Box::pin(async { panic!("resident entry reloaded") }), + None, + ) + .await + .unwrap(); + assert!(hit); + assert!(cache.get_resident(&key(2)).await.is_none()); + assert!(cache.size_bytes().await <= 160); + cache.clear().await; + assert_eq!(cache.num_entries().await, 0); + cache.insert(&key(0), Arc::new(0u64), 32, None).await; + assert!(cache.get_resident(&key(0)).await.is_some()); + } + + /// Fill a single-shard cache so that `victim` is its only cold entry, + /// then peek at its tier or get it, and insert one more entry. + async fn cold_victim_after(touch_with_get: bool) -> QuickCacheBackend { + const ENTRY_BYTES: usize = 100; + const HOT_ENTRIES: u8 = 9; + let key = |id| InternalCacheKey::from_bytes([id; 16]); + let entry_weight = key_footprint(&key(0)) + ENTRY_BYTES; + let capacity = entry_weight * (usize::from(HOT_ENTRIES) + 1); + let cache = QuickCacheBackend::with_capacity(capacity); + for id in 0..HOT_ENTRIES { + cache + .insert(&key(id), Arc::new(id), ENTRY_BYTES, None) + .await; + } + let victim = key(HOT_ENTRIES); + cache + .insert(&victim, Arc::new(0u8), ENTRY_BYTES, None) + .await; + if touch_with_get { + assert!(cache.get_resident(&victim).await.is_some()); + } else { + assert_eq!(cache.peek_tier(&victim).await, CacheTier::Resident); + } + cache + .insert(&key(HOT_ENTRIES + 1), Arc::new(0u8), ENTRY_BYTES, None) + .await; + cache + } + + #[tokio::test] + async fn tier_peek_leaves_eviction_order_unchanged() { + let victim = InternalCacheKey::from_bytes([9; 16]); + let peeked = cold_victim_after(false).await; + assert_eq!(peeked.peek_tier(&victim).await, CacheTier::Absent); + assert!(!peeked.cache.contains_key(&victim)); + + // A real access sets the reference bit, so the same victim survives. + let touched = cold_victim_after(true).await; + assert_eq!(touched.peek_tier(&victim).await, CacheTier::Resident); + } + + /// The priority tier evicts its oldest stamp first; a tier peek leaves + /// the stamp alone where a get refreshes it. + #[tokio::test] + async fn tier_peek_leaves_priority_order_unchanged() { + let key = |id| InternalCacheKey::from_bytes([id; 16]); + let codec = CacheCodec::new("test.priority", 1, |_, _| Ok(()), |_| Ok(Arc::new(()))) + .with_memory_priority(1); + for touch_with_get in [false, true] { + // Room for three 48-byte entries (32 bytes plus the 16-byte key). + let cache = QuickCacheBackend::with_capacity(160); + for id in 1..=3 { + cache + .insert(&key(id), Arc::new(u64::from(id)), 32, Some(codec)) + .await; + } + if touch_with_get { + assert!(cache.get_resident(&key(1)).await.is_some()); + } else { + assert_eq!(cache.peek_tier(&key(1)).await, CacheTier::Resident); + } + cache.insert(&key(4), Arc::new(4u64), 32, Some(codec)).await; + let (survivor, evicted) = if touch_with_get { + (key(1), key(2)) + } else { + (key(2), key(1)) + }; + assert_eq!( + cache.peek_tier(&evicted).await, + CacheTier::Absent, + "get={touch_with_get}" + ); + assert_eq!( + cache.peek_tier(&survivor).await, + CacheTier::Resident, + "get={touch_with_get}" + ); + } + } + #[test] fn entry_weight_includes_fixed_key() { let key = InternalCacheKey::from_bytes([0; 16]); From e003e44aafc2e12aaeac56d942ffe9bb0097bd45 Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Fri, 9 Oct 2026 21:28:43 +0800 Subject: [PATCH 2/3] feat(cache): add budgeted leases for resident entries --- rust/lance-core/src/cache/backend.rs | 62 +++ rust/lance-core/src/cache/mod.rs | 427 ++++++++++++++++++++- rust/lance-core/src/cache/moka.rs | 4 + rust/lance-core/src/cache/pin.rs | 504 ++++++++++++++++++++++++ rust/lance-core/src/cache/priority.rs | 79 +++- rust/lance-core/src/cache/quick.rs | 531 +++++++++++++++++++++++++- 6 files changed, 1585 insertions(+), 22 deletions(-) create mode 100644 rust/lance-core/src/cache/pin.rs diff --git a/rust/lance-core/src/cache/backend.rs b/rust/lance-core/src/cache/backend.rs index d2d57cfd53f..cc69a42697e 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 0028caafe0b..ea30f5af353 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, 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}, }; @@ -548,6 +553,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, @@ -808,6 +968,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()?, @@ -1222,6 +1439,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 9c297ce38f3..ddc85eb578d 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..9ae814939f5 --- /dev/null +++ b/rust/lance-core/src/cache/pin.rs @@ -0,0 +1,504 @@ +// 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 entry, or its next admission, 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); + } + } +} + +/// Where a pin's entry is admitted: the budget of the backend holding it, +/// the partition (shard) of its key and the bytes the backend charges. +#[derive(Debug, Clone)] +struct PinBinding { + budget: Arc, + partition: usize, + bytes: u64, +} + +impl PinBinding { + fn same_as(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.budget, &other.budget) + && self.partition == other.partition + && self.bytes == other.bytes + } +} + +#[derive(Debug, Default)] +struct PinState { + /// Live leases. + holders: usize, + /// Live admissions of the value in its backend. + records: usize, + binding: Option, + /// Whether the binding's budget holds the entry's bytes. + reserved: bool, + /// Whether the binding's budget counts the entry as leased. + counted_leased: bool, +} + +/// The pin of one cached value; see the [module docs](self). Leases and +/// admissions share it, so every admission of the value is pinned by the +/// same leases: a value evicted while leased (an overflow) and admitted again +/// is pinned again at once. +#[derive(Debug, Default)] +pub struct CachePin { + state: Mutex, + /// Mirror of `state.reserved`, which eviction reads without the lock. + 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 + /// entry reserves its bytes, or counts an overflow. + 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. A pin admitted to another budget, or with + /// another charge, moves to this one. + pub fn record( + pin: &Arc, + budget: &Arc, + partition: usize, + bytes: u64, + ) -> PinRecord { + let binding = PinBinding { + budget: budget.clone(), + partition, + bytes, + }; + let mut state = pin.lock(); + if !state + .binding + .as_ref() + .is_some_and(|bound| bound.same_as(&binding)) + { + Self::unbind(&mut state); + state.binding = Some(binding); + } + state.records += 1; + pin.sync(&mut state, true); + PinRecord { pin: pin.clone() } + } + + /// Whether a backend must keep the entry: admitted, leased, and within + /// its budget's cap. + 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().records + } + + /// Whether the entry is admitted and leased but found no room under its + /// budget's cap, so it stays evictable. + pub fn is_overflowed(&self) -> bool { + let state = self.lock(); + state.holders > 0 && state.records > 0 && state.binding.is_some() && !state.reserved + } + + 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) { + let mut state = self.lock(); + state.records = state.records.saturating_sub(1); + if state.records == 0 { + self.sync(&mut state, false); + } + } + + /// Release what the binding's budget holds for the entry. + fn unbind(state: &mut PinState) { + if let Some(binding) = &state.binding { + if state.counted_leased { + binding.budget.count_leased(binding.bytes, false); + } + if state.reserved { + binding.budget.release(binding.partition, binding.bytes); + } + } + state.counted_leased = false; + state.reserved = false; + } + + /// Bring the budget in line with the holders and admissions: count a + /// leased admitted entry as leased, release it once unleased or no + /// longer admitted, and, when `reserve`, reserve its bytes if it is not + /// yet pinned. + fn sync(&self, state: &mut PinState, reserve: bool) { + let active = state.holders > 0 && state.records > 0; + match state.binding.clone() { + Some(binding) if active => { + if !state.counted_leased { + binding.budget.count_leased(binding.bytes, true); + state.counted_leased = true; + } + if reserve && !state.reserved { + state.reserved = binding.budget.try_reserve(binding.partition, binding.bytes); + } + } + _ => Self::unbind(state), + } + self.pinned.store(state.reserved, 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, +} + +impl PinRecord { + /// Whether the backend must keep the entry; see [`CachePin::is_pinned`]. + pub fn is_pinned(&self) -> bool { + self.pin.is_pinned() + } + + /// 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(); + } +} + +/// 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 the leased entry 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); + } + + #[test] + fn admission_to_another_budget_moves_the_pin() { + let (first, second) = ( + Arc::new(PinBudget::new(1000, 1)), + Arc::new(PinBudget::new(1000, 1)), + ); + let pin = CachePin::new(); + let _lease = CachePin::lease(&pin); + let old = CachePin::record(&pin, &first, 0, 100); + assert_eq!(first.stats().pinned_bytes, 100); + let _new = CachePin::record(&pin, &second, 0, 100); + assert_eq!(first.stats().pinned_bytes, 0); + assert_eq!(second.stats().pinned_bytes, 100); + assert_eq!(pin.admissions(), 2); + drop(old); + // The pin is still admitted once, so it stays pinned. + assert_eq!(pin.admissions(), 1); + assert!(pin.is_pinned()); + } + + /// 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(); + assert_eq!(pin.admissions(), 0); + let record = Arc::new(CachePin::record(&pin, &budget, 0, 100)); + let copy = record.clone(); + assert_eq!(pin.admissions(), 1); + let replacement = CachePin::record(&pin, &budget, 0, 100); + assert_eq!(pin.admissions(), 2); + drop((record, copy)); + assert_eq!(pin.admissions(), 1); + drop(replacement); + assert_eq!(pin.admissions(), 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 1dd78dd63ff..cb0ad0753b1 100644 --- a/rust/lance-core/src/cache/quick.rs +++ b/rust/lance-core/src/cache/quick.rs @@ -5,17 +5,23 @@ //! whose hit path is one atomic bit — no read-op channel or inline //! housekeeping. Used for the session index and metadata caches; the index //! cache sees thousands of cache reads per query. +//! +//! 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; @@ -24,6 +30,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)] @@ -36,12 +67,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, } impl std::fmt::Debug for QuickCacheBackend { @@ -93,18 +134,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( @@ -127,19 +192,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 }; @@ -190,7 +274,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, @@ -227,6 +315,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)) @@ -238,6 +327,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 = { @@ -324,7 +488,7 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use super::*; - use crate::cache::{CacheKey, CacheTier, LanceCache}; + use crate::cache::{CacheKey, CacheLease, CacheTier, LanceCache}; struct TestKey { key: String, @@ -507,12 +671,349 @@ 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_capacity(PIN_TEST_CAPACITY)); + assert_eq!(backend.cache.num_shards(), 1); + (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); + } + #[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), From fb1cd3253ef0c42923462d775475b0ccd89ab2ca Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Sat, 10 Oct 2026 20:57:55 +0800 Subject: [PATCH 3/3] fix(cache): reserve pin budgets for each admission --- rust/lance-core/src/cache/pin.rs | 336 ++++++++++++++++++++--------- rust/lance-core/src/cache/quick.rs | 120 +++++++++++ 2 files changed, 360 insertions(+), 96 deletions(-) diff --git a/rust/lance-core/src/cache/pin.rs b/rust/lance-core/src/cache/pin.rs index 9ae814939f5..65bf3896288 100644 --- a/rust/lance-core/src/cache/pin.rs +++ b/rust/lance-core/src/cache/pin.rs @@ -14,7 +14,8 @@ //! 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 entry, or its next admission, tries again. +//! 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}; @@ -145,20 +146,43 @@ impl PinBudget { } } -/// Where a pin's entry is admitted: the budget of the backend holding it, -/// the partition (shard) of its key and the bytes the backend charges. -#[derive(Debug, Clone)] -struct PinBinding { +/// 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 PinBinding { - fn same_as(&self, other: &Self) -> bool { - Arc::ptr_eq(&self.budget, &other.budget) - && self.partition == other.partition - && self.bytes == other.bytes +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; + } } } @@ -166,23 +190,20 @@ impl PinBinding { struct PinState { /// Live leases. holders: usize, - /// Live admissions of the value in its backend. - records: usize, - binding: Option, - /// Whether the binding's budget holds the entry's bytes. - reserved: bool, - /// Whether the binding's budget counts the entry as leased. - counted_leased: bool, + /// 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 it, so every admission of the value is pinned by the -/// same leases: a value evicted while leased (an overflow) and admitted again -/// is pinned again at once. +/// 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, - /// Mirror of `state.reserved`, which eviction reads without the lock. + /// Whether any admission is pinned. Eviction uses the admission-specific + /// flag in its `PinRecord`, never this value-wide status. pinned: AtomicBool, } @@ -194,7 +215,8 @@ impl CachePin { /// 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 - /// entry reserves its bytes, or counts an overflow. + /// 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; @@ -207,35 +229,35 @@ impl CachePin { /// 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. A pin admitted to another budget, or with - /// another charge, moves to this one. + /// 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 binding = PinBinding { + 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(); - if !state - .binding - .as_ref() - .is_some_and(|bound| bound.same_as(&binding)) - { - Self::unbind(&mut state); - state.binding = Some(binding); + admission.sync(state.holders > 0); + state.admissions.push(admission); + pin.update_pinned(&state); + PinRecord { + pin: pin.clone(), + pinned, } - state.records += 1; - pin.sync(&mut state, true); - PinRecord { pin: pin.clone() } } - /// Whether a backend must keep the entry: admitted, leased, and within - /// its budget's cap. + /// 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) } @@ -251,14 +273,18 @@ impl CachePin { /// the value from the cache; one replaced by another admission of the /// same value does not. pub fn admissions(&self) -> usize { - self.lock().records + self.lock().admissions.len() } - /// Whether the entry is admitted and leased but found no room under its - /// budget's cap, so it stays evictable. + /// 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.records > 0 && state.binding.is_some() && !state.reserved + state.holders > 0 + && state + .admissions + .iter() + .any(|admission| !admission.pinned.load(Ordering::Relaxed)) } fn lock(&self) -> MutexGuard<'_, PinState> { @@ -273,47 +299,50 @@ impl CachePin { } } - fn drop_record(&self) { + fn drop_record(&self, pinned: &Arc) { let mut state = self.lock(); - state.records = state.records.saturating_sub(1); - if state.records == 0 { - self.sync(&mut state, false); + 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); } - /// Release what the binding's budget holds for the entry. - fn unbind(state: &mut PinState) { - if let Some(binding) = &state.binding { - if state.counted_leased { - binding.budget.count_leased(binding.bytes, false); - } - if state.reserved { - binding.budget.release(binding.partition, binding.bytes); - } - } - state.counted_leased = false; - state.reserved = false; - } - - /// Bring the budget in line with the holders and admissions: count a - /// leased admitted entry as leased, release it once unleased or no - /// longer admitted, and, when `reserve`, reserve its bytes if it is not - /// yet pinned. - fn sync(&self, state: &mut PinState, reserve: bool) { - let active = state.holders > 0 && state.records > 0; - match state.binding.clone() { - Some(binding) if active => { - if !state.counted_leased { - binding.budget.count_leased(binding.bytes, true); - state.counted_leased = true; - } - if reserve && !state.reserved { - state.reserved = binding.budget.try_reserve(binding.partition, binding.bytes); - } - } - _ => Self::unbind(state), + fn sync(&self, state: &mut PinState, is_leased: bool) { + for admission in &mut state.admissions { + admission.sync(is_leased); } - self.pinned.store(state.reserved, Ordering::Release); + 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, + ); } } @@ -324,12 +353,15 @@ impl CachePin { #[derive(Debug)] pub struct PinRecord { pin: Arc, + pinned: Arc, } impl PinRecord { - /// Whether the backend must keep the entry; see [`CachePin::is_pinned`]. + /// 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.pin.is_pinned() + self.pinned.load(Ordering::Acquire) } /// The pin of the recorded value. @@ -340,7 +372,7 @@ impl PinRecord { impl Drop for PinRecord { fn drop(&mut self) { - self.pin.drop_record(); + self.pin.drop_record(&self.pinned); } } @@ -357,7 +389,8 @@ impl CacheLease { &self.pin } - /// Whether the leased entry is pinned; see [`CachePin::is_pinned`]. + /// Whether any admission of the leased value is pinned; see + /// [`CachePin::is_pinned`]. pub fn is_pinned(&self) -> bool { self.pin.is_pinned() } @@ -465,24 +498,99 @@ mod tests { assert_eq!(budget.stats().pinned_bytes, 0); } - #[test] - fn admission_to_another_budget_moves_the_pin() { - let (first, second) = ( - Arc::new(PinBudget::new(1000, 1)), - Arc::new(PinBudget::new(1000, 1)), - ); + #[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 old = CachePin::record(&pin, &first, 0, 100); - assert_eq!(first.stats().pinned_bytes, 100); - let _new = CachePin::record(&pin, &second, 0, 100); - assert_eq!(first.stats().pinned_bytes, 0); - assert_eq!(second.stats().pinned_bytes, 100); + 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); - drop(old); - // The pin is still admitted once, so it stays pinned. + 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!(pin.is_pinned()); + 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. @@ -490,15 +598,51 @@ mod tests { 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/quick.rs b/rust/lance-core/src/cache/quick.rs index 9e94af5669b..9fefee9479a 100644 --- a/rust/lance-core/src/cache/quick.rs +++ b/rust/lance-core/src/cache/quick.rs @@ -1083,6 +1083,126 @@ mod tests { 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]);