[Breaking] Change in behavior of /status by introduction of component…

… based health

Introduce a component based health check system. Each type of component
should be able to opt into registering handlers implement some mechanical
checks of health. Health in this context is functionality expected to
work but runtime wise are semi no-op in terms of influencing the underlying
storage / rpc systems.

Opting in to the system requires for component to define a `HealthStatusIndicator`
and implement the `check_health()` function. Registration should automatically
be done by the existence of implementation and calling the `Store.register_health()`
function in component implementation. At the moment only `Store` based components
can register and a single health check is defined for `FilesystemStore`.

A global parameter `default_digest_size_health_check` has been introduced for
configuring the static seeded random bytes to fill data payload. The default
value is 1MB.

The `/status` endpoint has been updated to return the resulting string serialized
instances of `HealthStatus`, smart serialization of such objects is not implemented
at this time.

nativelink-config/src/cas_server.rs

@@ -629,6 +629,13 @@ 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.
+    ///
+    /// 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);