Introduce store_size metric.

Author: dkostyrev

nativelink-store/src/filesystem_store.rs

@@ -725,6 +725,10 @@ impl<Fe: FileEntry> FilesystemStore<Fe> {
             .ok_or_else(|| make_err!(Code::NotFound, "{digest} not found in filesystem store. This may indicate the file was evicted due to cache pressure. Consider increasing 'max_bytes' in your filesystem store's eviction_policy configuration."))
     }
 
+    pub fn get_len(&self) -> u64 {
+        self.evicting_map.len()
+    }
+
     async fn update_file(
         self: Pin<&Self>,
         mut entry: Fe,

nativelink-store/src/metrics_store.rs

@@ -4,9 +4,10 @@ use crate::memory_store::MemoryStore;
 use async_trait::async_trait;
 use nativelink_error::Error;
 use nativelink_metric::MetricsComponent;
+use nativelink_proto::build_event_stream::File;
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::health_utils::{HealthStatus, HealthStatusIndicator};
-use nativelink_util::metrics::{STORE_METRICS, StoreMetricAttrs, StoreType};
+use nativelink_util::metrics::{StoreMetricAttrs, StoreType, STORE_METRICS};
 use nativelink_util::store_trait::{
     RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
 };
@@ -25,19 +26,26 @@ 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()) {
+        if let Some(fs_store) = inner.downcast_ref::<FilesystemStore>(None) {
             #[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>> {
+                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() })) {
+            if let Err(e) = inner.register_remove_callback(Arc::new(EvictionCallback {
+                attrs: attrs.clone(),
+            })) {
                 tracing::error!("Failed to register remove callback: {:?}", e);
             }
+
+            STORE_METRICS.store_size.record(fs_store.get_len(), &attrs.store_size());
         }
 
         Arc::new(Self {
@@ -103,6 +111,10 @@ impl StoreDriver for MetricsStore {
                 .record(duration_ms as f64, &self.attrs.write_error());
         }
 
+        if let Some(fs_store) = self.inner.downcast_ref::<FilesystemStore>(None) {
+            STORE_METRICS.store_size.record(fs_store.get_len(), &self.attrs.store_size());
+        }
+
         result
     }
 
@@ -168,5 +180,4 @@ impl HealthStatusIndicator for MetricsStore {
 
 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/evicting_map.rs

@@ -108,6 +108,8 @@ struct State<
     #[metric(help = "Total size of all items in the store")]
     sum_store_size: u64,
 
+    store_len: u64,
+
     #[metric(help = "Number of bytes evicted from the store")]
     evicted_bytes: Counter,
     #[metric(help = "Number of items evicted from the store")]
@@ -148,6 +150,7 @@ impl<
             btree.remove(key);
         }
         self.sum_store_size -= eviction_item.data.len();
+        self.store_len -= 1;
         if replaced {
             self.replaced_items.inc();
             self.replaced_bytes.add(eviction_item.data.len());
@@ -234,6 +237,7 @@ where
                 lru: LruCache::unbounded(),
                 btree: None,
                 sum_store_size: 0,
+                store_len: 0,
                 evicted_bytes: Counter::default(),
                 evicted_items: CounterWithTime::default(),
                 replaced_bytes: Counter::default(),
@@ -546,6 +550,10 @@ where
             .await
     }
 
+    pub fn len(&self) -> u64 {
+        self.state.lock().store_len
+    }
+
     fn inner_insert_many<It>(
         &self,
         state: &mut State<K, Q, T, C>,
@@ -572,6 +580,7 @@ where
                 replaced_items.push(old_item);
             }
             state.sum_store_size += new_item_size;
+            state.store_len += 1;
             state.lifetime_inserted_bytes.add(new_item_size);
         }
 

nativelink-util/src/metrics.rs

@@ -1643,6 +1643,11 @@ pub static STORE_METRICS: LazyLock<StoreMetrics> = LazyLock::new(|| {
             .u64_counter("eviction_count")
             .with_description("Number of evictions")
             .build(),
+
+        store_size: meter
+            .u64_gauge("store_size")
+            .with_description("Number of items in the store")
+            .build(),
     }
 });
 
@@ -1654,6 +1659,8 @@ pub struct StoreMetrics {
     pub store_operations: metrics::Counter<u64>,
     /// Counter of evictions
     pub eviction_count: metrics::Counter<u64>,
+    /// Counter of items in the store
+    pub store_size: metrics::Gauge<u64>,
 }
 
 #[derive(Debug, Clone)]
@@ -1666,6 +1673,8 @@ pub struct StoreMetricAttrs {
     write_success: Vec<KeyValue>,
     write_error: Vec<KeyValue>,
     eviction: Vec<KeyValue>,
+    store_size: Vec<KeyValue>,
+
 }
 
 impl StoreMetricAttrs {
@@ -1675,11 +1684,12 @@ impl StoreMetricAttrs {
     /// type, instance ID).
     #[must_use]
     pub fn new_with_name(store_type: StoreType, name: &str) -> Self {
+        let base_attrs = vec![
+            KeyValue::new(STORE_TYPE, store_type.to_string()),
+            KeyValue::new(STORE_NAME, name.to_string()),
+        ];
         let make_attrs = |op: CacheOperationName, result: CacheOperationResult| {
-            let mut attrs = vec![
-                KeyValue::new(STORE_TYPE, store_type.to_string()),
-                KeyValue::new(STORE_NAME, name.to_string()),
-            ];
+            let mut attrs = base_attrs.clone();
             attrs.push(KeyValue::new(CACHE_OPERATION, op));
             attrs.push(KeyValue::new(CACHE_RESULT, result));
             attrs
@@ -1694,6 +1704,8 @@ impl StoreMetricAttrs {
             write_success: make_attrs(CacheOperationName::Write, CacheOperationResult::Success),
             write_error: make_attrs(CacheOperationName::Write, CacheOperationResult::Error),
             eviction: make_attrs(CacheOperationName::Evict, CacheOperationResult::Success),
+            store_size: base_attrs.clone(),
+
         }
     }
 
@@ -1726,4 +1738,8 @@ impl StoreMetricAttrs {
     pub fn eviction(&self) -> &[KeyValue] {
         &self.eviction
     }
+    #[must_use]
+    pub fn store_size(&self) -> &[KeyValue] {
+        &self.store_size
+    }
 }