From 5c0c521af1ca4b526733dc9e2fe01239c432b035 Mon Sep 17 00:00:00 2001 From: Sergey Troshkov Date: Thu, 10 Sep 2026 09:15:14 +0700 Subject: [PATCH] fix: avoid index cache capacity fragmentation --- rust/lance-core/src/cache/mod.rs | 2 +- rust/lance-core/src/cache/quick.rs | 312 ++++++++++++++++++++++++++++- rust/lance/src/session.rs | 159 ++++++++++++--- 3 files changed, 441 insertions(+), 32 deletions(-) diff --git a/rust/lance-core/src/cache/mod.rs b/rust/lance-core/src/cache/mod.rs index 44b5604bd1a..83e3f8685d6 100644 --- a/rust/lance-core/src/cache/mod.rs +++ b/rust/lance-core/src/cache/mod.rs @@ -62,7 +62,7 @@ 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 quick::{QuickCacheBackend, recommended_cache_shards}; +pub use quick::{QuickCacheBackend, QuickCacheShardPolicy, recommended_cache_shards}; pub use registry::{BackendBuildFn, BackendConfig, build_from_config, register_backend}; use std::any::TypeId; diff --git a/rust/lance-core/src/cache/quick.rs b/rust/lance-core/src/cache/quick.rs index 53e0d0e8870..c1554d994fe 100644 --- a/rust/lance-core/src/cache/quick.rs +++ b/rust/lance-core/src/cache/quick.rs @@ -1,10 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -//! [`CacheBackend`] backed by [quick_cache](https://crates.io/crates/quick_cache), -//! 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. +//! [`CacheBackend`] backed by [quick_cache](https://crates.io/crates/quick_cache). +//! A hit takes a shard read lock, clones the cached value, and marks it as +//! accessed. There is no read-operation channel or inline eviction work. use std::pin::Pin; @@ -37,6 +36,19 @@ pub struct QuickCacheBackend { cache: quick_cache::sync::Cache, } +/// Controls how a [`QuickCacheBackend`] divides its weight budget. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum QuickCacheShardPolicy { + /// Choose a shard count from the cache capacity and available parallelism. + #[default] + Recommended, + /// Give every entry access to one shared weight budget. + /// + /// This avoids capacity fragmentation for a small number of large, + /// unequal entries. Concurrent hits can still take the shard read lock. + Single, +} + impl std::fmt::Debug for QuickCacheBackend { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("QuickCacheBackend") @@ -54,10 +66,17 @@ const MIN_SHARD_SHARE: usize = 4 << 30; /// keeps each shard's budget >= 4 GiB so large entries stay admissible. /// Rounded down because quick_cache rounds requests up. pub fn recommended_cache_shards(capacity: usize) -> usize { - let by_cpu = std::thread::available_parallelism() + let available_parallelism = std::thread::available_parallelism() .map(|n| n.get()) - .unwrap_or(2) - / 2; + .unwrap_or(2); + recommended_cache_shards_for_parallelism(capacity, available_parallelism) +} + +fn recommended_cache_shards_for_parallelism( + capacity: usize, + available_parallelism: usize, +) -> usize { + let by_cpu = available_parallelism / 2; let shards = (capacity / MIN_SHARD_SHARE).min(by_cpu).max(1); let shards = if shards.is_power_of_two() { shards @@ -75,7 +94,35 @@ impl QuickCacheBackend { /// (weight = key footprint + declared size), sharded per /// [`recommended_cache_shards`]. pub fn with_capacity(capacity: usize) -> Self { - let shards = recommended_cache_shards(capacity); + Self::with_shard_policy(capacity, QuickCacheShardPolicy::Recommended) + } + + /// Create a weighted cache with an explicit shard policy. + /// + /// `capacity` bounds the sum of key footprints and declared entry sizes. + /// Each shard has an independent share of that bound. In addition, the + /// default Quick admission policy rejects an unpinned entry heavier than + /// approximately 97% of one shard's share. + /// + /// # Example + /// + /// ``` + /// use lance_core::cache::{QuickCacheBackend, QuickCacheShardPolicy}; + /// + /// let cache = QuickCacheBackend::with_shard_policy( + /// 8 << 30, + /// QuickCacheShardPolicy::Single, + /// ); + /// ``` + pub fn with_shard_policy(capacity: usize, shard_policy: QuickCacheShardPolicy) -> Self { + let shards = match shard_policy { + QuickCacheShardPolicy::Recommended => recommended_cache_shards(capacity), + QuickCacheShardPolicy::Single => 1, + }; + Self::with_shards(capacity, shards) + } + + fn with_shards(capacity: usize, shards: usize) -> Self { // Floor protects the shard count from quick_cache's items-per-shard // heuristic; ceiling bounds pre-allocation. let estimated_items = (capacity / ESTIMATED_AVG_ENTRY_BYTES).clamp(shards * 32, 1_000_000); @@ -94,6 +141,16 @@ impl QuickCacheBackend { ); Self { cache } } + + #[cfg(test)] + fn num_shards(&self) -> usize { + self.cache.num_shards() + } + + #[cfg(test)] + fn shard_index(&self, key: &InternalCacheKey) -> usize { + self.cache.shard_index(key) + } } #[async_trait] @@ -177,6 +234,8 @@ mod tests { use super::*; use crate::cache::{CacheKey, LanceCache}; + const TEST_CAPACITY: usize = 1_000; + struct TestKey { key: String, _phantom: PhantomData, @@ -214,6 +273,243 @@ mod tests { ); } + fn keys_in_shard( + cache: &QuickCacheBackend, + shard_index: usize, + count: usize, + ) -> Vec { + let mut keys = Vec::with_capacity(count); + for value in 0_u128.. { + let key = InternalCacheKey::from_bytes(value.to_le_bytes()); + if cache.shard_index(&key) == shard_index { + keys.push(key); + if keys.len() == count { + return keys; + } + } + } + unreachable!("the key space must contain enough keys for every shard") + } + + async fn load_declared_value( + cache: &QuickCacheBackend, + key: &InternalCacheKey, + value: usize, + size_bytes: usize, + loads: &AtomicUsize, + ) -> CacheEntry { + let (entry, _) = cache + .get_or_insert( + key, + Box::pin(async { + loads.fetch_add(1, Ordering::SeqCst); + Ok((Arc::new(value) as CacheEntry, size_bytes)) + }), + None, + ) + .await + .unwrap(); + assert_eq!(*entry.downcast_ref::().unwrap(), value); + entry + } + + #[test] + fn recommended_shards_cover_large_capacity_boundaries() { + assert_eq!(recommended_cache_shards_for_parallelism(8 << 30, 4), 2); + assert_eq!( + recommended_cache_shards_for_parallelism((8 << 30) - 1, 4), + 1 + ); + assert_eq!(recommended_cache_shards_for_parallelism(16 << 30, 8), 4); + assert_eq!( + recommended_cache_shards_for_parallelism((16 << 30) - 1, 8), + 2 + ); + } + + #[tokio::test] + async fn single_shard_avoids_fragmentation_for_unequal_entries() { + // Declared entry weights, including the 16-byte key, total 660. They + // fit the whole cache but not one 500-byte share of a two-shard cache. + let sizes = [168, 188, 256]; + let sharded = QuickCacheBackend::with_shards(TEST_CAPACITY, 2); + assert_eq!(sharded.num_shards(), 2); + let sharded_keys = keys_in_shard(&sharded, 0, sizes.len()); + let sharded_loads = AtomicUsize::new(0); + for _ in 0..3 { + for (value, (key, size_bytes)) in sharded_keys.iter().zip(sizes).enumerate() { + load_declared_value(&sharded, key, value, size_bytes, &sharded_loads).await; + } + } + assert!(sharded_loads.load(Ordering::SeqCst) > sizes.len()); + assert!(sharded.num_entries().await < sizes.len()); + assert!(sharded.size_bytes().await <= TEST_CAPACITY); + + let single = + QuickCacheBackend::with_shard_policy(TEST_CAPACITY, QuickCacheShardPolicy::Single); + let single_loads = AtomicUsize::new(0); + for _ in 0..3 { + for (value, (key, size_bytes)) in sharded_keys.iter().zip(sizes).enumerate() { + load_declared_value(&single, key, value, size_bytes, &single_loads).await; + } + } + assert_eq!(single_loads.load(Ordering::SeqCst), sizes.len()); + assert_eq!(single.num_entries().await, sizes.len()); + assert_eq!(single.size_bytes().await, 660); + } + + #[tokio::test] + async fn direct_insert_obeys_the_same_shard_budget() { + let sizes = [168, 188, 256]; + for (shards, expected_entries) in [(2, 2), (1, 3)] { + let cache = QuickCacheBackend::with_shards(TEST_CAPACITY, shards); + let keys = keys_in_shard(&cache, 0, sizes.len()); + for (value, (key, size_bytes)) in keys.iter().zip(sizes).enumerate() { + cache.insert(key, Arc::new(value), size_bytes, None).await; + } + assert_eq!(cache.num_entries().await, expected_entries); + assert!(cache.size_bytes().await <= TEST_CAPACITY); + } + } + + #[tokio::test] + async fn single_shard_preserves_capacity_and_hot_admission_limits() { + let cache = QuickCacheBackend::with_shards(TEST_CAPACITY, 1); + let keys = keys_in_shard(&cache, 0, 4); + + // The default hot target is 97% of capacity. Weight includes the key, + // so a declared size of 954 is admitted at weight 970, while 955 is not. + cache.insert(&keys[0], Arc::new(0_usize), 954, None).await; + assert!(cache.get(&keys[0], None).await.is_some()); + cache.clear().await; + cache.insert(&keys[1], Arc::new(1_usize), 955, None).await; + assert!(cache.get(&keys[1], None).await.is_none()); + + // Individually admissible entries whose working set exceeds the total + // budget must still reload and remain bounded. + let loads = AtomicUsize::new(0); + for _ in 0..2 { + for (value, key) in keys[1..].iter().enumerate() { + load_declared_value(&cache, key, value, 384, &loads).await; + } + } + assert!(loads.load(Ordering::SeqCst) > 3); + assert!(cache.size_bytes().await <= TEST_CAPACITY); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn concurrent_misses_share_a_load_and_retry_failure() { + let cache = Arc::new(QuickCacheBackend::with_shards(4096, 1)); + let key = InternalCacheKey::from_bytes([42; 16]); + let loads = Arc::new(AtomicUsize::new(0)); + let release = Arc::new(tokio::sync::Notify::new()); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + + let owner = { + let cache = cache.clone(); + let loads = loads.clone(); + let release = release.clone(); + tokio::spawn(async move { + cache + .get_or_insert( + &key, + Box::pin(async move { + loads.fetch_add(1, Ordering::SeqCst); + let _ = started_tx.send(()); + release.notified().await; + Err(crate::Error::timeout("test loader failed")) + }), + None, + ) + .await + }) + }; + started_rx.await.unwrap(); + + let contender = { + let cache = cache.clone(); + let loads = loads.clone(); + tokio::spawn(async move { + cache + .get_or_insert( + &key, + Box::pin(async move { + loads.fetch_add(1, Ordering::SeqCst); + Ok((Arc::new(7_usize) as CacheEntry, 8)) + }), + None, + ) + .await + }) + }; + tokio::task::yield_now().await; + assert_eq!(loads.load(Ordering::SeqCst), 1); + release.notify_one(); + assert!(owner.await.unwrap().is_err()); + let (value, is_hit) = contender.await.unwrap().unwrap(); + assert_eq!(*value.downcast_ref::().unwrap(), 7); + assert!(!is_hit); + assert_eq!(loads.load(Ordering::SeqCst), 2); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn cancelled_loader_releases_concurrent_miss() { + let cache = Arc::new(QuickCacheBackend::with_shards(4096, 1)); + let key = InternalCacheKey::from_bytes([24; 16]); + let loads = Arc::new(AtomicUsize::new(0)); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + + let owner = { + let cache = cache.clone(); + let loads = loads.clone(); + tokio::spawn(async move { + cache + .get_or_insert( + &key, + Box::pin(async move { + loads.fetch_add(1, Ordering::SeqCst); + let _ = started_tx.send(()); + std::future::pending::<()>().await; + Ok((Arc::new(1_usize) as CacheEntry, 8)) + }), + None, + ) + .await + }) + }; + started_rx.await.unwrap(); + + let contender = { + let cache = cache.clone(); + let loads = loads.clone(); + tokio::spawn(async move { + cache + .get_or_insert( + &key, + Box::pin(async move { + loads.fetch_add(1, Ordering::SeqCst); + Ok((Arc::new(2_usize) as CacheEntry, 8)) + }), + None, + ) + .await + }) + }; + tokio::task::yield_now().await; + assert_eq!(loads.load(Ordering::SeqCst), 1); + + owner.abort(); + assert!(owner.await.unwrap_err().is_cancelled()); + let (value, is_hit) = tokio::time::timeout(std::time::Duration::from_secs(1), contender) + .await + .expect("contender remained blocked after loader cancellation") + .unwrap() + .unwrap(); + assert_eq!(*value.downcast_ref::().unwrap(), 2); + assert!(!is_hit); + assert_eq!(loads.load(Ordering::SeqCst), 2); + } + #[tokio::test] async fn test_quick_backend_roundtrip_singleflight_and_eviction() { // Capacity must be large relative to one entry: quick_cache shards diff --git a/rust/lance/src/session.rs b/rust/lance/src/session.rs index 8328ba20b22..83cde73be48 100644 --- a/rust/lance/src/session.rs +++ b/rust/lance/src/session.rs @@ -4,7 +4,7 @@ use std::collections::HashMap; use std::sync::Arc; -use lance_core::cache::{CacheBackend, LanceCache, QuickCacheBackend}; +use lance_core::cache::{CacheBackend, LanceCache, QuickCacheBackend, QuickCacheShardPolicy}; use lance_core::deepsize::DeepSizeOf; use lance_core::{Error, Result}; use lance_index::IndexType; @@ -17,6 +17,10 @@ use crate::session::index_caches::GlobalIndexCache; use self::index_extension::IndexExtension; +/// Automatic index caches favor one shared budget because they commonly hold +/// a small number of large, unequal partition entries. +const AUTOMATIC_INDEX_CACHE_SHARD_POLICY: QuickCacheShardPolicy = QuickCacheShardPolicy::Single; + pub(crate) mod caches; pub mod index_caches; pub(crate) mod index_extension; @@ -26,7 +30,10 @@ pub(crate) mod index_extension; pub enum CacheSpec { /// Use the Lance-level default capacity for the tier. Default, - /// Use the default in-memory backend with this capacity. + /// Use the default in-memory backend with this weighted-entry capacity. + /// + /// This is not an RSS limit: active queries and in-progress loaders can + /// retain memory in addition to resident cache entries. Size(usize), /// Use an already constructed backend. Backend(Arc), @@ -106,8 +113,9 @@ impl Session { /// /// Parameters: /// - /// - ***index_cache_size***: the size of the index cache, backed by - /// [`QuickCacheBackend`]. + /// - ***index_cache_size***: the weighted-entry budget of the index cache, + /// backed by a single-shard [`QuickCacheBackend`]. The budget is shared + /// by all indices in the session and is not a process RSS limit. /// - ***metadata_cache_size***: the size of the metadata cache, backed by /// [`QuickCacheBackend`]. /// - ***store_registry***: the object store registry to use when opening @@ -119,12 +127,14 @@ impl Session { store_registry: Arc, ) -> Self { Self { - index_cache: GlobalIndexCache(LanceCache::with_backend(Arc::new( - QuickCacheBackend::with_capacity(index_cache_size), - ))), - metadata_cache: GlobalMetadataCache(LanceCache::with_backend(Arc::new( - QuickCacheBackend::with_capacity(metadata_cache_size), - ))), + index_cache: GlobalIndexCache(Self::build_automatic_quick_cache( + index_cache_size, + AUTOMATIC_INDEX_CACHE_SHARD_POLICY, + )), + metadata_cache: GlobalMetadataCache(Self::build_automatic_quick_cache( + metadata_cache_size, + QuickCacheShardPolicy::Recommended, + )), index_extensions: HashMap::new(), store_registry, spill_store: Arc::new(LocalSpillStore::default()), @@ -142,9 +152,10 @@ impl Session { ) -> Self { Self { index_cache: GlobalIndexCache(LanceCache::with_backend(index_cache_backend)), - metadata_cache: GlobalMetadataCache(LanceCache::with_backend(Arc::new( - QuickCacheBackend::with_capacity(metadata_cache_size), - ))), + metadata_cache: GlobalMetadataCache(Self::build_automatic_quick_cache( + metadata_cache_size, + QuickCacheShardPolicy::Recommended, + )), index_extensions: HashMap::new(), store_registry, spill_store: Arc::new(LocalSpillStore::default()), @@ -180,7 +191,7 @@ impl Session { /// /// Each [`CacheSpec`] controls one tier. [`CacheSpec::Default`] uses that /// tier's Lance-level default capacity, [`CacheSpec::Size`] uses the - /// default in-memory backend with an explicit capacity, and + /// default in-memory backend with an explicit weighted-entry budget, and /// [`CacheSpec::Backend`] uses a caller-provided backend. This keeps size /// and backend selection mutually exclusive. /// @@ -211,8 +222,16 @@ impl Session { metadata_cache: CacheSpec, store_registry: Arc, ) -> Self { - let index_cache = Self::build_cache(index_cache, DEFAULT_INDEX_CACHE_SIZE); - let metadata_cache = Self::build_cache(metadata_cache, DEFAULT_METADATA_CACHE_SIZE); + let index_cache = Self::build_cache( + index_cache, + DEFAULT_INDEX_CACHE_SIZE, + AUTOMATIC_INDEX_CACHE_SHARD_POLICY, + ); + let metadata_cache = Self::build_cache( + metadata_cache, + DEFAULT_METADATA_CACHE_SIZE, + QuickCacheShardPolicy::Recommended, + ); Self { index_cache: GlobalIndexCache(index_cache), metadata_cache: GlobalMetadataCache(metadata_cache), @@ -222,18 +241,28 @@ impl Session { } } - fn build_cache(spec: CacheSpec, default_size: usize) -> LanceCache { + fn build_cache( + spec: CacheSpec, + default_size: usize, + shard_policy: QuickCacheShardPolicy, + ) -> LanceCache { match spec { - CacheSpec::Default => { - LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(default_size))) - } - CacheSpec::Size(size) => { - LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(size))) - } + CacheSpec::Default => Self::build_automatic_quick_cache(default_size, shard_policy), + CacheSpec::Size(size) => Self::build_automatic_quick_cache(size, shard_policy), CacheSpec::Backend(backend) => LanceCache::with_backend(backend), } } + fn build_automatic_quick_cache( + capacity: usize, + shard_policy: QuickCacheShardPolicy, + ) -> LanceCache { + LanceCache::with_backend(Arc::new(QuickCacheBackend::with_shard_policy( + capacity, + shard_policy, + ))) + } + /// Register a new index extension. /// /// A name can only be registered once per type of index extension. @@ -332,6 +361,7 @@ impl Default for Session { mod tests { use super::*; use lance_core::cache::{CacheKey, UnsizedCacheKey}; + use lance_core::deepsize::Context; use lance_index::vector::VectorIndex; use std::borrow::Cow; use tokio::io::AsyncWriteExt; @@ -361,6 +391,50 @@ mod tests { } } + struct DeclaredWeight(usize); + + impl DeepSizeOf for DeclaredWeight { + fn deep_size_of_children(&self, _context: &mut Context) -> usize { + self.0 + } + } + + struct WeightedKey(u64); + + impl CacheKey for WeightedKey { + type ValueType = DeclaredWeight; + + fn key(&self) -> Cow<'_, str> { + self.0.to_string().into() + } + + fn type_name() -> &'static str { + "DeclaredWeight" + } + } + + async fn assert_fitting_working_set_is_resident(session: &Session, gibibytes: &[usize]) { + for (index, gibibytes) in gibibytes.iter().copied().enumerate() { + session + .index_cache + .insert_with_key( + &WeightedKey(index as u64), + Arc::new(DeclaredWeight(gibibytes << 30)), + ) + .await; + } + for index in 0..gibibytes.len() { + assert!( + session + .index_cache + .get_with_key(&WeightedKey(index as u64)) + .await + .is_some(), + "entry {index} should remain resident" + ); + } + } + #[tokio::test] async fn test_disable_index_cache() { let no_cache = Session::new(0, 0, Default::default()); @@ -373,6 +447,45 @@ mod tests { ); } + #[tokio::test] + async fn automatic_index_caches_use_one_shared_weight_budget() { + assert_eq!( + AUTOMATIC_INDEX_CACHE_SHARD_POLICY, + QuickCacheShardPolicy::Single + ); + + let session = Session::new(8 << 30, 0, Default::default()); + assert_fitting_working_set_is_resident(&session, &[3, 2, 2]).await; + + let session = Session::with_cache_backends( + CacheSpec::Size(8 << 30), + CacheSpec::Size(0), + Default::default(), + ); + assert_fitting_working_set_is_resident(&session, &[3, 2, 2]).await; + + let session = Session::with_cache_backends( + CacheSpec::Default, + CacheSpec::Size(0), + Default::default(), + ); + assert_fitting_working_set_is_resident(&session, &[2, 1, 1]).await; + + let session = Session::default(); + let shared_session = session.clone(); + session + .index_cache + .insert_with_key(&TestKey("shared-session"), Arc::new(vec![7])) + .await; + assert!( + shared_session + .index_cache + .get_with_key(&TestKey("shared-session")) + .await + .is_some() + ); + } + /// `with_cache_backends` should honor whichever tier the caller /// provided a backend for and fall back to that tier's default on the /// other tier.