Add Grpc, Memory & S3 store to health checker registry

Adds these stores to registry to be checked by the health checker if enabled.
TraceMachina · Feb 13, 2025 · a3556af · a3556af
1 parent 79c2357
commit a3556af
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs
@@ -40,7 +40,7 @@ use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
 use nativelink_util::connection_manager::ConnectionManager;
 use nativelink_util::digest_hasher::{default_digest_hasher_func, ACTIVE_HASHER_FUNC};
-use nativelink_util::health_utils::HealthStatusIndicator;
+use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatusIndicator};
 use nativelink_util::origin_context::ActiveOriginContext;
 use nativelink_util::proto_stream_utils::{
     FirstStream, WriteRequestStreamWrapper, WriteState, WriteStateWrapper,
@@ -765,6 +765,10 @@ impl StoreDriver for GrpcStore {
     fn as_any_arc(self: Arc<Self>) -> Arc<dyn std::any::Any + Sync + Send + 'static> {
         self
     }
+
+    fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) {
+        registry.register_indicator(self);
+    }
 }
 
 default_health_status_indicator!(GrpcStore);
diff --git a/nativelink-store/src/memory_store.rs b/nativelink-store/src/memory_store.rs
@@ -26,7 +26,7 @@ use nativelink_error::{Code, Error, ResultExt};
 use nativelink_metric::MetricsComponent;
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::evicting_map::{EvictingMap, LenEntry};
-use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator};
+use nativelink_util::health_utils::{default_health_status_indicator, HealthRegistryBuilder, HealthStatusIndicator};
 use nativelink_util::store_trait::{StoreDriver, StoreKey, StoreKeyBorrow, UploadSizeInfo};
 
 use crate::cas_utils::is_zero_digest;
@@ -194,6 +194,10 @@ impl StoreDriver for MemoryStore {
     fn as_any_arc(self: Arc<Self>) -> Arc<dyn std::any::Any + Sync + Send + 'static> {
         self
     }
+
+    fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) {
+        registry.register_indicator(self);
+    }
 }
 
 default_health_status_indicator!(MemoryStore);
diff --git a/nativelink-store/src/s3_store.rs b/nativelink-store/src/s3_store.rs
@@ -51,7 +51,7 @@ use nativelink_util::buf_channel::{
     make_buf_channel_pair, DropCloserReadHalf, DropCloserWriteHalf,
 };
 use nativelink_util::fs;
-use nativelink_util::health_utils::{HealthStatus, HealthStatusIndicator};
+use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatus, HealthStatusIndicator};
 use nativelink_util::instant_wrapper::InstantWrapper;
 use nativelink_util::retry::{Retrier, RetryResult};
 use nativelink_util::store_trait::{StoreDriver, StoreKey, UploadSizeInfo};
@@ -780,6 +780,10 @@ where
     fn as_any_arc(self: Arc<Self>) -> Arc<dyn std::any::Any + Sync + Send + 'static> {
         self
     }
+
+    fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) {
+        registry.register_indicator(self);
+    }
 }
 
 #[async_trait]