Introduce eviction count.

Author: dkostyrev

nativelink-store/src/default_store_factory.rs

@@ -24,7 +24,6 @@ use nativelink_error::Error;
 use nativelink_util::health_utils::HealthRegistryBuilder;
 use nativelink_util::metrics::StoreType;
 use nativelink_util::store_trait::{Store, StoreDriver};
-
 use crate::completeness_checking_store::CompletenessCheckingStore;
 use crate::compression_store::CompressionStore;
 use crate::dedup_store::DedupStore;
@@ -147,6 +146,7 @@ fn should_wrap_in_metrics_store(spec: &StoreSpec) -> bool {
     matches!(
         spec,
         StoreSpec::Memory(_)
+            | StoreSpec::Grpc(_)
             | StoreSpec::ExperimentalCloudObjectStore(_)
             | StoreSpec::ExperimentalMongo(_)
             | StoreSpec::Filesystem(_)

nativelink-store/src/metrics_store.rs

@@ -1,9 +1,12 @@
+use crate::callback_utils::RemoveItemCallbackHolder;
+use crate::filesystem_store::FilesystemStore;
+use crate::memory_store::MemoryStore;
 use async_trait::async_trait;
 use nativelink_error::Error;
 use nativelink_metric::MetricsComponent;
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::health_utils::{HealthStatus, HealthStatusIndicator};
-use nativelink_util::metrics::{StoreMetricAttrs, StoreType, STORE_METRICS};
+use nativelink_util::metrics::{STORE_METRICS, StoreMetricAttrs, StoreType};
 use nativelink_util::store_trait::{
     RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
 };
@@ -15,15 +18,31 @@ use std::time::Instant;
 #[derive(MetricsComponent, Debug)]
 pub struct MetricsStore {
     inner: Arc<Store>,
-    attrs: StoreMetricAttrs,
+    attrs: Arc<StoreMetricAttrs>,
 }
 
 impl MetricsStore {
     #[must_use]
     pub fn new(inner: Arc<Store>, name: &str, store_type: StoreType) -> Arc<Self> {
+        let attrs = Arc::new(StoreMetricAttrs::new_with_name(store_type, name));
+        if should_add_remove_callback(inner.clone()) {
+            #[derive(Debug)]
+            struct EvictionCallback {
+                attrs: Arc<StoreMetricAttrs>,
+            }
+            impl RemoveItemCallback for EvictionCallback {
+                fn callback<'a>(&'a self, store_key: StoreKey<'a>) -> Pin<Box<dyn Future<Output=()> + Send + 'a>> {
+                    Box::pin(async { STORE_METRICS.eviction_count.add(1, self.attrs.eviction()) })
+                }
+            }
+            if let Err(e) = inner.register_remove_callback(Arc::new(EvictionCallback { attrs: attrs.clone() })) {
+                tracing::error!("Failed to register remove callback: {:?}", e);
+            }
+        }
+
         Arc::new(Self {
             inner: inner.clone(),
-            attrs: StoreMetricAttrs::new_with_name(store_type, name),
+            attrs: attrs.clone(),
         })
     }
 }
@@ -76,7 +95,9 @@ impl StoreDriver for MetricsStore {
                 .store_operation_duration
                 .record(duration_ms as f64, &self.attrs.write_success());
         } else {
-            STORE_METRICS.store_operations.add(1, &self.attrs.write_error());
+            STORE_METRICS
+                .store_operations
+                .add(1, &self.attrs.write_error());
             STORE_METRICS
                 .store_operation_duration
                 .record(duration_ms as f64, &self.attrs.write_error());
@@ -144,3 +165,8 @@ impl HealthStatusIndicator for MetricsStore {
         self.inner.check_health(_namespace).await
     }
 }
+
+fn should_add_remove_callback(store: Arc<Store>) -> bool {
+    store.downcast_ref::<FilesystemStore>(None).is_some()
+        || store.downcast_ref::<MemoryStore>(None).is_some()
+}

nativelink-util/src/metrics.rs

@@ -1656,6 +1656,11 @@ pub static STORE_METRICS: LazyLock<StoreMetrics> = LazyLock::new(|| {
                 5000.0, // 5 seconds
             ])
             .build(),
+
+        eviction_count: meter
+            .u64_counter("eviction_count")
+            .with_description("Number of evictions")
+            .build(),
     }
 });
 
@@ -1665,6 +1670,8 @@ pub struct StoreMetrics {
     pub store_operation_duration: metrics::Histogram<f64>,
     /// Counter of store operations by type and result
     pub store_operations: metrics::Counter<u64>,
+    /// Counter of evictions
+    pub eviction_count: metrics::Counter<u64>,
 }
 
 #[derive(Debug, Clone)]
@@ -1676,6 +1683,7 @@ pub struct StoreMetricAttrs {
     read_error: Vec<KeyValue>,
     write_success: Vec<KeyValue>,
     write_error: Vec<KeyValue>,
+    eviction: Vec<KeyValue>,
 }
 
 impl StoreMetricAttrs {
@@ -1703,6 +1711,7 @@ impl StoreMetricAttrs {
             read_error: make_attrs(CacheOperationName::Read, CacheOperationResult::Error),
             write_success: make_attrs(CacheOperationName::Write, CacheOperationResult::Success),
             write_error: make_attrs(CacheOperationName::Write, CacheOperationResult::Error),
+            eviction: make_attrs(CacheOperationName::Evict, CacheOperationResult::Success),
         }
     }
 
@@ -1731,4 +1740,8 @@ impl StoreMetricAttrs {
     pub fn write_error(&self) -> &[KeyValue] {
         &self.write_error
     }
+    #[must_use]
+    pub fn eviction(&self) -> &[KeyValue] {
+        &self.eviction
+    }
 }