diff --git a/README.md b/README.md index 1cf6210b..fcd42aca 100644 --- a/README.md +++ b/README.md @@ -40,18 +40,21 @@ use std::sync::Arc; tokio_test::block_on(async { let cache = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(1); + // The identity names the source the data belongs to; reads under any + // other identity miss. One source here, so any constant will do. + let identity = 0; let values = Arc::new(UInt64Array::from(vec![10, 11, 12, 13, 14, 15])); // 1) insert - cache.insert(entry_id, values.clone()).await; + cache.insert(entry_id, identity, values.clone()).await; // 2) get - let all_rows = cache.get(&entry_id).await.expect("entry should exist"); + let all_rows = cache.get(&entry_id, identity).await.expect("entry should exist"); // 3) get filtered (selection pushdown): keep rows 0, 2, 4 let selection = BooleanBuffer::from(vec![true, false, true, false, true, false]); let selected_rows = cache - .get(&entry_id) + .get(&entry_id, identity) .with_selection(&selection) .await .expect("entry should exist"); @@ -63,7 +66,7 @@ tokio_test::block_on(async { Arc::new(Literal::new(ScalarValue::UInt64(Some(12)))), )); let predicate_mask = cache - .eval_predicate(&entry_id, &predicate) + .eval_predicate(&entry_id, identity, &predicate) .await .expect("entry should exist") .expect("predicate should be evaluated in cache"); diff --git a/benchmark/src/lib.rs b/benchmark/src/lib.rs index a72185b0..dec471f6 100644 --- a/benchmark/src/lib.rs +++ b/benchmark/src/lib.rs @@ -8,7 +8,7 @@ use fastrace::Span; use fastrace::future::FutureExt as _; use liquid_cache::cache::CacheStats; use liquid_cache::cache::{Evict, EvictionPolicy, TranscodeEvict}; -use liquid_cache_common::rpc::ExecutionMetricsResponse; +use liquid_cache_common::metrics::ExecutionMetricsResponse; use liquid_cache_datafusion_server::{ApiResponse, ExecutionStats}; use log::info; use serde::Serialize; diff --git a/examples/core.rs b/examples/core.rs index 2d84688e..f5823291 100644 --- a/examples/core.rs +++ b/examples/core.rs @@ -18,10 +18,8 @@ async fn main() -> Result<(), Box> { .await; let entry_id = EntryID::from(7); - // The identity names which source this data came from. A cache key packs - // its fields into fixed widths, so two sources can compute one key; the - // identity is what keeps a read from being served the other's data. One - // source here, so any constant will do. + // The identity names the source this data belongs to. One source here, + // so any constant will do. let identity = 0; let arrow_array = Arc::new(UInt64Array::from_iter_values(0..16)); storage diff --git a/src/common/Cargo.toml b/src/common/Cargo.toml index 582f4bc9..ba612de0 100644 --- a/src/common/Cargo.toml +++ b/src/common/Cargo.toml @@ -7,14 +7,18 @@ readme = "README.md" description = { workspace = true } repository = { workspace = true } +[features] +# Arrow Flight RPC message types used between the client and the server. +rpc = ["dep:arrow-flight", "dep:prost"] + [dependencies] -arrow-flight = { workspace = true } +arrow-flight = { workspace = true, optional = true } async-trait = { workspace = true } bytes = { workspace = true } chrono = "0.4.45" futures = { workspace = true } object_store = { workspace = true } -prost = { workspace = true } +prost = { workspace = true, optional = true } serde = { workspace = true } thiserror = "2.0.18" tokio = { workspace = true } @@ -22,3 +26,7 @@ url = { workspace = true } [dev-dependencies] tempfile = { workspace = true } +tokio = { workspace = true, features = ["macros"] } + +[package.metadata.docs.rs] +features = ["rpc"] diff --git a/src/common/README.md b/src/common/README.md index 502f4227..9a2875f5 100644 --- a/src/common/README.md +++ b/src/common/README.md @@ -3,3 +3,5 @@ Common types and utilities shared across LiquidCache components. This crate contains shared data structures, enums, and utility functions used throughout the LiquidCache ecosystem. It defines the core abstractions for cache modes, schema transformations, and RPC message types that enable communication between client and server components. + +The RPC message types (`rpc` module) are behind the `rpc` feature, which is off by default. The client and server crates enable it; in-process use through `liquid-cache-datafusion-local` does not need it, and so does not depend on Arrow Flight or gRPC. diff --git a/src/common/src/lib.rs b/src/common/src/lib.rs index 5b4b43b8..57593084 100644 --- a/src/common/src/lib.rs +++ b/src/common/src/lib.rs @@ -1,5 +1,7 @@ #![doc = include_str!("../README.md")] +pub mod metrics; pub mod mock_store; +#[cfg(feature = "rpc")] pub mod rpc; pub mod utils; diff --git a/src/common/src/metrics.rs b/src/common/src/metrics.rs new file mode 100644 index 00000000..00960a77 --- /dev/null +++ b/src/common/src/metrics.rs @@ -0,0 +1,18 @@ +//! Execution metrics reported for a query plan. + +#[derive(Clone, PartialEq, serde::Serialize, serde::Deserialize, Debug)] +pub struct ExecutionMetricsResponse { + pub pushdown_eval_time: u64, + pub cache_memory_usage: u64, + pub liquid_cache_usage: u64, +} + +impl ExecutionMetricsResponse { + pub fn zero() -> Self { + Self { + pushdown_eval_time: 0, + cache_memory_usage: 0, + liquid_cache_usage: 0, + } + } +} diff --git a/src/common/src/rpc.rs b/src/common/src/rpc.rs index 452d4ef5..0907085a 100644 --- a/src/common/src/rpc.rs +++ b/src/common/src/rpc.rs @@ -246,20 +246,3 @@ impl ProstMessageExt for FetchResults { } } } - -#[derive(Clone, PartialEq, serde::Serialize, serde::Deserialize, Debug)] -pub struct ExecutionMetricsResponse { - pub pushdown_eval_time: u64, - pub cache_memory_usage: u64, - pub liquid_cache_usage: u64, -} - -impl ExecutionMetricsResponse { - pub fn zero() -> Self { - Self { - pushdown_eval_time: 0, - cache_memory_usage: 0, - liquid_cache_usage: 0, - } - } -} diff --git a/src/core/README.md b/src/core/README.md index 4352a42d..26acdce7 100644 --- a/src/core/README.md +++ b/src/core/README.md @@ -24,14 +24,12 @@ let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(42); let arrow_array = Arc::new(UInt64Array::from_iter_values(0..1000)); -// Insert once; replacement/placement is handled by the cache policy -// `0` is the file identity: the unnarrowed name of the source this data -// belongs to. Cache keys pack their fields into fixed widths, so two sources -// can compute one key; the identity is what keeps a read from being served -// the other's data. One source here, so any constant will do. +// Insert once; replacement/placement is handled by the cache policy. +// `0` is the identity: the source this data belongs to. Reads under any +// other identity miss. One source here, so any constant will do. storage.insert(entry_id, 0, arrow_array.clone()).await; -assert!(storage.is_cached(&entry_id, 0)); +assert!(storage.contains(&entry_id, 0)); }); ``` diff --git a/src/core/src/cache/builders.rs b/src/core/src/cache/builders.rs index 46d98e0e..4a08b9be 100644 --- a/src/core/src/cache/builders.rs +++ b/src/core/src/cache/builders.rs @@ -175,8 +175,6 @@ pub fn default_max_memory_bytes() -> usize { pub struct Insert<'a> { pub(super) storage: &'a Arc, pub(super) entry_id: EntryID, - /// The unnarrowed name of the source this data belongs to, recorded with - /// the entry so a key collision cannot serve it to anyone else. pub(super) identity: u64, pub(super) batch: ArrayRef, pub(super) skip_gc: bool, @@ -223,11 +221,7 @@ impl<'a> Insert<'a> { } let batch = CacheEntry::memory_arrow(batch); self.storage - .insert_inner( - self.entry_id, - crate::cache::index::WriteIdentity::Owned(self.identity), - batch, - ) + .insert_inner(self.entry_id, self.identity, batch) .await } } diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 376993bc..c8c3bcbd 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -17,11 +17,7 @@ use super::{ }; use crate::cache::policies::{EvictionOutcome, EvictionPolicy}; use crate::cache::utils::arrow_to_bytes; -use crate::cache::{ - CacheExpression, LiquidExpr, - index::{ArtIndex, WriteIdentity}, - utils::EntryID, -}; +use crate::cache::{CacheExpression, LiquidExpr, index::ArtIndex, utils::EntryID}; use crate::cache::{CacheFull, CacheStats, EventTrace}; use crate::sync::Arc; @@ -29,6 +25,15 @@ use crate::sync::Arc; /// Cache storage for liquid cache. /// +/// # Identity +/// +/// Every entry is cached for an *identity*: a `u64` naming the source the data +/// belongs to. Callers that pack several fields into a fixed-width +/// [`EntryID`] can have two sources compute the same key; the identity is how +/// the cache tells them apart. A lookup under any identity other than the one +/// an entry was cached for is a miss, and an insert replaces whatever held the +/// key. Callers whose keys are already unique can pass any constant. +/// /// Example (async read): /// ```rust /// use liquid_cache::cache::{LiquidCacheBuilder, EntryID}; @@ -40,9 +45,8 @@ use crate::sync::Arc; /// /// let entry_id = EntryID::from(0); /// let arrow_array = Arc::new(UInt64Array::from_iter_values(0..32)); -/// // `0` is the file identity: cache keys pack their fields into fixed -/// // widths, so two sources can compute one key, and the identity is what -/// // keeps a read from being served the other's data. +/// // The identity names the source this data belongs to. One source here, +/// // so any constant will do. /// storage.insert(entry_id, 0, arrow_array.clone()).await; /// /// // Get the arrow array back asynchronously @@ -72,78 +76,10 @@ pub enum PrefetchResult { Absent, } -/// Disk an insert left for its caller to reclaim. -/// -/// A store object is addressed by entry id *and* the identity that wrote it, so -/// an object stops being reachable through the index the moment the key changes -/// hands or the write that produced it is dropped. Nothing else will ever delete -/// it or release its bytes: `release_disk` is driven by an index entry, and by -/// then no index entry names it. -#[derive(Default)] -struct DiskResidue { - /// `(identity, disk_bytes)` of a disk-resident entry this write displaced, - /// whose object is now unreachable and has to be deleted. - displaced: Option<(u64, usize)>, - /// Bytes of a displaced disk entry whose object this write *overwrote in - /// place*. The object is current and must be kept; only the superseded - /// entry's reservation is given back, or one object is charged twice. - superseded: Option, - /// The write itself did not land, so whatever the caller already wrote to - /// the store under its own identity is unreachable too. - dropped: bool, -} - -impl DiskResidue { - fn dropped() -> Self { - Self { - displaced: None, - superseded: None, - dropped: true, - } - } - - /// A displaced disk entry's object survives this insert unless the insert - /// put its own bytes over it. - /// - /// The store key is `(entry id, identity)`, so only a disk-resident entry - /// written under the identity that already held the key addresses the very - /// same object: there the put overwrote it, and reclaiming would delete the - /// bytes just written. Otherwise the object is left behind — a different - /// identity addresses a different key, and an entry that lives in memory - /// wrote nothing at all — and it becomes unreachable the moment the index - /// stops naming it. - fn displacing( - displaced: Option<&(u64, Arc)>, - writer: u64, - written: CachedBatchType, - ) -> Self { - let overwrites_in_place = matches!( - written, - CachedBatchType::DiskLiquid | CachedBatchType::DiskArrow - ); - let disk_bytes = displaced.and_then(|(identity, entry)| match entry.as_ref() { - CacheEntry::DiskLiquid { disk_bytes, .. } - | CacheEntry::DiskArrow { disk_bytes, .. } => Some((*identity, *disk_bytes)), - CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) => None, - }); - // Same identity and a disk-resident write means this put landed on the - // very object the displaced entry named: keep the object, give back only - // its reservation. Anything else leaves an object nothing can reach. - match disk_bytes { - Some((identity, bytes)) if overwrites_in_place && identity == writer => Self { - displaced: None, - superseded: Some(bytes), - dropped: false, - }, - other => Self { - displaced: other, - superseded: None, - dropped: false, - }, - } - } -} +/// Length of the identity trailer appended to every object written to disk. +const IDENTITY_TRAILER_LEN: usize = std::mem::size_of::(); +/// Builder returned by [`LiquidCache::insert`] for configuring cache writes. impl LiquidCache { /// Return current cache statistics: counts and resource usage. pub fn stats(&self) -> CacheStats { @@ -158,7 +94,7 @@ impl LiquidCache { let mut memory_arrow_bytes = 0usize; let mut memory_liquid_bytes = 0usize; - self.index.for_each(|_, _, batch| match batch { + self.index.for_each(|_, slot| match &slot.entry { CacheEntry::MemoryArrow(array) => { memory_arrow_entries += 1; memory_arrow_bytes += array.get_array_memory_size(); @@ -177,7 +113,6 @@ impl LiquidCache { CacheStats { total_entries, - identity_mismatches: self.index.identity_mismatches(), memory_arrow_entries, memory_liquid_entries, disk_liquid_entries, @@ -192,7 +127,7 @@ impl LiquidCache { } } - /// Insert a batch into the cache. + /// Insert a batch into the cache for `identity`. pub fn insert<'a>( self: &'a Arc, entry_id: EntryID, @@ -202,7 +137,7 @@ impl LiquidCache { Insert::new(self, entry_id, identity, batch_to_cache) } - /// Create a [`Get`] builder for the provided entry. + /// Create a [`Get`] builder for the provided entry, as cached for `identity`. pub fn get<'a>(&'a self, entry_id: &'a EntryID, identity: u64) -> Get<'a> { Get::new(self, entry_id, identity) } @@ -219,14 +154,12 @@ impl LiquidCache { /// Prefetch an entry into a memory-form snapshot without recording an access. pub async fn prefetch(&self, entry_id: &EntryID, identity: u64) -> PrefetchResult { - // Checked, not raw: a prefetch hands a snapshot to a caller, so an - // aliased key must read as absent rather than as someone else's rows. - let Some(entry) = self.index.get_checked(entry_id, identity) else { + let Some(slot) = self.index.get_for(entry_id, identity) else { return PrefetchResult::Absent; }; - match entry.as_ref() { - CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) => { - PrefetchResult::Snapshot(entry) + match &slot.entry { + entry @ (CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_)) => { + PrefetchResult::Snapshot(Arc::new(entry.clone())) } disk @ CacheEntry::DiskArrow { .. } => { let Some(array) = self.read_disk_arrow_array(entry_id, identity).await else { @@ -268,11 +201,11 @@ impl LiquidCache { ) -> Option { self.observer.on_try_read_liquid(); self.trace(InternalEvent::TryReadLiquid { entry: *entry_id }); - let batch = self.index.get_checked(entry_id, identity)?; + let slot = self.index.get_for(entry_id, identity)?; self.cache_policy - .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); + .notify_access(entry_id, CachedBatchType::from(&slot.entry)); - match batch.as_ref() { + match &slot.entry { CacheEntry::MemoryLiquid(array) => Some(array.clone()), entry @ CacheEntry::DiskLiquid { .. } => { let liquid = self.read_disk_liquid_array(entry_id, identity).await?; @@ -290,11 +223,12 @@ impl LiquidCache { } } - /// Iterate over all entries in the cache. + /// Iterate over all entries in the cache, whichever identity holds them. /// 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, u64, &CacheEntry)) { - self.index.for_each(&mut f); + pub fn for_each_entry(&self, mut f: impl FnMut(&EntryID, &CacheEntry)) { + self.index + .for_each(|entry_id, slot| f(entry_id, &slot.entry)); } /// Reset the cache. @@ -303,13 +237,19 @@ impl LiquidCache { self.budget.reset_usage(); } - /// Check whether the cache holds this entry *for this identity*. + /// Check whether the cache holds a batch under `entry_id` for `identity`. + pub fn contains(&self, entry_id: &EntryID, identity: u64) -> bool { + self.index.contains(entry_id, identity) + } + + /// How many lookups found their key held by a different identity, and + /// missed rather than serve it. /// - /// A key held by another identity reads as absent: it belongs to a source - /// that can no longer read it, and serving it here would hand one caller - /// another's rows. - pub fn is_cached(&self, entry_id: &EntryID, identity: u64) -> bool { - self.index.is_cached(entry_id, identity) + /// Zero unless two sources compute the same [`EntryID`]. Each such lookup + /// was served correctly; a rising count means those sources cannot share + /// the cache. + pub fn identity_mismatches(&self) -> u64 { + self.index.identity_mismatches() } /// Get the config of the cache. @@ -340,28 +280,25 @@ 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, identity, batch| { - entires.push((*entry_id, identity, batch.clone())); + self.index.for_each(|entry_id, slot| { + entires.push((*entry_id, slot.identity, slot.entry.clone())); }); - for (entry_id, flush_identity, batch) in entires { + for (entry_id, 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, flush_identity, &batch, bytes) + .write_batch_to_disk(entry_id, identity, &batch, bytes) .await { Ok(()) => { - let residue = 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"); - self.settle(entry_id, residue, Some((flush_identity, disk_bytes))) - .await; + self.try_insert( + entry_id, + identity, + CacheEntry::disk_arrow(array.data_type().clone(), disk_bytes), + ) + .expect("failed to insert disk arrow entry"); } Err(CacheFull) => self.drop_memory_entry(entry_id, &batch), } @@ -370,27 +307,19 @@ impl LiquidCache { let liquid_bytes = liquid_array.to_bytes(); let disk_bytes = liquid_bytes.len(); match self - .write_batch_to_disk( - entry_id, - flush_identity, - &batch, - Bytes::from(liquid_bytes), - ) + .write_batch_to_disk(entry_id, identity, &batch, Bytes::from(liquid_bytes)) .await { Ok(()) => { - let residue = self - .try_insert( - entry_id, - WriteIdentity::Rewrite(flush_identity), - CacheEntry::disk_liquid( - liquid_array.original_arrow_data_type(), - disk_bytes, - ), - ) - .expect("failed to insert disk liquid entry"); - self.settle(entry_id, residue, Some((flush_identity, disk_bytes))) - .await; + self.try_insert( + entry_id, + identity, + CacheEntry::disk_liquid( + liquid_array.original_arrow_data_type(), + disk_bytes, + ), + ) + .expect("failed to insert disk liquid entry"); } Err(CacheFull) => self.drop_memory_entry(entry_id, &batch), } @@ -450,19 +379,12 @@ impl LiquidCache { pub(crate) async fn insert_inner( &self, entry_id: EntryID, - identity: WriteIdentity, + identity: u64, mut batch_to_cache: CacheEntry, ) -> Result<(), CacheFull> { - // Set once this loop spills the entry to disk itself: those bytes are the - // caller's own write, so a rewrite dropped as stale has to reclaim them. - let mut wrote = None; loop { - let not_inserted = match self.try_insert(entry_id, identity, batch_to_cache) { - Ok(residue) => { - self.settle(entry_id, residue, wrote).await; - return Ok(()); - } - Err(not_inserted) => not_inserted, + let Err(not_inserted) = self.try_insert(entry_id, identity, batch_to_cache) else { + return Ok(()); }; self.trace(InternalEvent::InsertFailed { entry: entry_id, @@ -475,13 +397,8 @@ 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, identity.value(), not_inserted) + .write_in_memory_batch_to_disk(entry_id, identity, not_inserted) .await?; - if let CacheEntry::DiskLiquid { disk_bytes, .. } - | CacheEntry::DiskArrow { disk_bytes, .. } = &on_disk_batch - { - wrote = Some((identity.value(), *disk_bytes)); - } batch_to_cache = on_disk_batch; continue; } @@ -528,12 +445,12 @@ impl LiquidCache { fn try_insert( &self, entry_id: EntryID, - identity: WriteIdentity, + identity: u64, to_insert: CacheEntry, - ) -> Result { + ) -> Result<(), CacheEntry> { let new_memory_size = to_insert.memory_usage_bytes(); - let (cached_batch_type, outcome) = if let Some(entry) = self.index.get(&entry_id) { - let old_memory_size = entry.memory_usage_bytes(); + let cached_batch_type = if let Some(slot) = self.index.get(&entry_id) { + let old_memory_size = slot.entry.memory_usage_bytes(); if self .budget .try_update_memory_usage(old_memory_size, new_memory_size) @@ -542,29 +459,15 @@ impl LiquidCache { return Err(to_insert); } let batch_type = CachedBatchType::from(&to_insert); - let outcome = self.index.insert(&entry_id, identity, to_insert); - if !outcome.stored { - // A rewrite whose key changed hands since it was read. Give the - // reservation back rather than counting memory for an entry that - // was never stored, and tell the caller its disk write is now - // unreachable. - self.budget - .try_update_memory_usage(new_memory_size, old_memory_size) - .ok(); - return Ok(DiskResidue::dropped()); - } - (batch_type, outcome) + self.index.insert(&entry_id, identity, to_insert); + 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); - let outcome = self.index.insert(&entry_id, identity, to_insert); - if !outcome.stored { - self.budget.try_update_memory_usage(new_memory_size, 0).ok(); - return Ok(DiskResidue::dropped()); - } - (batch_type, outcome) + self.index.insert(&entry_id, identity, to_insert); + batch_type }; self.trace(InternalEvent::InsertSuccess { @@ -574,75 +477,7 @@ impl LiquidCache { self.cache_policy .notify_insert(&entry_id, cached_batch_type); - Ok(DiskResidue::displacing( - outcome.displaced.as_ref(), - identity.value(), - cached_batch_type, - )) - } - - /// Delete a store object nothing can reach any more and give its bytes back. - /// - /// Reached on the paths where an object outlives the index entry that named - /// it: a write dropped as stale after its bytes were already written, an - /// entry displaced by a write under a different identity, and a disk entry - /// replaced by a memory one — hydration, or a caller overwriting the value - /// — which puts nothing in the store and so leaves the old object whole. - /// - /// The deletion is conditional because a store key is - /// `(entry id, identity)` and an identity comes back: the file-id pool - /// hands a re-opened path its previous record, so a fresh write can occupy - /// this key between the moment a reclaim was decided and the moment it - /// runs. `holds_disk_entry` asks the index whether a live record names the - /// object right now; if one does, that record's write put the bytes that - /// are there and deleting them would leave it pointing at nothing. - /// - /// The check is a read, and the deletion that follows it is an await, so it - /// narrows the window rather than closing it: a put that has landed but - /// whose record has not yet been installed is still invisible here. Closing - /// it needs the store key to name the write, not just the writer. - async fn reclaim_orphaned_disk(&self, entry_id: EntryID, identity: u64, disk_bytes: usize) { - if !self.index.holds_disk_entry(&entry_id, identity) { - match self - .store - .remove(&entry_id_to_key(&entry_id, identity)) - .await - { - // `false` means the object was already gone, which is fine: the - // bytes still have to be given back either way. - Ok(_) | Err(t4::Error::NotFound) => {} - Err(error) => panic!("orphan remove failed: {error}"), - } - } - // The reservation belonged to the record that is gone, so it comes back - // whether or not the object did: an object left standing is one a newer - // record overwrote, and that record reserved its own bytes for it. - self.budget.release_disk(disk_bytes); - self.trace(InternalEvent::DiskEvict { - entry: entry_id, - bytes: disk_bytes, - }); - } - - /// Reclaim whatever an insert left unreachable, including the caller's own - /// write when it was dropped as stale. - /// - /// `wrote` is the (identity, bytes) the caller put in the store before the - /// insert, if any. - async fn settle(&self, entry_id: EntryID, residue: DiskResidue, wrote: Option<(u64, usize)>) { - if let Some((identity, bytes)) = residue.displaced { - self.reclaim_orphaned_disk(entry_id, identity, bytes).await; - } - if let Some(bytes) = residue.superseded { - // The object stays — this write overwrote it — so only the byte - // count the superseded entry held is returned. - self.budget.release_disk(bytes); - } - if residue.dropped - && let Some((identity, bytes)) = wrote - { - self.reclaim_orphaned_disk(entry_id, identity, bytes).await; - } + Ok(()) } fn drop_memory_entry(&self, entry_id: EntryID, _expected: &CacheEntry) { @@ -651,36 +486,36 @@ impl LiquidCache { }; assert!( matches!( - removed.as_ref(), + &removed.entry, CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) ), "flush should only drop memory entries" ); self.budget - .try_update_memory_usage(removed.memory_usage_bytes(), 0) + .try_update_memory_usage(removed.entry.memory_usage_bytes(), 0) .expect("memory release cannot fail"); self.cache_policy.notify_remove(&entry_id); } - async fn remove_disk_entry(&self, entry_id: EntryID, removed_identity: u64) { - // Checked: the caller read this entry earlier, and between then and now - // the key can change hands. Removing the new owner's record would strand - // its store object while releasing a byte count taken from the record - // just destroyed. - let Some(removed) = self.index.remove_checked(&entry_id, removed_identity) else { + async fn remove_disk_entry(&self, entry_id: EntryID) { + let Some(removed) = self.index.remove(&entry_id) else { return; }; - let disk_bytes = match removed.as_ref() { + let disk_bytes = match &removed.entry { CacheEntry::DiskLiquid { disk_bytes, .. } | CacheEntry::DiskArrow { disk_bytes, .. } => *disk_bytes, _ => panic!("remove_disk_entry called for non-disk entry"), }; - // Through the same reclaim as every other orphan: the record is gone, - // so the object is unreachable, and the gap before the delete is the - // gap a re-used identity can write into. - self.reclaim_orphaned_disk(entry_id, removed_identity, disk_bytes) - .await; + self.store + .remove(&entry_id_to_key(&entry_id)) + .await + .expect("disk remove failed"); + self.budget.release_disk(disk_bytes + IDENTITY_TRAILER_LEN); self.cache_policy.notify_remove(&entry_id); + self.trace(InternalEvent::DiskEvict { + entry: entry_id, + bytes: disk_bytes, + }); } /// Consume the trace of the cache, for testing only. @@ -719,46 +554,40 @@ impl LiquidCache { } async fn evict_victim_inner(&self, victim: EntryID) -> Result<(), CacheFull> { - // Read the identity alongside the entry: everything this loop writes - // back is a rewrite of what it just read, and must be dropped rather - // than relabelled if the key changes hands meanwhile. - let Some((identity, mut victim_entry)) = self.index.get_with_identity(&victim) else { + let Some(slot) = self.index.get(&victim) else { return Ok(()); }; + // The rewrite is made for the identity the victim was read under. If + // another identity takes the key over meanwhile, the rewrite still + // carries the old identity, so the new owner reads it as a miss. + let identity = slot.identity; + let mut victim_entry = slot.entry.clone(); self.trace(InternalEvent::EvictionVictim { entry: victim }); loop { - let outcome = self.eviction_policy.evict( - victim_entry.as_ref(), - self.metadata.lineage(&victim).as_deref(), - ); + let outcome = self + .eviction_policy + .evict(&victim_entry, self.metadata.lineage(&victim).as_deref()); match outcome { EvictionOutcome::Replace { entry: new_batch, bytes_to_write, } => { - // Remember what went to the store: if the rewrite is then - // dropped as stale, these bytes are unreachable and have to - // be reclaimed here. - let mut wrote = None; if let Some(bytes_to_write) = bytes_to_write { - let len = bytes_to_write.len(); self.write_batch_to_disk(victim, identity, &new_batch, bytes_to_write) .await?; - wrote = Some((identity, len)); } - match self.try_insert(victim, WriteIdentity::Rewrite(identity), new_batch) { - Ok(residue) => { - self.settle(victim, residue, wrote).await; + match self.try_insert(victim, identity, new_batch) { + Ok(()) => { break; } Err(batch) => { - victim_entry = Arc::new(batch); + victim_entry = batch; } } } EvictionOutcome::Remove => { - self.remove_disk_entry(victim, identity).await; + self.remove_disk_entry(victim).await; break; } } @@ -787,9 +616,7 @@ impl LiquidCache { cached: cached_type, new: new_type, }); - let _ = self - .insert_inner(*entry_id, WriteIdentity::Rewrite(identity), new_entry) - .await; + let _ = self.insert_inner(*entry_id, identity, new_entry).await; } } @@ -801,14 +628,19 @@ impl LiquidCache { expression: Option<&CacheExpression>, ) -> Option { self.observer.on_get(selection.is_some()); - let batch = self.index.get_checked(entry_id, identity)?; + let slot = self.index.get_for(entry_id, identity)?; self.cache_policy - .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); - self.read_entry_inner(entry_id, identity, batch.as_ref(), selection, expression) + .notify_access(entry_id, CachedBatchType::from(&slot.entry)); + self.read_entry_inner(entry_id, identity, &slot.entry, selection, expression) .await } - /// Read an already-looked-up cache entry. + /// Read an already-looked-up cache entry, cached for `identity`. + /// + /// `entry` is a value the caller already holds, so it is not checked + /// against `identity`: the caller must have obtained it for that identity, + /// e.g. from [`Self::prefetch`]. `identity` is checked when the entry's + /// data has to be read from disk. pub async fn read_entry( &self, entry_id: &EntryID, @@ -914,6 +746,13 @@ impl LiquidCache { } } + /// Write `bytes` to disk for `identity`, followed by the identity itself. + /// + /// Objects are keyed by entry id alone, so a write landing after the key + /// changed hands can leave one identity's bytes where another's entry + /// points. The trailer is what lets a read notice that and miss. It sits + /// after the payload so the payload keeps the alignment it is read with. + /// The disk budget is charged for the trailer as well as the payload. async fn write_batch_to_disk( &self, entry_id: EntryID, @@ -923,7 +762,11 @@ impl LiquidCache { ) -> Result<(), CacheFull> { let len = bytes.len(); loop { - if self.budget.try_reserve_disk(len).is_ok() { + if self + .budget + .try_reserve_disk(len + IDENTITY_TRAILER_LEN) + .is_ok() + { break; } let victims = self.cache_policy.find_disk_victim(8); @@ -931,11 +774,7 @@ impl LiquidCache { return Err(CacheFull); } for victim in victims { - // Each victim's object is addressed by the identity that wrote - // it, so look that up rather than assuming this writer's. - if let Some((victim_identity, _)) = self.index.get_with_identity(&victim) { - self.remove_disk_entry(victim, victim_identity).await; - } + self.remove_disk_entry(victim).await; } } self.trace(InternalEvent::IoWrite { @@ -943,19 +782,34 @@ impl LiquidCache { kind: CachedBatchType::from(batch), bytes: len, }); + let mut object = Vec::with_capacity(len + IDENTITY_TRAILER_LEN); + object.extend_from_slice(&bytes); + object.extend_from_slice(&identity.to_le_bytes()); self.store - .put(entry_id_to_key(&entry_id, identity), bytes.to_vec()) + .put(entry_id_to_key(&entry_id), object) .await .expect("write failed"); Ok(()) } - async fn read_disk_arrow_array(&self, entry_id: &EntryID, identity: u64) -> Option { - let bytes = match self.store.get(&entry_id_to_key(entry_id, identity)).await { - Ok(bytes) => bytes, + /// The payload stored under `entry_id`, if it was written for `identity`. + async fn read_disk_bytes(&self, entry_id: &EntryID, identity: u64) -> Option { + let mut object = match self.store.get(&entry_id_to_key(entry_id)).await { + Ok(object) => object, Err(t4::Error::NotFound) => return None, Err(error) => panic!("read failed: {error}"), }; + let payload_len = object.len().checked_sub(IDENTITY_TRAILER_LEN)?; + let trailer: [u8; IDENTITY_TRAILER_LEN] = object[payload_len..].try_into().unwrap(); + if u64::from_le_bytes(trailer) != identity { + return None; + } + object.truncate(payload_len); + Some(Bytes::from(object)) + } + + async fn read_disk_arrow_array(&self, entry_id: &EntryID, identity: u64) -> Option { + let bytes = self.read_disk_bytes(entry_id, identity).await?; let bytes_len = bytes.len(); let cursor = std::io::Cursor::new(bytes); let mut reader = @@ -974,17 +828,13 @@ impl LiquidCache { entry_id: &EntryID, identity: u64, ) -> Option { - let bytes = match self.store.get(&entry_id_to_key(entry_id, identity)).await { - Ok(bytes) => bytes, - Err(t4::Error::NotFound) => return None, - Err(error) => panic!("read failed: {error}"), - }; + let bytes = self.read_disk_bytes(entry_id, identity).await?; self.trace(InternalEvent::IoReadLiquid { entry: *entry_id, bytes: bytes.len(), }); Some(Arc::new(crate::liquid_array::LiquidArray::from_bytes( - Bytes::from(bytes), + bytes, ))) } @@ -996,20 +846,23 @@ impl LiquidCache { predicate: &LiquidExpr, ) -> Option { self.observer.on_eval_predicate(); - let batch = self.index.get_checked(entry_id, identity)?; + let slot = self.index.get_for(entry_id, identity)?; self.cache_policy - .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); + .notify_access(entry_id, CachedBatchType::from(&slot.entry)); self.eval_predicate_on_entry_inner( entry_id, identity, - batch.as_ref(), + &slot.entry, selection_opt, predicate, ) .await } - /// Evaluate a predicate on an already-looked-up cache entry. + /// Evaluate a predicate on an already-looked-up cache entry, cached for `identity`. + /// + /// As with [`Self::read_entry`], `entry` itself is not checked against + /// `identity`; only data read from disk is. pub async fn eval_predicate_on_entry( &self, entry_id: &EntryID, @@ -1153,547 +1006,6 @@ mod tests { } } - /// Hydrating a disk entry must not charge its bytes to the disk budget - /// twice. - /// - /// The read that materializes a `DiskArrow`/`DiskLiquid` entry replaces it - /// with a memory entry. The store object under `(entry id, identity)` and - /// its share of `used_disk_bytes` outlive that replacement, so the next - /// spill reserves the same byte count again for an object the put simply - /// overwrites. Over a fixed working set, repeated read/spill rounds make - /// `disk_usage_bytes` climb without a byte more being written. - #[tokio::test] - async fn hydrating_a_disk_entry_does_not_recharge_its_disk_bytes() { - let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; - let entry_id = EntryID::from(500usize); - let array = create_test_arrow_array(1024); - - store.insert(entry_id, 0, array.clone()).await.unwrap(); - store.flush_all_to_disk().await.unwrap(); - let charged = store.budget.disk_usage_bytes(); - assert!(charged > 0, "flush must have written bytes"); - - let mut usage = vec![charged]; - for _ in 0..3 { - // Reading a disk entry hydrates it back into memory ... - let read = store.get(&entry_id, 0).await.expect("present"); - assert_eq!(read.as_ref(), array.as_ref()); - assert!(matches!( - store.index().get(&entry_id).unwrap().as_ref(), - CacheEntry::MemoryArrow(_) - )); - // ... and the next flush spills the very same bytes again. - store.flush_all_to_disk().await.unwrap(); - usage.push(store.budget.disk_usage_bytes()); - } - - assert_eq!( - usage, - vec![charged; 4], - "one entry of a fixed size occupies the same disk across read/spill rounds" - ); - } - - /// A spill that overwrites an entry's own disk object in place must release - /// the copy it superseded. - /// - /// Reported by review. The store key is `(entry id, identity)`, so this - /// write's put landed on the very object the old entry named — deleting it - /// would destroy the bytes just written. But the old entry's reservation is - /// still counted, so one object ends up charged twice. - /// - /// The two steps are written out rather than reached through `insert`. - /// Which write displaces the disk entry is a consequence of budget - /// arithmetic — an Arrow batch is transcoded and kept in memory long - /// before it is spilled — so an `insert` that lands here today lands - /// somewhere else after any change to a policy or a size. Driving the two - /// steps a spill performs, bytes to the store and then the record, pins - /// the displacement this test is named for. - #[tokio::test] - async fn an_in_place_disk_overwrite_releases_the_copy_it_supersedes() { - let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; - let entry_id = EntryID::from(1usize); - - // Get the key onto disk first, so the next write displaces a disk entry. - store - .insert(entry_id, 7, create_test_arrow_array(512)) - .await - .unwrap(); - store.flush_all_to_disk().await.unwrap(); - let (named_once, charged_once) = charged_disk_bytes_match_the_index(&store); - assert!( - named_once > 0, - "the entry must be on disk for this to test anything" - ); - assert_eq!(named_once, charged_once, "baseline must be consistent"); - - // Same key, same identity, written to disk again over its own object. - let replacement = create_test_arrow_array(4096); - let bytes = arrow_to_bytes(&replacement).unwrap(); - let disk_bytes = bytes.len(); - let on_disk = CacheEntry::disk_arrow(replacement.data_type().clone(), disk_bytes); - store - .write_batch_to_disk(entry_id, 7, &on_disk, bytes) - .await - .unwrap(); - let residue = store - .try_insert(entry_id, WriteIdentity::Rewrite(7), on_disk) - .expect("the record swap must land: the key still holds identity 7"); - assert!( - residue.superseded.is_some(), - "this is the displacement the test exists to cover" - ); - store.settle(entry_id, residue, Some((7, disk_bytes))).await; - - let (named, charged) = charged_disk_bytes_match_the_index(&store); - assert_eq!( - charged, named, - "the superseded copy must be released: one object, one reservation" - ); - assert_eq!( - store - .get(&entry_id, 7) - .await - .expect("the overwrite must still be readable") - .as_ref(), - replacement.as_ref(), - "settling must keep the object this write put there" - ); - } - - /// `DiskResidue::displacing` decides whether a displaced entry's store - /// object survives the write that displaced it, and only one combination - /// means it did not: a disk-resident write under the identity that already - /// held the key addresses the very object the old record named, so the put - /// landed on it. Every other combination leaves an object behind that - /// nothing can reach. Getting it wrong either deletes bytes just written or - /// charges one object twice. - #[test] - fn only_a_same_identity_disk_write_lands_on_the_object_it_displaces() { - let on_disk = ( - 7u64, - Arc::new(CacheEntry::disk_arrow( - arrow::datatypes::DataType::Int64, - 512, - )), - ); - let in_memory = (7u64, Arc::new(create_test_array(8))); - let another = ( - 9u64, - Arc::new(CacheEntry::disk_arrow( - arrow::datatypes::DataType::Int64, - 512, - )), - ); - - for written in [CachedBatchType::DiskArrow, CachedBatchType::DiskLiquid] { - // The key records the identity, not the form, so a liquid write - // over an Arrow copy lands on the same object as an Arrow one. - let residue = DiskResidue::displacing(Some(&on_disk), 7, written); - assert_eq!( - residue.superseded, - Some(512), - "a {written:?} write under identity 7 overwrote 7's own object" - ); - assert_eq!(residue.displaced, None); - assert!(!residue.dropped); - - // A different identity is a different store key. - let residue = DiskResidue::displacing(Some(&another), 7, written); - assert_eq!(residue.displaced, Some((9, 512))); - assert_eq!(residue.superseded, None); - } - - // A memory write puts nothing in the store, so it cannot have - // overwritten anything. - let residue = DiskResidue::displacing(Some(&on_disk), 7, CachedBatchType::MemoryArrow); - assert_eq!(residue.displaced, Some((7, 512))); - assert_eq!(residue.superseded, None); - - // A displaced memory entry never had an object. - let residue = DiskResidue::displacing(Some(&in_memory), 7, CachedBatchType::DiskArrow); - assert_eq!(residue.displaced, None); - assert_eq!(residue.superseded, None); - - // Nothing displaced at all. - let residue = DiskResidue::displacing(None, 7, CachedBatchType::DiskArrow); - assert_eq!(residue.displaced, None); - assert_eq!(residue.superseded, None); - } - - /// Every byte counted against the disk budget must belong to an index - /// entry that is on disk. Anything else can never be released: - /// `release_disk` is only ever reached from an index entry. - fn charged_disk_bytes_match_the_index(cache: &LiquidCache) -> (usize, usize) { - let mut named = 0usize; - cache.for_each_entry(|_, _, entry| match entry { - CacheEntry::DiskLiquid { disk_bytes, .. } - | CacheEntry::DiskArrow { disk_bytes, .. } => named += *disk_bytes, - CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) => {} - }); - (named, cache.budget.disk_usage_bytes()) - } - - /// Overwriting a disk-resident entry under the identity that already holds - /// the key must not leave the superseded copy charged. - /// - /// The store key is `(entry id, identity)`, so the object the old entry - /// named is still there — but this insert wrote nothing to the store, so - /// it did not overwrite it, and the index no longer names it. - #[tokio::test] - async fn overwriting_a_disk_entry_releases_the_superseded_copy() { - let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; - let entry_id = EntryID::from(501usize); - - store - .insert(entry_id, 5, create_test_arrow_array(1024)) - .await - .unwrap(); - store.flush_all_to_disk().await.unwrap(); - assert!( - store.budget.disk_usage_bytes() > 0, - "flush must have written" - ); - - // Same identity, new value: the index entry becomes a memory one. - store - .insert(entry_id, 5, create_test_arrow_array(2048)) - .await - .unwrap(); - - let (named, charged) = charged_disk_bytes_match_the_index(&store); - assert_eq!( - charged, named, - "the superseded copy is still charged but no index entry names it" - ); - } - - /// An overwrite must never be readable as the value it replaced, and the - /// form recorded in the index must be the form the store actually holds. - #[tokio::test] - async fn an_overwritten_entry_never_reads_back_the_superseded_bytes() { - let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; - let entry_id = EntryID::from(502usize); - let first = create_test_arrow_array(1024); - let second: ArrayRef = Arc::new(arrow::array::Int64Array::from_iter_values( - (0..2048).map(|v| v + 1_000_000), - )); - - store.insert(entry_id, 5, first.clone()).await.unwrap(); - store.flush_all_to_disk().await.unwrap(); - assert!(matches!( - store.index().get(&entry_id).unwrap().as_ref(), - CacheEntry::DiskArrow { .. } - )); - - store.insert(entry_id, 5, second.clone()).await.unwrap(); - assert_eq!( - store.get(&entry_id, 5).await.expect("present").as_ref(), - second.as_ref(), - "the overwrite must be what a read returns" - ); - - // Spill again: whatever form the index records has to match the bytes - // the store now holds, or the next read decodes one as the other. - store.flush_all_to_disk().await.unwrap(); - assert!( - matches!( - store.index().get(&entry_id).unwrap().as_ref(), - CacheEntry::DiskArrow { .. } - ), - "an Arrow flush must be recorded as an Arrow copy" - ); - assert_eq!( - store.get(&entry_id, 5).await.expect("present").as_ref(), - second.as_ref(), - "reading the spilled copy must not return the superseded value" - ); - } - - /// A hydrated entry that is then transcoded and spilled must be recorded - /// as the form it was written in, not the form it was hydrated from. - #[tokio::test] - async fn a_hydrated_then_transcoded_entry_is_recorded_as_what_was_written() { - let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; - let entry_id = EntryID::from(503usize); - let array = create_test_arrow_array(1024); - - store.insert(entry_id, 5, array.clone()).await.unwrap(); - store.flush_all_to_disk().await.unwrap(); - - // Hydrate the Arrow copy back into memory ... - store.get(&entry_id, 5).await.expect("present"); - // ... transcode it to liquid, then spill it as liquid. - store - .evict_victim_inner(entry_id) - .await - .expect("transcode must fit"); - assert!(matches!( - store.index().get(&entry_id).unwrap().as_ref(), - CacheEntry::MemoryLiquid(_) - )); - store.flush_all_to_disk().await.unwrap(); - assert!( - matches!( - store.index().get(&entry_id).unwrap().as_ref(), - CacheEntry::DiskLiquid { .. } - ), - "a liquid flush must be recorded as a liquid copy" - ); - - assert_eq!( - store.get(&entry_id, 5).await.expect("present").as_ref(), - array.as_ref(), - "a liquid stub must not be decoded over Arrow IPC bytes" - ); - let (named, charged) = charged_disk_bytes_match_the_index(&store); - assert_eq!( - charged, named, - "the Arrow copy it was hydrated from is still charged" - ); - } - - /// A flush that cannot place an entry on disk drops it. Nothing of that - /// entry may stay charged against the disk budget afterwards. - #[tokio::test] - async fn a_flush_that_drops_an_entry_leaves_no_disk_charged_to_it() { - let array = create_test_arrow_array(1024); - let one_copy = arrow_to_bytes(&array).unwrap().len(); - let cache = LiquidCacheBuilder::new() - .with_max_memory_bytes(1 << 20) - .with_max_disk_bytes(one_copy) - .with_eviction_policy(Box::new(TranscodeEvict)) - .with_hydration_policy(Box::new(crate::cache::AlwaysHydrate::new())) - .with_cache_policy(Box::new(LiquidPolicy::new())) - .build() - .await; - let first = EntryID::from(504usize); - let second = EntryID::from(505usize); - - cache.insert(first, 0, array.clone()).await.unwrap(); - cache.flush_all_to_disk().await.unwrap(); - // Reading it brings it back into memory; the disk tier should now be - // empty, so the next flush has room for both entries. - cache.get(&first, 0).await.expect("present"); - - cache.insert(second, 0, array.clone()).await.unwrap(); - cache.flush_all_to_disk().await.unwrap(); - - let (named, charged) = charged_disk_bytes_match_the_index(&cache); - assert_eq!( - charged, named, - "disk charged to entries the flush dropped can never be released" - ); - } - - /// A takeover landing *during* a rewrite's `store.put` must leave the new - /// owner whole: its bytes, its record, and its share of the budget. - /// - /// The steps below are what `evict_victim_inner` does — read the entry with - /// the identity it holds, write it to the store, then swap the record — with - /// the takeover injected between the read and the write, which is the one - /// interleaving that ordering cannot be produced by calling it once. - /// - /// Three things hold it together, and the first is the decisive one: - /// `entry_id_to_key` puts the writer's identity in the store key, so the - /// stale put addresses its own object and can never reach the new owner's; - /// `WriteIdentity::Rewrite` makes the index refuse the record swap; and - /// `settle` reclaims the object the refused write had already put there. - #[tokio::test] - async fn a_takeover_during_a_rewrites_disk_write_leaves_the_new_owner_whole() { - let cache = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; - let entry_id = EntryID::from(600usize); - let theirs = create_test_arrow_array(1024); - let ours: ArrayRef = Arc::new(arrow::array::Int64Array::from_iter_values( - (0..512).map(|v| v + 7_000_000), - )); - - // Identity 7 caches an entry, and eviction picks it up. - cache.insert(entry_id, 7, theirs.clone()).await.unwrap(); - let (observed, _read) = cache.index().get_with_identity(&entry_id).unwrap(); - assert_eq!(observed, 7); - let stale_bytes = arrow_to_bytes(&theirs).unwrap(); - let stale_len = stale_bytes.len(); - let stale_rewrite = CacheEntry::disk_arrow(theirs.data_type().clone(), stale_len); - - // Identity 9 takes the key over and puts its own copy on disk. - cache.insert(entry_id, 9, ours.clone()).await.unwrap(); - cache.flush_all_to_disk().await.unwrap(); - let owner_bytes = match cache.index().get(&entry_id).unwrap().as_ref() { - CacheEntry::DiskArrow { disk_bytes, .. } => *disk_bytes, - other => panic!("expected the new owner on disk, found {other}"), - }; - - // Only now does the stale rewrite's write complete ... - cache - .write_batch_to_disk(entry_id, observed, &stale_rewrite, stale_bytes) - .await - .unwrap(); - // ... and reach the record swap. - let residue = cache - .try_insert(entry_id, WriteIdentity::Rewrite(observed), stale_rewrite) - .expect("a refused rewrite is not a failure to insert"); - cache - .settle(entry_id, residue, Some((observed, stale_len))) - .await; - - // The stale writer's object is gone; the new owner's is not. - assert!( - matches!( - cache.store.get(&entry_id_to_key(&entry_id, observed)).await, - Err(t4::Error::NotFound) - ), - "the refused rewrite must take its own write back" - ); - assert!( - cache - .store - .get(&entry_id_to_key(&entry_id, 9)) - .await - .is_ok(), - "the new owner's object must survive a stale writer" - ); - let (named, charged) = charged_disk_bytes_match_the_index(&cache); - assert_eq!(charged, named); - assert_eq!(charged, owner_bytes, "only the new owner's copy is charged"); - - // And the new owner's record still names bytes that decode to its rows. - assert!(matches!( - cache.index().get(&entry_id).unwrap().as_ref(), - CacheEntry::DiskArrow { .. } - )); - assert_eq!( - cache.get(&entry_id, 9).await.expect("present").as_ref(), - ours.as_ref(), - "the new owner must read its own rows, never the stale writer's" - ); - assert!( - cache.get(&entry_id, 7).await.is_none(), - "the displaced identity reads a miss" - ); - } - - /// A settlement must not delete an object a *later* write put at the same - /// store key. - /// - /// A store key is `(entry id, identity)` and an identity comes back: the - /// file-id pool hands a re-opened path its previous record, so the very key - /// a settlement decided to delete can be occupied again by a fresh write - /// before the delete runs. `settle` deletes behind an await, which is all - /// the room that needs. The steps below are what one insert does — take the - /// key over, then settle what that displaced — with the re-open and its - /// spill injected in between. - #[tokio::test] - async fn a_settlement_does_not_delete_an_object_written_after_it_was_decided() { - let cache = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; - let entry_id = EntryID::from(700usize); - let first = create_test_arrow_array(1024); - let takeover: ArrayRef = Arc::new(arrow::array::Int64Array::from_iter_values( - (0..512).map(|v| v + 3_000_000), - )); - let reborn: ArrayRef = Arc::new(arrow::array::Int64Array::from_iter_values( - (0..256).map(|v| v + 5_000_000), - )); - - // Identity 7 caches the key and spills it: an object at (E, 7). - cache.insert(entry_id, 7, first).await.unwrap(); - cache.flush_all_to_disk().await.unwrap(); - - // Identity 9 takes the key over. That displaces 7's disk entry, so this - // settlement is going to delete (E, 7) ... - let residue = cache - .try_insert( - entry_id, - WriteIdentity::Owned(9), - CacheEntry::memory_arrow(takeover), - ) - .expect("the takeover fits in memory"); - - // ... but before it runs, 7's path is re-opened, the pool hands the same - // identity back, and it caches this key again and spills it. That put - // lands on the very object the pending delete names. - cache.insert(entry_id, 7, reborn.clone()).await.unwrap(); - cache.flush_all_to_disk().await.unwrap(); - - // Only now does the settlement run. - cache.settle(entry_id, residue, None).await; - - assert_eq!( - cache - .get(&entry_id, 7) - .await - .expect("the re-cached entry must survive a settlement decided before it") - .as_ref(), - reborn.as_ref(), - "the settlement must not reach a write that came after it" - ); - let (named, charged) = charged_disk_bytes_match_the_index(&cache); - assert_eq!(charged, named, "one object, one reservation"); - } - - /// A rewrite that loses its key must not leave its disk write behind. - /// - /// The bytes were already in the store when the index refused the write, and - /// the store key carries the identity that wrote them, so nothing reachable - /// through the index names them afterwards: neither the object nor its share - /// of `used_disk_bytes` would ever come back. - #[tokio::test] - async fn a_dropped_rewrite_reclaims_the_disk_it_already_wrote() { - let store = create_cache_store(10 * 1024, Box::new(LiquidPolicy::new())).await; - let entry_id = EntryID::from(1usize); - - // The owner caches an entry and flushes it, so a disk object exists. - store - .insert(entry_id, 7, create_test_arrow_array(64)) - .await - .unwrap(); - store.flush_all_to_disk().await.unwrap(); - let after_flush = store.budget.disk_usage_bytes(); - assert!(after_flush > 0, "flush must have written bytes"); - - // Another identity takes the key over, so the earlier owner's rewrite is - // now stale. Its disk object is unreachable and must be reclaimed. - store - .insert(entry_id, 9, create_test_arrow_array(64)) - .await - .unwrap(); - assert_eq!( - store.budget.disk_usage_bytes(), - 0, - "the displaced owner's disk bytes must be released, not stranded" - ); - assert!( - store.get(&entry_id, 9).await.is_some(), - "the new owner must still read its own entry" - ); - assert!( - store.get(&entry_id, 7).await.is_none(), - "the displaced owner must not read the new owner's rows" - ); - } - - /// The removal path is identity-checked: a caller that read an entry earlier - /// must not destroy the record of whoever holds the key now. - #[tokio::test] - async fn removing_a_disk_entry_under_a_stale_identity_is_refused() { - let store = create_cache_store(10 * 1024, Box::new(LiquidPolicy::new())).await; - let entry_id = EntryID::from(2usize); - - store - .insert(entry_id, 3, create_test_arrow_array(64)) - .await - .unwrap(); - store.flush_all_to_disk().await.unwrap(); - - // A stale identity tries to evict it. Nothing of the current owner's may - // be touched. - store.remove_disk_entry(entry_id, 999).await; - assert!( - store.get(&entry_id, 3).await.is_some(), - "a stale remove must leave the current owner's entry readable" - ); - } - #[tokio::test] async fn test_basic_cache_operations() { // Test basic insert, get, and size tracking in one test @@ -1707,15 +1019,12 @@ 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, WriteIdentity::Owned(0), array1) - .await - .unwrap(); + store.insert_inner(entry_id1, 0, array1).await.unwrap(); // Verify budget usage and data correctness assert_eq!(store.budget.memory_usage_bytes(), size1); let retrieved1 = store.index().get(&entry_id1).unwrap(); - match retrieved1.as_ref() { + match &retrieved1.entry { CacheEntry::MemoryArrow(arr) => assert_eq!(arr.len(), 100), _ => panic!("Expected ArrowMemory"), } @@ -1723,19 +1032,13 @@ 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, WriteIdentity::Owned(0), array2) - .await - .unwrap(); + store.insert_inner(entry_id2, 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, WriteIdentity::Owned(0), array3) - .await - .unwrap(); + store.insert_inner(entry_id1, 0, array3).await.unwrap(); assert_eq!(store.budget.memory_usage_bytes(), size3 + size2); assert!(store.index().get(&EntryID::from(999)).is_none()); @@ -1755,19 +1058,19 @@ mod tests { let store = create_cache_store(8000, Box::new(advisor)).await; // Small budget to force advice store - .insert_inner(entry_id1, WriteIdentity::Owned(0), create_test_array(800)) + .insert_inner(entry_id1, 0, create_test_array(800)) .await .unwrap(); - match store.index().get(&entry_id1).unwrap().as_ref() { + match &store.index().get(&entry_id1).unwrap().entry { CacheEntry::MemoryArrow(_) => {} other => panic!("Expected ArrowMemory, got {other:?}"), } store - .insert_inner(entry_id2, WriteIdentity::Owned(0), create_test_array(800)) + .insert_inner(entry_id2, 0, create_test_array(800)) .await .unwrap(); - match store.index().get(&entry_id1).unwrap().as_ref() { + match &store.index().get(&entry_id1).unwrap().entry { CacheEntry::MemoryLiquid(_) => {} other => panic!("Expected LiquidMemory after eviction, got {other:?}"), } @@ -1883,17 +1186,64 @@ mod tests { store.flush_all_to_disk().await.unwrap(); { let entry = store.index().get(&entry_id).unwrap(); - assert!(matches!(entry.as_ref(), CacheEntry::DiskArrow { .. })); + assert!(matches!(&entry.entry, CacheEntry::DiskArrow { .. })); } 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(); - assert!(matches!(entry.as_ref(), CacheEntry::MemoryArrow(_))); + assert!(matches!(&entry.entry, CacheEntry::MemoryArrow(_))); } } + #[tokio::test] + async fn an_entry_is_served_only_to_its_identity() { + let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(321usize); + let array = create_test_arrow_array(8); + + store.insert(entry_id, 1, array.clone()).await.unwrap(); + assert!(store.contains(&entry_id, 1)); + assert!(!store.contains(&entry_id, 2)); + assert!(store.get(&entry_id, 2).await.is_none()); + assert!(matches!( + store.prefetch(&entry_id, 2).await, + PrefetchResult::Absent + )); + + store.flush_all_to_disk().await.unwrap(); + assert!(store.get(&entry_id, 2).await.is_none()); + let result = store.get(&entry_id, 1).await.expect("present"); + assert_eq!(result.as_ref(), array.as_ref()); + + let replacement = create_test_arrow_array(16); + store + .insert(entry_id, 2, replacement.clone()) + .await + .unwrap(); + assert!(!store.contains(&entry_id, 1)); + let result = store.get(&entry_id, 2).await.expect("present"); + assert_eq!(result.as_ref(), replacement.as_ref()); + } + + #[tokio::test] + async fn disk_bytes_written_for_another_identity_are_a_miss() { + let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(321usize); + let array = create_test_arrow_array(8); + + store.insert(entry_id, 1, array.clone()).await.unwrap(); + store.flush_all_to_disk().await.unwrap(); + let disk_entry = store.index().get(&entry_id).unwrap().entry.clone(); + assert!(matches!(disk_entry, CacheEntry::DiskArrow { .. })); + + // A record naming identity 2 over bytes identity 1 wrote: what a + // rewrite landing after the key changed hands leaves behind. + store.index().insert(&entry_id, 2, disk_entry); + assert!(store.get(&entry_id, 2).await.is_none()); + } + #[tokio::test] async fn missing_disk_blob_is_a_cache_miss() { let directory = tempfile::tempdir().unwrap(); @@ -1909,7 +1259,7 @@ mod tests { .await .unwrap(); cache.flush_all_to_disk().await.unwrap(); - store.remove(&entry_id_to_key(&id, 0)).await.unwrap(); + store.remove(&entry_id_to_key(&id)).await.unwrap(); assert!(cache.get(&id, 0).await.is_none()); } @@ -1923,24 +1273,20 @@ mod tests { Arc::new(crate::liquid_array::LiquidArray::from_arrow_array(&arrow_array).unwrap()); store - .insert_inner( - entry_id, - WriteIdentity::Owned(0), - CacheEntry::memory_liquid(liquid.clone()), - ) + .insert_inner(entry_id, 0, CacheEntry::memory_liquid(liquid.clone())) .await .unwrap(); store.flush_all_to_disk().await.unwrap(); { let entry = store.index().get(&entry_id).unwrap(); - assert!(matches!(entry.as_ref(), CacheEntry::DiskLiquid { .. })); + assert!(matches!(&entry.entry, CacheEntry::DiskLiquid { .. })); } 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(); - assert!(matches!(entry.as_ref(), CacheEntry::MemoryLiquid(_))); + assert!(matches!(&entry.entry, CacheEntry::MemoryLiquid(_))); } } @@ -1957,15 +1303,15 @@ mod tests { let err = cache.insert(EntryID::from(900usize), 0, array).await; assert_eq!(err, Err(CacheFull)); - assert!(!cache.is_cached(&EntryID::from(900usize), 0)); + assert!(!cache.contains(&EntryID::from(900usize), 0)); } #[tokio::test] async fn insert_until_disk_full_then_evicts_oldest_disk_entry() { let first_array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); let second_array: ArrayRef = Arc::new(Int32Array::from_iter_values(16..32)); - let first_bytes = arrow_to_bytes(&first_array).unwrap().len(); - let second_bytes = arrow_to_bytes(&second_array).unwrap().len(); + let first_bytes = arrow_to_bytes(&first_array).unwrap().len() + IDENTITY_TRAILER_LEN; + let second_bytes = arrow_to_bytes(&second_array).unwrap().len() + IDENTITY_TRAILER_LEN; let cache = LiquidCacheBuilder::new() .with_max_memory_bytes(1 << 20) .with_max_disk_bytes(first_bytes.max(second_bytes)) @@ -1978,14 +1324,14 @@ mod tests { let second = EntryID::from(911usize); cache.insert(first, 0, first_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(cache.is_cached(&first, 0)); + assert!(cache.contains(&first, 0)); cache.insert(second, 0, second_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.is_cached(&first, 0)); + assert!(!cache.contains(&first, 0)); assert!(matches!( - cache.index().get(&second).unwrap().as_ref(), + &cache.index().get(&second).unwrap().entry, CacheEntry::DiskArrow { .. } )); } @@ -2010,13 +1356,13 @@ mod tests { cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.is_cached(&first, 0) || !cache.is_cached(&second, 0)); + assert!(!cache.contains(&first, 0) || !cache.contains(&second, 0)); } #[tokio::test] async fn disk_eviction_releases_budget() { let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); - let disk_bytes = arrow_to_bytes(&array).unwrap().len(); + let disk_bytes = arrow_to_bytes(&array).unwrap().len() + IDENTITY_TRAILER_LEN; let cache = LiquidCacheBuilder::new() .with_max_memory_bytes(1 << 20) .with_max_disk_bytes(disk_bytes) @@ -2029,10 +1375,10 @@ mod tests { cache.flush_all_to_disk().await.unwrap(); let before = cache.stats().disk_usage_bytes; - cache.remove_disk_entry(entry, 0).await; + cache.remove_disk_entry(entry).await; assert_eq!(cache.stats().disk_usage_bytes, before - disk_bytes); - assert!(!cache.is_cached(&entry, 0)); + assert!(!cache.contains(&entry, 0)); } #[tokio::test] @@ -2050,6 +1396,6 @@ mod tests { let result = cache.flush_all_to_disk().await; assert_eq!(result, Ok(())); - assert!(!cache.is_cached(&entry_id, 0)); + assert!(!cache.contains(&entry_id, 0)); } } diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index 78e3bd63..2e5bd5f3 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -1,106 +1,70 @@ use congee::CongeeArc; use std::{ fmt::{Debug, Formatter}, - sync::atomic::{AtomicU64, AtomicUsize, Ordering}, + sync::{ + Arc, + atomic::{AtomicU64, AtomicUsize, Ordering}, + }, }; use crate::cache::{cached_batch::CacheEntry, utils::EntryID}; -use crate::sync::{Arc, RwLock}; +use crate::sync::RwLock; -/// The value stored in the ART. +/// What the index holds for a key: the entry, and the identity of the source +/// it was cached for. /// -/// `CongeeArc` frees a removed or replaced value through crossbeam-epoch's -/// deferred destruction: it clones the `Arc` and drops the clone only when a -/// later pin collects that epoch's garbage — up to 64 objects per thread wait -/// in a thread-local bag, and the global queue drains 8 bags per 128 pins. -/// With multi-megabyte arrays as values, that kept every evicted entry alive -/// for an unbounded, budget-invisible stretch (liquid-cache#43: a tier -/// reporting at its limit while the process held several times that). -/// -/// 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. -/// -/// 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>>, +/// A key does not name its source uniquely: callers pack fields into fixed +/// widths, so two sources can compute the same key. The identity is what tells +/// them apart, so a read under any other identity must treat the slot as +/// absent. +#[derive(Debug)] +pub(crate) struct Slot { + pub(crate) identity: u64, + pub(crate) entry: CacheEntry, } impl Slot { - fn new(identity: u64, entry: CacheEntry) -> Arc { - Arc::new(Self { - identity, - entry: RwLock::new(Some(Arc::new(entry))), - }) - } - - fn load(&self) -> Option> { - self.entry.read().unwrap().clone() - } - - fn take(&self) -> Option> { - self.entry.write().unwrap().take() + pub(crate) fn belongs_to(&self, identity: u64) -> bool { + self.identity == identity } } -/// 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), +/// The value stored in the tree: a [`Slot`] behind a cell the index empties +/// the moment it gives the slot up. +/// +/// `CongeeArc` frees a removed or replaced value through crossbeam-epoch's +/// deferred destruction: it clones the `Arc` and drops the clone only when a +/// later pin collects that epoch's garbage — up to 64 objects per thread wait +/// in a thread-local bag, and the global queue drains 8 bags per 128 pins. +/// With multi-megabyte arrays as values, that kept every evicted entry alive +/// for an unbounded, budget-invisible stretch: a tier reporting at its limit +/// while the process held several times that. +/// +/// Taking the slot out of the cell when the index gives it up means the +/// deferred drop reclaims only an empty shell, and the array dies with the +/// last caller-held reference. +struct Cell { + slot: RwLock>>, } -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, +impl Cell { + fn new(slot: Slot) -> Self { + Self { + slot: RwLock::new(Some(Arc::new(slot))), } } -} -/// What an [`ArtIndex::insert`] did. -/// -/// `displaced` is the entry this write replaced, with the identity that was -/// recorded against it. It is handed back rather than dropped here because a -/// store object is addressed by entry id *and* identity -/// (`entry_id_to_key`): once another identity holds the key, nothing reachable -/// through the index names the old object any more, so the caller has to delete -/// it and give its disk bytes back or both leak for the life of the process. -pub(crate) struct InsertOutcome { - /// Whether the batch was stored. False only for a stale rewrite. - pub(crate) stored: bool, - /// The entry this write replaced, and the identity that held it. - pub(crate) displaced: Option<(u64, Arc)>, -} + fn load(&self) -> Option> { + self.slot.read().unwrap().clone() + } -impl InsertOutcome { - fn dropped() -> Self { - Self { - stored: false, - displaced: None, - } + fn take(&self) -> Option> { + self.slot.write().unwrap().take() } } pub(crate) struct ArtIndex { - art: CongeeArc, + art: CongeeArc, entry_count: AtomicUsize, identity_mismatches: AtomicU64, } @@ -120,188 +84,64 @@ impl ArtIndex { } } - /// 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> { + /// The slot under `entry_id`, whichever identity holds it. + 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 - // tree lookup and the load. A remove reading as a miss is exactly as - // if it had won the race outright, but after a replace the key is - // still present with a new slot, so look the key up once more rather - // than report a cached entry as absent. - let slot = self.art.get(*entry_id, &guard)?; - if let Some(entry) = slot.load() { - return Some(entry); + // An empty cell means the slot was removed or replaced between the + // tree lookup and the load. A remove reading as a miss is as if it had + // won the race outright, but after a replace the key is still present + // under a new cell, so look it up once more rather than report a + // cached entry as absent. + let cell = self.art.get(*entry_id, &guard)?; + if let Some(slot) = cell.load() { + return Some(slot); } self.art.get(*entry_id, &guard)?.load() } - /// 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 { + /// The slot under `entry_id`, only if `identity` holds it. A slot held by + /// another identity is counted in [`Self::identity_mismatches`]. + pub(crate) fn get_for(&self, entry_id: &EntryID, identity: u64) -> Option> { + let slot = self.get(entry_id)?; + if !slot.belongs_to(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() + Some(slot) } - /// 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)) + /// How many lookups found their key held by another identity. + pub(crate) fn identity_mismatches(&self) -> u64 { + self.identity_mismatches.load(Ordering::Relaxed) } - pub(crate) fn is_cached(&self, entry_id: &EntryID, identity: u64) -> bool { - self.get_checked(entry_id, identity).is_some() + pub(crate) fn contains(&self, entry_id: &EntryID, identity: u64) -> bool { + self.get_for(entry_id, identity).is_some() } - /// Store `batch` under `entry_id`, reporting what the write did. - /// - /// See [`WriteIdentity`] for the two kinds of write and why they differ, and - /// [`InsertOutcome`] for why a displaced entry has to be handed back rather - /// than dropped here. - pub(crate) fn insert( - &self, - entry_id: &EntryID, - identity: WriteIdentity, - batch: CacheEntry, - ) -> InsertOutcome { + /// Store `batch` under `entry_id` for `identity`, replacing whatever held + /// the key, including a slot held by another identity. + pub(crate) fn insert(&self, entry_id: &EntryID, identity: u64, batch: CacheEntry) { 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 InsertOutcome::dropped(), + let slot = Slot { + identity, + entry: batch, }; let existing = self .art - .insert(*entry_id, Slot::new(identity, batch), &guard) + .insert(*entry_id, Arc::new(Cell::new(slot)), &guard) .expect("Insertion failed"); - let displaced = match existing { + match existing { Some(replaced) => { - let was = replaced.identity; - replaced.take().map(|entry| (was, entry)) + replaced.take(); } None => { self.entry_count.fetch_add(1, Ordering::Relaxed); - None } - }; - InsertOutcome { - stored: true, - displaced, } } - /// Remove an entry only if `identity` is the one that holds it. - /// - /// The unchecked [`Self::remove`] is for maintenance acting on whatever - /// occupies a key. A caller that read an entry earlier and then removes it - /// must come through here: between its read and its removal the key can - /// change hands, and removing the new owner's record would strand that - /// owner's store object while releasing a byte count taken from the record - /// it just destroyed. - /// - /// Deciding and removing are one operation, not a check followed by a - /// remove: `compute_if_present` decides against the value the tree - /// publishes, so a key that has changed hands is left alone rather than - /// destroyed, and no record is ever put back over a writer that arrived - /// meanwhile. - /// - /// The closure runs under an optimistic read and the tree upgrades to a - /// write only afterwards, so a failed upgrade runs it again — against - /// whatever the key holds on that attempt. Only the last run is the one - /// the tree acts on, so the verdict is recomputed on every run rather - /// than latched on the first: a run that found the record ours followed - /// by a retry that finds it someone else's must report a mismatch, not a - /// removal that never happened. - pub(crate) fn remove_checked( - &self, - entry_id: &EntryID, - identity: u64, - ) -> Option> { - let guard = self.art.pin(); - let mut was_ours = false; - let previous = self.art.compute_if_present( - *entry_id, - |slot| { - was_ours = slot.identity == identity; - if was_ours { - None - } else { - // Unchanged, and the tree short-circuits on an unchanged - // value, so nothing is republished over the holder. - Some(slot) - } - }, - &guard, - )?; - if !was_ours { - self.identity_mismatches.fetch_add(1, Ordering::Relaxed); - return None; - } - self.entry_count.fetch_sub(1, Ordering::Relaxed); - previous.take() - } - - /// Does a live record still name the store object at `(entry_id, identity)`? - /// - /// A store key is `(entry id, identity)` (`entry_id_to_key`), so this is - /// what a caller about to delete such an object has to ask: an identity - /// comes back — the file-id pool hands a re-opened path its previous - /// record — and a write under it can occupy the key again between the - /// moment a deletion was decided and the moment it runs. - pub(crate) fn holds_disk_entry(&self, entry_id: &EntryID, identity: u64) -> bool { - let Some((held, entry)) = self.get_with_identity(entry_id) else { - return false; - }; - held == identity - && matches!( - entry.as_ref(), - CacheEntry::DiskLiquid { .. } | CacheEntry::DiskArrow { .. } - ) - } - - pub(crate) fn remove(&self, entry_id: &EntryID) -> Option> { + pub(crate) fn remove(&self, entry_id: &EntryID) -> Option> { let guard = self.art.pin(); let removed = self.art.remove(*entry_id, &guard)?; self.entry_count.fetch_sub(1, Ordering::Relaxed); @@ -315,10 +155,10 @@ impl ArtIndex { self.entry_count.store(0, Ordering::Relaxed); } - pub(crate) fn for_each(&self, mut f: impl FnMut(&EntryID, u64, &CacheEntry)) { + pub(crate) fn for_each(&self, mut f: impl FnMut(&EntryID, &Slot)) { for id in self.art.keys() { - if let Some((identity, entry)) = self.get_with_identity(&id) { - f(&id, identity, &entry); + if let Some(slot) = self.get(&id) { + f(&id, &slot); } } } @@ -331,48 +171,38 @@ 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)] mod tests { use crate::cache::cached_batch::CacheEntry; use crate::cache::utils::create_test_array; - use crate::sync::thread; use super::*; #[test] - fn test_get_and_is_cached() { + fn test_get_and_contains() { let store = ArtIndex::new(); let entry_id1: EntryID = EntryID::from(1); let entry_id2: EntryID = EntryID::from(2); let array1 = create_test_array(100); // Initially, entries should not be cached - assert!(!store.is_cached(&entry_id1, 0)); - assert!(!store.is_cached(&entry_id2, 0)); + assert!(!store.contains(&entry_id1, 0)); + assert!(!store.contains(&entry_id2, 0)); assert!(store.get(&entry_id1).is_none()); // Insert an entry and verify it's cached { - store.insert(&entry_id1, WriteIdentity::Owned(0), array1.clone()); + store.insert(&entry_id1, 0, array1.clone()); } - assert!(store.is_cached(&entry_id1, 0)); - assert!(!store.is_cached(&entry_id2, 0)); + assert!(store.contains(&entry_id1, 0)); + assert!(!store.contains(&entry_id2, 0)); // Get should return the cached value - match store.get(&entry_id1) { - Some(batch) => match batch.as_ref() { + match store.get_for(&entry_id1, 0) { + Some(slot) => match &slot.entry { CacheEntry::MemoryArrow(arr) => assert_eq!(arr.len(), 100), _ => panic!("Expected ArrowMemory batch"), }, @@ -386,14 +216,14 @@ mod tests { let entry_id: EntryID = EntryID::from(1); let array = create_test_array(100); - store.insert(&entry_id, WriteIdentity::Owned(0), array.clone()); + store.insert(&entry_id, 0, array.clone()); let entry_id: EntryID = EntryID::from(1); - assert!(store.is_cached(&entry_id, 0)); + assert!(store.contains(&entry_id, 0)); store.reset(); let entry_id: EntryID = EntryID::from(1); - assert!(!store.is_cached(&entry_id, 0)); + assert!(!store.contains(&entry_id, 0)); } /// The array behind a removed or replaced entry must die with the last @@ -408,8 +238,8 @@ mod tests { unreachable!() }; let weak_first = Arc::downgrade(first_array); - store.insert(&id, WriteIdentity::Owned(0), first); - store.insert(&id, WriteIdentity::Owned(0), create_test_array(200)); + store.insert(&id, 0, first); + store.insert(&id, 0, create_test_array(200)); assert!( weak_first.upgrade().is_none(), "replaced entry still alive: held by the index's deferred drop" @@ -417,7 +247,7 @@ mod tests { let second = store.get(&id).unwrap(); let removed = store.remove(&id).unwrap(); - let CacheEntry::MemoryArrow(second_array) = removed.as_ref() else { + let CacheEntry::MemoryArrow(second_array) = &removed.entry else { unreachable!() }; let weak_second = Arc::downgrade(second_array); @@ -429,225 +259,20 @@ 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() { + fn a_slot_is_visible_only_to_its_identity() { let store = ArtIndex::new(); - let key: EntryID = EntryID::from(7); - - assert!( - store - .insert(&key, WriteIdentity::Owned(1), create_test_array(100)) - .stored - ); + let entry_id = EntryID::from(1); + store.insert(&entry_id, 7, create_test_array(10)); - // 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!(store.contains(&entry_id, 7)); + assert!(!store.contains(&entry_id, 8)); + assert!(store.get_for(&entry_id, 8).is_none()); 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)) - .stored - ); - 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)) - .stored - ); - 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)) - .stored - ); - - // The rewrite lands too late and must be dropped, not relabelled. - assert!( - !store - .insert( - &key, - WriteIdentity::Rewrite(observed), - create_test_array(100) - ) - .stored - ); - - 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}"), - } - } - - /// A checked removal must not overwrite or strand a record that arrived - /// while it ran. - /// - /// Three writers to one key, which is the shape the reviewer named A-B-C: - /// A removes under the identity it read, B takes the key over, and C takes - /// it over after B. Whatever order they land in, the index has to stay - /// self-consistent — every key it holds is counted, every record it holds - /// is one a writer installed paired with that writer's array, and a - /// removal only ever hands back the record it was entitled to. - /// - /// Deciding and removing as one operation is what makes that hold. A - /// removal that tested the identity, removed, then put back what it found - /// could publish its restore over C and lose C's record: the key would - /// hold B while `entry_count` counted C too. - /// - /// `takeovers` names the identities that race the removal, so a caller can - /// widen the race without restating the setup. - fn checked_removal_against_takeovers(takeovers: &[u64]) { - let index = Arc::new(ArtIndex::new()); - let key = EntryID::from(21); - index.insert(&key, WriteIdentity::Owned(1), create_test_array(10)); - - let remover = { - let index = Arc::clone(&index); - thread::spawn(move || index.remove_checked(&key, 1)) - }; - let takeovers: Vec<_> = takeovers - .iter() - .copied() - .map(|who| { - let index = Arc::clone(&index); - thread::spawn(move || { - index.insert( - &key, - WriteIdentity::Owned(who), - create_test_array(who as usize * 10), - ); - }) - }) - .collect(); - let removed = remover.join().unwrap(); - for takeover in takeovers { - takeover.join().unwrap(); - } - - if let Some(removed) = removed { - let CacheEntry::MemoryArrow(array) = removed.as_ref() else { - panic!("only memory entries were written") - }; - assert_eq!( - array.len(), - 10, - "a checked removal handed back a record it did not name" - ); - } - if let Some((identity, entry)) = index.get_with_identity(&key) { - let CacheEntry::MemoryArrow(array) = entry.as_ref() else { - panic!("only memory entries were written") - }; - assert_eq!( - array.len(), - identity as usize * 10, - "the key holds identity {identity} against another writer's array" - ); - } - assert_eq!( - index.entry_count(), - index.keys().len(), - "every key the index holds must be counted: an uncounted one is a \ - record nothing will ever release" - ); - } - - #[test] - fn concurrent_checked_removal_against_takeovers() { - checked_removal_against_takeovers(&[2, 3]); - } - - /// The same race as above, run until the tree makes a checked removal - /// retry. - /// - /// `compute_if_present` decides under an optimistic read and upgrades to a - /// write afterwards, so a removal whose upgrade loses runs its closure - /// again — and the key can belong to someone else by then. A single round - /// almost never reaches that interleaving, so the round above can hold a - /// verdict latched on the first run and still pass. More writers on the - /// key and enough rounds is what finds it. - #[test] - fn a_checked_removal_that_retries_reports_the_verdict_it_acted_on() { - for _ in 0..20_000 { - checked_removal_against_takeovers(&[2, 3, 4, 5]); - } - } - - /// The same three writers, every interleaving the model checker can reach. - #[cfg(feature = "shuttle")] - #[test] - fn shuttle_checked_removal_against_takeovers() { - crate::utils::shuttle_test(|| checked_removal_against_takeovers(&[2, 3])); - } - - /// 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)) - .stored - ); - assert!( - store - .insert(&key, WriteIdentity::Rewrite(5), create_test_array(20)) - .stored - ); - 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)) - .stored - ); - assert!(store.get(&key).is_none()); - assert_eq!(store.entry_count(), 0); + store.insert(&entry_id, 8, create_test_array(20)); + assert!(!store.contains(&entry_id, 7)); + assert!(store.contains(&entry_id, 8)); + assert_eq!(store.entry_count(), 1); } } diff --git a/src/core/src/cache/io_context.rs b/src/core/src/cache/io_context.rs index 12152a0a..7a7f71d3 100644 --- a/src/core/src/cache/io_context.rs +++ b/src/core/src/cache/io_context.rs @@ -26,21 +26,8 @@ pub trait EntryMetadata: Debug + Send + Sync { } /// Convert an [`EntryID`] to a t4 key (8-byte little-endian representation). -/// 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 index entry 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 agrees 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 +pub(crate) fn entry_id_to_key(entry_id: &EntryID) -> Vec { + usize::from(*entry_id).to_le_bytes().to_vec() } /// A default implementation of [`EntryMetadata`]. diff --git a/src/core/src/cache/observer/stats.rs b/src/core/src/cache/observer/stats.rs index e8844923..6c871583 100644 --- a/src/core/src/cache/observer/stats.rs +++ b/src/core/src/cache/observer/stats.rs @@ -110,12 +110,6 @@ define_runtime_stats! { pub struct CacheStats { /// Total number of entries in the cache. pub total_entries: usize, - /// How many lookups or writes found a key held by a different file. - /// - /// Expected to stay at zero. A non-zero value means two sources computed - /// the same `EntryID`, and each was served correctly only because the - /// identity check turned the collision into a miss. - pub identity_mismatches: u64, /// Number of in-memory Arrow entries. pub memory_arrow_entries: usize, /// Number of in-memory Liquid entries. diff --git a/src/core/src/cache/tests/policies.rs b/src/core/src/cache/tests/policies.rs index 928ee86f..e074adbd 100644 --- a/src/core/src/cache/tests/policies.rs +++ b/src/core/src/cache/tests/policies.rs @@ -104,7 +104,7 @@ async fn liquid_eviction_reads_memory_and_disk() { } let mut states = Vec::new(); - cache.for_each_entry(|_, _, entry| states.push(CachedBatchType::from(entry))); + cache.for_each_entry(|_, entry| states.push(CachedBatchType::from(entry))); assert!( states .iter() @@ -154,7 +154,7 @@ async fn liquid_eviction_reads_memory_and_disk() { cache.flush_all_to_disk().await.unwrap(); let mut disk_states = Vec::new(); - cache.for_each_entry(|_, _, entry| disk_states.push(CachedBatchType::from(entry))); + cache.for_each_entry(|_, entry| disk_states.push(CachedBatchType::from(entry))); assert!(disk_states.iter().all(|state| matches!( state, CachedBatchType::DiskLiquid | CachedBatchType::DiskArrow diff --git a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap index 5775f019..fc702d26 100644 --- a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap +++ b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap @@ -35,17 +35,14 @@ event=read entry=0 selection=false expr=None cached=DiskLiquid event=io_read_liquid entry=0 bytes=[bytes] event=hydrate entry=0 cached=DiskLiquid new=MemoryLiquid event=insert_success entry=0 kind=MemoryLiquid -event=disk_evict entry=0 bytes=[bytes] event=read entry=1 selection=false expr=None cached=DiskLiquid event=io_read_liquid entry=1 bytes=[bytes] event=hydrate entry=1 cached=DiskLiquid new=MemoryLiquid event=insert_success entry=1 kind=MemoryLiquid -event=disk_evict entry=1 bytes=[bytes] event=read entry=2 selection=false expr=None cached=DiskLiquid event=io_read_liquid entry=2 bytes=[bytes] event=hydrate entry=2 cached=DiskLiquid new=MemoryLiquid event=insert_success entry=2 kind=MemoryLiquid -event=disk_evict entry=2 bytes=[bytes] event=read entry=3 selection=false expr=None cached=MemoryLiquid event=read entry=4 selection=false expr=None cached=MemoryArrow ] diff --git a/src/datafusion-client/Cargo.toml b/src/datafusion-client/Cargo.toml index ddd207cd..42e04186 100644 --- a/src/datafusion-client/Cargo.toml +++ b/src/datafusion-client/Cargo.toml @@ -17,13 +17,13 @@ arrow-schema = { workspace = true } futures = { workspace = true } tonic = { workspace = true } uuid = { workspace = true } -liquid-cache-common = { workspace = true } +liquid-cache-common = { workspace = true, features = ["rpc"] } liquid-cache-datafusion = { workspace = true } fastrace = { workspace = true } fastrace-tonic = { workspace = true } tower = "0.5.3" tokio = { workspace = true } -object_store = { workspace = true } +object_store = { workspace = true, features = ["http"] } [dev-dependencies] insta = { workspace = true } diff --git a/src/datafusion-local/src/tests/batch_size_alignment.rs b/src/datafusion-local/src/tests/batch_size_alignment.rs index 609f1baa..ccb4f71b 100644 --- a/src/datafusion-local/src/tests/batch_size_alignment.rs +++ b/src/datafusion-local/src/tests/batch_size_alignment.rs @@ -1,9 +1,9 @@ -//! Regression tests for issue #13: the reader indexes the cache by batch id, and -//! the parquet fallback turns that id back into rows using the *cache* batch size. -//! Reading at the session batch size (`datafusion.execution.batch_size`) instead -//! therefore addressed the wrong rows whenever the two differed — the scan either -//! ran off the end of the row group and panicked, or silently returned rows from -//! the wrong offsets. +//! Regression tests for batch-size alignment: the reader indexes the cache by batch +//! id, and the parquet fallback turns that id back into rows using the *cache* batch +//! size. Reading at the session batch size (`datafusion.execution.batch_size`) +//! instead addresses the wrong rows whenever the two differ — the scan either runs +//! off the end of the row group and fails, or silently returns rows from the wrong +//! offsets. //! //! `LiquidCacheLocalBuilder::build` pins `execution.batch_size` to the cache batch //! size, which is why the default configuration never hit this. A host that sets a @@ -150,7 +150,7 @@ async fn narrowed_session_batch_size_with_limit_returns_aligned_rows() { assert_eq!(ids, (0..4290).collect::>()); } -/// Issue #13 question 2: the fully-cached path misaligned as well. The cold run +/// The fully-cached path misaligns as well. The cold run /// populates cache batches 0 and 1; the warm run reads only from the cache and /// used to emit rows 8192.. for query batch 1, because `read_arrow_array` applied /// the 2145-bit selection to the 8192-row stored chunk and `arrow::compute::filter` diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap index 212649c1..ffd9e6af 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap @@ -49,6 +49,6 @@ RuntimeStatsSnapshot: try_read_liquid_calls: 0 read_io_count: 3 write_io_count: 0 - disk_evictions: 3 + disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap index f9bee9ed..6309a155 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap @@ -48,6 +48,6 @@ RuntimeStatsSnapshot: try_read_liquid_calls: 0 read_io_count: 4 write_io_count: 0 - disk_evictions: 4 + disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap index 4a85ae85..6e881e54 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap @@ -64,6 +64,6 @@ RuntimeStatsSnapshot: try_read_liquid_calls: 0 read_io_count: 1 write_io_count: 0 - disk_evictions: 1 + disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap index 1923cf5a..5dcdb137 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap @@ -49,6 +49,6 @@ RuntimeStatsSnapshot: try_read_liquid_calls: 0 read_io_count: 3 write_io_count: 2 - disk_evictions: 3 + disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 diff --git a/src/datafusion-server/Cargo.toml b/src/datafusion-server/Cargo.toml index 804b48dc..a391c48f 100644 --- a/src/datafusion-server/Cargo.toml +++ b/src/datafusion-server/Cargo.toml @@ -21,7 +21,7 @@ tokio = { workspace = true } url = { workspace = true } liquid-cache-datafusion = { workspace = true } object_store = { workspace = true, features = ["aws", "http"] } -liquid-cache-common = { workspace = true } +liquid-cache-common = { workspace = true, features = ["rpc"] } tempfile = { workspace = true } axum = "0.8.9" serde = { workspace = true } diff --git a/src/datafusion-server/src/admin_server/handlers.rs b/src/datafusion-server/src/admin_server/handlers.rs index 029f928a..4614fe96 100644 --- a/src/datafusion-server/src/admin_server/handlers.rs +++ b/src/datafusion-server/src/admin_server/handlers.rs @@ -17,7 +17,7 @@ use datafusion::{ datasource::physical_plan::FileScanConfig, physical_plan::{ExecutionPlan, StatisticsArgs, StatisticsContext}, }; -use liquid_cache_common::rpc::ExecutionMetricsResponse; +use liquid_cache_common::metrics::ExecutionMetricsResponse; use liquid_cache_datafusion::LiquidParquetSource; use log::info; use serde::Serialize; diff --git a/src/datafusion-server/src/service.rs b/src/datafusion-server/src/service.rs index 04c67124..479daca0 100644 --- a/src/datafusion-server/src/service.rs +++ b/src/datafusion-server/src/service.rs @@ -7,7 +7,7 @@ use datafusion::{ }; use liquid_cache::{ByteCache, cache::EvictionPolicy}; use liquid_cache::{cache::HydrationPolicy, cache_policies::CachePolicy}; -use liquid_cache_common::rpc::ExecutionMetricsResponse; +use liquid_cache_common::metrics::ExecutionMetricsResponse; use liquid_cache_datafusion::{ cache::{ColumnLineages, LiquidCacheParquet, LiquidCacheParquetRef}, extract_execution_metrics, diff --git a/src/datafusion/Cargo.toml b/src/datafusion/Cargo.toml index 53e89258..76f1cbdf 100644 --- a/src/datafusion/Cargo.toml +++ b/src/datafusion/Cargo.toml @@ -17,7 +17,7 @@ futures = { workspace = true } tokio = { workspace = true } ahash = { workspace = true } log = { workspace = true } -object_store = { workspace = true, features = ["http"] } +object_store = { workspace = true } liquid-cache-common = { workspace = true } liquid-cache = { workspace = true } fastrace = { workspace = true } diff --git a/src/datafusion/src/cache/column.rs b/src/datafusion/src/cache/column.rs index 61bffe00..a0cdaba6 100644 --- a/src/datafusion/src/cache/column.rs +++ b/src/datafusion/src/cache/column.rs @@ -20,12 +20,13 @@ use std::sync::Arc; #[derive(Debug)] pub struct CachedColumn { cache_store: Arc, - /// The lease its keys are built from. Held so the id cannot be recycled - /// while this column can still compute keys from it; `identity()` is what - /// every cache read and write is checked against. - file_id: Arc, field: Arc, column_path: ColumnAccessPath, + /// The unnarrowed id of the file this column belongs to. `column_path` + /// keeps only its low 16 bits, so another file can compute the same + /// entry ids; the cache serves an entry only to the identity it was + /// cached for. + file_identity: u64, expression: Option>, snapshots: Arc, } @@ -58,8 +59,8 @@ impl CachedColumn { pub(crate) fn new( field: Arc, cache_store: Arc, - file_id: Arc, column_access_path: ColumnAccessPath, + file_identity: u64, expression: Option>, is_predicate_column: bool, snapshots: Arc, @@ -83,27 +84,25 @@ impl CachedColumn { Self { field, cache_store, - file_id, column_path: column_access_path, + file_identity, expression, snapshots, } } /// row_id must be on a batch boundary. - /// The unnarrowed name of the file this column reads, recorded with every - /// entry it caches and checked on every read. - pub(crate) fn identity(&self) -> u64 { - self.file_id.identity() - } - pub(crate) fn entry_id(&self, batch_id: BatchID) -> ParquetArrayID { self.column_path.entry_id(batch_id) } + pub(crate) fn file_identity(&self) -> u64 { + self.file_identity + } + pub(crate) fn contains(&self, batch_id: BatchID) -> bool { self.cache_store - .is_cached(&self.entry_id(batch_id).into(), self.identity()) + .contains(&self.entry_id(batch_id).into(), self.file_identity) } pub(crate) fn snapshot_entry(&self, batch_id: BatchID) -> Option> { @@ -144,7 +143,7 @@ impl CachedColumn { self.cache_store .eval_predicate_on_entry( &entry_id, - self.identity(), + self.file_identity, entry.as_ref(), Some(filter), &liquid_expr, @@ -153,7 +152,7 @@ impl CachedColumn { } None => { self.cache_store - .eval_predicate(&entry_id, self.identity(), &liquid_expr) + .eval_predicate(&entry_id, self.file_identity, &liquid_expr) .with_selection(filter) .await } @@ -205,7 +204,7 @@ impl CachedColumn { .cache_store .read_entry( &entry_id, - self.identity(), + self.file_identity, entry.as_ref(), Some(filter), self.expression.as_deref(), @@ -213,7 +212,7 @@ impl CachedColumn { .await; } self.cache_store - .get(&entry_id, self.identity()) + .get(&entry_id, self.file_identity) .with_selection(filter) .with_optional_expression_hint(self.expression()) .read() @@ -223,7 +222,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, self.identity()).await + self.cache_store.get(&entry_id, self.file_identity).await } /// Insert an array into the cache. @@ -237,7 +236,7 @@ impl CachedColumn { } self.cache_store - .insert(self.entry_id(batch_id).into(), self.identity(), array) + .insert(self.entry_id(batch_id).into(), self.file_identity, array) .await?; Ok(()) } @@ -254,7 +253,11 @@ impl CachedColumn { if self.snapshots.get(&entry_id).is_some() { return PrefetchOutcome::AlreadySnapshotted; } - match self.cache_store.prefetch(&entry_id, self.identity()).await { + match self + .cache_store + .prefetch(&entry_id, self.file_identity) + .await + { PrefetchResult::Snapshot(entry) => { self.snapshots.insert(entry_id, entry); PrefetchOutcome::Snapshotted diff --git a/src/datafusion/src/cache/file_id.rs b/src/datafusion/src/cache/file_id.rs deleted file mode 100644 index 94feea76..00000000 --- a/src/datafusion/src/cache/file_id.rs +++ /dev/null @@ -1,473 +0,0 @@ -//! 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)) - } - // A fresh id while the key's file field still has room, and only - // then the oldest released one. Recycling eagerly would be correct - // — the identity check turns an aliased read into a miss — but it - // costs the previous holder every entry it cached, since the new - // owner takes those keys over. Two files read one after another, - // or the two sides of a join, would evict each other for no reason - // while 65,000 ids sat unused. Reuse is what the narrow key field - // forces, not something to spend before it is needed. - None if inner.next <= u16::MAX as u64 => { - let id = inner.next; - inner.next += 1; - (id, None) - } - 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"); - - // The released id is available, but a fresh one is preferred while the - // key field has room: reuse costs the previous holder its entries, so - // it is spent only once there is nothing else to hand out. - let c = pool.acquire("c.parquet"); - assert_eq!(c.get(), 2); - assert_eq!(pool.live_count(), 2); - } - - #[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); - assert_eq!(pool.live_count(), 1, "the last holder released id 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(); - // Past the key width, so every acquire below takes the recycling path - // rather than a fresh id. Held, so the queue holds only what this test - // releases. - let _fillers = exhaust_key_width(&pool); - - 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" - ); - } - - /// Once the key field is exhausted there is nothing to hand out but - /// released ids, and the oldest goes first — giving the most recently - /// released file's entries the longer window before they are taken over. - #[test] - fn released_ids_are_reused_oldest_first() { - let pool = FileIdPool::new(); - let _fillers = exhaust_key_width(&pool); - let a = pool.acquire("a.parquet"); - let b = pool.acquire("b.parquet"); - let (a_id, b_id) = (a.get(), b.get()); - drop(a); - drop(b); - assert_eq!(pool.acquire("x.parquet").get(), a_id); - assert_eq!(pool.acquire("y.parquet").get(), b_id); - } - - /// Drive `next` past the 16-bit key field so the pool has no fresh ids to - /// hand out and must recycle. - /// - /// The fillers are returned rather than dropped, and the caller has to keep - /// them: dropping them here would leave 65,536 released records queued - /// ahead of whatever the test then releases, and FIFO reuse would hand back - /// a filler's id instead of the one under test. - #[must_use] - fn exhaust_key_width(pool: &Arc) -> Vec> { - (0..=(u16::MAX as u64)) - .map(|i| pool.acquire(&format!("filler_{i}.parquet"))) - .collect() - } - - #[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 b52379de..84674589 100644 --- a/src/datafusion/src/cache/id.rs +++ b/src/datafusion/src/cache/id.rs @@ -51,8 +51,11 @@ const _: () = assert!(std::mem::align_of::() == 8); impl ParquetArrayID { /// Creates a new CacheEntryID. + /// + /// Only the low 16 bits of `file_id` are kept, so distinct files can + /// produce the same id. The cache tells them apart by the full file id it + /// is given alongside the entry id. 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 +149,9 @@ pub struct ColumnAccessPath { impl ColumnAccessPath { /// Create a new instance of ColumnAccessPath. + /// + /// Only the low 16 bits of `file_id` are kept; 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 { @@ -226,9 +230,9 @@ mod tests { } #[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 test_cache_entry_id_new_keeps_low_16_bits_of_file_id() { + let wrapped = ParquetArrayID::new((u16::MAX as u64) + 6, 2, 3, BatchID::from_raw(4)); + assert_eq!(wrapped, ParquetArrayID::new(5, 2, 3, BatchID::from_raw(4))); } #[test] diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index 4501158f..7ef6a7a8 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -3,9 +3,7 @@ use crate::io::ParquetCacheMetadata; use crate::reader::{LiquidPredicate, extract_multi_column_or}; -use crate::sync::RwLock; -mod file_id; - +use crate::sync::{Mutex, RwLock}; use ahash::AHashMap; use arrow::array::{BooleanArray, RecordBatch, RecordBatchOptions}; use arrow::buffer::BooleanBuffer; @@ -23,6 +21,7 @@ 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 id; @@ -51,15 +50,6 @@ pub struct ParquetFileIdentity { } impl ParquetFileIdentity { - /// The key a file id is leased against. - /// - /// Both components, for the same reason the identity carries both: a path - /// is unique only within its object store, so keying on it alone would let - /// two stores share one lease and therefore one set of cache keys. - pub(crate) fn lease_key(&self) -> String { - format!("{}{}", self.object_store_url.as_str(), self.path) - } - /// Create an identity from an object store URL and an object path. pub fn new(object_store_url: ObjectStoreUrl, path: String) -> Self { Self { @@ -118,19 +108,18 @@ impl CachedRowGroup { fn new( cache_store: Arc, row_group_idx: u64, - file_id: Arc, + file_idx: u64, columns: &[CachedColumnSpec], snapshots: Arc, ) -> Self { let mut column_maps = ColumnMaps::default(); for (column_id, field, expression, is_predicate_column) in columns { - let column_access_path = - ColumnAccessPath::new(file_id.get(), row_group_idx, *column_id); + let column_access_path = ColumnAccessPath::new(file_idx, row_group_idx, *column_id); let column = Arc::new(CachedColumn::new( Arc::clone(field), Arc::clone(&cache_store), - Arc::clone(&file_id), column_access_path, + file_idx, expression.clone(), *is_predicate_column, Arc::clone(&snapshots), @@ -232,7 +221,7 @@ impl CachedRowGroup { Some(array) => Some(array), None => { self.cache_store - .try_read_liquid(&entry_id, column.identity()) + .try_read_liquid(&entry_id, column.file_identity()) .await } }; @@ -293,9 +282,7 @@ pub(crate) type CachedRowGroupRef = Arc; #[derive(Debug)] pub struct CachedFile { cache_store: Arc, - /// Held, not copied: the id returns to the pool when the last holder — this - /// file and everything derived from it — is dropped. - file_id: Arc, + file_id: u64, file_schema: SchemaRef, lineages: Arc, } @@ -303,7 +290,7 @@ pub struct CachedFile { impl CachedFile { fn new( cache_store: Arc, - file_id: Arc, + file_id: u64, file_schema: SchemaRef, lineages: Arc, ) -> Self { @@ -350,7 +337,7 @@ impl CachedFile { Arc::new(CachedRowGroup::new( self.cache_store.clone(), row_group_id, - Arc::clone(&self.file_id), + self.file_id, &columns, snapshots, )) @@ -374,12 +361,11 @@ pub(crate) type CachedFileRef = Arc; #[derive(Debug)] pub struct LiquidCacheParquet { /// Map object-store-qualified file identity to file id. - /// Leases the file ids that cache keys are built from, so the id space - /// tracks the files being read rather than every file ever read — see - /// [`file_id`]. - file_ids: Arc, + files: Mutex>, cache_store: Arc, + + current_file_id: AtomicU64, } /// A reference to the main cache structure. @@ -438,8 +424,9 @@ impl LiquidCacheParquet { .await; LiquidCacheParquet { - file_ids: file_id::FileIdPool::new(), + files: Mutex::new(AHashMap::new()), cache_store: cache_storage, + current_file_id: AtomicU64::new(0), } } @@ -460,34 +447,41 @@ impl LiquidCacheParquet { full_file_schema: SchemaRef, lineages: Arc, ) -> CachedFileRef { + let mut files = self.files.lock().unwrap(); + let file_id = *files + .entry(file_identity) + .or_insert_with(|| self.current_file_id.fetch_add(1, Ordering::Relaxed)); + drop(files); + Arc::new(CachedFile::new( self.cache_store.clone(), - self.file_ids.acquire(&file_identity.lease_key()), + file_id, full_file_schema, lineages, )) } - /// How many file ids are currently leased. + /// How many distinct files have been registered with this cache. /// - /// Bounded by the files being read, not by everything ever read. Rising - /// without bound means leases are being held longer than the reads that - /// need them. + /// File ids are never reused, so this counts every file ever registered, + /// not only those being read now. pub fn leased_file_ids(&self) -> usize { - self.file_ids.live_count() + self.files.lock().unwrap().len() } - /// How many ids have been handed out that do not fit the cache key's - /// 16-bit file field. Non-zero means more files are being read at once - /// than the key can name, and ids are being recycled under entries that - /// are still resident. + /// How many file ids have been handed out that do not fit the cache key's + /// 16-bit file field. Above zero, distinct files compute the same keys: + /// each is served correctly, but they cannot share the cache. pub fn file_ids_over_key_width(&self) -> u64 { - self.file_ids.over_key_width() + self.current_file_id + .load(Ordering::Relaxed) + .saturating_sub(1 << 16) } - /// How many cache lookups or writes found a key held by another file. + /// How many cache lookups found a key held by another file, and missed + /// rather than serve it. pub fn identity_mismatches(&self) -> u64 { - self.cache_store.stats().identity_mismatches + self.cache_store.identity_mismatches() } /// Get the batch size of the cache. @@ -536,7 +530,8 @@ impl LiquidCacheParquet { /// 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) { - self.file_ids.reset(); + let mut files = self.files.lock().unwrap(); + files.clear(); self.cache_store.reset(); } @@ -961,4 +956,70 @@ mod tests { let expected = BooleanBuffer::collect_bool(batch_size, |i| i == 1 || i == 3).into(); assert_eq!(result, expected); } + + /// A cache key holds 16 bits of file id, so the 65,537th file registered + /// computes the same keys as the first. Its reads must not be served the + /// first file's data, and its own data must still be cacheable. + #[tokio::test] + async fn files_whose_keys_collide_do_not_share_entries() { + let batch_size = 64; + let tmp_dir = tempfile::tempdir().unwrap(); + let store = t4::mount(tmp_dir.path().join("liquid_cache.t4")) + .await + .unwrap(); + let cache = LiquidCacheParquet::new( + batch_size, + usize::MAX, + usize::MAX, + store, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeEvict), + Box::new(AlwaysHydrate::new()), + ) + .await; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let identity = + |path: String| ParquetFileIdentity::new(ObjectStoreUrl::local_filesystem(), path); + let batch_id = BatchID::from_row_id(0, batch_size); + + let first = cache.register_or_get_file(identity("first".to_string()), schema.clone()); + let first_column = first.create_row_group(0, vec![]).get_column(0).unwrap(); + let first_data: ArrayRef = Arc::new(Int32Array::from_iter_values(0..batch_size as i32)); + first_column + .insert(batch_id, first_data.clone()) + .await + .unwrap(); + + for i in 1..=u16::MAX as u32 { + cache.register_or_get_file(identity(format!("filler-{i}")), schema.clone()); + } + + let colliding = cache.register_or_get_file(identity("colliding".to_string()), schema); + let colliding_column = colliding.create_row_group(0, vec![]).get_column(0).unwrap(); + assert_eq!( + colliding_column.entry_id(batch_id), + first_column.entry_id(batch_id), + "the 65,537th file should compute the first file's key" + ); + + assert!( + colliding_column + .get_arrow_array_test_only(batch_id) + .await + .is_none(), + "a file was served another file's cached data" + ); + + let colliding_data: ArrayRef = + Arc::new(Int32Array::from_iter_values((0..batch_size as i32).rev())); + colliding_column + .insert(batch_id, colliding_data.clone()) + .await + .expect("a file could not cache its own data"); + let read = colliding_column + .get_arrow_array_test_only(batch_id) + .await + .unwrap(); + assert_eq!(read.as_ref(), colliding_data.as_ref()); + } } diff --git a/src/datafusion/src/cache/stats.rs b/src/datafusion/src/cache/stats.rs index fadba50b..d1947fc6 100644 --- a/src/datafusion/src/cache/stats.rs +++ b/src/datafusion/src/cache/stats.rs @@ -123,36 +123,35 @@ 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::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::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, 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::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::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,21 +204,15 @@ mod tests { let mut row_start_id_sum = 0; let mut row_count_sum = 0; let mut memory_size_sum = 0; - // 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( - ParquetFileIdentity::new( - datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(), - format!("test_{file_no}.parquet"), - ), - schema.clone(), - ) - }) - .collect(); - for file in &files { + for file_no in 0..8 { + let file_name = format!("test_{file_no}.parquet"); + let file = cache.register_or_get_file( + ParquetFileIdentity::new( + datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(), + file_name, + ), + schema.clone(), + ); 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/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 3a0a9b27..7daebe62 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -654,7 +654,7 @@ mod tests { physical_plan::{ PhysicalExpr, collect, display::DisplayableExecutionPlan, filter_pushdown::PushedDown, }, - prelude::SessionContext, + prelude::{SessionConfig, SessionContext}, }; use liquid_cache::{ cache::{AlwaysHydrate, TranscodeEvict}, @@ -1012,13 +1012,21 @@ mod tests { writer.close().unwrap(); } - #[tokio::test] - async fn prunes_row_group_with_bloom_filter() { + /// Scan `a = 2` over a file whose two row groups both span `1..=4`, so + /// statistics keep both and only the bloom filter can rule one out. Returns + /// the rows produced and the bloom-filter pruning metric. + async fn scan_bloom_file( + bloom_filter_on_read: bool, + ) -> (usize, datafusion::physical_plan::metrics::PruningMetrics) { let tmp_dir = tempfile::tempdir().unwrap(); let parquet_path = tmp_dir.path().join("bloom.parquet"); write_bloom_file(&parquet_path); - let ctx = SessionContext::new(); + let config = SessionConfig::new().set_bool( + "datafusion.execution.parquet.bloom_filter_on_read", + bloom_filter_on_read, + ); + let ctx = SessionContext::new_with_config(config); ctx.register_parquet("t", parquet_path.to_str().unwrap(), Default::default()) .await .unwrap(); @@ -1034,7 +1042,7 @@ mod tests { let metrics = liquid_source(&rewritten).metrics().clone(); let batches = collect(rewritten, ctx.task_ctx()).await.unwrap(); - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + let rows = batches.iter().map(RecordBatch::num_rows).sum::(); let metric = metrics .clone_inner() @@ -1046,7 +1054,26 @@ mod tests { else { panic!("unexpected metric: {metric:?}"); }; + (rows, pruning_metrics) + } + + #[tokio::test] + async fn prunes_row_group_with_bloom_filter() { + let (rows, pruning_metrics) = scan_bloom_file(true).await; + assert_eq!(rows, 1); assert_eq!(pruning_metrics.pruned(), 1); + assert_eq!(pruning_metrics.matched(), 1); + } + + /// `bloom_filter_on_read = false` must keep the reader off the bloom + /// filters entirely: the row group the filter would have pruned is read, + /// and the predicate alone still produces the right answer. + #[tokio::test] + async fn skips_bloom_filter_when_bloom_filter_on_read_is_off() { + let (rows, pruning_metrics) = scan_bloom_file(false).await; + assert_eq!(rows, 1); + assert_eq!(pruning_metrics.pruned(), 0); + assert_eq!(pruning_metrics.matched(), 2); } /// Declining a scan costs it the cache, so the guard must key on what the scan diff --git a/src/datafusion/src/reader/plantime/morselizer.rs b/src/datafusion/src/reader/plantime/morselizer.rs index 8e9f1a3b..6e9ee546 100644 --- a/src/datafusion/src/reader/plantime/morselizer.rs +++ b/src/datafusion/src/reader/plantime/morselizer.rs @@ -67,6 +67,7 @@ pub(crate) struct LiquidMorselizer { pub(crate) parquet_file_reader_factory: Arc, pub(crate) object_store_url: ObjectStoreUrl, pub(crate) reorder_filters: bool, + pub(crate) enable_bloom_filter: bool, pub(crate) liquid_cache: LiquidCacheParquetRef, pub(crate) expr_adapter_factory: Arc, pub(crate) span: Option>, @@ -187,6 +188,7 @@ impl Morselizer for LiquidMorselizer { projection, predicate, reorder_filters: self.reorder_filters, + enable_bloom_filter: self.enable_bloom_filter, liquid_cache: self.liquid_cache.clone(), expr_adapter_factory: Arc::clone(&self.expr_adapter_factory), file_identity, @@ -236,6 +238,7 @@ struct PreparedLiquidOpen { projection: ProjectionExprs, predicate: Option>, reorder_filters: bool, + enable_bloom_filter: bool, liquid_cache: LiquidCacheParquetRef, expr_adapter_factory: Arc, file_identity: ParquetFileIdentity, @@ -466,7 +469,10 @@ fn prepare_and_prune_by_stats(mut loaded: MetadataLoadedLiquidOpen) -> Result
  • Result
  • Option { let mut kind = None; - cache.for_each_entry(|entry_id, _, entry| { + cache.for_each_entry(|entry_id, entry| { if entry_id == id { kind = Some(CachedBatchType::from(entry)); } diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index 69e0415b..85291e1f 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -48,6 +48,10 @@ impl LiquidParquetSource { self.table_parquet_options.global.reorder_filters } + fn bloom_filter_on_read(&self) -> bool { + self.table_parquet_options.global.bloom_filter_on_read + } + /// Set the span for the LiquidParquetSource pub fn with_span(&self, span: fastrace::Span) -> Self { Self { @@ -175,6 +179,7 @@ impl FileSource for LiquidParquetSource { parquet_file_reader_factory, object_store_url: base_config.object_store_url.clone(), reorder_filters: self.reorder_filters(), + enable_bloom_filter: self.bloom_filter_on_read(), expr_adapter_factory, span: execution_span.map(Arc::new), lineages: Arc::clone(&self.lineages), diff --git a/src/datafusion/src/utils.rs b/src/datafusion/src/utils.rs index e6d903cc..5d5b3aad 100644 --- a/src/datafusion/src/utils.rs +++ b/src/datafusion/src/utils.rs @@ -9,7 +9,7 @@ use datafusion::{ datasource::source::DataSourceExec, physical_plan::{ExecutionPlan, metrics::MetricValue}, }; -use liquid_cache_common::rpc::ExecutionMetricsResponse; +use liquid_cache_common::metrics::ExecutionMetricsResponse; use parquet::arrow::arrow_reader::RowSelector; use crate::cache::LiquidCacheParquetRef;