Add metrics to compression and existence cache store (#651)

TraceMachina · Feb 2, 2024 · 722c80b · 722c80b
1 parent 81b64f7
commit 722c80b
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/nativelink-store/src/compression_store.rs b/nativelink-store/src/compression_store.rs
@@ -26,6 +26,7 @@ use lz4_flex::block::{compress_into, decompress_into, get_maximum_output_size};
 use nativelink_error::{error_if, make_err, Code, Error, ResultExt};
 use nativelink_util::buf_channel::{make_buf_channel_pair, DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::{DigestInfo, JoinHandleDropGuard};
+use nativelink_util::metrics_utils::Registry;
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 use serde::{Deserialize, Serialize};
 
@@ -583,4 +584,9 @@ impl Store for CompressionStore {
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }
+
+    fn register_metrics(self: Arc<Self>, registry: &mut Registry) {
+        let inner_store_registry = registry.sub_registry_with_prefix("inner_store");
+        self.inner_store.clone().register_metrics(inner_store_registry);
+    }
 }
diff --git a/nativelink-store/src/existence_cache_store.rs b/nativelink-store/src/existence_cache_store.rs
@@ -23,6 +23,7 @@ use nativelink_error::{error_if, Error, ResultExt};
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
 use nativelink_util::evicting_map::{EvictingMap, LenEntry};
+use nativelink_util::metrics_utils::{CollectorState, MetricsComponent, Registry};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 
 #[derive(Clone, Debug)]
@@ -185,4 +186,15 @@ impl Store for ExistenceCacheStore {
     fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
         Box::new(self)
     }
+
+    fn register_metrics(self: Arc<Self>, registry: &mut Registry) {
+        let inner_store_registry = registry.sub_registry_with_prefix("inner_store");
+        self.inner_store.clone().register_metrics(inner_store_registry);
+    }
+}
+
+impl MetricsComponent for ExistenceCacheStore {
+    fn gather_metrics(&self, c: &mut CollectorState) {
+        self.existence_cache.gather_metrics(c)
+    }
 }