diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index b95f306b..376993bc 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -588,17 +588,35 @@ impl LiquidCache { /// 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) { - 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}"), + 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, @@ -657,16 +675,12 @@ impl LiquidCache { | CacheEntry::DiskArrow { disk_bytes, .. } => *disk_bytes, _ => panic!("remove_disk_entry called for non-disk entry"), }; - self.store - .remove(&entry_id_to_key(&entry_id, removed_identity)) - .await - .expect("disk remove failed"); - self.budget.release_disk(disk_bytes); + // 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.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. @@ -1184,18 +1198,23 @@ mod tests { /// the copy it superseded. /// /// Reported by review. The store key is `(entry id, identity)`, so this - /// insert's put landed on the very object the old entry named — deleting it + /// 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() { - // Tiny, so the second insert cannot stay in memory and — with only a - // disk entry present — has no memory victim to evict. `insert_inner` - // then spills the batch itself and re-inserts it over its own object. - let store = create_cache_store(64, Box::new(LiquidPolicy::new())).await; + 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 insert displaces a disk entry. + // 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 @@ -1209,16 +1228,97 @@ mod tests { 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 - .insert(entry_id, 7, create_test_arrow_array(4096)) + .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 @@ -1473,6 +1573,64 @@ mod tests { ); } + /// 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 diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index 54c64233..78e3bd63 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -240,35 +240,65 @@ impl ArtIndex { /// 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 slot = self.art.get(*entry_id, &guard)?; - if slot.identity != identity { - self.identity_mismatches.fetch_add(1, Ordering::Relaxed); - return None; - } - let removed = self.art.remove(*entry_id, &guard)?; - if removed.identity != identity { - // Lost the race between the check and the removal: put it back - // rather than destroying a record this caller has no claim on. - if let Some(entry) = removed.take() { - self.art - .insert( - *entry_id, - Slot::new(removed.identity, (*entry).clone()), - &guard, - ) - .expect("Insertion failed"); - } + 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); - removed.take() + 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> { @@ -316,6 +346,7 @@ impl ArtIndex { mod tests { use crate::cache::cached_batch::CacheEntry; use crate::cache::utils::create_test_array; + use crate::sync::thread; use super::*; @@ -485,6 +516,107 @@ mod tests { } } + /// 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