diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cdabc9a08..7b03ebe91 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -139,6 +139,13 @@ jobs: run: | cd src/core cargo test --features "shuttle" --release -- --test-threads=1 shuttle + # The parquet cache keeps its own concurrent state — the file id + # pool — in this package, and it has its own `shuttle` feature. + # Without this step the model checker never sees it. + - name: Run shuttle test (datafusion) + run: | + cd src/datafusion + cargo test --features "shuttle" --release -- --test-threads=1 shuttle address_san: name: Address Sanitizer diff --git a/benchmark/src/inprocess_runner.rs b/benchmark/src/inprocess_runner.rs index b45707cf6..03b6c8490 100644 --- a/benchmark/src/inprocess_runner.rs +++ b/benchmark/src/inprocess_runner.rs @@ -657,7 +657,7 @@ impl InProcessBenchmarkRunner { && let Some(cache) = &cache { unsafe { - cache.reset(); + cache.reset().await; } } diff --git a/examples/core.rs b/examples/core.rs index d172f601a..79644b639 100644 --- a/examples/core.rs +++ b/examples/core.rs @@ -18,13 +18,20 @@ async fn main() -> Result<(), Box> { .await; let entry_id = EntryID::from(7); + // Names whose data this is. Entry ids are packed and can alias between + // sources; the cache compares this and treats a mismatch as a miss, so a + // caller only ever reads back what it put in. + let identity = 1; let arrow_array = Arc::new(UInt64Array::from_iter_values(0..16)); - storage.insert(entry_id, arrow_array.clone()).await.unwrap(); + storage + .insert(entry_id, identity, arrow_array.clone()) + .await + .unwrap(); // Move data to disk so the read demonstrates async I/O storage.flush_all_to_disk().await.unwrap(); - let retrieved = storage.get(&entry_id).await.unwrap(); + let retrieved = storage.get(&entry_id, identity).await.unwrap(); assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); Ok(()) diff --git a/src/core/README.md b/src/core/README.md index d263aab51..e776a9ef7 100644 --- a/src/core/README.md +++ b/src/core/README.md @@ -22,12 +22,16 @@ tokio_test::block_on(async { let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(42); +// Names whose data this is. Entry ids are packed and can alias between +// sources, so the cache compares this on every read and treats a mismatch as +// a miss — a caller only ever reads back what it put in. +let identity = 1; let arrow_array = Arc::new(UInt64Array::from_iter_values(0..1000)); // Insert once; replacement/placement is handled by the cache policy -storage.insert(entry_id, arrow_array.clone()).await; +storage.insert(entry_id, identity, arrow_array.clone()).await; -assert!(storage.is_cached(&entry_id)); +assert!(storage.is_cached(&entry_id, identity)); }); ``` @@ -42,14 +46,15 @@ tokio_test::block_on(async { let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(7); +let identity = 1; let arrow_array = Arc::new(UInt64Array::from_iter_values(0..16)); -storage.insert(entry_id, arrow_array.clone()).await; +storage.insert(entry_id, identity, arrow_array.clone()).await; // Move data to disk so the read will demonstrate async I/O storage.flush_all_to_disk().await; // Read asynchronously -let retrieved = storage.get(&entry_id).await.unwrap(); +let retrieved = storage.get(&entry_id, identity).await.unwrap(); assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); }); ``` @@ -71,10 +76,11 @@ tokio_test::block_on(async { let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(8); +let identity = 1; let data = Arc::new(StringArray::from(vec![ Some("apple"), Some("banana"), None, Some("apple"), Some("cherry"), ])); -storage.insert(entry_id, data.clone()).await; +storage.insert(entry_id, identity, data.clone()).await; // Move data to disk so the read will demonstrate async I/O storage.flush_all_to_disk().await; @@ -94,7 +100,7 @@ let liquid_expr = liquid_cache::cache::LiquidExpr::try_new( // Read with predicate pushdown let mask = storage - .eval_predicate(&entry_id, &liquid_expr) + .eval_predicate(&entry_id, identity, &liquid_expr) .with_selection(&selection) .await .unwrap(); diff --git a/src/core/src/cache/builders.rs b/src/core/src/cache/builders.rs index af0572dac..e1db0e236 100644 --- a/src/core/src/cache/builders.rs +++ b/src/core/src/cache/builders.rs @@ -11,6 +11,7 @@ use super::core::LiquidCache; use super::io_context::{DefaultCacheMetadata, EntryMetadata}; use super::policies::{CachePolicy, HydrationPolicy, SqueezePolicy, TranscodeSqueezeEvict}; use super::{CacheExpression, CacheFull, EntryID, LiquidExpr, LiquidPolicy}; +use crate::cache::index::WriteIdentity; use crate::sync::Arc; /// Builder for [LiquidCache]. @@ -175,16 +176,23 @@ pub fn default_max_memory_bytes() -> usize { pub struct Insert<'a> { pub(super) storage: &'a Arc, pub(super) entry_id: EntryID, + pub(super) identity: u64, pub(super) batch: ArrayRef, pub(super) skip_gc: bool, pub(super) squeeze_hint: Option>, } impl<'a> Insert<'a> { - pub(super) fn new(storage: &'a Arc, entry_id: EntryID, batch: ArrayRef) -> Self { + pub(super) fn new( + storage: &'a Arc, + entry_id: EntryID, + identity: u64, + batch: ArrayRef, + ) -> Self { Self { storage, entry_id, + identity, batch, skip_gc: false, squeeze_hint: None, @@ -214,7 +222,9 @@ impl<'a> Insert<'a> { } let batch = CacheEntry::memory_arrow(batch); self.storage.supersede_disk_copy(self.entry_id).await; - self.storage.insert_inner(self.entry_id, batch).await + self.storage + .insert_inner(self.entry_id, WriteIdentity::Owned(self.identity), batch) + .await } } @@ -232,15 +242,17 @@ impl<'a> IntoFuture for Insert<'a> { pub struct Get<'a> { pub(super) storage: &'a LiquidCache, pub(super) entry_id: &'a EntryID, + pub(super) identity: u64, pub(super) selection: Option<&'a BooleanBuffer>, pub(super) expression_hint: Option>, } impl<'a> Get<'a> { - pub(super) fn new(storage: &'a LiquidCache, entry_id: &'a EntryID) -> Self { + pub(super) fn new(storage: &'a LiquidCache, entry_id: &'a EntryID, identity: u64) -> Self { Self { storage, entry_id, + identity, selection: None, expression_hint: None, } @@ -273,6 +285,7 @@ impl<'a> Get<'a> { self.storage .read_arrow_array( self.entry_id, + self.identity, self.selection, self.expression_hint.as_deref(), ) @@ -328,6 +341,7 @@ fn maybe_gc_view_arrays(array: &ArrayRef) -> Option { pub struct EvaluatePredicate<'a> { pub(super) storage: &'a LiquidCache, pub(super) entry_id: &'a EntryID, + pub(super) identity: u64, pub(super) predicate: &'a LiquidExpr, pub(super) selection: Option<&'a BooleanBuffer>, } @@ -336,11 +350,13 @@ impl<'a> EvaluatePredicate<'a> { pub(super) fn new( storage: &'a LiquidCache, entry_id: &'a EntryID, + identity: u64, predicate: &'a LiquidExpr, ) -> Self { Self { storage, entry_id, + identity, predicate, selection: None, } @@ -355,7 +371,7 @@ impl<'a> EvaluatePredicate<'a> { /// Evaluate the predicate against the cached data. pub async fn read(self) -> Option { self.storage - .eval_predicate_internal(self.entry_id, self.selection, self.predicate) + .eval_predicate_internal(self.entry_id, self.identity, self.selection, self.predicate) .await } } @@ -453,9 +469,9 @@ mod tests { let cache = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(123usize); - cache.insert(entry_id, root.clone()).await.unwrap(); + cache.insert(entry_id, 0, root.clone()).await.unwrap(); - let stored = cache.get(&entry_id).await.expect("array present"); + let stored = cache.get(&entry_id, 0).await.expect("array present"); let post_size = stored.get_array_memory_size(); // GC should have compacted the view arrays, reducing memory footprint. diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 2f85a1a7d..b6bc1c08a 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -18,7 +18,11 @@ use super::{ use crate::cache::DefaultSqueezeIo; use crate::cache::policies::{SqueezeOutcome, SqueezePolicy}; use crate::cache::utils::{LiquidCompressorStates, arrow_to_bytes}; -use crate::cache::{CacheExpression, LiquidExpr, index::ArtIndex, utils::EntryID}; +use crate::cache::{ + CacheExpression, LiquidExpr, + index::{ArtIndex, WriteIdentity}, + utils::EntryID, +}; use crate::cache::{CacheFull, CacheStats, EventTrace}; use crate::liquid_array::{ LiquidSqueezedArrayRef, SqueezeIoHandler, SqueezedBacking, SqueezedDate32Array, @@ -40,6 +44,11 @@ use std::collections::HashMap; /// (liquid-cache#43). #[derive(Debug, Clone, Copy)] struct DiskCopy { + /// Whose bytes these are. A key can change hands while a write to it is + /// in flight, and a declined rewrite leaves the object behind; without + /// this the next owner adopts it on kind and length alone and reads the + /// previous owner's rows as its own. + identity: u64, kind: DiskKind, bytes: usize, } @@ -53,22 +62,26 @@ enum DiskKind { impl DiskCopy { /// The store object an entry refers to: a disk stub's bytes, or the /// full serialisation a squeezed entry reads back through. - fn referenced_by(entry: &CacheEntry) -> Option { + fn referenced_by(identity: u64, entry: &CacheEntry) -> Option { match entry { CacheEntry::DiskLiquid { disk_bytes, .. } => Some(Self { + identity, kind: DiskKind::Liquid, bytes: *disk_bytes, }), CacheEntry::DiskArrow { disk_bytes, .. } => Some(Self { + identity, kind: DiskKind::Arrow, bytes: *disk_bytes, }), CacheEntry::MemorySqueezedLiquid(squeezed) => Some(match squeezed.disk_backing() { SqueezedBacking::Liquid(bytes) => Self { + identity, kind: DiskKind::Liquid, bytes, }, SqueezedBacking::Arrow(bytes) => Self { + identity, kind: DiskKind::Arrow, bytes, }, @@ -91,10 +104,10 @@ impl DiskCopy { /// /// let entry_id = EntryID::from(0); /// let arrow_array = Arc::new(UInt64Array::from_iter_values(0..32)); -/// storage.insert(entry_id, arrow_array.clone()).await; +/// storage.insert(entry_id, 0, arrow_array.clone()).await; /// /// // Get the arrow array back asynchronously -/// let retrieved = storage.get(&entry_id).await.unwrap(); +/// let retrieved = storage.get(&entry_id, 0).await.unwrap(); /// assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); /// }); /// ``` @@ -130,7 +143,7 @@ impl LiquidCache { let mut memory_liquid_bytes = 0usize; let mut memory_squeezed_liquid_bytes = 0usize; - self.index.for_each(|_, batch| match batch { + self.index.for_each(|_, _, batch| match batch { CacheEntry::MemoryArrow(array) => { memory_arrow_entries += 1; memory_arrow_bytes += array.get_array_memory_size(); @@ -161,6 +174,7 @@ impl LiquidCache { memory_arrow_bytes, memory_liquid_bytes, memory_squeezed_liquid_bytes, + identity_mismatches: self.index.identity_mismatches(), memory_usage_bytes, disk_usage_bytes, max_memory_bytes: self.config.max_memory_bytes(), @@ -173,23 +187,25 @@ impl LiquidCache { pub fn insert<'a>( self: &'a Arc, entry_id: EntryID, + identity: u64, batch_to_cache: ArrayRef, ) -> Insert<'a> { - Insert::new(self, entry_id, batch_to_cache) + Insert::new(self, entry_id, identity, batch_to_cache) } /// Create a [`Get`] builder for the provided entry. - pub fn get<'a>(&'a self, entry_id: &'a EntryID) -> Get<'a> { - Get::new(self, entry_id) + pub fn get<'a>(&'a self, entry_id: &'a EntryID, identity: u64) -> Get<'a> { + Get::new(self, entry_id, identity) } /// Create an [`EvaluatePredicate`] builder for evaluating predicates on cached data. pub fn eval_predicate<'a>( &'a self, entry_id: &'a EntryID, + identity: u64, predicate: &'a LiquidExpr, ) -> EvaluatePredicate<'a> { - EvaluatePredicate::new(self, entry_id, predicate) + EvaluatePredicate::new(self, entry_id, identity, predicate) } /// Try to read a liquid array from the cache. @@ -197,24 +213,31 @@ impl LiquidCache { pub async fn try_read_liquid( &self, entry_id: &EntryID, + identity: u64, ) -> Option { self.observer.on_try_read_liquid(); self.trace(InternalEvent::TryReadLiquid { entry: *entry_id }); - let batch = self.index.get(entry_id)?; + let batch = self.index.get_checked(entry_id, identity)?; self.cache_policy .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); match batch.as_ref() { CacheEntry::MemoryLiquid(array) => Some(array.clone()), entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) - .await; + let liquid = self.read_disk_liquid_array(entry_id, identity).await; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Liquid(&liquid), + None, + ) + .await; Some(liquid) } CacheEntry::MemorySqueezedLiquid(array) => match array.disk_backing() { SqueezedBacking::Liquid(_) => { - let liquid = self.read_disk_liquid_array(entry_id).await; + let liquid = self.read_disk_liquid_array(entry_id, identity).await; Some(liquid) } SqueezedBacking::Arrow(_) => None, @@ -226,20 +249,35 @@ impl LiquidCache { /// Iterate over all entries in the cache. /// No guarantees are made about the order of the entries. /// Isolation level: read-committed - pub fn for_each_entry(&self, mut f: impl FnMut(&EntryID, &CacheEntry)) { + pub fn for_each_entry(&self, mut f: impl FnMut(&EntryID, u64, &CacheEntry)) { self.index.for_each(&mut f); } /// Reset the cache. - pub fn reset(&self) { + /// + /// Deletes the store objects before forgetting the records that name them. + /// The store key carries the identity that wrote it, so everything after a + /// reset writes under new keys and nothing would ever overwrite these + /// again — dropping the records first would strand up to `max_disk_bytes` + /// per reset, unreachable and uncounted. + pub async fn reset(&self) { + let recorded: Vec<(EntryID, DiskCopy)> = { + let mut copies = self.disk_copies.lock().unwrap(); + copies.drain().collect() + }; + for (entry_id, copy) in recorded { + self.store + .remove(&entry_id_to_key(&entry_id, copy.identity)) + .await + .expect("disk remove failed"); + } self.index.reset(); self.budget.reset_usage(); - self.disk_copies.lock().unwrap().clear(); } /// Check if a batch is cached. - pub fn is_cached(&self, entry_id: &EntryID) -> bool { - self.index.is_cached(entry_id) + pub fn is_cached(&self, entry_id: &EntryID, identity: u64) -> bool { + self.index.is_cached(entry_id, identity) } /// Get the config of the cache. @@ -275,18 +313,27 @@ impl LiquidCache { /// Flush all entries to disk. pub async fn flush_all_to_disk(&self) -> Result<(), CacheFull> { let mut entires = Vec::new(); - self.for_each_entry(|entry_id, batch| { - entires.push((*entry_id, batch.clone())); + self.for_each_entry(|entry_id, identity, batch| { + entires.push((*entry_id, identity, batch.clone())); }); - for (entry_id, batch) in entires { + for (entry_id, flush_identity, batch) in entires { match &batch { CacheEntry::MemoryArrow(array) => { let bytes = arrow_to_bytes(array).expect("failed to convert arrow to bytes"); let disk_bytes = bytes.len(); - match self.write_batch_to_disk(entry_id, &batch, bytes).await { + match self + .write_batch_to_disk( + entry_id, + WriteIdentity::Rewrite(flush_identity), + &batch, + bytes, + ) + .await + { Ok(()) => { self.try_insert( entry_id, + WriteIdentity::Rewrite(flush_identity), CacheEntry::disk_arrow(array.data_type().clone(), disk_bytes), ) .expect("failed to insert disk arrow entry"); @@ -299,24 +346,35 @@ impl LiquidCache { if let Some(DiskCopy { kind: DiskKind::Liquid, bytes, - }) = self.disk_copy(&entry_id) + .. + }) = self.disk_copy(&entry_id, flush_identity) { // Hydrated from disk and never modified since: the // bytes are already there, flip the index rather // than re-serialising and rewriting them. - self.try_insert(entry_id, CacheEntry::disk_liquid(data_type, bytes)) - .expect("failed to insert disk liquid entry"); + self.try_insert( + entry_id, + WriteIdentity::Rewrite(flush_identity), + CacheEntry::disk_liquid(data_type, bytes), + ) + .expect("failed to insert disk liquid entry"); continue; } let liquid_bytes = liquid_array.to_bytes(); let disk_bytes = liquid_bytes.len(); match self - .write_batch_to_disk(entry_id, &batch, Bytes::from(liquid_bytes)) + .write_batch_to_disk( + entry_id, + WriteIdentity::Rewrite(flush_identity), + &batch, + Bytes::from(liquid_bytes), + ) .await { Ok(()) => { self.try_insert( entry_id, + WriteIdentity::Rewrite(flush_identity), CacheEntry::disk_liquid(data_type, disk_bytes), ) .expect("failed to insert disk liquid entry"); @@ -327,7 +385,7 @@ impl LiquidCache { CacheEntry::MemorySqueezedLiquid(array) => { // We don't have to do anything, because it's already on disk let disk_entry = Self::disk_entry_from_squeezed(array); - self.try_insert(entry_id, disk_entry) + self.try_insert(entry_id, WriteIdentity::Rewrite(flush_identity), disk_entry) .expect("failed to insert disk entry"); } CacheEntry::DiskArrow { .. } | CacheEntry::DiskLiquid { .. } => { @@ -344,6 +402,7 @@ impl LiquidCache { async fn write_in_memory_batch_to_disk( &self, entry_id: EntryID, + identity: WriteIdentity, batch: CacheEntry, ) -> Result { match &batch { @@ -351,6 +410,7 @@ impl LiquidCache { let squeeze_io: Arc = Arc::new(DefaultSqueezeIo::new( self.store.clone(), entry_id, + identity.value(), self.observer.clone(), )); let outcome = self.squeeze_policy.squeeze( @@ -367,7 +427,7 @@ impl LiquidCache { unreachable!("memory arrow squeeze cannot remove entry"); }; if let Some(bytes_to_write) = bytes_to_write { - self.write_batch_to_disk(entry_id, &new_batch, bytes_to_write) + self.write_batch_to_disk(entry_id, identity, &new_batch, bytes_to_write) .await?; } Ok(new_batch) @@ -377,13 +437,14 @@ impl LiquidCache { if let Some(DiskCopy { kind: DiskKind::Liquid, bytes, - }) = self.disk_copy(&entry_id) + .. + }) = self.disk_copy(&entry_id, identity.value()) { return Ok(CacheEntry::disk_liquid(data_type, bytes)); } let liquid_bytes = Bytes::from(liquid_array.to_bytes()); let disk_bytes = liquid_bytes.len(); - self.write_batch_to_disk(entry_id, &batch, liquid_bytes) + self.write_batch_to_disk(entry_id, identity, &batch, liquid_bytes) .await?; Ok(CacheEntry::disk_liquid(data_type, disk_bytes)) } @@ -406,10 +467,11 @@ impl LiquidCache { pub(crate) async fn insert_inner( &self, entry_id: EntryID, + identity: WriteIdentity, mut batch_to_cache: CacheEntry, ) -> Result<(), CacheFull> { loop { - let Err(not_inserted) = self.try_insert(entry_id, batch_to_cache) else { + let Err(not_inserted) = self.try_insert(entry_id, identity, batch_to_cache) else { return Ok(()); }; self.trace(InternalEvent::InsertFailed { @@ -423,7 +485,7 @@ impl LiquidCache { // this can happen if the entry to be inserted is too large, in that case, // we write it to disk let on_disk_batch = self - .write_in_memory_batch_to_disk(entry_id, not_inserted) + .write_in_memory_batch_to_disk(entry_id, identity, not_inserted) .await?; batch_to_cache = on_disk_batch; continue; @@ -469,7 +531,21 @@ impl LiquidCache { } } - fn disk_copy(&self, entry_id: &EntryID) -> Option { + /// The store object recorded for `entry_id`, but only if it belongs to + /// `identity`. A copy left by a previous owner reads as absent, so it is + /// never adopted by whoever holds the key now. + fn disk_copy(&self, entry_id: &EntryID, identity: u64) -> Option { + self.disk_copies + .lock() + .unwrap() + .get(entry_id) + .copied() + .filter(|copy| copy.identity == identity) + } + + /// The record regardless of owner, for paths that act on whatever object + /// is there — superseding it, discarding it, releasing its reservation. + fn any_disk_copy(&self, entry_id: &EntryID) -> Option { self.disk_copies.lock().unwrap().get(entry_id).copied() } @@ -484,7 +560,7 @@ impl LiquidCache { /// still land its result after the new one), so this covers the /// sequential case only. pub(crate) async fn supersede_disk_copy(&self, entry_id: EntryID) { - if self.disk_copy(&entry_id).is_none() { + if self.any_disk_copy(&entry_id).is_none() { return; } match self.index.get(&entry_id).as_deref() { @@ -512,7 +588,7 @@ impl LiquidCache { return; }; self.store - .remove(&entry_id_to_key(&entry_id)) + .remove(&entry_id_to_key(&entry_id, copy.identity)) .await .expect("disk remove failed"); self.budget.release_disk(copy.bytes); @@ -521,7 +597,12 @@ impl LiquidCache { /// If `outcome` demotes an entry to a form backed by a store object whose /// bytes are already there, drop the write and point the entry at the /// existing copy. - fn reuse_disk_copy(&self, entry_id: &EntryID, outcome: SqueezeOutcome) -> SqueezeOutcome { + fn reuse_disk_copy( + &self, + entry_id: &EntryID, + identity: u64, + outcome: SqueezeOutcome, + ) -> SqueezeOutcome { let (entry, bytes) = match outcome { SqueezeOutcome::Replace { entry, @@ -533,9 +614,10 @@ impl LiquidCache { entry, bytes_to_write: Some(bytes), }; - let (Some(copy), Some(wanted)) = - (self.disk_copy(entry_id), DiskCopy::referenced_by(&entry)) - else { + let (Some(copy), Some(wanted)) = ( + self.disk_copy(entry_id, identity), + DiskCopy::referenced_by(identity, &entry), + ) else { return keep_write(entry); }; if copy.kind != wanted.kind { @@ -568,7 +650,16 @@ impl LiquidCache { } } - fn try_insert(&self, entry_id: EntryID, to_insert: CacheEntry) -> Result<(), CacheEntry> { + /// A declined write is not an error: the entry being rewritten has since + /// been taken over or removed. Neither is worth retrying, so the + /// reservation is handed back and the call reports success with nothing + /// stored. See [`WriteIdentity`]. + fn try_insert( + &self, + entry_id: EntryID, + identity: WriteIdentity, + to_insert: CacheEntry, + ) -> Result<(), CacheEntry> { let new_memory_size = to_insert.memory_usage_bytes(); let cached_batch_type = if let Some(entry) = self.index.get(&entry_id) { let old_memory_size = entry.memory_usage_bytes(); @@ -580,14 +671,28 @@ impl LiquidCache { return Err(to_insert); } let batch_type = CachedBatchType::from(&to_insert); - self.index.insert(&entry_id, to_insert); + if !self.index.insert(&entry_id, identity, to_insert) { + // Restoring the reservation *grows* it again when the entry we + // were replacing was larger, so this can legitimately fail on + // a full cache. Nothing is stored either way; the budget is + // left under-counted rather than the process brought down. + let _ = self + .budget + .try_update_memory_usage(new_memory_size, old_memory_size); + return Ok(()); + } batch_type } else { if self.budget.try_reserve_memory(new_memory_size).is_err() { return Err(to_insert); } let batch_type = CachedBatchType::from(&to_insert); - self.index.insert(&entry_id, to_insert); + if !self.index.insert(&entry_id, identity, to_insert) { + self.budget + .try_update_memory_usage(new_memory_size, 0) + .expect("memory release cannot fail"); + return Ok(()); + } batch_type }; @@ -633,11 +738,15 @@ impl LiquidCache { | CacheEntry::DiskArrow { disk_bytes, .. } => *disk_bytes, _ => panic!("remove_disk_entry called for non-disk entry"), }; - self.store - .remove(&entry_id_to_key(&entry_id)) - .await - .expect("disk remove failed"); - self.disk_copies.lock().unwrap().remove(&entry_id); + // Take the record first: it names the owner whose object this is, and + // the key needs it. + let removed_copy = self.disk_copies.lock().unwrap().remove(&entry_id); + if let Some(copy) = removed_copy { + self.store + .remove(&entry_id_to_key(&entry_id, copy.identity)) + .await + .expect("disk remove failed"); + } self.budget.release_disk(disk_bytes); self.cache_policy.notify_remove(&entry_id); self.trace(InternalEvent::DiskEvict { @@ -682,7 +791,9 @@ impl LiquidCache { } async fn squeeze_victim_inner(&self, to_squeeze: EntryID) -> Result<(), CacheFull> { - let Some(mut to_squeeze_batch) = self.index.get(&to_squeeze) else { + let Some((squeezed_identity, mut to_squeeze_batch)) = + self.index.get_with_identity(&to_squeeze) + else { return Ok(()); }; self.trace(InternalEvent::SqueezeVictim { entry: to_squeeze }); @@ -692,6 +803,7 @@ impl LiquidCache { let squeeze_io: Arc = Arc::new(DefaultSqueezeIo::new( self.store.clone(), to_squeeze, + squeezed_identity, self.observer.clone(), )); @@ -707,7 +819,7 @@ impl LiquidCache { squeeze_hint, &squeeze_io, ); - let outcome = self.reuse_disk_copy(&to_squeeze, outcome); + let outcome = self.reuse_disk_copy(&to_squeeze, squeezed_identity, outcome); match outcome { SqueezeOutcome::Replace { @@ -715,10 +827,19 @@ impl LiquidCache { bytes_to_write, } => { if let Some(bytes_to_write) = bytes_to_write { - self.write_batch_to_disk(to_squeeze, &new_batch, bytes_to_write) - .await?; + self.write_batch_to_disk( + to_squeeze, + WriteIdentity::Rewrite(squeezed_identity), + &new_batch, + bytes_to_write, + ) + .await?; } - match self.try_insert(to_squeeze, new_batch) { + match self.try_insert( + to_squeeze, + WriteIdentity::Rewrite(squeezed_identity), + new_batch, + ) { Ok(()) => { break; } @@ -747,6 +868,7 @@ impl LiquidCache { async fn maybe_hydrate( &self, entry_id: &EntryID, + identity: u64, cached: &CacheEntry, materialized: MaterializedEntry<'_>, expression: Option<&CacheExpression>, @@ -766,19 +888,22 @@ impl LiquidCache { cached: cached_type, new: new_type, }); - let _ = self.insert_inner(*entry_id, new_entry).await; + let _ = self + .insert_inner(*entry_id, WriteIdentity::Rewrite(identity), new_entry) + .await; } } pub(crate) async fn read_arrow_array( &self, entry_id: &EntryID, + identity: u64, selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, ) -> Option { use arrow::array::BooleanArray; - let batch = self.index.get(entry_id)?; + let batch = self.index.get_checked(entry_id, identity)?; self.cache_policy .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); self.trace(InternalEvent::Read { @@ -801,11 +926,11 @@ impl LiquidCache { None => Some(array.to_arrow_array()), }, CacheEntry::DiskArrow { .. } | CacheEntry::DiskLiquid { .. } => { - self.read_disk_array(batch.as_ref(), entry_id, expression, selection) + self.read_disk_array(batch.as_ref(), entry_id, identity, expression, selection) .await } CacheEntry::MemorySqueezedLiquid(array) => { - self.read_squeezed_array(array, entry_id, expression, selection) + self.read_squeezed_array(array, entry_id, identity, expression, selection) .await } } @@ -815,6 +940,7 @@ impl LiquidCache { &self, entry: &CacheEntry, entry_id: &EntryID, + identity: u64, expression: Option<&CacheExpression>, selection: Option<&BooleanBuffer>, ) -> Option { @@ -825,9 +951,10 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let full_array = self.read_disk_arrow_array(entry_id).await; + let full_array = self.read_disk_arrow_array(entry_id, identity).await; self.maybe_hydrate( entry_id, + identity, entry, MaterializedEntry::Arrow(&full_array), expression, @@ -847,9 +974,10 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let liquid = self.read_disk_liquid_array(entry_id).await; + let liquid = self.read_disk_liquid_array(entry_id, identity).await; self.maybe_hydrate( entry_id, + identity, entry, MaterializedEntry::Liquid(&liquid), expression, @@ -868,6 +996,7 @@ impl LiquidCache { &self, array: &LiquidSqueezedArrayRef, entry_id: &EntryID, + identity: u64, expression: Option<&CacheExpression>, selection: Option<&BooleanBuffer>, ) -> Option { @@ -881,7 +1010,7 @@ impl LiquidCache { } if let Some(array) = self - .try_read_squeezed_variant_array(array, entry_id, expression, selection) + .try_read_squeezed_variant_array(array, entry_id, identity, expression, selection) .await { self.observer.on_get_squeezed_success(); @@ -926,6 +1055,7 @@ impl LiquidCache { &self, array: &LiquidSqueezedArrayRef, entry_id: &EntryID, + identity: u64, expression: Option<&CacheExpression>, selection: Option<&BooleanBuffer>, ) -> Option { @@ -940,9 +1070,10 @@ impl LiquidCache { let full_array = if !all_paths_present { let batch = CacheEntry::MemorySqueezedLiquid(array.clone()); self.observer.on_get_squeezed_needs_io(); - let full_array = self.read_disk_arrow_array(entry_id).await; + let full_array = self.read_disk_arrow_array(entry_id, identity).await; self.maybe_hydrate( entry_id, + identity, &batch, MaterializedEntry::Arrow(&full_array), expression, @@ -968,6 +1099,7 @@ impl LiquidCache { async fn write_batch_to_disk( &self, entry_id: EntryID, + identity: WriteIdentity, batch: &CacheEntry, bytes: Bytes, ) -> Result<(), CacheFull> { @@ -989,8 +1121,26 @@ impl LiquidCache { kind: CachedBatchType::from(batch), bytes: len, }); + // A *rewrite* replays an entry read earlier, so a takeover in the + // meantime makes it stale: its object would be unreachable, and + // superseding on its behalf below would delete the live owner's. Drop + // it, and hand back the reservation taken above — nothing records + // those bytes, so nothing would ever release them. + // + // An *owned* write is the caller taking the key, and the index has not + // caught up yet by construction. Dropping it would leave the caller's + // own index entry pointing at bytes that were never written, and the + // next read of it panics. + if let WriteIdentity::Rewrite(rewriting) = identity + && let Some((current, _)) = self.index.get_with_identity(&entry_id) + && current != rewriting + { + self.budget.release_disk(len); + return Ok(()); + } + let identity = identity.value(); self.store - .put(entry_id_to_key(&entry_id), bytes.to_vec()) + .put(entry_id_to_key(&entry_id, identity), bytes.to_vec()) .await .expect("write failed"); // `bytes` is whatever `batch` serialises to: Arrow IPC for an arrow @@ -1003,23 +1153,40 @@ impl LiquidCache { }, CacheEntry::DiskLiquid { .. } | CacheEntry::MemoryLiquid(_) => DiskKind::Liquid, }; - let previous = self - .disk_copies - .lock() - .unwrap() - .insert(entry_id, DiskCopy { kind, bytes: len }); + let previous = self.disk_copies.lock().unwrap().insert( + entry_id, + DiskCopy { + identity, + kind, + bytes: len, + }, + ); if let Some(previous) = previous { - // The put replaced the object under this key, so the previous - // copy's reservation goes with it. + // Same owner: the put replaced that object, so its reservation + // goes with it and there is nothing left to delete. + // + // Different owner: this is the current owner superseding one that + // has let the key go (a stale writer never reaches here). The key + // carries the identity, so the put landed somewhere else and the + // previous object is still there — with no record naming it and + // nothing that would ever reach it. Releasing its reservation + // without removing it would leave disk held by a blob the budget + // has stopped counting. + if previous.identity != identity { + self.store + .remove(&entry_id_to_key(&entry_id, previous.identity)) + .await + .expect("disk remove failed"); + } self.budget.release_disk(previous.bytes); } Ok(()) } - async fn read_disk_arrow_array(&self, entry_id: &EntryID) -> ArrayRef { + async fn read_disk_arrow_array(&self, entry_id: &EntryID, identity: u64) -> ArrayRef { let bytes = self .store - .get(&entry_id_to_key(entry_id)) + .get(&entry_id_to_key(entry_id, identity)) .await .expect("read failed"); let bytes_len = bytes.len(); @@ -1038,10 +1205,11 @@ impl LiquidCache { async fn read_disk_liquid_array( &self, entry_id: &EntryID, + identity: u64, ) -> crate::liquid_array::LiquidArrayRef { let bytes = self .store - .get(&entry_id_to_key(entry_id)) + .get(&entry_id_to_key(entry_id, identity)) .await .expect("read failed"); self.trace(InternalEvent::IoReadLiquid { @@ -1060,13 +1228,14 @@ impl LiquidCache { pub(crate) async fn eval_predicate_internal( &self, entry_id: &EntryID, + identity: u64, selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, ) -> Option { use arrow::array::BooleanArray; self.observer.on_eval_predicate(); - let batch = self.index.get(entry_id)?; + let batch = self.index.get_checked(entry_id, identity)?; self.cache_policy .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); self.trace(InternalEvent::EvalPredicate { @@ -1083,23 +1252,27 @@ impl LiquidCache { owned.as_ref().unwrap() }); let selection_array = BooleanArray::new(selection.clone(), None); - let filtered = arrow::compute::filter(array, &selection_array) - .expect("selection must match array length"); - Some(self.eval_predicate_on_array(filtered, predicate)) + let filtered = arrow::compute::filter(array, &selection_array).ok()?; + self.eval_predicate_on_array(filtered, predicate) } entry @ CacheEntry::DiskArrow { .. } => { - let array = self.read_disk_arrow_array(entry_id).await; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Arrow(&array), None) - .await; + let array = self.read_disk_arrow_array(entry_id, identity).await; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Arrow(&array), + None, + ) + .await; let mut owned = None; let selection = selection_opt.unwrap_or_else(|| { owned = Some(BooleanBuffer::new_set(array.len())); owned.as_ref().unwrap() }); let selection_array = BooleanArray::new(selection.clone(), None); - let filtered = arrow::compute::filter(&array, &selection_array) - .expect("selection must match array length"); - Some(self.eval_predicate_on_array(filtered, predicate)) + let filtered = arrow::compute::filter(&array, &selection_array).ok()?; + self.eval_predicate_on_array(filtered, predicate) } CacheEntry::MemoryLiquid(array) => { let mut owned = None; @@ -1107,18 +1280,24 @@ impl LiquidCache { owned = Some(BooleanBuffer::new_set(array.len())); owned.as_ref().unwrap() }); - Some(array.try_eval_predicate(predicate, selection)) + array.try_eval_predicate(predicate, selection) } entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) - .await; + let liquid = self.read_disk_liquid_array(entry_id, identity).await; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Liquid(&liquid), + None, + ) + .await; let mut owned = None; let selection = selection_opt.unwrap_or_else(|| { owned = Some(BooleanBuffer::new_set(liquid.len())); owned.as_ref().unwrap() }); - Some(liquid.try_eval_predicate(predicate, selection)) + liquid.try_eval_predicate(predicate, selection) } CacheEntry::MemorySqueezedLiquid(array) => { self.eval_predicate_on_squeezed(array, selection_opt, predicate) @@ -1138,25 +1317,25 @@ impl LiquidCache { owned = Some(BooleanBuffer::new_set(array.len())); owned.as_ref().unwrap() }); - Some(array.try_eval_predicate(predicate, selection).await) + array.try_eval_predicate(predicate, selection).await } - fn eval_predicate_on_array(&self, array: ArrayRef, predicate: &LiquidExpr) -> BooleanArray { + /// `None` when the cached array cannot answer the predicate. See the + /// free function of the same name in `liquid_array`. + fn eval_predicate_on_array( + &self, + array: ArrayRef, + predicate: &LiquidExpr, + ) -> Option { let schema = Arc::new(Schema::new(vec![Field::new( "liquid_predicate_col", array.data_type().clone(), true, )])); - let record_batch = - RecordBatch::try_new(schema, vec![array]).expect("single-column predicate batch"); - let result = predicate - .physical_expr() - .evaluate(&record_batch) - .expect("validated LiquidExpr must evaluate"); - let boolean_array = result - .into_array(record_batch.num_rows()) - .expect("predicate output must be an array"); - boolean_array.as_boolean().clone() + let record_batch = RecordBatch::try_new(schema, vec![array]).ok()?; + let result = predicate.physical_expr().evaluate(&record_batch).ok()?; + let boolean_array = result.into_array(record_batch.num_rows()).ok()?; + Some(boolean_array.as_boolean().clone()) } } @@ -1218,7 +1397,10 @@ mod tests { let entry_id1: EntryID = EntryID::from(1); let array1 = create_test_array(100); let size1 = array1.memory_usage_bytes(); - store.insert_inner(entry_id1, array1).await.unwrap(); + store + .insert_inner(entry_id1, WriteIdentity::Owned(0), array1) + .await + .unwrap(); // Verify budget usage and data correctness assert_eq!(store.budget.memory_usage_bytes(), size1); @@ -1231,13 +1413,19 @@ mod tests { let entry_id2: EntryID = EntryID::from(2); let array2 = create_test_array(200); let size2 = array2.memory_usage_bytes(); - store.insert_inner(entry_id2, array2).await.unwrap(); + store + .insert_inner(entry_id2, WriteIdentity::Owned(0), array2) + .await + .unwrap(); assert_eq!(store.budget.memory_usage_bytes(), size1 + size2); let array3 = create_test_array(150); let size3 = array3.memory_usage_bytes(); - store.insert_inner(entry_id1, array3).await.unwrap(); + store + .insert_inner(entry_id1, WriteIdentity::Owned(0), array3) + .await + .unwrap(); assert_eq!(store.budget.memory_usage_bytes(), size3 + size2); assert!(store.index().get(&EntryID::from(999)).is_none()); @@ -1256,6 +1444,7 @@ mod tests { store .insert_inner( entry_id, + WriteIdentity::Owned(0), CacheEntry::memory_squeezed_liquid(squeezed.clone()), ) .await @@ -1263,7 +1452,7 @@ mod tests { let expr = Arc::new(CacheExpression::extract_date32(Date32Field::Year)); let result = store - .get(&entry_id) + .get(&entry_id, 0) .with_expression_hint(expr) .read() .await @@ -1294,7 +1483,7 @@ mod tests { let store = create_cache_store(8000, Box::new(advisor)).await; // Small budget to force advice store - .insert_inner(entry_id1, create_test_array(800)) + .insert_inner(entry_id1, WriteIdentity::Owned(0), create_test_array(800)) .await .unwrap(); match store.index().get(&entry_id1).unwrap().as_ref() { @@ -1303,7 +1492,7 @@ mod tests { } store - .insert_inner(entry_id2, create_test_array(800)) + .insert_inner(entry_id2, WriteIdentity::Owned(0), create_test_array(800)) .await .unwrap(); match store.index().get(&entry_id1).unwrap().as_ref() { @@ -1353,7 +1542,7 @@ mod tests { let unique_id = thread_id * ops_per_thread + i; let entry_id: EntryID = EntryID::from(unique_id); let array = create_test_arrow_array(100); - store.insert(entry_id, array).await.unwrap(); + store.insert(entry_id, 0, array).await.unwrap(); } }); })); @@ -1387,8 +1576,14 @@ mod tests { // Insert two small batches let arr1: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64)); let arr2: ArrayRef = Arc::new(Int32Array::from_iter_values(0..128)); - storage.insert(EntryID::from(1usize), arr1).await.unwrap(); - storage.insert(EntryID::from(2usize), arr2).await.unwrap(); + storage + .insert(EntryID::from(1usize), 0, arr1) + .await + .unwrap(); + storage + .insert(EntryID::from(2usize), 0, arr2) + .await + .unwrap(); // Stats after insert: 2 entries, memory usage > 0, disk usage == 0 let s = storage.stats(); @@ -1412,14 +1607,14 @@ mod tests { let entry_id = EntryID::from(321usize); let array = create_test_arrow_array(8); - store.insert(entry_id, array.clone()).await.unwrap(); + store.insert(entry_id, 0, array.clone()).await.unwrap(); store.flush_all_to_disk().await.unwrap(); { let entry = store.index().get(&entry_id).unwrap(); assert!(matches!(entry.as_ref(), CacheEntry::DiskArrow { .. })); } - let result = store.get(&entry_id).await.expect("present"); + let result = store.get(&entry_id, 0).await.expect("present"); assert_eq!(result.as_ref(), array.as_ref()); { let entry = store.index().get(&entry_id).unwrap(); @@ -1436,7 +1631,11 @@ mod tests { let liquid = transcode_liquid_inner(&arrow_array, &compressor).unwrap(); store - .insert_inner(entry_id, CacheEntry::memory_liquid(liquid.clone())) + .insert_inner( + entry_id, + WriteIdentity::Owned(0), + CacheEntry::memory_liquid(liquid.clone()), + ) .await .unwrap(); store.flush_all_to_disk().await.unwrap(); @@ -1445,7 +1644,7 @@ mod tests { assert!(matches!(entry.as_ref(), CacheEntry::DiskLiquid { .. })); } - let result = store.get(&entry_id).await.expect("present"); + let result = store.get(&entry_id, 0).await.expect("present"); assert_eq!(result.as_ref(), arrow_array.as_ref()); { let entry = store.index().get(&entry_id).unwrap(); @@ -1453,6 +1652,245 @@ mod tests { } } + /// An owned write is the caller taking the key, and the index has not + /// caught up yet by construction — the write happens first, the index + /// entry second. Dropping it as "stale" would leave the caller's own entry + /// pointing at bytes that were never written, and the next read panics. + #[tokio::test] + async fn an_owned_write_is_not_dropped_because_the_index_lags() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(51usize); + + // Identity 1 holds the key. + store + .insert_inner(entry_id, WriteIdentity::Owned(1), create_test_array(100)) + .await + .unwrap(); + + // Identity 2 takes it over. Its bytes land before its index entry + // does, so the index still names 1 at write time. + let taking_over = create_test_array(200); + let CacheEntry::MemoryArrow(array) = &taking_over else { + unreachable!("create_test_array builds an arrow entry") + }; + let bytes = arrow_to_bytes(array).unwrap(); + store + .write_batch_to_disk(entry_id, WriteIdentity::Owned(2), &taking_over, bytes) + .await + .unwrap(); + + // The bytes must actually be there, or the entry installed next reads + // a missing object. + let read_back = store.read_disk_arrow_array(&entry_id, 2).await; + assert_eq!( + read_back.len(), + 200, + "an owned write was dropped, so its entry would point at nothing" + ); + } + + /// A reset must delete the objects it forgets. The store key carries the + /// identity that wrote it, so everything written after a reset lands under + /// new keys — nothing would overwrite the old objects, and with their + /// records gone nothing would ever find them either. + #[tokio::test] + async fn reset_deletes_the_store_objects_it_forgets() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(61usize); + + store + .insert_inner(entry_id, WriteIdentity::Owned(1), create_test_array(100)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + assert!(store.budget.disk_usage_bytes() > 0, "something spilled"); + assert!( + store + .store + .get(&crate::cache::io_context::entry_id_to_key(&entry_id, 1)) + .await + .is_ok(), + "the object is there before the reset" + ); + + store.reset().await; + + assert!( + store + .store + .get(&crate::cache::io_context::entry_id_to_key(&entry_id, 1)) + .await + .is_err(), + "reset left an object nothing can reach and nothing counts" + ); + assert_eq!(store.budget.disk_usage_bytes(), 0); + } + + /// A dropped write must hand back the disk it reserved. Nothing records + /// those bytes — no `DiskCopy` names them — so no later path would ever + /// release them, and repeated takeovers during squeezes would walk the + /// disk tally up to its limit while holding nothing. + #[tokio::test] + async fn a_dropped_stale_write_releases_its_reservation() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(41usize); + + // Identity 2 owns the key. + store + .insert_inner(entry_id, WriteIdentity::Owned(2), create_test_array(100)) + .await + .unwrap(); + let disk_before = store.budget.disk_usage_bytes(); + + // A rewrite for an identity that has since lost the key is dropped. + let stale_entry = create_test_array(50); + let CacheEntry::MemoryArrow(stale_array) = &stale_entry else { + unreachable!("create_test_array builds an arrow entry") + }; + let stale_bytes = arrow_to_bytes(stale_array).unwrap(); + store + .write_batch_to_disk( + entry_id, + WriteIdentity::Rewrite(1), + &stale_entry, + stale_bytes, + ) + .await + .unwrap(); + + assert_eq!( + store.budget.disk_usage_bytes(), + disk_before, + "a dropped write must not keep the disk it reserved" + ); + } + + /// Taking a key over must not strand the previous owner's object. + /// + /// Once the identity is part of the store key, a new owner's write lands + /// somewhere else rather than on top — so the old object survives its own + /// record. Releasing its reservation without deleting it leaves disk held + /// by a blob nothing can reach and the budget has stopped counting. + #[tokio::test] + async fn taking_a_key_over_removes_the_previous_owner_s_object() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(31usize); + + store + .insert_inner(entry_id, WriteIdentity::Owned(1), create_test_array(100)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + let disk_after_first = store.budget.disk_usage_bytes(); + assert!(disk_after_first > 0, "the first owner spilled to disk"); + + // A different owner takes the key and spills too. + store + .insert_inner(entry_id, WriteIdentity::Owned(2), create_test_array(100)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + + // The first owner's object is gone, not merely unaccounted. + assert!( + store + .store + .get(&crate::cache::io_context::entry_id_to_key(&entry_id, 1)) + .await + .is_err(), + "the previous owner's object outlived its record" + ); + assert_eq!( + store.budget.disk_usage_bytes(), + disk_after_first, + "disk accounting should track one object, not two" + ); + } + + /// The store object has to be owned too, not just the record naming it. + /// + /// A write is issued before the index rewrite that would have declined it, + /// so it can land *after* another owner has taken the key over and become + /// disk-backed. Addressed by the packed id alone, that write overwrites + /// bytes the new owner's index entry and `DiskCopy` record both agree are + /// its own — silent wrong data, past every check. The identity belongs in + /// the store key so the two never address one object. + #[tokio::test] + async fn a_late_write_for_a_previous_owner_cannot_reach_the_current_one() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(21usize); + + // Identity 2 owns the key and is disk-backed. + store + .insert_inner(entry_id, WriteIdentity::Owned(2), create_test_array(200)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + let before = store.read_disk_arrow_array(&entry_id, 2).await; + assert_eq!(before.len(), 200); + + // Identity 1's write lands late, carrying a different array. + let stale_entry = create_test_array(37); + let CacheEntry::MemoryArrow(stale_array) = &stale_entry else { + unreachable!("create_test_array builds an arrow entry") + }; + let stale_bytes = arrow_to_bytes(stale_array).unwrap(); + store + .write_batch_to_disk( + entry_id, + WriteIdentity::Rewrite(1), + &stale_entry, + stale_bytes, + ) + .await + .unwrap(); + + // The current owner still reads its own rows. + let after = store.read_disk_arrow_array(&entry_id, 2).await; + assert_eq!( + after.len(), + 200, + "a write for a previous owner reached the current owner's object" + ); + } + + /// A rewrite that is declined has already written its bytes, so the store + /// object and its record outlive the entry they were built for. If the + /// next owner of the key could see that record it would adopt the object + /// on kind and length alone and read the previous owner's rows as its own. + #[tokio::test] + async fn a_disk_copy_is_invisible_to_whoever_holds_the_key_next() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(5usize); + + // Identity 1 caches and spills, recording a store object for this key. + store + .insert_inner(entry_id, WriteIdentity::Owned(1), create_test_array(100)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + assert!( + store.disk_copy(&entry_id, 1).is_some(), + "the owner sees the object it wrote" + ); + + // Identity 2 takes the key over. The object is still on disk, and the + // record still names identity 1. + store + .insert_inner(entry_id, WriteIdentity::Owned(2), create_test_array(200)) + .await + .unwrap(); + + assert!( + store.disk_copy(&entry_id, 2).is_none(), + "the new owner must not adopt the object the previous one left" + ); + assert!( + store.any_disk_copy(&entry_id).is_some(), + "the record is still there for the paths that reclaim it" + ); + } + #[tokio::test] async fn insert_returns_cache_full_when_memory_and_disk_are_saturated() { let cache = LiquidCacheBuilder::new() @@ -1463,10 +1901,10 @@ mod tests { .await; let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); - let err = cache.insert(EntryID::from(900usize), array).await; + let err = cache.insert(EntryID::from(900usize), 0, array).await; assert_eq!(err, Err(CacheFull)); - assert!(!cache.is_cached(&EntryID::from(900usize))); + assert!(!cache.is_cached(&EntryID::from(900usize), 0)); } #[tokio::test] @@ -1485,14 +1923,14 @@ mod tests { let first = EntryID::from(910usize); let second = EntryID::from(911usize); - cache.insert(first, first_array).await.unwrap(); + cache.insert(first, 0, first_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(cache.is_cached(&first)); + assert!(cache.is_cached(&first, 0)); - cache.insert(second, second_array).await.unwrap(); + cache.insert(second, 0, second_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.is_cached(&first)); + assert!(!cache.is_cached(&first, 0)); assert!(matches!( cache.index().get(&second).unwrap().as_ref(), CacheEntry::DiskArrow { .. } @@ -1513,13 +1951,13 @@ mod tests { .await; let first = EntryID::from(912usize); let second = EntryID::from(913usize); - cache.insert(first, first_array).await.unwrap(); + cache.insert(first, 0, first_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - cache.insert(second, second_array).await.unwrap(); + cache.insert(second, 0, second_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.is_cached(&first) || !cache.is_cached(&second)); + assert!(!cache.is_cached(&first, 0) || !cache.is_cached(&second, 0)); } #[tokio::test] @@ -1534,14 +1972,14 @@ mod tests { .build() .await; let entry = EntryID::from(914usize); - cache.insert(entry, array).await.unwrap(); + cache.insert(entry, 0, array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); let before = cache.stats().disk_usage_bytes; cache.remove_disk_entry(entry).await; assert_eq!(cache.stats().disk_usage_bytes, before - disk_bytes); - assert!(!cache.is_cached(&entry)); + assert!(!cache.is_cached(&entry, 0)); } #[tokio::test] @@ -1554,12 +1992,12 @@ mod tests { .await; let entry_id = EntryID::from(901usize); let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); - cache.insert(entry_id, array).await.unwrap(); + cache.insert(entry_id, 0, array).await.unwrap(); let result = cache.flush_all_to_disk().await; assert_eq!(result, Ok(())); - assert!(!cache.is_cached(&entry_id)); + assert!(!cache.is_cached(&entry_id, 0)); } async fn hydrating_cache() -> Arc { @@ -1594,15 +2032,15 @@ mod tests { let v1: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); let v2: ArrayRef = Arc::new(Int32Array::from_iter_values(100..164)); - cache.insert(id, v1).await.unwrap(); + cache.insert(id, 0, v1).await.unwrap(); let v1_disk_bytes = demote_to_disk(&cache, id).await; assert_eq!(cache.budget().disk_usage_bytes(), v1_disk_bytes); - cache.insert(id, v2.clone()).await.unwrap(); + cache.insert(id, 0, v2.clone()).await.unwrap(); let disk_after_overwrite = cache.budget().disk_usage_bytes(); let v2_disk_bytes = demote_to_disk(&cache, id).await; - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), v2.as_ref(), "read back the superseded value"); assert_eq!( disk_after_overwrite, 0, @@ -1621,9 +2059,9 @@ mod tests { let v1: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); let v2: ArrayRef = Arc::new(Int32Array::from_iter_values(100..164)); - cache.insert(id, v1.clone()).await.unwrap(); + cache.insert(id, 0, v1.clone()).await.unwrap(); let v1_disk_bytes = demote_to_disk(&cache, id).await; - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), v1.as_ref()); assert!(matches!( cache.index().get(&id).unwrap().as_ref(), @@ -1631,11 +2069,11 @@ mod tests { )); assert_eq!(cache.budget().disk_usage_bytes(), v1_disk_bytes); - cache.insert(id, v2.clone()).await.unwrap(); + cache.insert(id, 0, v2.clone()).await.unwrap(); let disk_after_overwrite = cache.budget().disk_usage_bytes(); let v2_disk_bytes = demote_to_disk(&cache, id).await; - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), v2.as_ref(), "read back the superseded value"); assert_eq!(disk_after_overwrite, 0); assert_eq!(cache.budget().disk_usage_bytes(), v2_disk_bytes); @@ -1651,13 +2089,13 @@ mod tests { let id = EntryID::from(922usize); let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64)); - cache.insert(id, array.clone()).await.unwrap(); + cache.insert(id, 0, array.clone()).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); assert!(matches!( cache.index().get(&id).unwrap().as_ref(), CacheEntry::DiskArrow { .. } )); - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), array.as_ref()); assert!(matches!( cache.index().get(&id).unwrap().as_ref(), @@ -1666,7 +2104,7 @@ mod tests { let disk_bytes = demote_to_disk(&cache, id).await; assert_eq!(cache.budget().disk_usage_bytes(), disk_bytes); - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), array.as_ref()); } @@ -1686,7 +2124,7 @@ mod tests { let expr = Arc::new(CacheExpression::extract_date32(Date32Field::Year)); cache - .insert(id, dates.clone()) + .insert(id, 0, dates.clone()) .with_squeeze_hint(expr.clone()) .await .unwrap(); @@ -1702,7 +2140,7 @@ mod tests { // Drain the IO counters so the count below covers only the re-eviction. let _ = cache.observer().runtime_snapshot(); - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), dates.as_ref()); assert!(matches!( cache.index().get(&id).unwrap().as_ref(), @@ -1725,7 +2163,7 @@ mod tests { assert_eq!(cache.budget().disk_usage_bytes(), disk_bytes); let years = cache - .get(&id) + .get(&id, 0) .with_expression_hint(expr) .read() .await @@ -1766,7 +2204,7 @@ mod tests { let expr = expr.clone(); async move { cache - .insert(id, dates) + .insert(id, 0, dates) .with_squeeze_hint(expr) .await .unwrap(); @@ -1797,11 +2235,11 @@ mod tests { // Too big for memory, and the disk tier is full with no victims. let too_big: ArrayRef = Arc::new(Int32Array::from_iter_values(0..(1 << 16))); - let result = cache.insert(id, too_big).await; + let result = cache.insert(id, 0, too_big).await; assert_eq!(result, Err(CacheFull)); - assert!(!cache.is_cached(&id)); - assert!(cache.get(&id).await.is_none()); + assert!(!cache.is_cached(&id, 0)); + assert!(cache.get(&id, 0).await.is_none()); assert_eq!(cache.budget().disk_usage_bytes(), 0); assert_eq!(cache.budget().memory_usage_bytes(), 0); } @@ -1823,9 +2261,9 @@ mod tests { .await; let id = EntryID::from(925usize); - cache.insert(id, array.clone()).await.unwrap(); + cache.insert(id, 0, array.clone()).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), array.as_ref()); assert!(matches!( cache.index().get(&id).unwrap().as_ref(), @@ -1837,7 +2275,7 @@ mod tests { // that is full with the entry's own copy, so the entry is dropped. cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.is_cached(&id)); + assert!(!cache.is_cached(&id, 0)); assert_eq!( cache.budget().disk_usage_bytes(), 0, diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index cf8881ab5..4ec252484 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -1,7 +1,7 @@ use congee::CongeeArc; use std::{ fmt::{Debug, Formatter}, - sync::atomic::{AtomicUsize, Ordering}, + sync::atomic::{AtomicU64, AtomicUsize, Ordering}, }; use crate::cache::{cached_batch::CacheEntry, utils::EntryID}; @@ -20,25 +20,65 @@ use crate::sync::{Arc, RwLock}; /// So the tree stores a small slot and the payload is taken out of it the /// moment the index gives the entry up. The deferred drop then reclaims only /// an empty shell, and the array dies with the last caller-held reference. -struct Slot(RwLock>>); +/// +/// The slot also records the identity of what it holds. `EntryID` is a packed +/// integer whose fields are narrower than the values they encode, so two +/// distinct sources can compute the same key; the key alone therefore cannot +/// answer "is this the entry I asked for?". `identity` is the caller's +/// unnarrowed name for the data, compared on every read through +/// [`ArtIndex::get_checked`], which turns such aliasing into a miss rather +/// than a wrong answer. +struct Slot { + identity: u64, + entry: RwLock>>, +} impl Slot { - fn new(entry: CacheEntry) -> Arc { - Arc::new(Self(RwLock::new(Some(Arc::new(entry))))) + fn new(identity: u64, entry: CacheEntry) -> Arc { + Arc::new(Self { + identity, + entry: RwLock::new(Some(Arc::new(entry))), + }) } fn load(&self) -> Option> { - self.0.read().unwrap().clone() + self.entry.read().unwrap().clone() } fn take(&self) -> Option> { - self.0.write().unwrap().take() + self.entry.write().unwrap().take() + } +} + +/// Whose data a write carries, and on what terms. +#[derive(Debug, Clone, Copy)] +pub(crate) enum WriteIdentity { + /// A caller storing its own data. Takes the key over if another identity + /// holds it: that identity belongs to a source that cannot read this key + /// any more, so leaving its entry there would cost the key to both. + Owned(u64), + /// Maintenance rewriting an entry it read earlier — transcode, squeeze, + /// hydrate, spill. It carries the identity the entry was read under and + /// lands only if the key still holds it. Adopting whatever is there + /// instead would relabel one source's data with another's whenever a + /// takeover lands between the read and the write, and the new owner would + /// then read those rows as its own. + Rewrite(u64), +} + +impl WriteIdentity { + /// The identity this write carries, whichever kind it is. + pub(crate) fn value(&self) -> u64 { + match self { + Self::Owned(id) | Self::Rewrite(id) => *id, + } } } pub(crate) struct ArtIndex { art: CongeeArc, entry_count: AtomicUsize, + identity_mismatches: AtomicU64, } impl Debug for ArtIndex { @@ -52,9 +92,16 @@ impl ArtIndex { Self { art: CongeeArc::new(), entry_count: AtomicUsize::new(0), + identity_mismatches: AtomicU64::new(0), } } + /// Look up an entry without checking whose it is. + /// + /// This is for maintenance that acts on whatever currently occupies a key — + /// eviction, squeezing, disk supersession, iteration for stats. A read + /// serving a caller must use [`Self::get_checked`] instead, so that a key + /// collision cannot return one caller another's data. pub(crate) fn get(&self, entry_id: &EntryID) -> Option> { let guard = self.art.pin(); // An empty slot means the entry was removed or replaced between the @@ -69,15 +116,79 @@ impl ArtIndex { self.art.get(*entry_id, &guard)?.load() } - pub(crate) fn is_cached(&self, entry_id: &EntryID) -> bool { - self.get(entry_id).is_some() + /// Look up an entry, returning it only if it is the one `identity` names. + /// + /// A mismatch reads as a miss, so the caller re-reads from its source and + /// gets correct data. It is also counted: the tally is expected to stay at + /// zero, and a non-zero value means two sources are computing the same + /// `EntryID`. + pub(crate) fn get_checked(&self, entry_id: &EntryID, identity: u64) -> Option> { + let guard = self.art.pin(); + let slot = self.art.get(*entry_id, &guard)?; + if slot.identity != identity { + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + return None; + } + if let Some(entry) = slot.load() { + return Some(entry); + } + // Re-read as in `get`: a replace leaves the key present under a new + // slot, which carries its own identity and must be checked again. + let slot = self.art.get(*entry_id, &guard)?; + if slot.identity != identity { + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + return None; + } + slot.load() + } + + /// Look up an entry together with the identity recorded against it, for + /// maintenance that must rewrite it under the identity it observed. + pub(crate) fn get_with_identity(&self, entry_id: &EntryID) -> Option<(u64, Arc)> { + let guard = self.art.pin(); + let slot = self.art.get(*entry_id, &guard)?; + let identity = slot.identity; + if let Some(entry) = slot.load() { + return Some((identity, entry)); + } + let slot = self.art.get(*entry_id, &guard)?; + let identity = slot.identity; + slot.load().map(|entry| (identity, entry)) + } + + pub(crate) fn is_cached(&self, entry_id: &EntryID, identity: u64) -> bool { + self.get_checked(entry_id, identity).is_some() } - pub(crate) fn insert(&self, entry_id: &EntryID, batch: CacheEntry) { + /// Store `batch` under `entry_id`, returning whether it was stored. + /// + /// See [`WriteIdentity`] for the two kinds of write and why they differ. + pub(crate) fn insert( + &self, + entry_id: &EntryID, + identity: WriteIdentity, + batch: CacheEntry, + ) -> bool { let guard = self.art.pin(); + let existing_identity = self.art.get(*entry_id, &guard).map(|slot| slot.identity); + let identity = match (identity, existing_identity) { + (WriteIdentity::Owned(new), Some(old)) => { + if new != old { + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + } + new + } + (WriteIdentity::Owned(new), None) => new, + // The key still holds what this rewrite was built from. + (WriteIdentity::Rewrite(expected), Some(old)) if expected == old => expected, + // It does not: the entry was taken over or removed while this + // rewrite was in flight, so the payload belongs to a source that + // no longer owns the key. Drop it. + (WriteIdentity::Rewrite(_), _) => return false, + }; let existing = self .art - .insert(*entry_id, Slot::new(batch), &guard) + .insert(*entry_id, Slot::new(identity, batch), &guard) .expect("Insertion failed"); match existing { Some(replaced) => drop(replaced.take()), @@ -85,6 +196,7 @@ impl ArtIndex { self.entry_count.fetch_add(1, Ordering::Relaxed); } } + true } pub(crate) fn remove(&self, entry_id: &EntryID) -> Option> { @@ -101,10 +213,10 @@ impl ArtIndex { self.entry_count.store(0, Ordering::Relaxed); } - pub(crate) fn for_each(&self, mut f: impl FnMut(&EntryID, &CacheEntry)) { + pub(crate) fn for_each(&self, mut f: impl FnMut(&EntryID, u64, &CacheEntry)) { for id in self.art.keys() { - if let Some(entry) = self.get(&id) { - f(&id, &entry); + if let Some((identity, entry)) = self.get_with_identity(&id) { + f(&id, identity, &entry); } } } @@ -117,6 +229,15 @@ impl ArtIndex { pub(crate) fn entry_count(&self) -> usize { self.entry_count.load(Ordering::Relaxed) } + + /// How many lookups or inserts found a key held by a different identity. + /// + /// Expected to stay at zero. A non-zero value means two sources compute the + /// same `EntryID`, and every one of them was served correctly only because + /// the check turned it into a miss. + pub(crate) fn identity_mismatches(&self) -> u64 { + self.identity_mismatches.load(Ordering::Relaxed) + } } #[cfg(test)] @@ -134,17 +255,17 @@ mod tests { let array1 = create_test_array(100); // Initially, entries should not be cached - assert!(!store.is_cached(&entry_id1)); - assert!(!store.is_cached(&entry_id2)); + assert!(!store.is_cached(&entry_id1, 0)); + assert!(!store.is_cached(&entry_id2, 0)); assert!(store.get(&entry_id1).is_none()); // Insert an entry and verify it's cached { - store.insert(&entry_id1, array1.clone()); + store.insert(&entry_id1, WriteIdentity::Owned(0), array1.clone()); } - assert!(store.is_cached(&entry_id1)); - assert!(!store.is_cached(&entry_id2)); + assert!(store.is_cached(&entry_id1, 0)); + assert!(!store.is_cached(&entry_id2, 0)); // Get should return the cached value match store.get(&entry_id1) { @@ -162,14 +283,14 @@ mod tests { let entry_id: EntryID = EntryID::from(1); let array = create_test_array(100); - store.insert(&entry_id, array.clone()); + store.insert(&entry_id, WriteIdentity::Owned(0), array.clone()); let entry_id: EntryID = EntryID::from(1); - assert!(store.is_cached(&entry_id)); + assert!(store.is_cached(&entry_id, 0)); store.reset(); let entry_id: EntryID = EntryID::from(1); - assert!(!store.is_cached(&entry_id)); + assert!(!store.is_cached(&entry_id, 0)); } /// The array behind a removed or replaced entry must die with the last @@ -184,8 +305,8 @@ mod tests { unreachable!() }; let weak_first = Arc::downgrade(first_array); - store.insert(&id, first); - store.insert(&id, create_test_array(200)); + store.insert(&id, WriteIdentity::Owned(0), first); + store.insert(&id, WriteIdentity::Owned(0), create_test_array(200)); assert!( weak_first.upgrade().is_none(), "replaced entry still alive: held by the index's deferred drop" @@ -204,4 +325,93 @@ mod tests { ); assert_eq!(store.entry_count(), 0); } + + /// Two files whose ids narrow to the same `EntryID` must not read each + /// other's data. Before the identity check this returned the incumbent's + /// array, which is a wrong answer whenever the two happen to share a type. + #[test] + fn an_entry_is_never_served_to_a_different_identity() { + let store = ArtIndex::new(); + let key: EntryID = EntryID::from(7); + + assert!(store.insert(&key, WriteIdentity::Owned(1), create_test_array(100))); + + // The colliding file asks for the same key and is told nothing is there. + assert!(store.get_checked(&key, 2).is_none()); + assert!(!store.is_cached(&key, 2)); + assert_eq!(store.identity_mismatches(), 2); + + // The owner still reads its own entry. + assert!(store.get_checked(&key, 1).is_some()); + + // The colliding file takes the key over. It has to: it cannot read + // what is there, so leaving it would cost the key to both of them. + assert!(store.insert(&key, WriteIdentity::Owned(2), create_test_array(200))); + assert!( + store.get_checked(&key, 1).is_none(), + "the displaced file reads a miss, never the other file's rows" + ); + match store.get_checked(&key, 2).unwrap().as_ref() { + CacheEntry::MemoryArrow(array) => assert_eq!(array.len(), 200), + other => panic!("expected the new owner's array, found {other}"), + } + } + + /// A rewrite is built from an entry read earlier, and the key can be taken + /// over in between — a squeeze reads, awaits a disk write, then stores. If + /// the rewrite adopted whatever identity held the key by then, it would + /// relabel the old file's data as the new owner's, and the new owner would + /// read those rows as a hit. Carrying the identity it read under makes the + /// stale write drop instead. + #[test] + fn a_rewrite_does_not_land_on_a_key_taken_over_since_it_was_read() { + let store = ArtIndex::new(); + let key: EntryID = EntryID::from(11); + + // File A caches, and something begins rewriting that entry. + assert!(store.insert(&key, WriteIdentity::Owned(1), create_test_array(100))); + let (observed, _read) = store.get_with_identity(&key).unwrap(); + assert_eq!(observed, 1); + + // File B takes the key over while that rewrite is in flight. + assert!(store.insert(&key, WriteIdentity::Owned(2), create_test_array(200))); + + // The rewrite lands too late and must be dropped, not relabelled. + assert!(!store.insert( + &key, + WriteIdentity::Rewrite(observed), + create_test_array(100) + )); + + match store.get_checked(&key, 2).unwrap().as_ref() { + CacheEntry::MemoryArrow(array) => assert_eq!( + array.len(), + 200, + "the new owner must still read its own rows, not the rewrite's" + ), + other => panic!("expected the new owner's array, found {other}"), + } + } + + /// Maintenance rewrites a key in place and must neither change whose the + /// entry is nor bring back one that has been removed — a stale reader that + /// misses goes on to insert what it read, and that write must not land + /// under a key nobody owns any more. + #[test] + fn maintenance_preserves_identity_and_cannot_resurrect_a_removed_key() { + let store = ArtIndex::new(); + let key: EntryID = EntryID::from(9); + + assert!(store.insert(&key, WriteIdentity::Owned(5), create_test_array(10))); + assert!(store.insert(&key, WriteIdentity::Rewrite(5), create_test_array(20))); + assert!( + store.get_checked(&key, 5).is_some(), + "rewriting in place kept the identity" + ); + + store.remove(&key); + assert!(!store.insert(&key, WriteIdentity::Rewrite(5), create_test_array(30))); + assert!(store.get(&key).is_none()); + assert_eq!(store.entry_count(), 0); + } } diff --git a/src/core/src/cache/io_context.rs b/src/core/src/cache/io_context.rs index a9c8974fe..1709fec66 100644 --- a/src/core/src/cache/io_context.rs +++ b/src/core/src/cache/io_context.rs @@ -38,9 +38,21 @@ pub trait EntryMetadata: Debug + Send + Sync { fn get_compressor(&self, entry_id: &EntryID) -> Arc; } -/// Convert an [`EntryID`] to a t4 key (8-byte little-endian representation). -pub(crate) fn entry_id_to_key(entry_id: &EntryID) -> Vec { - usize::from(*entry_id).to_le_bytes().to_vec() +/// Convert an [`EntryID`] and the identity that owns it to a t4 key. +/// +/// Both halves, not just the entry id. `EntryID` is a packed integer whose +/// fields are narrower than the values they encode, so two sources can compute +/// one id — and a store object addressed by that id alone is *shared*. Scoping +/// only the `DiskCopy` record is not enough: a write in flight for one owner +/// can land after another has taken the key over and installed its own disk +/// entry, overwriting bytes the new owner's index entry and record both agree +/// are its own. With the identity in the key the two address different +/// objects, so a late write cannot reach the other's bytes at all. +pub(crate) fn entry_id_to_key(entry_id: &EntryID, identity: u64) -> Vec { + let mut key = Vec::with_capacity(16); + key.extend_from_slice(&usize::from(*entry_id).to_le_bytes()); + key.extend_from_slice(&identity.to_le_bytes()); + key } /// A default implementation of [`EntryMetadata`]. @@ -84,15 +96,23 @@ impl EntryMetadata for DefaultCacheMetadata { pub struct DefaultSqueezeIo { store: t4::Store, entry_id: EntryID, + /// The owner whose object this reads. See [`entry_id_to_key`]. + identity: u64, observer: Arc, } impl DefaultSqueezeIo { /// Create a new instance of [DefaultSqueezeIo]. - pub fn new(store: t4::Store, entry_id: EntryID, observer: Arc) -> Self { + pub fn new( + store: t4::Store, + entry_id: EntryID, + identity: u64, + observer: Arc, + ) -> Self { Self { store, entry_id, + identity, observer, } } @@ -101,7 +121,7 @@ impl DefaultSqueezeIo { #[async_trait::async_trait] impl SqueezeIoHandler for DefaultSqueezeIo { async fn read(&self, range: Option>) -> std::io::Result { - let key = entry_id_to_key(&self.entry_id); + let key = entry_id_to_key(&self.entry_id, self.identity); let bytes = match range { Some(range) => { let len = range.end - range.start; diff --git a/src/core/src/cache/observer/stats.rs b/src/core/src/cache/observer/stats.rs index fa0c3d9ae..bd827b752 100644 --- a/src/core/src/cache/observer/stats.rs +++ b/src/core/src/cache/observer/stats.rs @@ -142,6 +142,13 @@ pub struct CacheStats { pub memory_liquid_bytes: usize, /// Total size of in-memory Squeezed-Liquid entries in bytes. pub memory_squeezed_liquid_bytes: usize, + /// Lookups and inserts that found a key held by a different identity. + /// + /// Expected to stay at zero. A non-zero value means two sources compute + /// the same `EntryID` — each was served correctly, because the check + /// turns the collision into a miss, but the cache is not holding what + /// either of them could use. + pub identity_mismatches: u64, /// Total memory usage of the cache. pub memory_usage_bytes: usize, /// Total disk usage of the cache. diff --git a/src/core/src/cache/tests/policies.rs b/src/core/src/cache/tests/policies.rs index 29c849902..b0e159606 100644 --- a/src/core/src/cache/tests/policies.rs +++ b/src/core/src/cache/tests/policies.rs @@ -18,12 +18,12 @@ async fn default_policies() { for i in 0..5 { let entry_id = EntryID::from(i); - cache.insert(entry_id, test_array.clone()).await.unwrap(); + cache.insert(entry_id, 0, test_array.clone()).await.unwrap(); } for i in 0..5 { let entry_id = EntryID::from(i); - let array = cache.get(&entry_id).read().await.unwrap(); + let array = cache.get(&entry_id, 0).read().await.unwrap(); assert_eq!(array.len(), test_array.len()); } @@ -44,17 +44,17 @@ async fn insert_wont_fit_cache() { .build() .await; cache - .insert(EntryID::from(0), test_array.clone()) + .insert(EntryID::from(0), 0, test_array.clone()) .await .unwrap(); let array_3x = arrow::compute::concat(&[&test_array, &test_array, &test_array]).unwrap(); let array_9x = arrow::compute::concat(&[&array_3x, &array_3x, &array_3x]).unwrap(); let array_27x = arrow::compute::concat(&[&array_9x, &array_9x, &array_9x]).unwrap(); cache - .insert(EntryID::from(1), array_27x.clone()) + .insert(EntryID::from(1), 0, array_27x.clone()) .await .unwrap(); - cache.get(&EntryID::from(1)).read().await.unwrap(); + cache.get(&EntryID::from(1), 0).read().await.unwrap(); let trace = cache.consume_event_trace(); let json_trace = serde_json::to_string(&trace).unwrap(); diff --git a/src/core/src/cache/tests/squeezed.rs b/src/core/src/cache/tests/squeezed.rs index 3d7c152e7..0f3e75e2f 100644 --- a/src/core/src/cache/tests/squeezed.rs +++ b/src/core/src/cache/tests/squeezed.rs @@ -41,7 +41,7 @@ async fn read_squeezed_date_time() { for i in 0..4 { let entry_id = EntryID::from(i); cache - .insert(entry_id, array.clone()) + .insert(entry_id, 0, array.clone()) .with_squeeze_hint(expression.clone()) .await .unwrap(); @@ -50,14 +50,14 @@ async fn read_squeezed_date_time() { for i in 0..4 { let entry_id = EntryID::from(i); let array = cache - .get(&entry_id) + .get(&entry_id, 0) .with_expression_hint(expression.clone()) .await .unwrap(); assert_eq!(array.len(), array.len()); } cache - .get(&EntryID::from(1)) + .get(&EntryID::from(1), 0) .with_expression_hint(Arc::new(CacheExpression::extract_date32( Date32Field::Month, ))) @@ -111,14 +111,14 @@ async fn read_squeezed_variant_path() { for i in 0..3 { let entry_id = EntryID::from(i); cache - .insert(entry_id, variant_array.clone()) + .insert(entry_id, 0, variant_array.clone()) .with_squeeze_hint(name_expr.clone()) .await .unwrap(); } let squeezed = cache - .get(&EntryID::from(0)) + .get(&EntryID::from(0), 0) .with_expression_hint(name_expr.clone()) .read() .await @@ -126,13 +126,13 @@ async fn read_squeezed_variant_path() { assert_eq!(squeezed.len(), variant_array.len()); cache - .get(&EntryID::from(0)) + .get(&EntryID::from(0), 0) .with_expression_hint(age_expr.clone()) .read() .await .unwrap(); cache - .get(&EntryID::from(1)) + .get(&EntryID::from(1), 0) .with_expression_hint(zipcode_expr.clone()) .read() .await @@ -171,19 +171,22 @@ async fn read_squeezed_int64_array() { let entry_id = EntryID::from(i); if i % 2 == 0 { cache - .insert(entry_id, int64_array.clone()) + .insert(entry_id, 0, int64_array.clone()) .with_squeeze_hint(expression.clone()) .await .unwrap(); } else { - cache.insert(entry_id, int64_array.clone()).await.unwrap(); + cache + .insert(entry_id, 0, int64_array.clone()) + .await + .unwrap(); } } for i in 0..4 { let entry_id = EntryID::from(i); let array = cache - .get(&entry_id) + .get(&entry_id, 0) .with_expression_hint(expression.clone()) .read() .await diff --git a/src/core/src/liquid_array/byte_view_array/mod.rs b/src/core/src/liquid_array/byte_view_array/mod.rs index 1d484148b..0be3fcf6b 100644 --- a/src/core/src/liquid_array/byte_view_array/mod.rs +++ b/src/core/src/liquid_array/byte_view_array/mod.rs @@ -354,11 +354,15 @@ impl LiquidArray for LiquidByteViewArray { Arc::new(dict) } - fn try_eval_predicate(&self, expr: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + fn try_eval_predicate( + &self, + expr: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { let filtered = helpers::filter_inner(self, filter); helpers::try_eval_predicate_in_memory(expr.physical_expr(), &filtered) - .unwrap_or_else(|| eval_predicate_on_array(filtered.to_arrow_array(), expr)) + .or_else(|| eval_predicate_on_array(filtered.to_arrow_array(), expr)) } fn to_bytes(&self) -> Vec { @@ -487,13 +491,17 @@ impl LiquidSqueezedArray for LiquidByteViewArray { /// /// Note that the filter is a boolean buffer, not a boolean array, i.e., filter can't be nullable. /// The returned boolean mask is nullable if the the original array is nullable. - async fn try_eval_predicate(&self, expr: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + async fn try_eval_predicate( + &self, + expr: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { // Reuse generic filter path first to reduce input rows if any let filtered = helpers::filter_inner(self, filter); if let Some(mask) = helpers::try_eval_predicate_on_disk(expr.physical_expr(), &filtered).await { - mask + Some(mask) } else { eval_predicate_on_array(filtered.to_arrow_array().await, expr) } diff --git a/src/core/src/liquid_array/decimal_array.rs b/src/core/src/liquid_array/decimal_array.rs index 374ff9cc7..4e7b7e43b 100644 --- a/src/core/src/liquid_array/decimal_array.rs +++ b/src/core/src/liquid_array/decimal_array.rs @@ -545,7 +545,7 @@ impl LiquidSqueezedArray for LiquidDecimalQuantizedArray { &self, liquid_expr: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { let filtered = self.filter_inner(filter); let expr = if let Some(expr) = unwrap_dynamic_filter(liquid_expr.physical_expr()) { @@ -575,7 +575,7 @@ impl LiquidSqueezedArray for LiquidDecimalQuantizedArray { match filtered.try_eval_predicate_inner(&op, literal) { Ok(Some(mask)) => { self.io.trace_io_saved(); - return mask; + return Some(mask); } Ok(None) => { let fallback = self.filter(filter).await; @@ -588,8 +588,7 @@ impl LiquidSqueezedArray for LiquidDecimalQuantizedArray { let full = self.hydrate_full_arrow().await; let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); + let filtered_arr = arrow::compute::filter(&full, &selection_array).ok()?; let filtered_len = filtered_arr.len(); let lhs = ColumnarValue::Array(filtered_arr); @@ -606,12 +605,11 @@ impl LiquidSqueezedArray for LiquidDecimalQuantizedArray { return eval_predicate_on_array(fallback, liquid_expr); } }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone() + // A comparison that cannot evaluate means this array is not what the + // predicate was built for. Report it as unanswerable rather than + // asserting: the caller falls back to the source. + let result = result.ok()?; + Some(result.into_array(filtered_len).ok()?.as_boolean().clone()) } } @@ -685,7 +683,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = BooleanArray::from(vec![Some(true), Some(true), None, Some(true)]); assert_eq!(got, expected); assert_eq!(io.reads(), 0); diff --git a/src/core/src/liquid_array/float_array.rs b/src/core/src/liquid_array/float_array.rs index 683dece88..6c4a2b1a2 100644 --- a/src/core/src/liquid_array/float_array.rs +++ b/src/core/src/liquid_array/float_array.rs @@ -993,7 +993,7 @@ where &self, liquid_expr: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { // Apply selection first to reduce input rows let filtered = self.filter_inner(filter); let expr = liquid_expr.physical_expr(); @@ -1005,7 +1005,7 @@ where let supported_op = Operator::from_datafusion(op); if let Some(supported_op) = supported_op { match filtered.try_eval_predicate_inner(&supported_op, literal) { - Ok(Some(mask)) => return mask, + Ok(Some(mask)) => return Some(mask), Ok(None) => { let fallback = self.filter(filter).await; return eval_predicate_on_array(fallback, liquid_expr); @@ -1018,8 +1018,7 @@ where let full = self.hydrate_full_arrow().await; let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); + let filtered_arr = arrow::compute::filter(&full, &selection_array).ok()?; let filtered_len = filtered_arr.len(); let lhs = ColumnarValue::Array(filtered_arr); @@ -1036,12 +1035,10 @@ where return eval_predicate_on_array(fallback, liquid_expr); } }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - return result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone(); + // Unanswerable rather than fatal: this array is not what the + // predicate was built for, so the caller falls back. + let result = result.ok()?; + return Some(result.into_array(filtered_len).ok()?.as_boolean().clone()); } } let fallback = self.filter(filter).await; @@ -1291,7 +1288,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1323,7 +1321,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1386,7 +1385,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1418,7 +1418,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { diff --git a/src/core/src/liquid_array/hybrid_primitive_array.rs b/src/core/src/liquid_array/hybrid_primitive_array.rs index f2d0845bb..65563cdc1 100644 --- a/src/core/src/liquid_array/hybrid_primitive_array.rs +++ b/src/core/src/liquid_array/hybrid_primitive_array.rs @@ -338,7 +338,7 @@ where &self, liquid_expr: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { // Apply selection first to reduce input rows let filtered = self.filter_inner(filter); @@ -374,7 +374,7 @@ where match filtered.try_eval_predicate_inner(&supported_op, literal) { Ok(Some(mask)) => { self.io.trace_io_saved(); - return mask; + return Some(mask); } Ok(None) => { let fallback = self.filter(filter).await; @@ -389,8 +389,7 @@ where let full = self.hydrate_full_arrow().await; let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); + let filtered_arr = arrow::compute::filter(&full, &selection_array).ok()?; let filtered_len = filtered_arr.len(); let lhs_array = match lhs_kind { PredicateLhs::Plain => filtered_arr, @@ -415,12 +414,11 @@ where return eval_predicate_on_array(fallback, liquid_expr); } }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone() + // A comparison that cannot evaluate means this array is not what the + // predicate was built for. Report it as unanswerable rather than + // asserting: the caller falls back to the source. + let result = result.ok()?; + Some(result.into_array(filtered_len).ok()?.as_boolean().clone()) } } @@ -701,7 +699,7 @@ where &self, liquid_expr: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { // Apply selection first to reduce input rows let filtered = self.filter_inner(filter); @@ -737,7 +735,7 @@ where match filtered.try_eval_predicate_inner(&supported_op, literal) { Ok(Some(mask)) => { self.io.trace_io_saved(); - return mask; + return Some(mask); } Ok(None) => { let fallback = self.filter(filter).await; @@ -752,8 +750,7 @@ where let full = self.hydrate_full_arrow().await; let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); + let filtered_arr = arrow::compute::filter(&full, &selection_array).ok()?; let filtered_len = filtered_arr.len(); let lhs_array = match lhs_kind { PredicateLhs::Plain => filtered_arr, @@ -778,12 +775,11 @@ where return eval_predicate_on_array(fallback, liquid_expr); } }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone() + // A comparison that cannot evaluate means this array is not what the + // predicate was built for. Report it as unanswerable rather than + // asserting: the caller falls back to the source. + let result = result.ok()?; + Some(result.into_array(filtered_len).ok()?.as_boolean().clone()) } } @@ -994,7 +990,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = expected_for(op, k); assert_eq!(io.reads(), 0); assert_eq!(got, expected); @@ -1015,7 +1012,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = expected_for(op, k); assert!(io.reads() > 0); assert_eq!(got, expected); @@ -1081,7 +1079,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = expected_for(op, k); assert_eq!(io.reads(), 0); assert_eq!(got, expected); @@ -1101,7 +1100,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = expected_for(op, k); assert!(io.reads() > 0); assert_eq!(got, expected); @@ -1143,7 +1143,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1175,7 +1176,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1225,7 +1227,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1257,7 +1260,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { diff --git a/src/core/src/liquid_array/linear_integer_array.rs b/src/core/src/liquid_array/linear_integer_array.rs index 4a8bc14ed..bbbcee7ef 100644 --- a/src/core/src/liquid_array/linear_integer_array.rs +++ b/src/core/src/liquid_array/linear_integer_array.rs @@ -351,7 +351,11 @@ where filter::filter(&arr, &selection).unwrap() } - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + fn try_eval_predicate( + &self, + predicate: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { let arr = self.filter(filter); eval_predicate_on_array(arr, predicate) } diff --git a/src/core/src/liquid_array/mod.rs b/src/core/src/liquid_array/mod.rs index 7776c7640..cb5b3d00c 100644 --- a/src/core/src/liquid_array/mod.rs +++ b/src/core/src/liquid_array/mod.rs @@ -124,7 +124,11 @@ pub trait LiquidArray: std::fmt::Debug + Send + Sync { /// /// Note that the filter is a boolean buffer, not a boolean array, i.e., filter can't be nullable. /// The returned boolean mask is nullable if the the original array is nullable. - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + fn try_eval_predicate( + &self, + predicate: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { let filtered = self.filter(filter); eval_predicate_on_array(filtered, predicate) } @@ -252,7 +256,7 @@ pub trait LiquidSqueezedArray: std::fmt::Debug + Send + Sync { &self, predicate: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { let filtered = self.filter(filter).await; eval_predicate_on_array(filtered, predicate) } @@ -262,21 +266,26 @@ pub trait LiquidSqueezedArray: std::fmt::Debug + Send + Sync { fn disk_backing(&self) -> SqueezedBacking; } -pub(crate) fn eval_predicate_on_array(array: ArrayRef, predicate: &LiquidExpr) -> BooleanArray { +/// Evaluate `predicate` against a one-column batch built from `array`. +/// +/// Returns `None` when the array cannot answer the predicate — a data type the +/// expression was not built for, or a batch it cannot be evaluated against. +/// That is reachable whenever a cached entry is not the one the predicate was +/// built for, so it must not be an assertion: the caller treats `None` as "the +/// cache cannot answer", materializes from the source and evaluates there. +pub(crate) fn eval_predicate_on_array( + array: ArrayRef, + predicate: &LiquidExpr, +) -> Option { let schema = Arc::new(Schema::new(vec![Field::new( "liquid_predicate_col", array.data_type().clone(), true, )])); - let record_batch = RecordBatch::try_new(schema, vec![array]).expect("predicate input batch"); - let result = predicate - .physical_expr() - .evaluate(&record_batch) - .expect("validated LiquidExpr must evaluate"); - let boolean_array = result - .into_array(record_batch.num_rows()) - .expect("predicate output must be an array"); - boolean_array.as_boolean().clone() + let record_batch = RecordBatch::try_new(schema, vec![array]).ok()?; + let result = predicate.physical_expr().evaluate(&record_batch).ok()?; + let boolean_array = result.into_array(record_batch.num_rows()).ok()?; + Some(boolean_array.as_boolean().clone()) } /// A trait to read the backing bytes of a squeezed array from disk. diff --git a/src/core/src/liquid_array/primitive_array.rs b/src/core/src/liquid_array/primitive_array.rs index c651112b5..eb0b2eb78 100644 --- a/src/core/src/liquid_array/primitive_array.rs +++ b/src/core/src/liquid_array/primitive_array.rs @@ -373,7 +373,11 @@ where arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap() } - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + fn try_eval_predicate( + &self, + predicate: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { let filtered = self.filter(filter); eval_predicate_on_array(filtered, predicate) } @@ -573,7 +577,11 @@ where arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap() } - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + fn try_eval_predicate( + &self, + predicate: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { let filtered = self.filter(filter); eval_predicate_on_array(filtered, predicate) } diff --git a/src/core/src/liquid_array/squeezed_date32_array.rs b/src/core/src/liquid_array/squeezed_date32_array.rs index e2497e556..75b817340 100644 --- a/src/core/src/liquid_array/squeezed_date32_array.rs +++ b/src/core/src/liquid_array/squeezed_date32_array.rs @@ -480,7 +480,7 @@ impl LiquidSqueezedArray for SqueezedDate32Array { &self, predicate: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { let filtered = self.filter(filter).await; eval_predicate_on_array(filtered, predicate) } diff --git a/src/core/src/liquid_array/tests.rs b/src/core/src/liquid_array/tests.rs index de3475737..ac04d8c9c 100644 --- a/src/core/src/liquid_array/tests.rs +++ b/src/core/src/liquid_array/tests.rs @@ -185,8 +185,9 @@ mod byte_view_tests { let expr: Arc = Arc::new(BinaryExpr::new(col, Operator::Eq, lit)); let liquid = make_byte_view(&input); - let result = - liquid.try_eval_predicate(&crate::cache::LiquidExpr::new_unchecked(expr), &mask); + let result = liquid + .try_eval_predicate(&crate::cache::LiquidExpr::new_unchecked(expr), &mask) + .expect("predicate must evaluate in this test"); let expected = BooleanArray::from(vec![ Some(true), None, diff --git a/src/core/study/cache_storage.rs b/src/core/study/cache_storage.rs index e4d6533a3..b4c4c5a91 100644 --- a/src/core/study/cache_storage.rs +++ b/src/core/study/cache_storage.rs @@ -92,7 +92,7 @@ fn main() { continue; }; if storage - .eval_predicate(id, &liquid_expr) + .eval_predicate(id, 0, &liquid_expr) .with_selection(&selection) .await .is_some() @@ -145,7 +145,7 @@ fn load_and_insert_referer( let id = EntryID::from(idx); ids.push(id); total_size += array.get_array_memory_size(); - storage.insert(id, array).await.unwrap(); + storage.insert(id, 0, array).await.unwrap(); idx += 1; } diff --git a/src/core/study/squeeze_integer.rs b/src/core/study/squeeze_integer.rs index cfdb89cc1..a19a99d3f 100644 --- a/src/core/study/squeeze_integer.rs +++ b/src/core/study/squeeze_integer.rs @@ -400,8 +400,10 @@ fn try_eval_or_fetch( ) -> (BooleanArray, usize) { io.reset_bytes_read(); let maybe_expr = LiquidExpr::try_new(expr.clone(), &hybrid.original_arrow_data_type(), None); - if let Some(liquid_expr) = maybe_expr { - let mask = futures::executor::block_on(hybrid.try_eval_predicate(&liquid_expr, filter)); + if let Some(liquid_expr) = maybe_expr + && let Some(mask) = + futures::executor::block_on(hybrid.try_eval_predicate(&liquid_expr, filter)) + { return (mask, io.bytes_read()); } // Not supported in hybrid form: materialize from full bytes and compute via Arrow. diff --git a/src/core/tests/memory_footprint.rs b/src/core/tests/memory_footprint.rs index 836255880..63da0eb00 100644 --- a/src/core/tests/memory_footprint.rs +++ b/src/core/tests/memory_footprint.rs @@ -103,7 +103,7 @@ fn make_entry(seed: u64, rows: usize) -> ArrayRef { fn indexed_bytes(cache: &LiquidCache) -> usize { let mut sum = 0; - cache.for_each_entry(|_, e| sum += e.memory_usage_bytes()); + cache.for_each_entry(|_, _, e| sum += e.memory_usage_bytes()); sum } @@ -111,7 +111,7 @@ fn indexed_bytes(cache: &LiquidCache) -> usize { /// index rather than the budget, so the budget can be checked against it. fn indexed_disk_bytes(cache: &LiquidCache) -> usize { let mut sum = 0; - cache.for_each_entry(|_, e| { + cache.for_each_entry(|_, _, e| { sum += match e { CacheEntry::DiskLiquid { disk_bytes, .. } | CacheEntry::DiskArrow { disk_bytes, .. } => *disk_bytes, @@ -168,7 +168,7 @@ async fn heap_footprint_tracks_budget_for_oversized_working_set() { // column batch as arrow, dropping the caller's copy right after. for i in 0..ENTRIES { let arr = make_entry(i as u64, ROWS); - cache.insert(EntryID::from(i), arr).await.unwrap(); + cache.insert(EntryID::from(i), 0, arr).await.unwrap(); } report(&cache, "after fill", baseline); let idle_after_fill = live() - baseline; @@ -189,7 +189,7 @@ async fn heap_footprint_tracks_budget_for_oversized_working_set() { reset_peak(); for _pass in 0..2 { for i in 0..ENTRIES { - let arr = cache.get(&EntryID::from(i)).await.unwrap(); + let arr = cache.get(&EntryID::from(i), 0).await.unwrap(); assert_eq!(arr.len(), ROWS); drop(arr); } @@ -249,7 +249,7 @@ async fn heap_footprint_tracks_budget_for_oversized_working_set() { // What survives once the index is emptied is held by the store, the // policy, or the compressor state — not by indexed entries. - cache.reset(); + cache.reset().await; report(&cache, "after reset", baseline); let idle_after_reset = live() - baseline; assert!( diff --git a/src/datafusion-server/src/admin_server/handlers.rs b/src/datafusion-server/src/admin_server/handlers.rs index 1be5b0be7..f3164ba02 100644 --- a/src/datafusion-server/src/admin_server/handlers.rs +++ b/src/datafusion-server/src/admin_server/handlers.rs @@ -51,7 +51,7 @@ pub(crate) async fn reset_cache_handler(State(state): State>) -> J info!("Resetting cache..."); let cache = state.liquid_cache.cache(); unsafe { - cache.reset(); + cache.reset().await; } Json(ApiResponse { diff --git a/src/datafusion/src/cache/column.rs b/src/datafusion/src/cache/column.rs index 1f582fd86..f0c2592fd 100644 --- a/src/datafusion/src/cache/column.rs +++ b/src/datafusion/src/cache/column.rs @@ -10,7 +10,7 @@ use parquet::arrow::arrow_reader::ArrowPredicate; use crate::{ LiquidPredicate, - cache::{BatchID, ColumnAccessPath, ParquetArrayID}, + cache::{BatchID, ColumnAccessPath, ParquetArrayID, file_id::FileId}, }; use std::sync::Arc; @@ -20,6 +20,15 @@ pub struct CachedColumn { cache_store: Arc, field: Arc, column_path: ColumnAccessPath, + /// The file id before it is narrowed into `column_path`. Two files whose + /// ids differ only in the bits `ColumnAccessPath` drops share every + /// `EntryID` this column computes; the cache compares this value to tell + /// them apart and treat the other file's data as a miss. + /// + /// Held as a lease rather than copied: a row group outlives the + /// `CachedFile` it came from, and the id must stay allocated for as long + /// as anything can still compute a key from it. + file_id: Arc, expression: Option>, } @@ -46,6 +55,7 @@ impl CachedColumn { field: Arc, cache_store: Arc, column_access_path: ColumnAccessPath, + file_id: Arc, expression: Option>, is_predicate_column: bool, ) -> Self { @@ -69,6 +79,7 @@ impl CachedColumn { field, cache_store, column_path: column_access_path, + file_id, expression, } } @@ -78,8 +89,14 @@ impl CachedColumn { self.column_path.entry_id(batch_id) } + /// The never-reused name of the file this column belongs to. + pub(crate) fn identity(&self) -> u64 { + self.file_id.identity() + } + pub(crate) fn is_cached(&self, batch_id: BatchID) -> bool { - self.cache_store.is_cached(&self.entry_id(batch_id).into()) + self.cache_store + .is_cached(&self.entry_id(batch_id).into(), self.identity()) } /// Returns the Arrow field metadata for this cached column. @@ -92,9 +109,12 @@ impl CachedColumn { self.expression.clone() } - fn array_to_record_batch(&self, array: ArrayRef) -> RecordBatch { + /// `None` when the array does not match this column's field — the cache + /// returned something built for a different column. The caller treats that + /// as "cannot answer from cache" and reads the source instead. + fn array_to_record_batch(&self, array: ArrayRef) -> Option { let schema = Arc::new(Schema::new(vec![self.field.clone()])); - RecordBatch::try_new(schema, vec![array]).unwrap() + RecordBatch::try_new(schema, vec![array]).ok() } /// Evaluates a predicate on a cached column. @@ -114,7 +134,7 @@ impl CachedColumn { if let Some(liquid_expr) = liquid_expr && let Some(boolean_array) = self .cache_store - .eval_predicate(&entry_id, &liquid_expr) + .eval_predicate(&entry_id, self.identity(), &liquid_expr) .with_selection(filter) .await { @@ -126,7 +146,7 @@ impl CachedColumn { } let array = self.get_arrow_array_with_filter(batch_id, filter).await?; - let record_batch = self.array_to_record_batch(array); + let record_batch = self.array_to_record_batch(array)?; let boolean_array = match predicate.evaluate(record_batch) { Ok(arr) => arr, Err(err) => return Some(Err(err)), @@ -160,7 +180,7 @@ impl CachedColumn { ) -> Option { let entry_id = self.entry_id(batch_id).into(); self.cache_store - .get(&entry_id) + .get(&entry_id, self.identity()) .with_selection(filter) .with_optional_expression_hint(self.expression()) .read() @@ -170,7 +190,7 @@ impl CachedColumn { #[cfg(test)] pub(crate) async fn get_arrow_array_test_only(&self, batch_id: BatchID) -> Option { let entry_id = self.entry_id(batch_id).into(); - self.cache_store.get(&entry_id).await + self.cache_store.get(&entry_id, self.identity()).await } /// Insert an array into the cache. @@ -184,7 +204,7 @@ impl CachedColumn { } self.cache_store - .insert(self.entry_id(batch_id).into(), array) + .insert(self.entry_id(batch_id).into(), self.identity(), array) .await?; Ok(()) } diff --git a/src/datafusion/src/cache/file_id.rs b/src/datafusion/src/cache/file_id.rs new file mode 100644 index 000000000..e0f96a376 --- /dev/null +++ b/src/datafusion/src/cache/file_id.rs @@ -0,0 +1,438 @@ +//! Allocation of the file ids that name cached data. +//! +//! An id is the part of a cache key that says which file an entry came from. +//! It is narrowed into 16 bits by [`crate::cache::ColumnAccessPath`], so the +//! supply of *distinct* keys is finite while the number of files a process +//! opens is not. Handing ids out from a counter that only ever climbs means a +//! long-lived process eventually reuses a key while its previous owner's data +//! is still cached. +//! +//! So an id is a lease rather than a permanent assignment. It is held by +//! everything that can still compute a key from it — the file handle, the row +//! groups and columns derived from it — and returns to the pool when the last +//! of them is dropped. The live id count is then bounded by what is actually +//! being read, not by everything that has ever been read. +//! +//! Cache *entries* deliberately do not hold a lease. An id can be reused while +//! entries keyed from it are still resident, and each entry records the +//! identity of the file it came from, so the new owner's reads miss rather +//! than returning the previous owner's rows. +//! +//! Its writes are not refused, though. A key held by another identity belongs +//! to a file that has already let its id go, so nothing can read that entry +//! any more and the new owner takes the key over +//! (`liquid_cache::cache::ArtIndex::insert`). Refusing instead would leave the +//! key occupied by data nobody can use, and on a cache below its budget +//! nothing evicts it — the new owner would never cache that key again. +//! +//! The alternative, releasing ids from inside index removal, would take a +//! process-wide lock underneath a crossbeam-epoch pin and deadlock against +//! `reset`. Keeping id lifetime and entry lifetime separate is what avoids +//! that, and the identity check is what makes the overlap safe. + +use std::collections::VecDeque; +use std::sync::atomic::{AtomicU64, Ordering}; + +use ahash::AHashMap; +// Through `crate::sync`, not `std::sync`: under the shuttle test feature this +// resolves to shuttle's primitives, which is what lets the model checker +// explore interleavings across `acquire` and `release`. A `std::sync::Mutex` +// is opaque to it, so the pool would be excluded from the very job that is +// meant to cover it. +use crate::sync::{Arc, Mutex, Weak}; + +/// A leased file id. The id returns to its pool when this is dropped. +/// +/// Two numbers, because they answer different questions and only one of them +/// can be recycled: +/// +/// * `id` goes into the cache key, whose file field is 16 bits wide. It has to +/// be recycled or a long-lived process runs out. +/// * `identity` names *which file* an entry came from, and is never reused. +/// It cannot be the recycled id: a file that inherits id 0 from a file that +/// has finished would otherwise be indistinguishable from it, and would read +/// the entries it left behind — the exact aliasing the identity exists to +/// catch. +#[derive(Debug)] +pub(crate) struct FileId { + id: u64, + identity: u64, + path: String, + pool: Arc, +} + +impl FileId { + /// The narrow, recycled id the cache key is built from. + pub(crate) fn get(&self) -> u64 { + self.id + } + + /// The wide, never-reused name for this file, recorded alongside every + /// entry so a recycled key cannot serve one file another's data. + pub(crate) fn identity(&self) -> u64 { + self.identity + } +} + +impl Drop for FileId { + fn drop(&mut self) { + self.pool.release(&self.path, self.id, self.identity); + } +} + +/// Hands out file ids and takes them back. +#[derive(Debug, Default)] +pub(crate) struct FileIdPool { + inner: Mutex, + /// Ids ever allocated that did not fit the key's 16-bit file field. A + /// non-zero count means keys are aliasing and the cache is refusing to + /// serve entries across the alias, which is correct but costs hit rate. + over_key_width: AtomicU64, +} + +#[derive(Debug, Default)] +struct PoolInner { + /// Live leases by path, so concurrent readers of one file share an id. + /// Entries are weak: the map never keeps a file alive on its own, and a + /// path is removed when its lease is released. + leases: AHashMap>, + /// Released ids, reused oldest-first. FIFO rather than LIFO on purpose: a + /// just-released id is the one whose entries are most likely still + /// resident, and reusing it last gives them the longest window to be + /// evicted before anything keys over them. + /// + /// Each carries the path that released it and the identity it had. If the + /// same path comes back it keeps that identity, so its cached entries are + /// still its own and still readable — a file read twice is a cache hit, + /// not a collision. Any other path gets a fresh identity, so it cannot + /// read what the previous holder left behind. + free: VecDeque, + next: u64, + /// Only ever climbs. A `u64` of these is not a resource worth reclaiming: + /// at one a microsecond it outlasts the hardware. + next_identity: u64, +} + +#[derive(Debug)] +struct Released { + id: u64, + path: String, + identity: u64, +} + +impl FileIdPool { + pub(crate) fn new() -> Arc { + Arc::new(Self::default()) + } + + /// The lease for `path`, shared with any reader already holding one. + pub(crate) fn acquire(self: &Arc, path: &str) -> Arc { + let mut inner = self.inner.lock().unwrap(); + if let Some(existing) = inner.leases.get(path).and_then(Weak::upgrade) { + return existing; + } + // Prefer this path's own released record, wherever it sits in the + // queue. Matching only the front would restore an identity just when + // release order happens to match acquire order — release order is + // stream completion order and acquire order is partition open order, + // so for any scan over more than one file they diverge and every + // re-read would orphan the entries it cached last time. + let mine = inner.free.iter().position(|r| r.path == path); + let (id, reusable_identity) = match mine { + Some(at) => { + let released = inner.free.remove(at).expect("index came from the queue"); + (released.id, Some(released.identity)) + } + None => match inner.free.pop_front() { + Some(released) => (released.id, None), + None => { + let id = inner.next; + inner.next += 1; + (id, None) + } + }, + }; + if id > u16::MAX as u64 { + self.over_key_width.fetch_add(1, Ordering::Relaxed); + } + let identity = match reusable_identity { + Some(identity) => identity, + None => { + let identity = inner.next_identity; + inner.next_identity += 1; + identity + } + }; + let lease = Arc::new(FileId { + id, + identity, + path: path.to_string(), + pool: Arc::clone(self), + }); + inner + .leases + .insert(path.to_string(), Arc::downgrade(&lease)); + lease + } + + fn release(&self, path: &str, id: u64, self_identity: u64) { + let Ok(mut inner) = self.inner.lock() else { + // A poisoned pool means some other thread panicked holding it. + // Losing one id is better than panicking again inside a drop. + return; + }; + // Only drop the path if it still points at the lease being released. + // A new lease for the same path may already have replaced it, and + // removing that one would hand the same file two live ids. + if inner + .leases + .get(path) + .is_some_and(|weak| weak.strong_count() == 0) + { + inner.leases.remove(path); + } + inner.free.push_back(Released { + id, + path: path.to_string(), + identity: self_identity, + }); + } + + /// Ids currently leased. Bounded by what is being read, which is what + /// keeps the 16-bit key field from running out. + pub(crate) fn live_count(&self) -> usize { + self.inner.lock().map(|i| i.leases.len()).unwrap_or(0) + } + + /// Ids handed out that do not fit the key's file field. Expected to stay + /// at zero. + pub(crate) fn over_key_width(&self) -> u64 { + self.over_key_width.load(Ordering::Relaxed) + } + + /// Forget every lease and start ids from zero again. + /// + /// Only valid when nothing holds a lease; callers that still do would keep + /// computing keys from ids this pool is free to hand out again. + pub(crate) fn reset(&self) { + if let Ok(mut inner) = self.inner.lock() { + inner.leases.clear(); + inner.free.clear(); + inner.next = 0; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A file re-opened after other files have come and gone must still find + /// its own record. Release order is stream completion order and acquire + /// order is partition open order, so the two rarely line up; matching only + /// the front of the queue would hand a re-read a fresh identity and orphan + /// everything it cached before. + #[test] + fn a_reopened_path_finds_its_record_anywhere_in_the_queue() { + let pool = FileIdPool::new(); + + let a = pool.acquire("a.parquet"); + let b = pool.acquire("b.parquet"); + let (a_id, a_identity) = (a.get(), a.identity()); + let (b_id, b_identity) = (b.get(), b.identity()); + + // Released a-then-b, so b's record sits behind a's. + drop(a); + drop(b); + + // Re-open b first: the queue front is a's record, not b's. + let b_again = pool.acquire("b.parquet"); + assert_eq!(b_again.get(), b_id, "b should get its own id back"); + assert_eq!( + b_again.identity(), + b_identity, + "b should keep its identity, or its cached entries are orphaned" + ); + + let a_again = pool.acquire("a.parquet"); + assert_eq!(a_again.get(), a_id); + assert_eq!(a_again.identity(), a_identity); + } + + /// Two leases alive at the same time must never share an id, and never + /// share an identity. That is the property the whole scheme rests on: + /// a shared id means two files computing one key, and a shared identity + /// means the check that catches it cannot tell them apart. + /// + /// Run under the model checker because `acquire` and `release` race by + /// construction — a lease is released from `Drop`, on whatever thread + /// happened to hold it last. + fn concurrent_leases_stay_distinct() { + let pool = FileIdPool::new(); + let mut threads = Vec::new(); + + for t in 0..3 { + let pool = Arc::clone(&pool); + threads.push(crate::sync::thread::spawn(move || { + for i in 0..3 { + let mine = pool.acquire(&format!("f{t}-{i}.parquet")); + + // Held at the same time, so they cannot be the same file. + let probe = pool.acquire("probe.parquet"); + assert_ne!(mine.get(), probe.get(), "two live leases shared an id"); + assert_ne!( + mine.identity(), + probe.identity(), + "two live leases shared an identity" + ); + drop(probe); + + // The same path always resolves to the same lease. + let again = pool.acquire(&format!("f{t}-{i}.parquet")); + assert_eq!(mine.get(), again.get()); + assert_eq!(mine.identity(), again.identity()); + } + })); + } + + for thread in threads { + thread.join().unwrap(); + } + } + + #[test] + fn concurrent_leases_stay_distinct_single_threaded() { + concurrent_leases_stay_distinct(); + } + + #[cfg(feature = "shuttle")] + #[test] + fn shuttle_concurrent_leases_stay_distinct() { + let mut runner = shuttle::PortfolioRunner::new(true, Default::default()); + let cores = std::thread::available_parallelism().unwrap().get().min(4); + for _ in 0..cores { + runner.add(shuttle::scheduler::PctScheduler::new(10, 1_000)); + } + runner.run(concurrent_leases_stay_distinct); + } + + #[test] + fn concurrent_readers_of_one_path_share_a_lease() { + let pool = FileIdPool::new(); + let first = pool.acquire("a.parquet"); + let second = pool.acquire("a.parquet"); + assert_eq!(first.get(), second.get()); + assert_eq!(pool.live_count(), 1); + } + + #[test] + fn an_id_returns_to_the_pool_when_its_last_holder_drops() { + let pool = FileIdPool::new(); + let a = pool.acquire("a.parquet"); + let b = pool.acquire("b.parquet"); + assert_eq!((a.get(), b.get()), (0, 1)); + assert_eq!(pool.live_count(), 2); + + drop(a); + assert_eq!(pool.live_count(), 1, "the released path is forgotten"); + + // Reused rather than climbing to 2: the supply tracks what is being + // read, which is the whole point. + let c = pool.acquire("c.parquet"); + assert_eq!(c.get(), 0); + } + + #[test] + fn a_second_holder_keeps_the_id_alive() { + let pool = FileIdPool::new(); + let first = pool.acquire("a.parquet"); + let second = pool.acquire("a.parquet"); + drop(first); + assert_eq!(pool.live_count(), 1); + // Bound, not a temporary: an unheld lease is released the moment the + // expression ends, which would put its id back before the next call. + let other = pool.acquire("b.parquet"); + assert_eq!(other.get(), 1, "id 0 is still leased"); + drop(second); + let recycled = pool.acquire("c.parquet"); + assert_eq!(recycled.get(), 0); + } + + /// The two numbers have to move independently. Reusing an id is how the + /// key space stays bounded; reusing an *identity* for a different file is + /// how one file reads another's entries. Re-opening the same path must + /// keep its identity, or every lease boundary silently empties the cache. + #[test] + fn identity_follows_the_path_while_the_id_is_recycled() { + let pool = FileIdPool::new(); + + let first = pool.acquire("a.parquet"); + let (a_id, a_identity) = (first.get(), first.identity()); + drop(first); + + // Same file again: same id and the same name, so its cached entries + // are still its own. + let reopened = pool.acquire("a.parquet"); + assert_eq!(reopened.get(), a_id); + assert_eq!( + reopened.identity(), + a_identity, + "re-opening a file must keep its identity, or its cache is dead" + ); + drop(reopened); + + // A different file inherits the id but must not inherit the name. + let other = pool.acquire("b.parquet"); + assert_eq!(other.get(), a_id, "the id is recycled"); + assert_ne!( + other.identity(), + a_identity, + "a different file must not be able to read what the last one left" + ); + } + + #[test] + fn released_ids_are_reused_oldest_first() { + let pool = FileIdPool::new(); + let a = pool.acquire("a.parquet"); + let b = pool.acquire("b.parquet"); + drop(a); + drop(b); + // 0 was released first, so it is handed out first — giving b's entries + // the longer eviction window. + assert_eq!(pool.acquire("x.parquet").get(), 0); + assert_eq!(pool.acquire("y.parquet").get(), 1); + } + + #[test] + fn ids_beyond_the_key_width_are_counted() { + let pool = FileIdPool::new(); + { + let mut inner = pool.inner.lock().unwrap(); + inner.next = u16::MAX as u64; + } + let _fits = pool.acquire("fits.parquet"); + assert_eq!(pool.over_key_width(), 0); + let _over = pool.acquire("over.parquet"); + assert_eq!(pool.over_key_width(), 1); + } + + /// A path re-registered while its old lease is being dropped must not lose + /// the new lease's entry in the map — that would give one file two live + /// ids and split its cache. + #[test] + fn releasing_a_stale_lease_leaves_a_newer_one_alone() { + let pool = FileIdPool::new(); + let first = pool.acquire("a.parquet"); + let first_id = first.get(); + drop(first); + let second = pool.acquire("a.parquet"); + assert_eq!(second.get(), first_id, "the id came back round"); + assert_eq!(pool.live_count(), 1); + assert_eq!( + pool.acquire("a.parquet").get(), + second.get(), + "the live lease is still the one the map points at" + ); + } +} diff --git a/src/datafusion/src/cache/id.rs b/src/datafusion/src/cache/id.rs index b52379dea..11f90a07f 100644 --- a/src/datafusion/src/cache/id.rs +++ b/src/datafusion/src/cache/id.rs @@ -51,8 +51,19 @@ const _: () = assert!(std::mem::align_of::() == 8); impl ParquetArrayID { /// Creates a new CacheEntryID. + /// + /// `file_id` is narrowed to the width the packed key gives it, and that is + /// deliberately not an assertion: a process that outlives 65,536 distinct + /// files goes on working, because the cache compares each entry's recorded + /// identity and treats an aliased key as a miss. Asserting would panic in + /// debug builds on a case release builds handle, leaving the shipped + /// behaviour untestable. `LiquidCacheParquet` counts ids that will not fit. + /// + /// The other two keep their assertion, because that identity is the file + /// id *alone*: it cannot tell row group 0 from row group 65,536, or column + /// 0 from column 65,536, within one file. Narrowing those is unguarded and + /// silently wrong, so it stays an assertion rather than a handled case. pub fn new(file_id: u64, row_group_id: u64, column_id: u64, batch_id: BatchID) -> Self { - debug_assert!(file_id <= u16::MAX as u64); debug_assert!(row_group_id <= u16::MAX as u64); debug_assert!(column_id <= u16::MAX as u64); Self { @@ -146,8 +157,11 @@ pub struct ColumnAccessPath { impl ColumnAccessPath { /// Create a new instance of ColumnAccessPath. + /// + /// `file_id` narrows unguarded and is handled by the identity check; the + /// other two assert, because the identity cannot distinguish them — see + /// [`ParquetArrayID::new`]. pub fn new(file_id: u64, row_group_id: u64, column_id: u64) -> Self { - debug_assert!(file_id <= u16::MAX as u64); debug_assert!(row_group_id <= u16::MAX as u64); debug_assert!(column_id <= u16::MAX as u64); Self { @@ -225,21 +239,34 @@ mod tests { assert_eq!(entry_id.batch_id_inner(), *batch_id as u64); } + /// `file_id` wraps rather than panicking, in every build. The consequence + /// — two sources computing one key — is caught by the identity the cache + /// records alongside each entry, not here. This pins the wrap down so the + /// aliasing it produces stays a known, reproducible condition rather than + /// a debug-only assertion the shipped binary never evaluates. #[test] - #[should_panic] - fn test_cache_entry_id_new_panic_file_id() { - ParquetArrayID::new((u16::MAX as u64) + 1, 0, 0, BatchID::from_raw(0)); + fn a_file_id_wraps_rather_than_panicking_above_its_width() { + let wrapped = ParquetArrayID::new((u16::MAX as u64) + 1, 1, 2, BatchID::from_raw(0)); + assert_eq!(wrapped.file_id_inner(), 0); + + // Which is exactly the aliasing the identity check exists to absorb. + let first = ParquetArrayID::new(0, 1, 2, BatchID::from_raw(0)); + assert_eq!(usize::from(wrapped), usize::from(first)); } + /// The other two fields keep their assertion. The identity recorded + /// against an entry is the file id alone, so it cannot tell row group 0 + /// from row group 65,536 within one file — narrowing those is unguarded + /// and silently wrong, not absorbed. #[test] - #[should_panic] - fn test_cache_entry_id_new_panic_row_group_id() { + #[should_panic(expected = "row_group_id")] + fn an_over_width_row_group_still_asserts() { ParquetArrayID::new(0, (u16::MAX as u64) + 1, 0, BatchID::from_raw(0)); } #[test] - #[should_panic] - fn test_cache_entry_id_new_panic_column_id() { + #[should_panic(expected = "column_id")] + fn an_over_width_column_still_asserts() { ParquetArrayID::new(0, 0, (u16::MAX as u64) + 1, BatchID::from_raw(0)); } diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index 814d0d324..404093ed9 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -3,7 +3,6 @@ use crate::io::ParquetCacheMetadata; use crate::reader::{LiquidPredicate, extract_multi_column_or}; -use crate::sync::Mutex; use ahash::AHashMap; use arrow::array::{BooleanArray, RecordBatch, RecordBatchOptions}; use arrow::buffer::BooleanBuffer; @@ -16,12 +15,14 @@ use parquet::arrow::arrow_reader::ArrowPredicate; use std::collections::HashMap; use std::path::Path; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; mod column; +mod file_id; mod id; mod stats; +use file_id::{FileId, FileIdPool}; + pub(crate) use column::InsertArrowArrayError; pub use column::{CachedColumn, CachedColumnRef}; pub(crate) use id::ColumnAccessPath; @@ -58,16 +59,18 @@ impl CachedRowGroup { fn new( cache_store: Arc, row_group_idx: u64, - file_idx: u64, + file_id: Arc, columns: &[CachedColumnSpec], ) -> Self { let mut column_maps = ColumnMaps::default(); for (column_id, field, expression, is_predicate_column) in columns { - let column_access_path = ColumnAccessPath::new(file_idx, row_group_idx, *column_id); + let column_access_path = + ColumnAccessPath::new(file_id.get(), row_group_idx, *column_id); let column = Arc::new(CachedColumn::new( Arc::clone(field), Arc::clone(&cache_store), column_access_path, + Arc::clone(&file_id), expression.clone(), *is_predicate_column, )); @@ -138,7 +141,10 @@ impl CachedRowGroup { } }; let entry_id = column.entry_id(batch_id).into(); - let liquid_array = self.cache_store.try_read_liquid(&entry_id).await; + let liquid_array = self + .cache_store + .try_read_liquid(&entry_id, column.identity()) + .await; let liquid_array = match liquid_array { None => { combined_buffer = None; @@ -146,7 +152,15 @@ impl CachedRowGroup { } Some(array) => array, }; - let buffer = liquid_array.try_eval_predicate(&liquid_expr, selection); + // Leave the loop rather than the function, as the two + // arms above do: an array that cannot answer the predicate + // does not mean the column is unreadable, and the arrow + // fallback below may still serve it from the cache. + let Some(buffer) = liquid_array.try_eval_predicate(&liquid_expr, selection) + else { + combined_buffer = None; + break; + }; combined_buffer = Some(match combined_buffer { None => buffer, @@ -190,7 +204,9 @@ pub(crate) type CachedRowGroupRef = Arc; #[derive(Debug)] pub struct CachedFile { cache_store: Arc, - file_id: u64, + /// Held, not copied: the id stays allocated for as long as anything can + /// still compute a cache key from it. + file_id: Arc, file_schema: SchemaRef, squeeze_hints: Arc, } @@ -198,7 +214,7 @@ pub struct CachedFile { impl CachedFile { fn new( cache_store: Arc, - file_id: u64, + file_id: Arc, file_schema: SchemaRef, squeeze_hints: Arc, ) -> Self { @@ -236,11 +252,17 @@ impl CachedFile { Arc::new(CachedRowGroup::new( self.cache_store.clone(), row_group_id, - self.file_id, + Arc::clone(&self.file_id), &columns, )) } + /// The leased id this file's cache keys are built from. + #[cfg(test)] + pub(crate) fn file_id(&self) -> u64 { + self.file_id.get() + } + /// Return the configured cache batch size. pub fn batch_size(&self) -> usize { self.cache_store.config().batch_size() @@ -258,12 +280,12 @@ pub(crate) type CachedFileRef = Arc; /// The main cache structure. #[derive(Debug)] pub struct LiquidCacheParquet { - /// Map file path to file id. - files: Mutex>, + /// Leases the file ids that name cached data. Ids come back when nothing + /// is reading the file any more, so the number in use tracks what is being + /// read rather than everything ever read — see [`file_id`]. + file_ids: Arc, cache_store: Arc, - - current_file_id: AtomicU64, } /// A reference to the main cache structure. @@ -322,9 +344,8 @@ impl LiquidCacheParquet { .await; LiquidCacheParquet { - files: Mutex::new(AHashMap::new()), + file_ids: FileIdPool::new(), cache_store: cache_storage, - current_file_id: AtomicU64::new(0), } } @@ -345,15 +366,9 @@ impl LiquidCacheParquet { full_file_schema: SchemaRef, squeeze_hints: Arc, ) -> CachedFileRef { - let mut files = self.files.lock().unwrap(); - let file_id = *files - .entry(file_path.clone()) - .or_insert_with(|| self.current_file_id.fetch_add(1, Ordering::Relaxed)); - drop(files); - Arc::new(CachedFile::new( self.cache_store.clone(), - file_id, + self.file_ids.acquire(&file_path), full_file_schema, squeeze_hints, )) @@ -384,6 +399,34 @@ impl LiquidCacheParquet { self.cache_store.budget().disk_usage_bytes() } + /// How many file ids are currently leased. + /// + /// This tracks the files being read, not the files ever read. It is the + /// number that has to stay under the cache key's 16-bit file field, so it + /// is worth watching: rising without bound means leases are being held by + /// something that should have let go. + pub fn leased_file_ids(&self) -> usize { + self.file_ids.live_count() + } + + /// How many ids have been handed out that do not fit the cache key's file + /// field. + /// + /// Expected to stay at zero. Above zero, distinct files are computing the + /// same keys — served correctly, because each entry records which file it + /// came from, but unable to share the cache. + pub fn file_ids_over_key_width(&self) -> u64 { + self.file_ids.over_key_width() + } + + /// How many cache lookups or writes found a key held by another file. + /// + /// The consequence of the counter above, and the one that proves the + /// aliasing is being caught rather than served. + pub fn identity_mismatches(&self) -> u64 { + self.cache_store.stats().identity_mismatches + } + /// Flush the cache trace to a file. pub fn flush_trace(&self, to_file: impl AsRef) { self.cache_store.observer().flush_cache_trace(to_file); @@ -404,10 +447,9 @@ impl LiquidCacheParquet { /// # Safety /// This is unsafe because resetting the cache while other threads are using the cache may cause undefined behavior. /// You should only call this when no one else is using the cache. - pub unsafe fn reset(&self) { - let mut files = self.files.lock().unwrap(); - files.clear(); - self.cache_store.reset(); + pub async unsafe fn reset(&self) { + self.file_ids.reset(); + self.cache_store.reset().await; } /// Flush all memory-based entries to disk while preserving their format. @@ -437,7 +479,7 @@ mod tests { use super::*; use crate::cache::{CachedRowGroupRef, LiquidCacheParquet}; use crate::reader::FilterCandidateBuilder; - use arrow::array::{Array, Int32Array}; + use arrow::array::{Array, ArrayRef, Int32Array}; use arrow::buffer::BooleanBuffer; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; @@ -470,6 +512,189 @@ mod tests { file.create_row_group(0, vec![]) } + /// Recycling a file id must not recycle a file's *name*. + /// + /// The id is narrow and reused so the key space cannot run out. If the + /// identity recorded against each entry were that same id, the next file + /// to inherit it would be indistinguishable from the one that gave it + /// back, and would read the entries it left behind — reintroducing the + /// aliasing the identity exists to catch, at every lease boundary rather + /// than only past 65,536 files. + #[tokio::test] + async fn a_file_inheriting_a_recycled_id_does_not_read_its_predecessors_data() { + let batch_size = 8; + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let tmp_dir = tempfile::tempdir().unwrap(); + let store = crate::test_utils::mount_test_store(tmp_dir.path()).await; + let cache = LiquidCacheParquet::new( + batch_size, + usize::MAX, + usize::MAX, + store, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeSqueezeEvict), + Box::new(AlwaysHydrate::new()), + ) + .await; + + let batch_id = BatchID::from_row_id(0, batch_size); + let filter = BooleanBuffer::new_set(batch_size); + let first_data: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8])); + + let first_id = { + let first = + cache.register_or_get_file("first.parquet".to_string(), Arc::clone(&schema)); + let column = first.create_row_group(0, vec![]).get_column(0).unwrap(); + column + .insert(batch_id, Arc::clone(&first_data)) + .await + .unwrap(); + first.file_id() + }; // lease dropped here, so the id goes back to the pool + + let second = cache.register_or_get_file("second.parquet".to_string(), schema); + assert_eq!( + second.file_id(), + first_id, + "the id must actually be recycled, or this test proves nothing" + ); + + let column = second.create_row_group(0, vec![]).get_column(0).unwrap(); + assert!(!column.is_cached(batch_id)); + assert!( + column + .get_arrow_array_with_filter(batch_id, &filter) + .await + .is_none(), + "inheriting an id must not inherit the entries keyed from it" + ); + + // It must also be able to cache. The predecessor's entries are keyed + // where this file's belong and nobody can read them any more, so they + // give way — otherwise a cache under its budget, where nothing is ever + // evicted, would leave this file permanently uncacheable. + let second_data: ArrayRef = Arc::new(Int32Array::from(vec![9, 9, 9, 9, 9, 9, 9, 9])); + column + .insert(batch_id, Arc::clone(&second_data)) + .await + .expect("the inheriting file must be able to cache"); + let got = column + .get_arrow_array_with_filter(batch_id, &filter) + .await + .expect("the new owner reads back its own rows"); + assert_eq!(got.as_ref(), second_data.as_ref()); + } + + /// What part of the fix is actually for: a process that reads far more + /// files than it holds open at once must not exhaust the key's 16-bit file + /// field. Before ids were leased this counter only ever climbed, so a + /// long-lived instance wrapped it purely by having *seen* enough files. + #[tokio::test] + async fn reading_files_one_after_another_does_not_consume_the_id_space() { + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let tmp_dir = tempfile::tempdir().unwrap(); + let store = crate::test_utils::mount_test_store(tmp_dir.path()).await; + let cache = LiquidCacheParquet::new( + 8, + usize::MAX, + usize::MAX, + store, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeSqueezeEvict), + Box::new(AlwaysHydrate::new()), + ) + .await; + + // Well past the ceiling in total, but only ever one open at a time. + for i in 0..(u16::MAX as usize + 1_000) { + let file = cache.register_or_get_file(format!("scan-{i}.parquet"), Arc::clone(&schema)); + assert_eq!( + file.file_id(), + 0, + "each file should reuse the id the previous one gave back" + ); + } + + // And a file opened now still fits the key field. + let after = cache.register_or_get_file("after.parquet".to_string(), schema); + assert!(after.file_id() <= u16::MAX as u64); + } + + /// The bug in its real shape, walked through the actual registration path. + /// + /// `ColumnAccessPath` narrows the file id to 16 bits, so the 65,537th + /// distinct file a process registers is keyed identically to the first. + /// Before entries recorded their identity, the newcomer read the + /// incumbent's data — a panic when the column types differed, silently + /// wrong rows when they matched. + #[tokio::test] + async fn a_file_past_the_key_ceiling_does_not_read_the_first_file_s_data() { + let batch_size = 8; + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let tmp_dir = tempfile::tempdir().unwrap(); + let store = crate::test_utils::mount_test_store(tmp_dir.path()).await; + let cache = LiquidCacheParquet::new( + batch_size, + usize::MAX, + usize::MAX, + store, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeSqueezeEvict), + Box::new(AlwaysHydrate::new()), + ) + .await; + + let batch_id = BatchID::from_row_id(0, batch_size); + let filter = BooleanBuffer::new_set(batch_size); + + // File id 0, with data in the cache. + let first = cache.register_or_get_file("first.parquet".to_string(), Arc::clone(&schema)); + let first_column = first.create_row_group(0, vec![]).get_column(0).unwrap(); + let first_data: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8])); + first_column + .insert(batch_id, Arc::clone(&first_data)) + .await + .unwrap(); + + // Burn the rest of the 16-bit id space. The handles are held: ids are + // leased, so files that are opened and closed hand their id straight + // back and the ceiling is only reachable with this many files open at + // once. + let _fillers: Vec<_> = (1..=u16::MAX as usize) + .map(|i| cache.register_or_get_file(format!("filler-{i}.parquet"), Arc::clone(&schema))) + .collect(); + + // Id 65536, which narrows to 0. + let wrapped = cache.register_or_get_file("wrapped.parquet".to_string(), schema); + let wrapped_column = wrapped.create_row_group(0, vec![]).get_column(0).unwrap(); + + assert_eq!( + usize::from(wrapped_column.entry_id(batch_id)), + usize::from(first_column.entry_id(batch_id)), + "the packed keys must actually collide, or this test proves nothing" + ); + + // The newcomer must not be handed the incumbent's rows. + assert!(!wrapped_column.is_cached(batch_id)); + assert!( + wrapped_column + .get_arrow_array_with_filter(batch_id, &filter) + .await + .is_none(), + "a colliding key must read as a miss, not as the other file's data" + ); + + // And the incumbent still reads its own. + let got = first_column + .get_arrow_array_with_filter(batch_id, &filter) + .await + .expect("the owner's entry is still there"); + assert_eq!(got.as_ref(), first_data.as_ref()); + } + /// Issue #19: `NOT (s = s)` simplifies to `s IS NULL AND NULL`, so a conjunct /// that reads no column reaches the row filter. It has to survive candidate /// building and then evaluate against the selection's row count — an diff --git a/src/datafusion/src/cache/stats.rs b/src/datafusion/src/cache/stats.rs index 474abcb2d..1eb8863ce 100644 --- a/src/datafusion/src/cache/stats.rs +++ b/src/datafusion/src/cache/stats.rs @@ -123,37 +123,38 @@ impl LiquidCacheParquet { /// Write the stats of the cache to a parquet file. pub fn write_stats(&self, parquet_file_path: impl AsRef) -> Result<(), ParquetError> { let mut writer = StatsWriter::new(parquet_file_path)?; - self.cache_store.for_each_entry(|entry_id, cached_batch| { - let memory_size = cached_batch.memory_usage_bytes(); - let row_count = match cached_batch { - CacheEntry::MemoryArrow(array) => Some(array.len() as u64), - CacheEntry::MemoryLiquid(array) => Some(array.len() as u64), - CacheEntry::MemorySqueezedLiquid(array) => Some(array.len() as u64), - CacheEntry::DiskLiquid { .. } => None, - CacheEntry::DiskArrow { .. } => None, // We'd need to read it to get the count - }; - let cache_type = match cached_batch { - CacheEntry::MemoryArrow(_) => "InMemory", - CacheEntry::MemoryLiquid(_) => "LiquidMemory", - CacheEntry::MemorySqueezedLiquid(_) => "LiquidSqueezed", - CacheEntry::DiskLiquid { .. } => "OnDiskLiquid", - CacheEntry::DiskArrow { .. } => "OnDiskArrow", - }; - let reference_count = cached_batch.reference_count(); - let entry_id = ParquetArrayID::from(*entry_id); - writer - .append_entry( - &entry_id.display_path(), - entry_id.row_group_id_inner(), - entry_id.column_id_inner(), - entry_id.batch_id_inner() * self.batch_size() as u64, - row_count, - memory_size as u64, - cache_type, - reference_count as u64, - ) - .unwrap(); - }); + self.cache_store + .for_each_entry(|entry_id, _identity, cached_batch| { + let memory_size = cached_batch.memory_usage_bytes(); + let row_count = match cached_batch { + CacheEntry::MemoryArrow(array) => Some(array.len() as u64), + CacheEntry::MemoryLiquid(array) => Some(array.len() as u64), + CacheEntry::MemorySqueezedLiquid(array) => Some(array.len() as u64), + CacheEntry::DiskLiquid { .. } => None, + CacheEntry::DiskArrow { .. } => None, // We'd need to read it to get the count + }; + let cache_type = match cached_batch { + CacheEntry::MemoryArrow(_) => "InMemory", + CacheEntry::MemoryLiquid(_) => "LiquidMemory", + CacheEntry::MemorySqueezedLiquid(_) => "LiquidSqueezed", + CacheEntry::DiskLiquid { .. } => "OnDiskLiquid", + CacheEntry::DiskArrow { .. } => "OnDiskArrow", + }; + let reference_count = cached_batch.reference_count(); + let entry_id = ParquetArrayID::from(*entry_id); + writer + .append_entry( + &entry_id.display_path(), + entry_id.row_group_id_inner(), + entry_id.column_id_inner(), + entry_id.batch_id_inner() * self.batch_size() as u64, + row_count, + memory_size as u64, + cache_type, + reference_count as u64, + ) + .unwrap(); + }); writer.finish()?; Ok(()) @@ -205,9 +206,15 @@ mod tests { let mut row_start_id_sum = 0; let mut row_count_sum = 0; let mut memory_size_sum = 0; - for file_no in 0..8 { - let file_name = format!("test_{file_no}.parquet"); - let file = cache.register_or_get_file(file_name, schema.clone()); + // Held for the whole loop, not per iteration: a file id is leased and + // comes back when its handle drops, so releasing each file before + // opening the next would hand them all the same id. + let files: Vec<_> = (0..8) + .map(|file_no| { + cache.register_or_get_file(format!("test_{file_no}.parquet"), schema.clone()) + }) + .collect(); + for file in &files { for rg in 0..8 { let row_group = file.create_row_group(rg, vec![]); for col in 0..8 { diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index fadc8e70d..778717b30 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -385,9 +385,10 @@ impl LiquidCacheReaderInner { arrays.push(array); } - Ok(Some( - RecordBatch::try_new(self.schema.clone(), arrays).unwrap(), - )) + // A batch that does not match the declared schema is a cache that + // handed back something other than what was asked for. Report it; + // unwinding here aborts the stream mid-flight with no error to show. + Ok(Some(RecordBatch::try_new(self.schema.clone(), arrays)?)) } async fn read_parquet_batch_and_fill_cache(