Existence cache remove callbacks

.bazelrc

@@ -104,12 +104,14 @@ build --@rules_rust//:clippy_flag=-Dclippy::alloc_instead_of_core
 build --@rules_rust//:clippy_flag=-Dclippy::as_underscore
 build --@rules_rust//:clippy_flag=-Wclippy::dbg_macro
 build --@rules_rust//:clippy_flag=-Wclippy::decimal_literal_representation
+build --@rules_rust//:clippy_flag=-Dclippy::elidable_lifetime_names
 build --@rules_rust//:clippy_flag=-Aclippy::get_unwrap
 build --@rules_rust//:clippy_flag=-Dclippy::missing_const_for_fn
 build --@rules_rust//:clippy_flag=-Aclippy::missing_docs_in_private_items
 build --@rules_rust//:clippy_flag=-Wclippy::print_stdout
+build --@rules_rust//:clippy_flag=-Dclippy::semicolon_if_nothing_returned
 build --@rules_rust//:clippy_flag=-Dclippy::std_instead_of_core
-build --@rules_rust//:clippy_flag=-Wclippy::todo
+build --@rules_rust//:clippy_flag=-Dclippy::todo
 build --@rules_rust//:clippy_flag=-Aclippy::too_long_first_doc_paragraph
 build --@rules_rust//:clippy_flag=-Wclippy::unimplemented
 build --@rules_rust//:clippy_flag=-Aclippy::unwrap_in_result
@@ -154,10 +156,13 @@ build:nightly --@rules_rust//rust/toolchain/channel=nightly
 # AddressSanitizer to detect memory leaks.
 build:asan --config=nightly
 build:asan --@rules_rust//:extra_rustc_flags=-Zsanitizer=address
+# asan explicitly needs to _not_ use the Nix toolchains and so this
+# is needed to load the rust toolchains in Bazel
+build:asan --extra_toolchains=@rust_toolchains//:all
 
 # ThreadSanitizer to detect data races. Tests under tsan shouldn't be cached
 # since they tend to produce irreproducible false negatives.
-build:tsan --config=nightly -c dbg
+build:tsan --config=nightly --config=debug
 build:tsan --@rules_rust//:extra_rustc_flags=-Zsanitizer=thread
 build:tsan --cache_test_results=no
 build:tsan --run_under=//tools:tsan.sh

.github/workflows/sanitizers.yaml

@@ -34,6 +34,14 @@ jobs:
         uses: >- # v4.2.2
           actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
 
+      - name: Free disk space
+        uses: >- # v2.0.0
+          endersonmenezes/free-disk-space@3f9ec39ebae520864ac93467ee395f5237585c21
+        with:
+          remove_android: false # Takes too long.
+          remove_dotnet: true
+          remove_haskell: true
+
       - name: Setup Bazel
         uses: >- # v0.13.0
           bazel-contrib/setup-bazel@663f88d97adf17db2523a5b385d9407a562e5551

Cargo.toml

@@ -136,15 +136,17 @@ pedantic = { level = "warn", priority = -1 }
 # Restriction Denies with default priority
 alloc-instead-of-core = "deny"
 as-underscore = "deny"
+elidable-lifetime-names = "deny"
+semicolon-if-nothing-returned = "deny"
 std-instead-of-core = "deny"
+todo = "deny"
 
 # Restriction Warnings with default priority
 dbg-macro = "warn"
 decimal-literal-representation = "warn"
 get-unwrap = "allow"                    # TODO(jhpratt) Flip
 missing-docs-in-private-items = "allow" # TODO(jhpratt) Flip
 print-stdout = "warn"
-todo = "warn"
 unimplemented = "warn"
 unwrap-in-result = "allow"              # TODO(jhpratt) Flip
 unwrap-used = "allow"                   # TODO(jhpratt) Flip

nativelink-error/src/lib.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use core::convert::Into;
+use std::sync::{MutexGuard, PoisonError};
 
 use nativelink_metric::{
     MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent,
@@ -187,6 +188,12 @@ impl From<tokio::task::JoinError> for Error {
     }
 }
 
+impl<T> From<PoisonError<MutexGuard<'_, T>>> for Error {
+    fn from(err: PoisonError<MutexGuard<'_, T>>) -> Self {
+        make_err!(Code::Internal, "{}", err.to_string())
+    }
+}
+
 impl From<serde_json5::Error> for Error {
     fn from(err: serde_json5::Error) -> Self {
         match err {

nativelink-scheduler/src/memory_awaited_action_db.rs

@@ -303,7 +303,8 @@ impl SortedAwaitedActions {
 pub struct AwaitedActionDbImpl<I: InstantWrapper, NowFn: Fn() -> I> {
     /// A lookup table to lookup the state of an action by its client operation id.
     #[metric(group = "client_operation_ids")]
-    client_operation_to_awaited_action: EvictingMap<OperationId, Arc<ClientAwaitedAction>, I>,
+    client_operation_to_awaited_action:
+        EvictingMap<OperationId, OperationId, Arc<ClientAwaitedAction>, I>,
 
     /// A lookup table to lookup the state of an action by its worker operation id.
     #[metric(group = "operation_ids")]

nativelink-store/BUILD.bazel

@@ -11,6 +11,7 @@ rust_library(
     name = "nativelink-store",
     srcs = [
         "src/ac_utils.rs",
+        "src/callback_utils.rs",
         "src/cas_utils.rs",
         "src/common_s3_utils.rs",
         "src/completeness_checking_store.rs",

nativelink-store/Cargo.toml

@@ -74,7 +74,7 @@ mongodb = { version = "3", features = [
   "rustls-tls",
 ], default-features = false }
 opentelemetry = { version = "0.29.1", default-features = false }
-parking_lot = "0.12.3"
+parking_lot = { version = "0.12.3", features = ["arc_lock", "send_guard"] }
 patricia_tree = { version = "0.9.0", default-features = false }
 prost = { version = "0.13.5", default-features = false }
 rand = { version = "0.9.0", default-features = false, features = [

nativelink-store/src/ac_utils.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 // @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
-// TODO(palfrey): IMPORTANT TODO: IMPORTING THIS SOMETMIES BREAKS
+// TODO(palfrey): IMPORTANT TODO: IMPORTING THIS SOMETIMES BREAKS
 //                    THREADSAFETY. FIGURE OUT WHY AND MOVE IT TO UTILS.
 // @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
 

nativelink-store/src/callback_utils.rs

@@ -0,0 +1,41 @@
+// Copyright 2025 The NativeLink Authors. All rights reserved.
+//
+// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    See LICENSE file for details
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use nativelink_util::evicting_map::RemoveStateCallback;
+use nativelink_util::store_trait::{RemoveItemCallback, StoreKey};
+use tonic::async_trait;
+
+// Generic struct to hold a RemoveItemCallback ref for the purposes
+// of a RemoveStateCallback call
+#[derive(Debug)]
+pub struct RemoveItemCallbackHolder {
+    callback_fn: Arc<Box<dyn RemoveItemCallback>>,
+}
+
+impl RemoveItemCallbackHolder {
+    pub fn new(callback: &Arc<Box<dyn RemoveItemCallback>>) -> Self {
+        Self {
+            callback_fn: callback.clone(),
+        }
+    }
+}
+
+#[async_trait]
+impl RemoveStateCallback<StoreKey<'static>> for RemoveItemCallbackHolder {
+    async fn callback(&self, key: &StoreKey<'static>) {
+        self.callback_fn.callback(key).await;
+    }
+}

nativelink-store/src/completeness_checking_store.rs

@@ -28,7 +28,9 @@ 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::metrics_utils::CounterWithTime;
-use nativelink_util::store_trait::{Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo};
+use nativelink_util::store_trait::{
+    RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
+};
 use parking_lot::Mutex;
 use tokio::sync::Notify;
 use tracing::warn;
@@ -387,6 +389,15 @@ impl StoreDriver for CompletenessCheckingStore {
     fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
         self
     }
+
+    fn register_remove_callback(
+        self: Arc<Self>,
+        callback: &Arc<Box<dyn RemoveItemCallback>>,
+    ) -> Result<(), Error> {
+        self.ac_store.register_remove_callback(callback)?;
+        self.cas_store.register_remove_callback(callback)?;
+        Ok(())
+    }
 }
 
 default_health_status_indicator!(CompletenessCheckingStore);

nativelink-store/src/compression_store.rs

@@ -30,7 +30,9 @@ use nativelink_util::buf_channel::{
 };
 use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator};
 use nativelink_util::spawn;
-use nativelink_util::store_trait::{Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo};
+use nativelink_util::store_trait::{
+    RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
+};
 use serde::{Deserialize, Serialize};
 
 use crate::cas_utils::is_zero_digest;
@@ -639,6 +641,13 @@ impl StoreDriver for CompressionStore {
     fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
         self
     }
+
+    fn register_remove_callback(
+        self: Arc<Self>,
+        callback: &Arc<Box<dyn RemoveItemCallback>>,
+    ) -> Result<(), Error> {
+        self.inner_store.register_remove_callback(callback)
+    }
 }
 
 default_health_status_indicator!(CompressionStore);

nativelink-store/src/dedup_store.rs

@@ -26,7 +26,9 @@ use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
 use nativelink_util::fastcdc::FastCDC;
 use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator};
-use nativelink_util::store_trait::{Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo};
+use nativelink_util::store_trait::{
+    RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
+};
 use serde::{Deserialize, Serialize};
 use tokio_util::codec::FramedRead;
 use tokio_util::io::StreamReader;
@@ -376,6 +378,15 @@ impl StoreDriver for DedupStore {
     fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
         self
     }
+
+    fn register_remove_callback(
+        self: Arc<Self>,
+        callback: &Arc<Box<dyn RemoveItemCallback>>,
+    ) -> Result<(), Error> {
+        self.index_store.register_remove_callback(callback)?;
+        self.content_store.register_remove_callback(callback)?;
+        Ok(())
+    }
 }
 
 default_health_status_indicator!(DedupStore);

nativelink-store/src/existence_cache_store.rs

@@ -14,7 +14,7 @@
 
 use core::pin::Pin;
 use std::borrow::Cow;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 use std::time::SystemTime;
 
 use async_trait::async_trait;
@@ -26,7 +26,11 @@ use nativelink_util::common::DigestInfo;
 use nativelink_util::evicting_map::{EvictingMap, LenEntry};
 use nativelink_util::health_utils::{HealthStatus, HealthStatusIndicator};
 use nativelink_util::instant_wrapper::InstantWrapper;
-use nativelink_util::store_trait::{Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo};
+use nativelink_util::store_trait::{
+    RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
+};
+use parking_lot::Mutex;
+use tracing::{debug, info, trace};
 
 #[derive(Clone, Debug)]
 struct ExistenceItem(u64);
@@ -47,7 +51,13 @@ impl LenEntry for ExistenceItem {
 pub struct ExistenceCacheStore<I: InstantWrapper> {
     #[metric(group = "inner_store")]
     inner_store: Store,
-    existence_cache: EvictingMap<DigestInfo, ExistenceItem, I>,
+    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: Arc<Mutex<Option<Vec<StoreKey<'static>>>>>,
 }
 
 impl ExistenceCacheStore<SystemTime> {
@@ -56,6 +66,39 @@ impl ExistenceCacheStore<SystemTime> {
     }
 }
 
+#[async_trait]
+impl<I: InstantWrapper> RemoveItemCallback for ExistenceCacheStore<I> {
+    async fn callback(&self, store_key: &StoreKey<'_>) {
+        debug!(?store_key, "Removing item from cache due to callback");
+        let new_key = store_key.borrow();
+        let deleted_key = self.existence_cache.remove(&new_key.into_digest()).await;
+        if !deleted_key {
+            info!(?store_key, "Failed to delete key from cache on callback");
+        }
+    }
+}
+
+#[derive(Debug)]
+struct ExistenceCacheCallback<I: InstantWrapper> {
+    cache: Weak<ExistenceCacheStore<I>>,
+}
+
+#[async_trait]
+impl<I: InstantWrapper> RemoveItemCallback for ExistenceCacheCallback<I> {
+    async fn callback(&self, store_key: &StoreKey<'_>) {
+        let cache = self.cache.upgrade();
+        if let Some(local_cache) = cache {
+            if let Some(callbacks) = &mut *local_cache.pause_remove_callbacks.lock_arc() {
+                callbacks.push(store_key.borrow().into_owned());
+            } else {
+                local_cache.callback(store_key).await;
+            }
+        } else {
+            debug!("Cache dropped, so not doing callback");
+        }
+    }
+}
+
 impl<I: InstantWrapper> ExistenceCacheStore<I> {
     pub fn new_with_time(
         spec: &ExistenceCacheSpec,
@@ -64,10 +107,19 @@ impl<I: InstantWrapper> ExistenceCacheStore<I> {
     ) -> Arc<Self> {
         let empty_policy = EvictionPolicy::default();
         let eviction_policy = spec.eviction_policy.as_ref().unwrap_or(&empty_policy);
-        Arc::new(Self {
+        let existence_cache_store = Arc::new(Self {
             inner_store,
             existence_cache: EvictingMap::new(eviction_policy, anchor_time),
-        })
+            pause_remove_callbacks: Arc::new(Mutex::new(None)),
+        });
+        let other_ref = Arc::downgrade(&existence_cache_store);
+        existence_cache_store
+            .inner_store
+            .register_remove_callback(&Arc::new(Box::new(ExistenceCacheCallback {
+                cache: other_ref,
+            })))
+            .expect("Register remove callback should work");
+        existence_cache_store
     }
 
     pub async fn exists_in_cache(&self, digest: &DigestInfo) -> bool {
@@ -184,15 +236,31 @@ impl<I: InstantWrapper> StoreDriver for ExistenceCacheStore<I> {
                 .err_tip(|| "In ExistenceCacheStore::update")?;
             return Ok(());
         }
+        {
+            let mut locked_callbacks = self.pause_remove_callbacks.lock_arc();
+            if locked_callbacks.is_none() {
+                locked_callbacks.replace(vec![]);
+            }
+        }
+        trace!(?digest, "Inserting into inner cache");
         let result = self.inner_store.update(digest, reader, size_info).await;
         if result.is_ok() {
+            trace!(?digest, "Inserting into existence cache");
             if let UploadSizeInfo::ExactSize(size) = size_info {
                 let _ = self
                     .existence_cache
                     .insert(digest, ExistenceItem(size))
                     .await;
             }
         }
+        {
+            let mut locked_callbacks = self.pause_remove_callbacks.lock_arc();
+            if let Some(callbacks) = locked_callbacks.take() {
+                for store_key in callbacks {
+                    self.callback(&store_key).await;
+                }
+            }
+        }
         result
     }
 
@@ -228,6 +296,13 @@ impl<I: InstantWrapper> StoreDriver for ExistenceCacheStore<I> {
     fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
         self
     }
+
+    fn register_remove_callback(
+        self: Arc<Self>,
+        callback: &Arc<Box<dyn RemoveItemCallback>>,
+    ) -> Result<(), Error> {
+        self.inner_store.register_remove_callback(callback)
+    }
 }
 
 #[async_trait]