fix(cache): leave a disk victim that a read hydrated

src/core/src/cache/core.rs

@@ -497,14 +497,28 @@ impl LiquidCache {
         self.cache_policy.notify_remove(&entry_id);
     }
 
+    /// Remove `entry_id` if the index still holds it on disk.
+    ///
+    /// Callers pick the victim from advice read earlier: the policy's disk
+    /// queue, or a slot loaded before the eviction policy ran. A read can have
+    /// hydrated the key back into memory since, under the same identity. That
+    /// record holds memory, not disk, and its writer files it with the policy,
+    /// so it is left alone. Its disk reservation is not released here either:
+    /// the hydrating insert replaced the disk record, and with it the only
+    /// byte count this call could release.
     async fn remove_disk_entry(&self, entry_id: EntryID) {
-        let Some(removed) = self.index.remove(&entry_id) else {
+        let Some(removed) = self.index.remove_if(&entry_id, |slot| {
+            matches!(
+                slot.entry,
+                CacheEntry::DiskLiquid { .. } | CacheEntry::DiskArrow { .. }
+            )
+        }) else {
             return;
         };
         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"),
+            _ => unreachable!("remove_if removed a memory entry"),
         };
         self.store
             .remove(&entry_id_to_key(&entry_id))
@@ -1398,4 +1412,118 @@ mod tests {
         assert_eq!(result, Ok(()));
         assert!(!cache.contains(&entry_id, 0));
     }
+
+    type DiskVictimHook = Box<dyn FnOnce(&[EntryID]) + Send>;
+
+    /// A [`LiquidPolicy`] that runs a hook once, after it hands out disk
+    /// victims and before the cache acts on them.
+    struct HookAfterDiskVictims {
+        inner: LiquidPolicy,
+        hook: Arc<std::sync::Mutex<Option<DiskVictimHook>>>,
+    }
+
+    impl std::fmt::Debug for HookAfterDiskVictims {
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            f.debug_struct("HookAfterDiskVictims").finish()
+        }
+    }
+
+    impl CachePolicy for HookAfterDiskVictims {
+        fn find_memory_victim(&self, cnt: usize) -> Vec<EntryID> {
+            self.inner.find_memory_victim(cnt)
+        }
+
+        fn find_disk_victim(&self, cnt: usize) -> Vec<EntryID> {
+            let victims = self.inner.find_disk_victim(cnt);
+            let hook = self.hook.lock().unwrap().take();
+            if let Some(hook) = hook {
+                hook(&victims);
+            }
+            victims
+        }
+
+        fn notify_insert(&self, entry_id: &EntryID, batch_type: CachedBatchType) {
+            self.inner.notify_insert(entry_id, batch_type);
+        }
+
+        fn notify_access(&self, entry_id: &EntryID, batch_type: CachedBatchType) {
+            self.inner.notify_access(entry_id, batch_type);
+        }
+
+        fn notify_remove(&self, entry_id: &EntryID) {
+            self.inner.notify_remove(entry_id);
+        }
+    }
+
+    /// The memory budget equals what the index's memory entries hold.
+    fn assert_memory_budget_matches_index(cache: &LiquidCache) {
+        let mut held = 0;
+        cache
+            .index()
+            .for_each(|_, slot| held += slot.entry.memory_usage_bytes());
+        assert_eq!(cache.budget.memory_usage_bytes(), held);
+    }
+
+    /// A disk victim is picked from the policy's disk queue, and a read can
+    /// hydrate that entry back into memory before the eviction removes it.
+    /// The eviction must then leave the memory entry alone: it is no longer a
+    /// disk entry, and its record is the hydrating read's.
+    #[test]
+    fn disk_victim_hydrated_before_its_removal_is_left_in_memory() {
+        let hydrated = EntryID::from(920usize);
+        let flushed = EntryID::from(921usize);
+        let hydrated_array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16));
+        let flushed_array: ArrayRef = Arc::new(Int32Array::from_iter_values(16..32));
+        let disk_bytes = arrow_to_bytes(&hydrated_array).unwrap().len() + IDENTITY_TRAILER_LEN;
+        let hook = Arc::new(std::sync::Mutex::new(None::<DiskVictimHook>));
+        let cache = block_on(
+            LiquidCacheBuilder::new()
+                .with_max_memory_bytes(1 << 20)
+                .with_max_disk_bytes(disk_bytes)
+                .with_eviction_policy(Box::new(TranscodeEvict))
+                .with_cache_policy(Box::new(HookAfterDiskVictims {
+                    inner: LiquidPolicy::new(),
+                    hook: hook.clone(),
+                }))
+                .build(),
+        );
+
+        block_on(async {
+            cache
+                .insert(hydrated, 0, hydrated_array.clone())
+                .await
+                .unwrap();
+            cache.flush_all_to_disk().await.unwrap();
+            cache.insert(flushed, 0, flushed_array).await.unwrap();
+        });
+        assert!(matches!(
+            &cache.index().get(&hydrated).unwrap().entry,
+            CacheEntry::DiskArrow { .. }
+        ));
+
+        // Between the advice and the removal, a read hydrates the victim.
+        let reader = cache.clone();
+        let expected = hydrated_array.clone();
+        *hook.lock().unwrap() = Some(Box::new(move |victims: &[EntryID]| {
+            assert_eq!(victims, &[hydrated]);
+            std::thread::spawn(move || {
+                let read = block_on(async { reader.get(&hydrated, 0).await });
+                assert_eq!(read.expect("present").as_ref(), expected.as_ref());
+            })
+            .join()
+            .unwrap();
+        }));
+
+        // Flushing the second entry needs the first one's disk space.
+        assert_eq!(block_on(cache.flush_all_to_disk()), Ok(()));
+
+        assert!(hook.lock().unwrap().is_none(), "the hook never ran");
+        assert!(matches!(
+            &cache.index().get(&hydrated).unwrap().entry,
+            CacheEntry::MemoryArrow(_)
+        ));
+        let read = block_on(async { cache.get(&hydrated, 0).await });
+        assert_eq!(read.expect("present").as_ref(), hydrated_array.as_ref());
+        assert_memory_budget_matches_index(&cache);
+    }
 }

src/core/src/cache/index.rs

@@ -148,6 +148,35 @@ impl ArtIndex {
         removed.take()
     }
 
+    /// Remove the slot under `entry_id` only if `pred` holds for it.
+    ///
+    /// The check and the removal are one step: a slot installed after `pred`
+    /// looked is not removed unseen, since the tree rejects the removal and
+    /// `pred` is asked again about the slot now there.
+    pub(crate) fn remove_if(
+        &self,
+        entry_id: &EntryID,
+        pred: impl Fn(&Slot) -> bool,
+    ) -> Option<Arc<Slot>> {
+        let guard = self.art.pin();
+        let mut matched = false;
+        let removed = self.art.compute_if_present(
+            *entry_id,
+            |cell| {
+                // The tree calls this again when its node changed under it, so
+                // the last call decides.
+                matched = cell.load().is_some_and(|slot| pred(&slot));
+                if matched { None } else { Some(cell) }
+            },
+            &guard,
+        )?;
+        if !matched {
+            return None;
+        }
+        self.entry_count.fetch_sub(1, Ordering::Relaxed);
+        removed.take()
+    }
+
     pub(crate) fn reset(&self) {
         for k in self.art.keys() {
             self.remove(&k);
@@ -275,4 +304,27 @@ mod tests {
         assert!(store.contains(&entry_id, 8));
         assert_eq!(store.entry_count(), 1);
     }
+
+    #[test]
+    fn remove_if_removes_only_a_matching_slot() {
+        let store = ArtIndex::new();
+        let entry_id = EntryID::from(1);
+        store.insert(&entry_id, 7, create_test_array(10));
+
+        assert!(
+            store
+                .remove_if(&entry_id, |slot| slot.identity == 8)
+                .is_none()
+        );
+        assert!(store.contains(&entry_id, 7));
+        assert_eq!(store.entry_count(), 1);
+
+        let removed = store
+            .remove_if(&entry_id, |slot| slot.identity == 7)
+            .expect("matching slot is removed");
+        assert_eq!(removed.identity, 7);
+        assert!(store.get(&entry_id).is_none());
+        assert_eq!(store.entry_count(), 0);
+        assert!(store.remove_if(&entry_id, |_| true).is_none());
+    }
 }