existence_cache: refcount pause_remove_callbacks; fix insert-after-evict resurrection

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};
@@ -49,17 +50,108 @@ 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<StoreKey<'static>>,
+}
+
+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<StoreKey<'static>> {
+        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. If the guard is dropped
+/// while it is the last holder with a non-empty queue (only reachable on
+/// 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<I>>,
+}
+
+impl<I: InstantWrapper> 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.
+        let Some(store) = self.store.take() else {
+            return;
+        };
+        let keys = store.pause_remove_callbacks.lock().release();
+        store.drain_pending_removals(keys).await;
+    }
+}
+
+impl<I: InstantWrapper> Drop for PauseGuard<'_, I> {
+    fn drop(&mut self) {
+        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;
+            });
+        }
+    }
+}
+
 #[derive(Debug, MetricsComponent)]
 pub struct ExistenceCacheStore<I: InstantWrapper> {
     #[metric(group = "inner_store")]
     inner_store: Store,
     existence_cache: EvictingMap<DigestInfo, DigestInfo, ExistenceItem, I>,
 
-    // 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<Option<Vec<StoreKey<'static>>>>,
+    // We need to pause remove callbacks temporarily while operating on the
+    // inner store; see `PauseState`.
+    pause_remove_callbacks: Mutex<PauseState>,
+
+    // 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<Weak<Self>>,
 }
 
 impl ExistenceCacheStore<SystemTime> {
@@ -98,9 +190,11 @@ impl<I: InstantWrapper> RemoveItemCallback for ExistenceCacheCallback<I> {
     ) -> Pin<Box<dyn Future<Output = ()> + 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,8 +218,10 @@ impl<I: InstantWrapper> ExistenceCacheStore<I> {
         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()),
+            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
@@ -146,6 +242,30 @@ impl<I: InstantWrapper> ExistenceCacheStore<I> {
         self.existence_cache.remove(digest).await;
     }
 
+    /// 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) }
+    }
+
+    /// 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<StoreKey<'static>>) {
+        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],
@@ -166,12 +286,25 @@ impl<I: InstantWrapper> ExistenceCacheStore<I> {
             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(&not_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.
         {
@@ -187,6 +320,10 @@ impl<I: InstantWrapper> ExistenceCacheStore<I> {
             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();
@@ -246,12 +383,7 @@ impl<I: InstantWrapper> StoreDriver for ExistenceCacheStore<I> {
         // 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![]);
-            }
-        }
+        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 {
@@ -261,16 +393,11 @@ impl<I: InstantWrapper> StoreDriver for ExistenceCacheStore<I> {
                 .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. If this future is cancelled
+        // before here, `pause`'s Drop still releases the hold.
+        pause.resume().await;
         result
     }
 
@@ -282,6 +409,12 @@ impl<I: InstantWrapper> StoreDriver for ExistenceCacheStore<I> {
         length: Option<u64>,
     ) -> 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.
+        let pause = self.pause_remove_callbacks();
         let result = self
             .inner_store
             .get_part(digest, writer, offset, length)
@@ -292,6 +425,7 @@ impl<I: InstantWrapper> StoreDriver for ExistenceCacheStore<I> {
                 .insert(digest, ExistenceItem(digest.size_bytes()))
                 .await;
         }
+        pause.resume().await;
         result
     }