From 91351b119f421ef5a90ab222e10f3e0579004392 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 00:55:53 -0400 Subject: [PATCH 1/4] existence_cache: refcount pause_remove_callbacks; fix insert-after-evict resurrection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ExistenceCacheStore pauses inner-store remove callbacks while it writes a blob, so that an eviction fired inline during the write is applied only after the key is inserted into the existence cache. Two defects made this mechanism racy — both rooted in the pause being a single, non-refcounted Option> slot shared by all concurrent operations: 1. Non-refcounted pause. Concurrent update() calls all shared one pause slot. The first update to finish unconditionally take()s it, ending the pause window for every other in-flight update. A remove callback (eviction under cache pressure) delivered during a still-running update's inner write then ran immediately as a no-op — the key was not yet in the cache — after which that update inserted a stale "exists" entry for a blob the inner store had already dropped. The cache then permanently claimed existence, and later re-uploads were drained and silently discarded. 2. Insert-after-evict ordering in get_part. get_part inserted into the existence cache without pausing removals at all. An inner store that evicts a key while a read of it is still in flight (e.g. a filesystem store streaming from an open file handle) fired the removal before the insert, so get_part resurrected a zombie exists-entry the inner store no longer held. Fix: replace the Option slot with a refcounted PauseState { pause_count, pending }. Each operation calls pause_remove_callbacks() on entry and resume_remove_callbacks() after its cache insert; queued removals drain only when the LAST concurrent holder releases (count -> 0), always after every holder has inserted. Because the insert happens before the release, a key evicted during its own operation is removed after it is inserted and cannot be left resurrected. get_part now uses the same pause guard around its read + insert. Includes three reproductions (fail unfixed, pass fixed): - concurrent_update_pause_refcount_race_test: first finisher steals the shared pause, losing an eviction during a concurrent update. - concurrent_update_unpause_race_stale_existence_test: stolen pause window leaves a stale exists entry for an oversized/evicted blob. - get_part_races_inner_eviction_and_resurrects_zombie_entry: get_part resurrects an entry evicted mid-read. Provenance: adversarial review fleet. Co-Authored-By: Claude Fable 5 --- nativelink-store/src/existence_cache_store.rs | 92 +++- .../tests/existence_store_test.rs | 481 +++++++++++++++++- 2 files changed, 546 insertions(+), 27 deletions(-) diff --git a/nativelink-store/src/existence_cache_store.rs b/nativelink-store/src/existence_cache_store.rs index 48b8e05ac04..f9ba225630b 100644 --- a/nativelink-store/src/existence_cache_store.rs +++ b/nativelink-store/src/existence_cache_store.rs @@ -49,17 +49,32 @@ impl LenEntry for ExistenceItem { } } +/// Tracks in-flight operations that must defer remove callbacks. +/// +/// `pause_count` is the number of concurrent operations (e.g. `update()` +/// calls) currently holding the pause. While it is non-zero, remove +/// callbacks from the inner store are queued in `pending` instead of being +/// applied immediately. Only when the *last* holder releases the pause +/// (count drops to zero) is the queue drained, after every holder has +/// inserted its key into the existence cache. This guarantees that an +/// eviction notification for a key delivered during that key's own +/// operation is applied *after* the insert, so the cache cannot be left +/// claiming existence of a blob the inner store dropped. +#[derive(Debug, Default)] +struct PauseState { + pause_count: usize, + pending: Vec>, +} + #[derive(Debug, MetricsComponent)] pub struct ExistenceCacheStore { #[metric(group = "inner_store")] inner_store: Store, existence_cache: EvictingMap, - // We need to pause them temporarily when inserting into the inner store - // as if it immediately expires them, we should only apply the remove callbacks - // afterwards. If this is None, we're not pausing; if it's Some it's the location to - // store them in temporarily - pause_remove_callbacks: Mutex>>>, + // We need to pause remove callbacks temporarily while operating on the + // inner store; see `PauseState`. + pause_remove_callbacks: Mutex, } impl ExistenceCacheStore { @@ -98,9 +113,11 @@ impl RemoveItemCallback for ExistenceCacheCallback { ) -> Pin + Send + 'a>> { let cache = self.cache.upgrade(); if let Some(local_cache) = cache { - if let Some(callbacks) = local_cache.pause_remove_callbacks.lock().as_mut() { - callbacks.push(store_key.into_owned()); + let mut state = local_cache.pause_remove_callbacks.lock(); + if state.pause_count > 0 { + state.pending.push(store_key.into_owned()); } else { + drop(state); let store_key = store_key.into_owned(); return Box::pin(async move { local_cache.callback(store_key).await; @@ -124,7 +141,7 @@ impl ExistenceCacheStore { let existence_cache_store = Arc::new(Self { inner_store, existence_cache: EvictingMap::new(eviction_policy, anchor_time), - pause_remove_callbacks: Mutex::new(None), + pause_remove_callbacks: Mutex::new(PauseState::default()), }); let other_ref = Arc::downgrade(&existence_cache_store); existence_cache_store @@ -146,6 +163,37 @@ impl ExistenceCacheStore { self.existence_cache.remove(digest).await; } + /// Begin deferring remove callbacks for the duration of an operation. + /// Must be paired with exactly one call to `resume_remove_callbacks`. + fn pause_remove_callbacks(&self) { + self.pause_remove_callbacks.lock().pause_count += 1; + } + + /// Release one hold on the pause. If this was the last concurrent + /// holder, drain the queued removals now — after every holder has + /// finished inserting into the existence cache — so evictions that + /// fired during the paused window take effect (including for keys + /// that were just inserted). + async fn resume_remove_callbacks(&self) { + let keys = { + let mut state = self.pause_remove_callbacks.lock(); + state.pause_count = state + .pause_count + .checked_sub(1) + .expect("resume_remove_callbacks called without matching pause"); + if state.pause_count == 0 { + core::mem::take(&mut state.pending) + } else { + Vec::new() + } + }; + let mut callbacks: FuturesUnordered<_> = keys + .into_iter() + .map(|store_key| self.callback(store_key)) + .collect(); + while callbacks.next().await.is_some() {} + } + async fn inner_has_with_results( self: Pin<&Self>, keys: &[DigestInfo], @@ -246,12 +294,7 @@ impl StoreDriver for ExistenceCacheStore { // wrong (the tier under it lost the blob and the remove callback did // not arrive), discarding the upload is what turns a stale answer // into a blob nothing holds. The entry is refreshed on success. - { - let mut locked_callbacks = self.pause_remove_callbacks.lock(); - if locked_callbacks.is_none() { - locked_callbacks.replace(vec![]); - } - } + self.pause_remove_callbacks(); trace!(?digest, "Inserting into inner cache"); let result = self.inner_store.update(digest, reader, size_info).await; if let Ok(size) = &result { @@ -261,16 +304,10 @@ impl StoreDriver for ExistenceCacheStore { .insert(digest, ExistenceItem(*size)) .await; } - { - let maybe_keys = self.pause_remove_callbacks.lock().take(); - if let Some(keys) = maybe_keys { - let mut callbacks: FuturesUnordered<_> = keys - .into_iter() - .map(|store_key| self.callback(store_key)) - .collect(); - while callbacks.next().await.is_some() {} - } - } + // Note: the insert above happens *before* releasing the pause, so a + // removal for `digest` queued during the inner update is applied + // after the insert and cannot be lost. + self.resume_remove_callbacks().await; result } @@ -282,6 +319,12 @@ impl StoreDriver for ExistenceCacheStore { length: Option, ) -> Result<(), Error> { let digest = key.into_digest(); + // Pause remove callbacks over the read + insert: the inner store may + // evict `digest` while a read of it is still in flight (e.g. a + // filesystem store keeps streaming from an open file handle). The + // eviction is queued and replayed after our insert, so the cache + // cannot resurrect an entry the inner store dropped mid-read. + self.pause_remove_callbacks(); let result = self .inner_store .get_part(digest, writer, offset, length) @@ -292,6 +335,7 @@ impl StoreDriver for ExistenceCacheStore { .insert(digest, ExistenceItem(digest.size_bytes())) .await; } + self.resume_remove_callbacks().await; result } diff --git a/nativelink-store/tests/existence_store_test.rs b/nativelink-store/tests/existence_store_test.rs index 7f0eff95343..2a6eda54d94 100644 --- a/nativelink-store/tests/existence_store_test.rs +++ b/nativelink-store/tests/existence_store_test.rs @@ -13,27 +13,36 @@ // limitations under the License. use core::pin::Pin; +use core::sync::atomic::{AtomicBool, Ordering}; use core::time::Duration; +use std::borrow::Cow; use std::sync::Arc; use async_trait::async_trait; +use bytes::Bytes; use mock_instant::thread_local::MockClock; use nativelink_config::stores::{ EvictionPolicy, ExistenceCacheSpec, MemorySpec, NoopSpec, StoreSpec, }; -use nativelink_error::{Error, ResultExt}; +use nativelink_error::{Code, Error, ResultExt, make_err}; use nativelink_macro::nativelink_test; -use nativelink_metric::MetricsComponent; +use nativelink_metric::{ + MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, +}; use nativelink_store::existence_cache_store::ExistenceCacheStore; use nativelink_store::memory_store::MemoryStore; use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use nativelink_util::common::DigestInfo; -use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; +use nativelink_util::health_utils::{ + HealthStatus, HealthStatusIndicator, default_health_status_indicator, +}; use nativelink_util::instant_wrapper::MockInstantWrapped; use nativelink_util::store_trait::{ RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, }; +use parking_lot::Mutex; use pretty_assertions::assert_eq; +use tokio::sync::Notify; const VALID_HASH1: &str = "0123456789abcdef000000000000000000010000000000000123456789abcdef"; @@ -366,3 +375,469 @@ async fn an_aged_positive_asks_the_inner_store_again() -> Result<(), Error> { ); Ok(()) } + +// Reproduces the pause_remove_callbacks race: the pause slot is a single +// non-refcounted Option shared by all concurrent update() calls. The first +// update to finish take()s it, ending the pause for every other in-flight +// update. A remove callback (eviction under pressure) firing inside the +// second update's inner_store.update await then runs immediately, BEFORE +// that update inserts into the existence cache, leaving a permanent stale +// exists-entry for a blob the inner store does not hold. Subsequent uploads +// of that key are drained and silently discarded. +#[nativelink_test] +async fn concurrent_update_pause_refcount_race_test() -> Result<(), Error> { + use core::pin::Pin; + use std::collections::HashMap; + use std::sync::Arc; + + use async_trait::async_trait; + use nativelink_metric::MetricsComponent; + use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; + use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; + use nativelink_util::store_trait::{RemoveCallback, StoreDriver, StoreKey, UploadSizeInfo}; + use tokio::sync::Notify; + + const VALID_HASH2: &str = "0223456789abcdef000000000000000000010000000000000123456789abcdef"; + + type GateMap = HashMap, Arc)>; + + #[derive(MetricsComponent)] + struct GatedStore { + inner: Store, + mem: Arc, + // Captured ExistenceCacheCallback registered by ExistenceCacheStore. + callback: Mutex>, + // (entered, release) per digest. Each gate fires at most once; once + // consumed the digest is removed so a later re-upload of the same key + // passes straight through to the inner store. + gates: Mutex, + // Key whose insertion should trigger an inline eviction (i.e. the + // inner store evicts the just-inserted entry under cache pressure + // and fires the remove callback while still inside update()). + evict_key: DigestInfo, + } + + impl core::fmt::Debug for GatedStore { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str("GatedStore") + } + } + + #[async_trait] + impl StoreDriver for GatedStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + digests: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + self.inner.has_with_results(digests, results).await + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + let digest = key.borrow().into_digest(); + let gate = self.gates.lock().remove(&digest); + let was_gated = gate.is_some(); + if let Some((entered, release)) = gate { + entered.notify_one(); + release.notified().await; + } + let result = self.inner.update(key, reader, size_info).await; + // Only the first (gated) upload of evict_key simulates the + // eviction-under-pressure; a later re-upload must persist. + if was_gated && digest == self.evict_key { + // Simulate eviction-under-pressure of the just-inserted + // entry: the evicting map removes it and fires the remove + // callback inline, inside this awaited update (see + // filesystem_store.rs "Insertion can unref this entry + // immediately under cache pressure"). + self.mem.remove_entry(digest.into()).await; + let callback = self.callback.lock().clone(); + if let Some(callback) = callback { + callback.callback(digest.into()).await; + } + } + result + } + + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + self.inner.get_part(key, writer, offset, length).await + } + + fn inner_store(&self, _digest: Option) -> &dyn StoreDriver { + self + } + + fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback( + self: Arc, + callback: RemoveCallback, + ) -> Result<(), Error> { + *self.callback.lock() = Some(callback); + Ok(()) + } + } + default_health_status_indicator!(GatedStore); + + let digest_x = DigestInfo::try_new(VALID_HASH1, 2).unwrap(); + let digest_y = DigestInfo::try_new(VALID_HASH2, 2).unwrap(); + + let entered_x = Arc::new(Notify::new()); + let release_x = Arc::new(Notify::new()); + let entered_y = Arc::new(Notify::new()); + let release_y = Arc::new(Notify::new()); + + let mem = MemoryStore::new(&MemorySpec::default()); + let gated = Arc::new(GatedStore { + inner: Store::new(mem.clone()), + mem: mem.clone(), + callback: Mutex::new(None), + gates: Mutex::new(HashMap::from([ + (digest_x, (entered_x.clone(), release_x.clone())), + (digest_y, (entered_y.clone(), release_y.clone())), + ])), + evict_key: digest_y, + }); + + let spec = ExistenceCacheSpec { + backend: StoreSpec::Noop(NoopSpec::default()), + eviction_policy: Option::default(), + }; + let store = ExistenceCacheStore::new(&spec, Store::new(gated.clone())); + + // Update A (key X): reaches inner update, having installed the pause. + let store_a = store.clone(); + let task_a = tokio::spawn(async move { store_a.update_oneshot(digest_x, "aa".into()).await }); + entered_x.notified().await; + + // Update B (key Y): sees the pause already installed and relies on it. + let store_b = store.clone(); + let task_b = tokio::spawn(async move { store_b.update_oneshot(digest_y, "bb".into()).await }); + entered_y.notified().await; + + // A finishes completely: inserts X and take()s the shared pause slot, + // ending the pause window for B as well. + release_x.notify_one(); + task_a.await.unwrap()?; + + // B finishes its inner update; the just-inserted Y is evicted under + // pressure and the remove callback fires inline (during B's await), + // then B inserts Y into the existence cache. + release_y.notify_one(); + task_b.await.unwrap()?; + + // The inner store does not hold Y (it was evicted). + assert_eq!(mem.has(digest_y).await?, None, "Y was evicted from inner"); + + // Correct behavior: the eviction notification for Y, delivered during + // B's update, must not be lost - the existence cache must not claim Y + // exists. With the non-refcounted pause slot, A's take() ended the + // pause, the removal ran before B's insert, and the cache is stale. + assert!( + !store.exists_in_cache(&digest_y).await, + "existence cache claims Y exists but inner store does not hold it" + ); + + // Demonstrate the harm: a re-upload of Y is drained and discarded. + store.update_oneshot(digest_y, "bb".into()).await?; + assert_eq!( + mem.has(digest_y).await?, + Some(2), + "re-upload of Y was silently discarded due to stale existence cache entry" + ); + Ok(()) +} + +const VALID_HASH2: &str = "aa23456789abcdef000000000000000000010000000000000123456789abcdef"; + +// Repro for the pause_remove_callbacks race: it is an Option-as-boolean with no +// refcount. Two concurrent update() calls share one pause window; the first +// finisher unconditionally take()s the flag, unpausing the second update while +// its inner-store write is still in flight. If the inner store then fires a +// remove callback for the in-flight key (eviction/oversized skip), the callback +// runs immediately as a no-op (key not yet cached), and the second update then +// inserts a stale "exists" entry for a blob the inner store does not hold. +#[nativelink_test] +async fn concurrent_update_unpause_race_stale_existence_test() -> Result<(), Error> { + let spec = ExistenceCacheSpec { + backend: StoreSpec::Noop(NoopSpec::default()), + eviction_policy: Option::default(), + }; + // max_bytes small enough that digest_b's write is skipped (and remove + // callbacks fired) by the inner memory store. + let inner_store = Store::new(MemoryStore::new(&MemorySpec { + eviction_policy: Some(EvictionPolicy { + max_bytes: 150, + ..Default::default() + }), + })); + let store = ExistenceCacheStore::new_with_time( + &spec, + inner_store.clone(), + MockInstantWrapped::default(), + ); + + let digest_a = DigestInfo::try_new(VALID_HASH1, 10).unwrap(); + let digest_b = DigestInfo::try_new(VALID_HASH2, 200).unwrap(); + + // Task B: starts update(), sets the pause flag, then parks inside the inner + // store's update awaiting stream data. + let (mut tx_b, rx_b) = nativelink_util::buf_channel::make_buf_channel_pair(); + let store_b = store.clone(); + let task_b = nativelink_util::spawn!("update_b", async move { + store_b + .update(digest_b, rx_b, UploadSizeInfo::ExactSize(200)) + .await + }); + // Single-threaded runtime: let B run until it awaits the reader. + for _ in 0..16 { + tokio::task::yield_now().await; + } + + // Task A: full update while B is in flight. A finds the pause flag Some + // (set by B), and at the end unconditionally take()s it -- stealing B's + // pause window. + store + .update_oneshot(digest_a, vec![0u8; 10].into()) + .await + .err_tip(|| "update A failed")?; + + // Now finish B. The inner memory store drains it, refuses to store it + // (200 >= max_bytes) and fires remove callbacks for digest_b -- but the + // pause flag is already None, so the removal runs immediately as a no-op. + // B then inserts digest_b into the existence cache: stale entry. + tx_b.send(Bytes::from(vec![0u8; 200])) + .await + .err_tip(|| "send b")?; + tx_b.send_eof().err_tip(|| "eof b")?; + task_b + .await + .expect("join b") + .err_tip(|| "update B failed")?; + + assert_eq!( + inner_store.has(digest_b).await?, + None, + "precondition: inner store must not hold digest_b" + ); + assert_eq!( + store.has(digest_b).await.err_tip(|| "has b")?, + None, + "existence cache claims digest_b exists but the inner store evicted it \ + (pause window stolen by concurrent update)" + ); + Ok(()) +} + +/// Inner store standing in for a filesystem store: a read that is in-flight +/// keeps streaming successfully even if the entry is evicted mid-read +/// (filesystem reads hold the `FileEntry` Arc). The test controls exactly +/// where the read is suspended via `Notify` rendezvous points. +#[derive(Debug)] +struct PausingReadStore { + present: AtomicBool, + read_started: Notify, + allow_read_finish: Notify, + remove_callbacks: Mutex>, + data: Bytes, +} + +impl PausingReadStore { + fn new(data: Bytes) -> Arc { + Arc::new(Self { + present: AtomicBool::new(true), + read_started: Notify::new(), + allow_read_finish: Notify::new(), + remove_callbacks: Mutex::new(vec![]), + data, + }) + } + + /// Simulate the evicting map evicting `key`: drop residency and fire the + /// registered remove callbacks, exactly like `EvictingMap` does. + async fn evict(&self, key: StoreKey<'static>) { + self.present.store(false, Ordering::SeqCst); + let callbacks: Vec = self.remove_callbacks.lock().clone(); + for callback in callbacks { + callback.callback(key.borrow()).await; + } + } +} + +impl MetricsComponent for PausingReadStore { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } +} + +#[async_trait] +impl StoreDriver for PausingReadStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + _keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + let value = if self.present.load(Ordering::SeqCst) { + Some(self.data.len() as u64) + } else { + None + }; + for result in results.iter_mut() { + *result = value; + } + Ok(()) + } + + async fn update( + self: Pin<&Self>, + _key: StoreKey<'_>, + mut reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + let size = reader.drain().await?; + self.present.store(true, Ordering::SeqCst); + Ok(size) + } + + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + if !self.present.load(Ordering::SeqCst) { + return Err(make_err!(Code::NotFound, "Not found in PausingReadStore")); + } + // Rendezvous: tell the test the read is parked at a real await point, + // then wait until the test lets the (still valid) stream complete. + self.read_started.notify_one(); + self.allow_read_finish.notified().await; + writer.send(self.data.clone()).await?; + writer.send_eof()?; + Ok(()) + } + + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + + fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback(self: Arc, callback: RemoveCallback) -> Result<(), Error> { + self.remove_callbacks.lock().push(callback); + Ok(()) + } +} + +#[async_trait] +impl HealthStatusIndicator for PausingReadStore { + fn get_name(&self) -> &'static str { + "PausingReadStore" + } + + async fn check_health(&self, namespace: Cow<'static, str>) -> HealthStatus { + StoreDriver::check_health(Pin::new(self), namespace).await + } +} + +// Repro: get_part() re-inserts the existence entry after the inner read +// succeeds without participating in the pause_remove_callbacks protocol. +// If the inner store evicts the digest while the read is in flight (legal: +// filesystem reads hold the FileEntry Arc and stream to completion), the +// eviction's remove-callback purges the existence cache first, and get_part +// then resurrects a zombie exists-entry for data the inner store no longer +// has. A subsequent update() sees has()==true and silently discards the +// re-upload: data loss. +#[nativelink_test] +async fn get_part_races_inner_eviction_and_resurrects_zombie_entry() -> Result<(), Error> { + const VALUE: &str = "123"; + let inner = PausingReadStore::new(Bytes::from_static(VALUE.as_bytes())); + let store = ExistenceCacheStore::new_with_time( + &ExistenceCacheSpec { + backend: StoreSpec::Noop(NoopSpec::default()), + eviction_policy: Option::default(), + }, + Store::new(inner.clone()), + MockInstantWrapped::default(), + ); + let digest = DigestInfo::try_new(VALID_HASH1, VALUE.len() as u64)?; + + // Task A: client read. Parks inside inner get_part at a genuine await + // point (stand-in for the filesystem store's read/send loop). + let store_clone = store.clone(); + let read_task = + tokio::spawn(async move { store_clone.get_part_unchunked(digest, 0, None).await }); + inner.read_started.notified().await; + + // Task B: cache pressure evicts D from the inner store while the read is + // suspended. The remove callback purges the existence cache (correctly). + inner.evict(digest.into()).await; + assert!( + !store.exists_in_cache(&digest).await, + "Eviction callback should have purged the existence cache" + ); + + // Task A resumes; the in-flight read still completes Ok. + inner.allow_read_finish.notify_one(); + let read_result = read_task.await.expect("read task panicked")?; + assert_eq!(read_result, Bytes::from_static(VALUE.as_bytes())); + + // BUG: get_part re-inserted the existence entry after the eviction + // notification was already consumed. The inner store no longer has D. + assert!( + !store.exists_in_cache(&digest).await, + "Zombie exists-entry: existence cache claims digest is present but \ + the inner store evicted it" + ); + + // Consequence: a re-upload of D is drained and discarded by update(), + // so the data is lost while has() keeps reporting true. + store.update_oneshot(digest, VALUE.into()).await?; + let mut inner_result = [None]; + Pin::new(&*inner) + .has_with_results(&[digest.into()], &mut inner_result) + .await?; + assert_eq!( + inner_result[0], + Some(VALUE.len() as u64), + "Re-upload was silently discarded due to zombie existence entry" + ); + + Ok(()) +} From b947771c3c4d8e59cfded39423bc6b393a435e85 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 19:48:04 -0400 Subject: [PATCH 2/4] existence_cache: drop-safe pause (RAII guard) so a cancelled update can't wedge invalidation (N6) Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/existence_cache_store.rs | 99 +++++++--- .../tests/existence_store_test.rs | 174 ++++++++++++++++++ 2 files changed, 247 insertions(+), 26 deletions(-) diff --git a/nativelink-store/src/existence_cache_store.rs b/nativelink-store/src/existence_cache_store.rs index f9ba225630b..b04adc98e63 100644 --- a/nativelink-store/src/existence_cache_store.rs +++ b/nativelink-store/src/existence_cache_store.rs @@ -66,6 +66,59 @@ struct PauseState { pending: Vec>, } +impl PauseState { + /// Release one hold. Returns the queued removals to drain iff this was the + /// last holder; otherwise returns an empty vec and leaves the queue for a + /// later holder to drain. + fn release(&mut self) -> Vec> { + self.pause_count = self + .pause_count + .checked_sub(1) + .expect("PauseState::release called without a matching pause"); + if self.pause_count == 0 { + core::mem::take(&mut self.pending) + } else { + Vec::new() + } + } +} + +/// RAII hold on the remove-callback pause. +/// +/// Dropping the guard always decrements the pause count, so a cancelled +/// operation cannot leave invalidation globally suspended (finding N6). The +/// normal path calls [`PauseGuard::resume`] instead, which additionally drains +/// any queued removals once the last holder is gone. If the guard is dropped +/// while it is the last holder with a non-empty queue (only reachable on +/// cancellation), those removals stay queued and are drained by the next +/// operation to release the pause — the async drain cannot run from `Drop`. +struct PauseGuard<'a, I: InstantWrapper> { + store: Option<&'a ExistenceCacheStore>, +} + +impl<'a, I: InstantWrapper> PauseGuard<'a, I> { + /// Release the hold and drain queued removals if this was the last holder. + async fn resume(mut self) { + // Take the store so `Drop` becomes a no-op; we handle release here. + let Some(store) = self.store.take() else { + return; + }; + let keys = store.pause_remove_callbacks.lock().release(); + store.drain_pending_removals(keys).await; + } +} + +impl Drop for PauseGuard<'_, I> { + fn drop(&mut self) { + if let Some(store) = self.store.take() { + // Cancellation path: release the count so the pause mechanism is + // never left permanently engaged. Any queued removals remain for + // the next holder's `resume` to drain. + let _ = store.pause_remove_callbacks.lock().release(); + } + } +} + #[derive(Debug, MetricsComponent)] pub struct ExistenceCacheStore { #[metric(group = "inner_store")] @@ -163,30 +216,23 @@ impl ExistenceCacheStore { self.existence_cache.remove(digest).await; } - /// Begin deferring remove callbacks for the duration of an operation. - /// Must be paired with exactly one call to `resume_remove_callbacks`. - fn pause_remove_callbacks(&self) { + /// Begin deferring remove callbacks for the duration of an operation and + /// return an RAII guard that releases the hold when dropped. + /// + /// The guard's `Drop` decrements the pause count unconditionally, so a + /// cancelled operation (e.g. an `update`/`get_part` future dropped mid-await + /// by a timeout) can never leave the pause count permanently elevated and + /// disable invalidation globally. On the normal path, call + /// [`PauseGuard::resume`] to release the hold *and* drain any queued + /// removals once the last holder is gone. + fn pause_remove_callbacks(&self) -> PauseGuard<'_, I> { self.pause_remove_callbacks.lock().pause_count += 1; + PauseGuard { store: Some(self) } } - /// Release one hold on the pause. If this was the last concurrent - /// holder, drain the queued removals now — after every holder has - /// finished inserting into the existence cache — so evictions that - /// fired during the paused window take effect (including for keys - /// that were just inserted). - async fn resume_remove_callbacks(&self) { - let keys = { - let mut state = self.pause_remove_callbacks.lock(); - state.pause_count = state - .pause_count - .checked_sub(1) - .expect("resume_remove_callbacks called without matching pause"); - if state.pause_count == 0 { - core::mem::take(&mut state.pending) - } else { - Vec::new() - } - }; + /// Drain the queued removals when the last holder has released. Called by + /// [`PauseGuard::resume`]; the count has already been decremented there. + async fn drain_pending_removals(&self, keys: Vec>) { let mut callbacks: FuturesUnordered<_> = keys .into_iter() .map(|store_key| self.callback(store_key)) @@ -294,7 +340,7 @@ impl StoreDriver for ExistenceCacheStore { // wrong (the tier under it lost the blob and the remove callback did // not arrive), discarding the upload is what turns a stale answer // into a blob nothing holds. The entry is refreshed on success. - self.pause_remove_callbacks(); + let pause = self.pause_remove_callbacks(); trace!(?digest, "Inserting into inner cache"); let result = self.inner_store.update(digest, reader, size_info).await; if let Ok(size) = &result { @@ -306,8 +352,9 @@ impl StoreDriver for ExistenceCacheStore { } // Note: the insert above happens *before* releasing the pause, so a // removal for `digest` queued during the inner update is applied - // after the insert and cannot be lost. - self.resume_remove_callbacks().await; + // after the insert and cannot be lost. If this future is cancelled + // before here, `pause`'s Drop still releases the hold. + pause.resume().await; result } @@ -324,7 +371,7 @@ impl StoreDriver for ExistenceCacheStore { // filesystem store keeps streaming from an open file handle). The // eviction is queued and replayed after our insert, so the cache // cannot resurrect an entry the inner store dropped mid-read. - self.pause_remove_callbacks(); + let pause = self.pause_remove_callbacks(); let result = self .inner_store .get_part(digest, writer, offset, length) @@ -335,7 +382,7 @@ impl StoreDriver for ExistenceCacheStore { .insert(digest, ExistenceItem(digest.size_bytes())) .await; } - self.resume_remove_callbacks().await; + pause.resume().await; result } diff --git a/nativelink-store/tests/existence_store_test.rs b/nativelink-store/tests/existence_store_test.rs index 2a6eda54d94..05e7fbd27ae 100644 --- a/nativelink-store/tests/existence_store_test.rs +++ b/nativelink-store/tests/existence_store_test.rs @@ -841,3 +841,177 @@ async fn get_part_races_inner_eviction_and_resurrects_zombie_entry() -> Result<( Ok(()) } + +/// Inner store whose `update` parks at a real await point until released, so a +/// test can drop the outer `update` future while the pause is held. `evict` +/// fires the registered remove callbacks like `EvictingMap` does. +#[derive(Debug)] +struct BlockingUpdateStore { + present: AtomicBool, + update_started: Notify, + allow_update_finish: Notify, + remove_callbacks: Mutex>, + len: u64, +} + +impl BlockingUpdateStore { + fn new(len: u64) -> Arc { + Arc::new(Self { + present: AtomicBool::new(true), + update_started: Notify::new(), + allow_update_finish: Notify::new(), + remove_callbacks: Mutex::new(vec![]), + len, + }) + } + + async fn evict(&self, key: StoreKey<'static>) { + self.present.store(false, Ordering::SeqCst); + let callbacks: Vec = self.remove_callbacks.lock().clone(); + for callback in callbacks { + callback.callback(key.borrow()).await; + } + } +} + +impl MetricsComponent for BlockingUpdateStore { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } +} + +#[async_trait] +impl StoreDriver for BlockingUpdateStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + _keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + let value = if self.present.load(Ordering::SeqCst) { + Some(self.len) + } else { + None + }; + for result in results.iter_mut() { + *result = value; + } + Ok(()) + } + + async fn update( + self: Pin<&Self>, + _key: StoreKey<'_>, + mut reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + // Park at a genuine await point with the pause held, then wait for + // the test. The test drops the future before ever releasing this. + self.update_started.notify_one(); + self.allow_update_finish.notified().await; + let size = reader.drain().await?; + self.present.store(true, Ordering::SeqCst); + Ok(size) + } + + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + _writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + Err(make_err!(Code::NotFound, "not exercised")) + } + + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + + fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback(self: Arc, callback: RemoveCallback) -> Result<(), Error> { + self.remove_callbacks.lock().push(callback); + Ok(()) + } +} + +#[async_trait] +impl HealthStatusIndicator for BlockingUpdateStore { + fn get_name(&self) -> &'static str { + "BlockingUpdateStore" + } + + async fn check_health(&self, namespace: Cow<'static, str>) -> HealthStatus { + StoreDriver::check_health(Pin::new(self), namespace).await + } +} + +/// Regression for N6: a cancelled `update` must not leave remove callbacks +/// paused forever. Start an `update` that parks with the pause held, drop its +/// future (as a timeout would), then evict a cached key through the inner +/// store. With the RAII pause guard the drop releases the hold, so the +/// eviction callback runs immediately and the existence cache stops reporting +/// the key. Without the guard the pause count stays elevated and the cache +/// keeps claiming the evicted key exists. +#[nativelink_test] +async fn cancelled_update_does_not_leave_callbacks_paused() -> Result<(), Error> { + const LEN: u64 = 2; + let inner = BlockingUpdateStore::new(LEN); + let store = ExistenceCacheStore::new_with_time( + &ExistenceCacheSpec { + backend: StoreSpec::Noop(NoopSpec::default()), + eviction_policy: Option::default(), + }, + Store::new(inner.clone()), + MockInstantWrapped::default(), + ); + + let digest_a = DigestInfo::try_new(VALID_HASH1, LEN)?; + let digest_b = DigestInfo::try_new(VALID_HASH2, LEN)?; + + // Populate A's existence entry via a has() that the inner store answers + // positively (present == true). + assert_eq!(store.has(digest_a).await?, Some(LEN)); + assert!(store.exists_in_cache(&digest_a).await); + + // Start an update of B that parks inside the inner store holding the pause. + let store_b = store.clone(); + let mut update_b = Box::pin(async move { + store_b + .update_oneshot(digest_b, vec![0u8; LEN as usize].into()) + .await + }); + // Drive it until it parks at the inner update await, then cancel it by + // dropping the future. + tokio::select! { + _ = &mut update_b => panic!("update should be parked, not complete"), + () = inner.update_started.notified() => {} + } + drop(update_b); + + // Evict A through the inner store; the remove callback must take effect + // now that the cancelled update released its pause. + inner.evict(digest_a.into()).await; + + assert!( + !store.exists_in_cache(&digest_a).await, + "existence cache still claims A exists: a cancelled update left \ + remove callbacks permanently paused (N6)" + ); + + Ok(()) +} From 5f272943c223b2c984d2fe53847e5e695aa89ce9 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 19:59:13 -0400 Subject: [PATCH 3/4] existence_cache: fix N6 test harness + clippy (elide lifetime, per-key presence) The regression test short-circuited: BlockingUpdateStore reported a single global `present` flag, so update(B)'s inner_has_with_results saw B as present and returned early, never reaching the parking inner update. Track presence per-key (HashSet) so A is present (populates its existence entry) while B is absent (update proceeds to the parking await). Elide the needless explicit lifetime on PauseGuard, replace let-underscore drop with drop(), and use usize::try_from for the reader length. Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/existence_cache_store.rs | 4 +- .../tests/existence_store_test.rs | 44 ++++++++++--------- 2 files changed, 26 insertions(+), 22 deletions(-) diff --git a/nativelink-store/src/existence_cache_store.rs b/nativelink-store/src/existence_cache_store.rs index b04adc98e63..076ab4155fe 100644 --- a/nativelink-store/src/existence_cache_store.rs +++ b/nativelink-store/src/existence_cache_store.rs @@ -96,7 +96,7 @@ struct PauseGuard<'a, I: InstantWrapper> { store: Option<&'a ExistenceCacheStore>, } -impl<'a, I: InstantWrapper> PauseGuard<'a, I> { +impl PauseGuard<'_, I> { /// Release the hold and drain queued removals if this was the last holder. async fn resume(mut self) { // Take the store so `Drop` becomes a no-op; we handle release here. @@ -114,7 +114,7 @@ impl Drop for PauseGuard<'_, I> { // Cancellation path: release the count so the pause mechanism is // never left permanently engaged. Any queued removals remain for // the next holder's `resume` to drain. - let _ = store.pause_remove_callbacks.lock().release(); + drop(store.pause_remove_callbacks.lock().release()); } } } diff --git a/nativelink-store/tests/existence_store_test.rs b/nativelink-store/tests/existence_store_test.rs index 05e7fbd27ae..94149f106bc 100644 --- a/nativelink-store/tests/existence_store_test.rs +++ b/nativelink-store/tests/existence_store_test.rs @@ -16,6 +16,7 @@ use core::pin::Pin; use core::sync::atomic::{AtomicBool, Ordering}; use core::time::Duration; use std::borrow::Cow; +use std::collections::HashSet; use std::sync::Arc; use async_trait::async_trait; @@ -847,7 +848,7 @@ async fn get_part_races_inner_eviction_and_resurrects_zombie_entry() -> Result<( /// fires the registered remove callbacks like `EvictingMap` does. #[derive(Debug)] struct BlockingUpdateStore { - present: AtomicBool, + present: Mutex>>, update_started: Notify, allow_update_finish: Notify, remove_callbacks: Mutex>, @@ -855,9 +856,9 @@ struct BlockingUpdateStore { } impl BlockingUpdateStore { - fn new(len: u64) -> Arc { + fn new(len: u64, present: &[StoreKey<'static>]) -> Arc { Arc::new(Self { - present: AtomicBool::new(true), + present: Mutex::new(present.iter().cloned().collect()), update_started: Notify::new(), allow_update_finish: Notify::new(), remove_callbacks: Mutex::new(vec![]), @@ -866,7 +867,7 @@ impl BlockingUpdateStore { } async fn evict(&self, key: StoreKey<'static>) { - self.present.store(false, Ordering::SeqCst); + self.present.lock().remove(&key); let callbacks: Vec = self.remove_callbacks.lock().clone(); for callback in callbacks { callback.callback(key.borrow()).await; @@ -892,32 +893,33 @@ impl StoreDriver for BlockingUpdateStore { async fn has_with_results( self: Pin<&Self>, - _keys: &[StoreKey<'_>], + keys: &[StoreKey<'_>], results: &mut [Option], ) -> Result<(), Error> { - let value = if self.present.load(Ordering::SeqCst) { - Some(self.len) - } else { - None - }; - for result in results.iter_mut() { - *result = value; + let present = self.present.lock(); + for (key, result) in keys.iter().zip(results.iter_mut()) { + *result = if present.contains(&key.borrow().into_owned()) { + Some(self.len) + } else { + None + }; } Ok(()) } async fn update( self: Pin<&Self>, - _key: StoreKey<'_>, + key: StoreKey<'_>, mut reader: DropCloserReadHalf, _size_info: UploadSizeInfo, ) -> Result { // Park at a genuine await point with the pause held, then wait for // the test. The test drops the future before ever releasing this. + let key = key.into_owned(); self.update_started.notify_one(); self.allow_update_finish.notified().await; let size = reader.drain().await?; - self.present.store(true, Ordering::SeqCst); + self.present.lock().insert(key); Ok(size) } @@ -970,7 +972,12 @@ impl HealthStatusIndicator for BlockingUpdateStore { #[nativelink_test] async fn cancelled_update_does_not_leave_callbacks_paused() -> Result<(), Error> { const LEN: u64 = 2; - let inner = BlockingUpdateStore::new(LEN); + let digest_a = DigestInfo::try_new(VALID_HASH1, LEN)?; + let digest_b = DigestInfo::try_new(VALID_HASH2, LEN)?; + + // Only A is present in the inner store; B is absent so its update proceeds + // to the parking inner `update` (rather than short-circuiting on `has`). + let inner = BlockingUpdateStore::new(LEN, &[digest_a.into()]); let store = ExistenceCacheStore::new_with_time( &ExistenceCacheSpec { backend: StoreSpec::Noop(NoopSpec::default()), @@ -980,11 +987,8 @@ async fn cancelled_update_does_not_leave_callbacks_paused() -> Result<(), Error> MockInstantWrapped::default(), ); - let digest_a = DigestInfo::try_new(VALID_HASH1, LEN)?; - let digest_b = DigestInfo::try_new(VALID_HASH2, LEN)?; - // Populate A's existence entry via a has() that the inner store answers - // positively (present == true). + // positively (A is present). assert_eq!(store.has(digest_a).await?, Some(LEN)); assert!(store.exists_in_cache(&digest_a).await); @@ -992,7 +996,7 @@ async fn cancelled_update_does_not_leave_callbacks_paused() -> Result<(), Error> let store_b = store.clone(); let mut update_b = Box::pin(async move { store_b - .update_oneshot(digest_b, vec![0u8; LEN as usize].into()) + .update_oneshot(digest_b, vec![0u8; usize::try_from(LEN).unwrap()].into()) .await }); // Drive it until it parks at the inner update await, then cancel it by From a74c1c200ae6478d07795743c0974c9d51e609bb Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 20:36:58 -0400 Subject: [PATCH 4/4] existence_cache: drain removals queued before a cancelled last holder (C20/N6 drop residual) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The RAII PauseGuard's Drop decremented the pause count (fixing N6's permanently-wedged pause) but discarded the pending-removal vector that `release()` returns when the last holder goes away. On the evict-before-cancel ordering — evict cached A while upload B holds the sole pause, then cancel B — A's queued invalidation was dropped on the floor and the existence cache kept reporting A forever even though the inner store had evicted it. Drop now drains those queued removals. Since Drop cannot await and the removal is async, it spawns a detached background task (via a weak self-reference) to run drain_pending_removals. New regression `cancelled_update_drains_removals_queued_before_cancel` uses the evict-*before*-cancel ordering and fails on the prior Drop (removals discarded); the existing N6 test `cancelled_update_does_not_leave_callbacks_paused` (evict-after-cancel) is kept. Also closes the C13 has-path residual: inner_has_with_results now pauses remove callbacks around the inner `has` + the positive-result insert, so an eviction delivered mid-`has` lands after the insert instead of being a no-op before it. Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/existence_cache_store.rs | 65 +++++++++++--- .../tests/existence_store_test.rs | 85 +++++++++++++++++++ 2 files changed, 139 insertions(+), 11 deletions(-) diff --git a/nativelink-store/src/existence_cache_store.rs b/nativelink-store/src/existence_cache_store.rs index 076ab4155fe..17ed8b374fc 100644 --- a/nativelink-store/src/existence_cache_store.rs +++ b/nativelink-store/src/existence_cache_store.rs @@ -23,6 +23,7 @@ use futures::stream::FuturesUnordered; use nativelink_config::stores::{EvictionPolicy, ExistenceCacheSpec}; use nativelink_error::{Error, ResultExt, error_if}; use nativelink_metric::MetricsComponent; +use nativelink_util::background_spawn; use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use nativelink_util::common::DigestInfo; use nativelink_util::evicting_map::{EvictingMap, LenEntry}; @@ -87,11 +88,12 @@ impl PauseState { /// /// Dropping the guard always decrements the pause count, so a cancelled /// operation cannot leave invalidation globally suspended (finding N6). The -/// normal path calls [`PauseGuard::resume`] instead, which additionally drains -/// any queued removals once the last holder is gone. If the guard is dropped +/// normal path calls [`PauseGuard::resume`] instead. If the guard is dropped /// while it is the last holder with a non-empty queue (only reachable on -/// cancellation), those removals stay queued and are drained by the next -/// operation to release the pause — the async drain cannot run from `Drop`. +/// cancellation), those removals must still be applied — otherwise the cache +/// keeps claiming a blob exists after the inner store evicted it (C20/N6 drop +/// residual). Since `Drop` cannot await and the removal is async, `Drop` spawns +/// a detached task (via the weak self-reference) to drain them. struct PauseGuard<'a, I: InstantWrapper> { store: Option<&'a ExistenceCacheStore>, } @@ -110,11 +112,27 @@ impl PauseGuard<'_, I> { impl Drop for PauseGuard<'_, I> { fn drop(&mut self) { - if let Some(store) = self.store.take() { - // Cancellation path: release the count so the pause mechanism is - // never left permanently engaged. Any queued removals remain for - // the next holder's `resume` to drain. - drop(store.pause_remove_callbacks.lock().release()); + let Some(store) = self.store.take() else { + return; + }; + // Cancellation path: release the count so the pause mechanism is never + // left permanently engaged (finding N6). + let keys = store.pause_remove_callbacks.lock().release(); + if keys.is_empty() { + return; + } + // We were the last holder and there are removals that were queued + // *before* this operation was cancelled (C20/N6 drop residual). These + // must still be applied, otherwise the cache can keep claiming a blob + // exists after the inner store evicted it. `Drop` cannot await, and the + // removal is async, so spawn a detached task that drains them. The + // weak self-reference keeps this `'static` without extending the + // store's lifetime past its last strong owner. + let weak = store.self_ref.lock().clone(); + if let Some(store) = weak.upgrade() { + background_spawn!("existence_cache_drain_pending_removals", async move { + store.drain_pending_removals(keys).await; + }); } } } @@ -128,6 +146,12 @@ pub struct ExistenceCacheStore { // We need to pause remove callbacks temporarily while operating on the // inner store; see `PauseState`. pause_remove_callbacks: Mutex, + + // A weak self-reference, set once at construction. The cancellation path in + // `PauseGuard::drop` needs an owned (`'static`) handle to the store so it + // can spawn the async drain of any removals queued before cancellation; a + // borrowed `&self` cannot outlive the `Drop` call. See `PauseGuard::drop`. + self_ref: Mutex>, } impl ExistenceCacheStore { @@ -195,7 +219,9 @@ impl ExistenceCacheStore { inner_store, existence_cache: EvictingMap::new(eviction_policy, anchor_time), pause_remove_callbacks: Mutex::new(PauseState::default()), + self_ref: Mutex::new(Weak::new()), }); + *existence_cache_store.self_ref.lock() = Arc::downgrade(&existence_cache_store); let other_ref = Arc::downgrade(&existence_cache_store); existence_cache_store .inner_store @@ -260,12 +286,25 @@ impl ExistenceCacheStore { return Ok(()); } + // Pause remove callbacks over the inner `has` + the insert below + // (finding C13). The inner store may evict one of `not_cached_keys` + // between answering `has` positively and our insert; without the pause + // that eviction's callback would run before the insert and be a no-op, + // leaving the cache claiming existence of a blob the inner store + // dropped. Queuing it until after the insert lets the drain remove it. + let pause = self.pause_remove_callbacks(); + // Now query only the items not found in the cache. let mut inner_results = vec![None; not_cached_keys.len()]; - self.inner_store + let has_result = self + .inner_store .has_with_results(¬_cached_keys, &mut inner_results) .await - .err_tip(|| "In ExistenceCacheStore::inner_has_with_results")?; + .err_tip(|| "In ExistenceCacheStore::inner_has_with_results"); + if let Err(err) = has_result { + pause.resume().await; + return Err(err); + } // Insert found from previous query into our cache. { @@ -281,6 +320,10 @@ impl ExistenceCacheStore { drop(self.existence_cache.insert_many(inserts).await); } + // Insert is done; release the pause and drain any eviction queued + // during the inner `has` so it lands after our insert (C13). + pause.resume().await; + // Merge the results from the cache and the query. { let mut inner_results_iter = inner_results.into_iter(); diff --git a/nativelink-store/tests/existence_store_test.rs b/nativelink-store/tests/existence_store_test.rs index 94149f106bc..c1b65f41ae1 100644 --- a/nativelink-store/tests/existence_store_test.rs +++ b/nativelink-store/tests/existence_store_test.rs @@ -1019,3 +1019,88 @@ async fn cancelled_update_does_not_leave_callbacks_paused() -> Result<(), Error> Ok(()) } + +/// Regression for the C20/N6 drop residual: a removal queued *before* the last +/// pause holder is cancelled must still be applied. This is the +/// evict-*before*-cancel ordering, which the N6 test above does not cover (it +/// evicts after the drop, so there is nothing queued to lose). +/// +/// Schedule: +/// 1. Populate A's positive existence entry via `has()`. +/// 2. Start an `update` of B that parks in the inner store holding the sole +/// pause guard. +/// 3. Evict A through the inner store *while the pause is held*: the remove +/// callback is queued in `pending` instead of applied, so the outer cache +/// still reports A. +/// 4. Drop (cancel) B's future. The last holder releases the pause. +/// +/// With the buggy `Drop` that discards the vector returned by `release()`, the +/// queued removal of A is lost and the cache keeps claiming A exists forever. +/// A correct `Drop` must drain those queued removals, so A disappears from the +/// cache after cancellation. +#[nativelink_test] +async fn cancelled_update_drains_removals_queued_before_cancel() -> Result<(), Error> { + const LEN: u64 = 2; + let digest_a = DigestInfo::try_new(VALID_HASH1, LEN)?; + let digest_b = DigestInfo::try_new(VALID_HASH2, LEN)?; + + // Only A is present in the inner store; B is absent so its update proceeds + // to the parking inner `update` (rather than short-circuiting on `has`). + let inner = BlockingUpdateStore::new(LEN, &[digest_a.into()]); + let store = ExistenceCacheStore::new_with_time( + &ExistenceCacheSpec { + backend: StoreSpec::Noop(NoopSpec::default()), + eviction_policy: Option::default(), + }, + Store::new(inner.clone()), + MockInstantWrapped::default(), + ); + + // Populate A's existence entry via a has() the inner store answers + // positively (A is present). + assert_eq!(store.has(digest_a).await?, Some(LEN)); + assert!(store.exists_in_cache(&digest_a).await); + + // Start an update of B that parks inside the inner store holding the pause. + let store_b = store.clone(); + let mut update_b = Box::pin(async move { + store_b + .update_oneshot(digest_b, vec![0u8; usize::try_from(LEN).unwrap()].into()) + .await + }); + tokio::select! { + _ = &mut update_b => panic!("update should be parked, not complete"), + () = inner.update_started.notified() => {} + } + + // Evict A through the inner store *while the pause is held*. The removal is + // queued in `pending`; the outer cache still reports A for now. + inner.evict(digest_a.into()).await; + assert!( + store.exists_in_cache(&digest_a).await, + "removal should be queued (not applied) while the pause is held" + ); + + // Cancel B by dropping its future. The last holder releases the pause; the + // queued removal of A must still be drained. `Drop` cannot await, so the + // drain is spawned; give it a bounded number of scheduler turns to run. + drop(update_b); + + let mut drained = false; + for _ in 0..1000 { + if !store.exists_in_cache(&digest_a).await { + drained = true; + break; + } + tokio::task::yield_now().await; + } + + assert!( + drained, + "existence cache still claims A exists: a removal queued before the \ + last pause holder was cancelled was discarded instead of drained \ + (C20/N6 drop residual)" + ); + + Ok(()) +}