diff --git a/rust/lance-core/src/cache/backend.rs b/rust/lance-core/src/cache/backend.rs index 013f4e262e7..cae4cc119ec 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 e9e3a9bce59..b6b22aea548 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 @@ -368,17 +376,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; }; @@ -454,7 +536,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); @@ -590,6 +672,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 + } + /// Weighted capacity in bytes, if the cache is alive and its backend /// reports one. pub fn capacity_bytes(&self) -> Option { @@ -1082,6 +1223,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 4a09b2147fa..92b32a5534d 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 028c1cb2023..20ce7b7766a 100644 --- a/rust/lance-core/src/cache/quick.rs +++ b/rust/lance-core/src/cache/quick.rs @@ -5,7 +5,10 @@ //! A hit takes a shard read lock, clones the cached value, and marks it as //! accessed. There is no read-operation channel or inline eviction work. +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; @@ -33,8 +36,11 @@ impl quick_cache::Weighter for EntryWeighter { } pub struct QuickCacheBackend { - cache: quick_cache::sync::Cache, capacity: usize, + generation: AtomicU64, + priority_active: AtomicBool, + priority: Mutex>, + cache: quick_cache::sync::Cache, } /// Controls how a [`QuickCacheBackend`] divides its weight budget. @@ -140,7 +146,51 @@ impl QuickCacheBackend { Default::default(), Default::default(), ); - Self { cache, capacity } + 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); } #[cfg(test)] @@ -156,7 +206,36 @@ impl QuickCacheBackend { #[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) } @@ -165,40 +244,86 @@ 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 capacity_bytes(&self) -> Option { @@ -207,10 +332,20 @@ impl CacheBackend for QuickCacheBackend { 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( @@ -218,14 +353,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::(), ) } } @@ -237,7 +385,7 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use super::*; - use crate::cache::{CacheKey, LanceCache}; + use crate::cache::{CacheKey, CacheTier, LanceCache}; const TEST_CAPACITY: usize = 1_000; @@ -265,6 +413,172 @@ 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()); + } + + #[rstest::rstest] + #[case::recommended(QuickCacheShardPolicy::Recommended)] + #[case::single(QuickCacheShardPolicy::Single)] + #[tokio::test] + async fn priority_budget_retains_metadata_and_singleflight_entries( + #[case] shard_policy: QuickCacheShardPolicy, + ) { + let cache = QuickCacheBackend::with_shard_policy(160, shard_policy); + 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()); + // Priority admission disables Quick's ordinary budget, but callers + // must still see the configured budget for the active priority tier. + assert_eq!(cache.capacity_bytes(), Some(160)); + 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); + assert_eq!(cache.capacity_bytes(), Some(160)); + 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]);