GH-768 Store config uniqueness enforced by StoreManager

nativelink-config/src/stores.rs

@@ -432,6 +432,21 @@ pub enum StoreSpec {
     noop(NoopSpec),
 }
 
+impl StoreSpec {
+    // To enforce no duplicate connection configs for a store, add it to the matcher and implement
+    // disallow_duplicates_digest() on it. Returns `None` for stores that are not being enforced unique.
+    pub fn disallow_duplicates_digest(&self) -> Option<String> {
+        match self {
+            Self::memory(spec) => Some(spec.disallow_duplicates_digest()),
+            Self::experimental_s3_store(spec) => Some(spec.disallow_duplicates_digest()),
+            Self::filesystem(spec) => Some(spec.disallow_duplicates_digest()),
+            Self::grpc(spec) => Some(spec.disallow_duplicates_digest()),
+            Self::redis_store(spec) => Some(spec.disallow_duplicates_digest()),
+            _ => None,
+        }
+    }
+}
+
 /// Configuration for an individual shard of the store.
 #[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(deny_unknown_fields)]
@@ -514,6 +529,12 @@ pub struct FilesystemSpec {
     pub block_size: u64,
 }
 
+impl FilesystemSpec {
+    fn disallow_duplicates_digest(&self) -> String {
+        format!("{}{}", self.content_path, self.temp_path)
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(deny_unknown_fields)]
 pub struct FastSlowSpec {
@@ -535,6 +556,12 @@ pub struct MemorySpec {
     pub eviction_policy: Option<EvictionPolicy>,
 }
 
+impl MemorySpec {
+    pub fn disallow_duplicates_digest(&self) -> String {
+        "InMemoryStore".into()
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(deny_unknown_fields)]
 pub struct DedupSpec {
@@ -787,6 +814,14 @@ pub struct S3Spec {
     pub disable_http2: bool,
 }
 
+impl S3Spec {
+    pub fn disallow_duplicates_digest(&self) -> String {
+        let key_prefix = self.key_prefix.as_deref().unwrap_or_default();
+
+        format!("{}{}{}", self.region, self.bucket, key_prefix)
+    }
+}
+
 #[allow(non_camel_case_types)]
 #[derive(Serialize, Deserialize, Debug, Clone, Copy)]
 pub enum StoreType {
@@ -852,6 +887,22 @@ pub struct GrpcSpec {
     pub connections_per_endpoint: usize,
 }
 
+impl GrpcSpec {
+    // todo: could improve duplication detection to individual endpoints to disallow accidental re-use
+    fn disallow_duplicates_digest(&self) -> String {
+        format!(
+            "{}{}",
+            self.instance_name,
+            self.endpoints
+                .clone()
+                .into_iter()
+                .map(|endpoint| endpoint.address)
+                .collect::<Vec<String>>()
+                .join(",")
+        )
+    }
+}
+
 /// The possible error codes that might occur on an upstream request.
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
 pub enum ErrorCode {
@@ -993,6 +1044,12 @@ pub struct RedisSpec {
     pub retry: Retry,
 }
 
+impl RedisSpec {
+    fn disallow_duplicates_digest(&self) -> String {
+        format!("{}{}", self.addresses.clone().join(","), self.key_prefix)
+    }
+}
+
 #[derive(Debug, Default, Deserialize, Serialize, Clone, PartialEq, Eq)]
 #[serde(rename_all = "lowercase")]
 pub enum RedisMode {

nativelink-service/BUILD.bazel

@@ -52,6 +52,7 @@ rust_test_suite(
         "tests/bep_server_test.rs",
         "tests/bytestream_server_test.rs",
         "tests/cas_server_test.rs",
+        "tests/store_overlap_rules_test.rs",
         "tests/worker_api_server_test.rs",
     ],
     proc_macro_deps = [

nativelink-service/tests/ac_server_test.rs

@@ -12,20 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::env;
 use std::pin::Pin;
 use std::sync::Arc;
 
 use bytes::BytesMut;
 use maplit::hashmap;
-use nativelink_config::stores::{MemorySpec, StoreSpec};
+use nativelink_config::stores::{FilesystemSpec, MemorySpec, StoreSpec};
 use nativelink_error::Error;
 use nativelink_macro::nativelink_test;
 use nativelink_proto::build::bazel::remote::execution::v2::action_cache_server::ActionCache;
 use nativelink_proto::build::bazel::remote::execution::v2::{
     digest_function, ActionResult, Digest, GetActionResultRequest, UpdateActionResultRequest,
 };
 use nativelink_service::ac_server::AcServer;
-use nativelink_store::default_store_factory::store_factory;
+use nativelink_store::default_store_factory::make_and_add_store_to_manager;
 use nativelink_store::store_manager::StoreManager;
 use nativelink_util::common::DigestInfo;
 use nativelink_util::store_trait::StoreLike;
@@ -48,29 +49,46 @@ async fn insert_into_store<T: Message>(
     let data_len = store_data.len();
     let digest = DigestInfo::try_new(hash, action_size)?;
     store.update_oneshot(digest, store_data.freeze()).await?;
+
     Ok(data_len.try_into().unwrap())
 }
 
 async fn make_store_manager() -> Result<Arc<StoreManager>, Error> {
     let store_manager = Arc::new(StoreManager::new());
-    store_manager.add_store(
+    make_and_add_store_to_manager(
         "main_cas",
-        store_factory(
-            &StoreSpec::memory(MemorySpec::default()),
-            &store_manager,
-            None,
-        )
-        .await?,
-    );
-    store_manager.add_store(
+        &StoreSpec::memory(MemorySpec::default()),
+        &store_manager,
+        None,
+    )
+    .await?;
+
+    let current_dir = env::current_dir().expect("Failed to get current directory");
+
+    let default_filesystem_spec = FilesystemSpec::default();
+
+    make_and_add_store_to_manager(
         "main_ac",
-        store_factory(
-            &StoreSpec::memory(MemorySpec::default()),
-            &store_manager,
-            None,
-        )
-        .await?,
-    );
+        &StoreSpec::filesystem(FilesystemSpec {
+            content_path: current_dir
+                .join("testing_data/ac/content_path")
+                .into_os_string()
+                .into_string()
+                .unwrap(),
+            temp_path: current_dir
+                .join("testing_data/ac/tmp_path")
+                .into_os_string()
+                .into_string()
+                .unwrap(),
+            read_buffer_size: default_filesystem_spec.read_buffer_size,
+            eviction_policy: default_filesystem_spec.eviction_policy,
+            block_size: default_filesystem_spec.block_size,
+        }),
+        &store_manager,
+        None,
+    )
+    .await?;
+
     Ok(store_manager)
 }
 

nativelink-service/tests/bep_server_test.rs

@@ -35,7 +35,7 @@ use nativelink_proto::google::devtools::build::v1::{
     PublishBuildToolEventStreamRequest, PublishLifecycleEventRequest, StreamId,
 };
 use nativelink_service::bep_server::BepServer;
-use nativelink_store::default_store_factory::store_factory;
+use nativelink_store::default_store_factory::make_and_add_store_to_manager;
 use nativelink_store::store_manager::StoreManager;
 use nativelink_util::buf_channel::make_buf_channel_pair;
 use nativelink_util::channel_body_for_tests::ChannelBody;
@@ -52,15 +52,14 @@ const BEP_STORE_NAME: &str = "main_bep";
 /// Utility function to construct a [`StoreManager`]
 async fn make_store_manager() -> Result<Arc<StoreManager>, Error> {
     let store_manager = Arc::new(StoreManager::new());
-    store_manager.add_store(
+    make_and_add_store_to_manager(
         BEP_STORE_NAME,
-        store_factory(
-            &StoreSpec::memory(MemorySpec::default()),
-            &store_manager,
-            None,
-        )
-        .await?,
-    );
+        &StoreSpec::memory(MemorySpec::default()),
+        &store_manager,
+        None,
+    )
+    .await?;
+
     Ok(store_manager)
 }
 

nativelink-service/tests/bytestream_server_test.rs

@@ -34,7 +34,7 @@ use nativelink_proto::google::bytestream::{
     QueryWriteStatusRequest, QueryWriteStatusResponse, ReadRequest, WriteRequest, WriteResponse,
 };
 use nativelink_service::bytestream_server::ByteStreamServer;
-use nativelink_store::default_store_factory::store_factory;
+use nativelink_store::default_store_factory::make_and_add_store_to_manager;
 use nativelink_store::store_manager::StoreManager;
 use nativelink_util::channel_body_for_tests::ChannelBody;
 use nativelink_util::common::{encode_stream_proto, DigestInfo};
@@ -58,15 +58,14 @@ const HASH1: &str = "0123456789abcdef000000000000000000000000000000000123456789a
 
 async fn make_store_manager() -> Result<Arc<StoreManager>, Error> {
     let store_manager = Arc::new(StoreManager::new());
-    store_manager.add_store(
+    make_and_add_store_to_manager(
         "main_cas",
-        store_factory(
-            &StoreSpec::memory(MemorySpec::default()),
-            &store_manager,
-            None,
-        )
-        .await?,
-    );
+        &StoreSpec::memory(MemorySpec::default()),
+        &store_manager,
+        None,
+    )
+    .await?;
+
     Ok(store_manager)
 }
 

nativelink-service/tests/cas_server_test.rs

@@ -30,7 +30,7 @@ use nativelink_proto::build::bazel::remote::execution::v2::{
 use nativelink_proto::google::rpc::Status as GrpcStatus;
 use nativelink_service::cas_server::CasServer;
 use nativelink_store::ac_utils::serialize_and_upload_message;
-use nativelink_store::default_store_factory::store_factory;
+use nativelink_store::default_store_factory::make_and_add_store_to_manager;
 use nativelink_store::store_manager::StoreManager;
 use nativelink_util::common::DigestInfo;
 use nativelink_util::digest_hasher::DigestHasherFunc;
@@ -47,15 +47,14 @@ const BAD_HASH: &str = "BAD_HASH";
 
 async fn make_store_manager() -> Result<Arc<StoreManager>, Error> {
     let store_manager = Arc::new(StoreManager::new());
-    store_manager.add_store(
+    make_and_add_store_to_manager(
         "main_cas",
-        store_factory(
-            &StoreSpec::memory(MemorySpec::default()),
-            &store_manager,
-            None,
-        )
-        .await?,
-    );
+        &StoreSpec::memory(MemorySpec::default()),
+        &store_manager,
+        None,
+    )
+    .await?;
+
     Ok(store_manager)
 }
 

nativelink-service/tests/store_overlap_rules_test.rs

@@ -0,0 +1,69 @@
+// Copyright 2024 The NativeLink Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use nativelink_config::stores::{MemorySpec, StoreSpec};
+use nativelink_error::Error;
+use nativelink_macro::nativelink_test;
+use nativelink_store::default_store_factory::make_and_add_store_to_manager;
+use nativelink_store::store_manager::StoreManager;
+
+#[nativelink_test]
+async fn same_datasource_disallowed_simple() -> Result<(), Error> {
+    let store_manager = Arc::new(StoreManager::new());
+    assert!(make_and_add_store_to_manager(
+        "main_cas",
+        &StoreSpec::memory(MemorySpec::default()),
+        &store_manager,
+        None,
+    )
+    .await
+    .is_ok());
+
+    assert!(make_and_add_store_to_manager(
+        "main_ac",
+        &StoreSpec::memory(MemorySpec::default()),
+        &store_manager,
+        None,
+    )
+    .await
+    .is_err());
+
+    Ok(())
+}
+
+#[nativelink_test]
+async fn same_datasource_disallowed_complex() -> Result<(), Error> {
+    let store_manager = Arc::new(StoreManager::new());
+    assert!(make_and_add_store_to_manager(
+        "main_cas",
+        &StoreSpec::memory(MemorySpec::default()),
+        &store_manager,
+        None,
+    )
+    .await
+    .is_ok());
+
+    assert!(make_and_add_store_to_manager(
+        "main_ac",
+        &StoreSpec::memory(MemorySpec::default()),
+        &store_manager,
+        None,
+    )
+    .await
+    .is_err());
+
+    Ok(())
+}

nativelink-store/src/default_store_factory.rs

@@ -42,7 +42,24 @@ use crate::verify_store::VerifyStore;
 
 type FutureMaybeStore<'a> = Box<dyn Future<Output = Result<Store, Error>> + 'a>;
 
-pub fn store_factory<'a>(
+pub async fn make_and_add_store_to_manager<'a>(
+    name: &'a str,
+    backend: &'a StoreSpec,
+    store_manager: &'a Arc<StoreManager>,
+    maybe_health_registry_builder: Option<&'a mut HealthRegistryBuilder>,
+) -> Result<(), Error> {
+    if let Some(digest) = backend.disallow_duplicates_digest() {
+        store_manager.digest_not_already_present(&digest)?;
+        store_manager.config_digest_add(digest);
+    }
+
+    let store = store_factory(backend, store_manager, maybe_health_registry_builder).await?;
+    store_manager.add_store(name, store)?;
+
+    Ok(())
+}
+
+fn store_factory<'a>(
     backend: &'a StoreSpec,
     store_manager: &'a Arc<StoreManager>,
     maybe_health_registry_builder: Option<&'a mut HealthRegistryBuilder>,