From 9524196fb37fb3f8ce47e259b0873dfa1fdfa2f7 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 24 Sep 2026 15:03:55 +0530 Subject: [PATCH] fix(cache): stop a reclaim deleting an object written after it was decided MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 its cached entries stay readable. `settle` decides to delete such a key and does it behind an await, and in between the key can be occupied again — by that same identity's next spill, or by the identity itself returning. The delete then destroyed a live record's bytes and left it pointing at nothing. `reclaim_orphaned_disk` now asks the index whether a live record names the object before deleting, as late as it can, and releases the byte count either way: an object left standing is one a newer record overwrote, and that record reserved its own bytes for it. `remove_disk_entry` goes through the same reclaim, since the gap between its removal and its delete is the same gap. That narrows the window rather than closing it. A put that has landed but whose record is not yet installed is invisible to the check. Closing it needs the store key to name the write and not just the writer, which means a generation in the key — and with one, a disk write never lands on the object it displaces, so there is no in-place overwrite left to keep. That is a larger change than this, and it is written up rather than made here. `ArtIndex::remove_checked` is closed outright. It tested the identity, removed, then put back what it found if the test had gone stale — and the restore was an unconditional insert, so a writer that took the key over in between lost its record while `entry_count` went on counting it. Deciding and removing are now one `compute_if_present`, so a key that changed hands is left alone instead. The verdict is recomputed on every run of that closure rather than latched on the first. `compute_if_present` decides under an optimistic read and upgrades to a write afterwards, so a removal whose upgrade loses runs the closure again — against whatever the key holds by then. A first run that finds the record ours followed by a retry that finds it someone else's acts on the retry, and must report the mismatch: latching the first verdict would decrement `entry_count` for a record still in the tree and hand the caller the new owner's entry to release. Tests: - `a_settlement_does_not_delete_an_object_written_after_it_was_decided` stages the A-B-A order: identity 7 spills, 9 takes the key over, 7 comes back and spills again, and only then does the settlement 9 displaced run. Without the check the settlement deletes 7's new object and the read misses. - `checked_removal_against_takeovers` takes the identities that race the removal, so the three tests over it state only what they vary. It runs as one round under threads, as every interleaving under shuttle, and as `a_checked_removal_that_retries_reports_the_verdict_it_acted_on`: more writers and enough rounds to reach a removal whose upgrade loses, which one round reaches too rarely to rely on. Against a latched verdict the last fails every run while the single round passes. - `an_in_place_disk_overwrite_releases_the_copy_it_supersedes` could not fail. It asserted through `insert`, which reached the branch it names only as a consequence of budget arithmetic, and the whole crate passed with the superseded handling deleted. It now writes out the two steps a spill performs — bytes to the store, then the record — and asserts the branch was taken. Its memory budget goes back to an ordinary one, since it no longer needs to provoke a particular eviction. - `only_a_same_identity_disk_write_lands_on_the_object_it_displaces` covers `DiskResidue::displacing` directly, for each combination of written form and displaced form. --- src/core/src/cache/core.rs | 208 +++++++++++++++++++++++++++++++----- src/core/src/cache/index.rs | 170 +++++++++++++++++++++++++---- 2 files changed, 334 insertions(+), 44 deletions(-) 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