Introduce component based health

nativelink-config/src/cas_server.rs

@@ -629,6 +629,15 @@ pub struct GlobalConfig {
     ///
     /// Default: ConfigDigestHashFunction::sha256
     pub default_digest_hash_function: Option<ConfigDigestHashFunction>,
+
+    /// Default digest size to use for health check when running
+    /// diagnostics checks. Health checks are expected to use this
+    /// size for filling a buffer that is used for creation of
+    /// digest.
+    ///
+    /// Default: 1024*1024 (1MiB)
+    #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")]
+    pub default_digest_size_health_check: usize,
 }
 
 #[derive(Deserialize, Debug)]

nativelink-service/tests/ac_server_test.rs

@@ -55,6 +55,7 @@ async fn make_store_manager() -> Result<Arc<StoreManager>, Error> {
             &nativelink_config::stores::StoreConfig::memory(nativelink_config::stores::MemoryStore::default()),
             &store_manager,
             Some(&mut <Registry>::default()),
+            None,
         )
         .await?,
     );
@@ -64,6 +65,7 @@ async fn make_store_manager() -> Result<Arc<StoreManager>, Error> {
             &nativelink_config::stores::StoreConfig::memory(nativelink_config::stores::MemoryStore::default()),
             &store_manager,
             Some(&mut <Registry>::default()),
+            None,
         )
         .await?,
     );

nativelink-service/tests/bytestream_server_test.rs

@@ -41,6 +41,7 @@ async fn make_store_manager() -> Result<Arc<StoreManager>, Error> {
             &nativelink_config::stores::StoreConfig::memory(nativelink_config::stores::MemoryStore::default()),
             &store_manager,
             Some(&mut <Registry>::default()),
+            None,
         )
         .await?,
     );

nativelink-service/tests/cas_server_test.rs

@@ -41,6 +41,7 @@ async fn make_store_manager() -> Result<Arc<StoreManager>, Error> {
             &nativelink_config::stores::StoreConfig::memory(nativelink_config::stores::MemoryStore::default()),
             &store_manager,
             Some(&mut <Registry>::default()),
+            None,
         )
         .await?,
     );

nativelink-store/BUILD.bazel

@@ -58,6 +58,7 @@ rust_library(
         "@crates//:serde",
         "@crates//:sha2",
         "@crates//:shellexpand",
+        "@crates//:tempfile",
         "@crates//:tokio",
         "@crates//:tokio-stream",
         "@crates//:tokio-util",

nativelink-store/Cargo.toml

@@ -30,6 +30,7 @@ rand = "0.8.5"
 serde = "1.0.193"
 sha2 = "0.10.8"
 shellexpand = "3.1.0"
+tempfile = "3.9.0"
 tokio = { version = "1.35.1" }
 tokio-stream = { version = "0.1.14", features = ["fs"] }
 tokio-util = { version = "0.7.10" }

nativelink-store/src/completeness_checking_store.rs

@@ -25,6 +25,7 @@ use nativelink_proto::build::bazel::remote::execution::v2::{
 };
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
+use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 use parking_lot::Mutex;
 use tokio::sync::Notify;
@@ -361,3 +362,5 @@ impl Store for CompletenessCheckingStore {
         Box::new(self)
     }
 }
+
+default_health_status_indicator!(CompletenessCheckingStore);

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::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::metrics_utils::Registry;
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 use serde::{Deserialize, Serialize};
@@ -590,3 +591,5 @@ impl Store for CompressionStore {
         self.inner_store.clone().register_metrics(inner_store_registry);
     }
 }
+
+default_health_status_indicator!(CompressionStore);

nativelink-store/src/dedup_store.rs

@@ -24,6 +24,7 @@ use nativelink_error::{make_err, Code, Error, ResultExt};
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf, StreamReader};
 use nativelink_util::common::DigestInfo;
 use nativelink_util::fastcdc::FastCDC;
+use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 use serde::{Deserialize, Serialize};
 use tokio_util::codec::FramedRead;
@@ -349,3 +350,5 @@ impl Store for DedupStore {
         Box::new(self)
     }
 }
+
+default_health_status_indicator!(DedupStore);

nativelink-store/src/default_store_factory.rs

@@ -19,6 +19,7 @@ use futures::stream::FuturesOrdered;
 use futures::{Future, TryStreamExt};
 use nativelink_config::stores::StoreConfig;
 use nativelink_error::Error;
+use nativelink_util::health_utils::HealthRegistryBuilder;
 use nativelink_util::metrics_utils::Registry;
 use nativelink_util::store_trait::Store;
 
@@ -44,51 +45,52 @@ pub fn store_factory<'a>(
     backend: &'a StoreConfig,
     store_manager: &'a Arc<StoreManager>,
     maybe_store_metrics: Option<&'a mut Registry>,
+    maybe_health_registry_builder: Option<&'a mut HealthRegistryBuilder>,
 ) -> Pin<FutureMaybeStore<'a>> {
     Box::pin(async move {
         let store: Arc<dyn Store> = match backend {
             StoreConfig::memory(config) => Arc::new(MemoryStore::new(config)),
             StoreConfig::experimental_s3_store(config) => Arc::new(S3Store::new(config).await?),
             StoreConfig::verify(config) => Arc::new(VerifyStore::new(
                 config,
-                store_factory(&config.backend, store_manager, None).await?,
+                store_factory(&config.backend, store_manager, None, None).await?,
             )),
             StoreConfig::compression(config) => Arc::new(CompressionStore::new(
                 *config.clone(),
-                store_factory(&config.backend, store_manager, None).await?,
+                store_factory(&config.backend, store_manager, None, None).await?,
             )?),
             StoreConfig::dedup(config) => Arc::new(DedupStore::new(
                 config,
-                store_factory(&config.index_store, store_manager, None).await?,
-                store_factory(&config.content_store, store_manager, None).await?,
+                store_factory(&config.index_store, store_manager, None, None).await?,
+                store_factory(&config.content_store, store_manager, None, None).await?,
             )),
             StoreConfig::existence_cache(config) => Arc::new(ExistenceCacheStore::new(
                 config,
-                store_factory(&config.backend, store_manager, None).await?,
+                store_factory(&config.backend, store_manager, None, None).await?,
             )),
             StoreConfig::completeness_checking(config) => Arc::new(CompletenessCheckingStore::new(
-                store_factory(&config.backend, store_manager, None).await?,
-                store_factory(&config.cas_store, store_manager, None).await?,
+                store_factory(&config.backend, store_manager, None, None).await?,
+                store_factory(&config.cas_store, store_manager, None, None).await?,
             )),
             StoreConfig::fast_slow(config) => Arc::new(FastSlowStore::new(
                 config,
-                store_factory(&config.fast, store_manager, None).await?,
-                store_factory(&config.slow, store_manager, None).await?,
+                store_factory(&config.fast, store_manager, None, None).await?,
+                store_factory(&config.slow, store_manager, None, None).await?,
             )),
             StoreConfig::filesystem(config) => Arc::new(<FilesystemStore>::new(config).await?),
             StoreConfig::ref_store(config) => Arc::new(RefStore::new(config, Arc::downgrade(store_manager))),
             StoreConfig::size_partitioning(config) => Arc::new(SizePartitioningStore::new(
                 config,
-                store_factory(&config.lower_store, store_manager, None).await?,
-                store_factory(&config.upper_store, store_manager, None).await?,
+                store_factory(&config.lower_store, store_manager, None, None).await?,
+                store_factory(&config.upper_store, store_manager, None, None).await?,
             )),
             StoreConfig::grpc(config) => Arc::new(GrpcStore::new(config).await?),
             StoreConfig::noop => Arc::new(NoopStore::new()),
             StoreConfig::shard(config) => {
                 let stores = config
                     .stores
                     .iter()
-                    .map(|store_config| store_factory(&store_config.store, store_manager, None))
+                    .map(|store_config| store_factory(&store_config.store, store_manager, None, None))
                     .collect::<FuturesOrdered<_>>()
                     .try_collect::<Vec<_>>()
                     .await?;
@@ -98,6 +100,11 @@ pub fn store_factory<'a>(
         if let Some(store_metrics) = maybe_store_metrics {
             store.clone().register_metrics(store_metrics);
         }
+
+        if let Some(health_registry_builder) = maybe_health_registry_builder {
+            store.clone().register_health(health_registry_builder);
+        }
+
         Ok(store)
     })
 }

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::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::metrics_utils::{CollectorState, MetricsComponent, Registry};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 
@@ -198,3 +199,5 @@ impl MetricsComponent for ExistenceCacheStore {
         self.existence_cache.gather_metrics(c)
     }
 }
+
+default_health_status_indicator!(ExistenceCacheStore);

nativelink-store/src/fast_slow_store.rs

@@ -22,6 +22,7 @@ use futures::{join, FutureExt};
 use nativelink_error::{make_err, Code, Error, ResultExt};
 use nativelink_util::buf_channel::{make_buf_channel_pair, DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
+use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::metrics_utils::Registry;
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 
@@ -268,3 +269,5 @@ impl Store for FastSlowStore {
         self.slow_store.clone().register_metrics(slow_store_registry);
     }
 }
+
+default_health_status_indicator!(FastSlowStore);

nativelink-store/src/filesystem_store.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::borrow::Cow;
 use std::ffi::OsString;
 use std::fmt::{Debug, Formatter};
 use std::pin::Pin;
@@ -29,6 +30,7 @@ use nativelink_error::{make_err, make_input_err, Code, Error, ResultExt};
 use nativelink_util::buf_channel::{make_buf_channel_pair, DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::{fs, DigestInfo};
 use nativelink_util::evicting_map::{EvictingMap, LenEntry};
+use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatus, HealthStatusIndicator};
 use nativelink_util::metrics_utils::{Collector, CollectorState, MetricsComponent, Registry};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt, SeekFrom};
@@ -750,6 +752,10 @@ impl<Fe: FileEntry> Store for FilesystemStore<Fe> {
     fn register_metrics(self: Arc<Self>, registry: &mut Registry) {
         registry.register_collector(Box::new(Collector::new(&self)));
     }
+
+    fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) {
+        registry.register_indicator(self);
+    }
 }
 
 impl<Fe: FileEntry> MetricsComponent for FilesystemStore<Fe> {
@@ -777,3 +783,14 @@ impl<Fe: FileEntry> MetricsComponent for FilesystemStore<Fe> {
         c.publish("evicting_map", self.evicting_map.as_ref(), "");
     }
 }
+
+#[async_trait]
+impl<Fe: FileEntry> HealthStatusIndicator for FilesystemStore<Fe> {
+    fn get_name(&self) -> &'static str {
+        "FilesystemStore"
+    }
+
+    async fn check_health(&self, namespace: Cow<'static, str>) -> HealthStatus {
+        Store::check_health(Pin::new(self), namespace).await
+    }
+}

nativelink-store/src/grpc_store.rs

@@ -37,11 +37,12 @@ use nativelink_proto::google::bytestream::{
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
 use nativelink_util::grpc_utils::ConnectionManager;
+use nativelink_util::health_utils::HealthStatusIndicator;
 use nativelink_util::resource_info::ResourceInfo;
 use nativelink_util::retry::{Retrier, RetryResult};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
-use nativelink_util::tls_utils;
 use nativelink_util::write_request_stream_wrapper::WriteRequestStreamWrapper;
+use nativelink_util::{default_health_status_indicator, tls_utils};
 use parking_lot::Mutex;
 use prost::Message;
 use rand::rngs::OsRng;
@@ -844,3 +845,5 @@ impl Store for GrpcStore {
         Box::new(self)
     }
 }
+
+default_health_status_indicator!(GrpcStore);

nativelink-store/src/memory_store.rs

@@ -23,6 +23,7 @@ use nativelink_error::{Code, Error, ResultExt};
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
 use nativelink_util::evicting_map::{EvictingMap, LenEntry};
+use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::metrics_utils::{Collector, CollectorState, MetricsComponent, Registry};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 
@@ -170,3 +171,5 @@ impl MetricsComponent for MemoryStore {
         c.publish("evicting_map", &self.evicting_map, "");
     }
 }
+
+default_health_status_indicator!(MemoryStore);

nativelink-store/src/noop_store.rs

@@ -19,6 +19,7 @@ use async_trait::async_trait;
 use nativelink_error::{make_err, Code, Error, ResultExt};
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
+use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 
 #[derive(Default)]
@@ -71,3 +72,5 @@ impl Store for NoopStore {
         Box::new(self)
     }
 }
+
+default_health_status_indicator!(NoopStore);

nativelink-store/src/ref_store.rs

@@ -20,6 +20,7 @@ use async_trait::async_trait;
 use nativelink_error::{make_err, make_input_err, Code, Error, ResultExt};
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
+use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 use tracing::error;
 
@@ -137,3 +138,5 @@ impl Store for RefStore {
         Box::new(self)
     }
 }
+
+default_health_status_indicator!(RefStore);

nativelink-store/src/s3_store.rs

@@ -40,6 +40,7 @@ use hyper_rustls::{HttpsConnector, MaybeHttpsStream};
 use nativelink_error::{error_if, make_err, make_input_err, Code, Error, ResultExt};
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
+use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::retry::{Retrier, RetryResult};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 use rand::rngs::OsRng;
@@ -534,3 +535,5 @@ impl Store for S3Store {
         Box::new(self)
     }
 }
+
+default_health_status_indicator!(S3Store);

nativelink-store/src/shard_store.rs

@@ -21,6 +21,7 @@ use futures::stream::{FuturesUnordered, TryStreamExt};
 use nativelink_error::{error_if, Error, ResultExt};
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
+use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::metrics_utils::Registry;
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 
@@ -191,3 +192,5 @@ impl Store for ShardStore {
         }
     }
 }
+
+default_health_status_indicator!(ShardStore);

nativelink-store/src/size_partitioning_store.rs

@@ -19,6 +19,7 @@ use async_trait::async_trait;
 use nativelink_error::{Error, ResultExt};
 use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
+use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 use tokio::join;
 
@@ -128,3 +129,5 @@ impl Store for SizePartitioningStore {
         Box::new(self)
     }
 }
+
+default_health_status_indicator!(SizePartitioningStore);

nativelink-store/src/verify_store.rs

@@ -22,6 +22,7 @@ use nativelink_error::{make_input_err, Error, ResultExt};
 use nativelink_util::buf_channel::{make_buf_channel_pair, DropCloserReadHalf, DropCloserWriteHalf};
 use nativelink_util::common::DigestInfo;
 use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc};
+use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator};
 use nativelink_util::metrics_utils::{Collector, CollectorState, CounterWithTime, MetricsComponent, Registry};
 use nativelink_util::store_trait::{Store, UploadSizeInfo};
 
@@ -199,3 +200,5 @@ impl MetricsComponent for VerifyStore {
         );
     }
 }
+
+default_health_status_indicator!(VerifyStore);