add support for

asr2003 · Dec 15, 2024 · e54c1a9 · e54c1a9
1 parent f672af7
commit e54c1a9
Show file tree

Hide file tree

Showing 8 changed files with 1,109 additions and 2 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -70,6 +70,7 @@ opentelemetry = { version = "0.27.1", default-features = false }
 prometheus = { version = "0.13.4", default-features = false }
 opentelemetry-prometheus = "0.27.0"
 serde_json = "1.0.133"
+google-cloud-storage = "0.23.0"
 
 [workspace.cargo-features-manager.keep]
 async-lock = ["std"]

diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs
@@ -54,6 +54,9 @@ pub enum StoreSpec {
     ///
     memory(MemorySpec),
 
+    /// TODO(asr2003): Add documentation.
+    experimental_gcs_store(GCSSpec),
+
     /// S3 store will use Amazon's S3 service as a backend to store
     /// the files. This configuration can be used to share files
     /// across multiple instances.
@@ -724,6 +727,26 @@ pub struct EvictionPolicy {
     pub max_count: u64,
 }
 
+#[derive(Serialize, Deserialize, Debug, Default, Clone)]
+#[serde(deny_unknown_fields)]
+pub struct GCSSpec {
+    /// Region name for GCS objects.
+    #[serde(default, deserialize_with = "convert_string_with_shellexpand")]
+    pub region: String,
+
+    /// Optional key prefix for GCS objects.
+    #[serde(default)]
+    pub key_prefix: Option<String>,
+
+    /// Bucket name to use as the backend.
+    #[serde(default, deserialize_with = "convert_string_with_shellexpand")]
+    pub bucket: String,
+
+    /// Retry configuration to use when a network request fails.
+    #[serde(default)]
+    pub retry: Retry,
+}
+
 #[derive(Serialize, Deserialize, Debug, Default, Clone)]
 #[serde(deny_unknown_fields)]
 pub struct S3Spec {

diff --git a/nativelink-store/BUILD.bazel b/nativelink-store/BUILD.bazel
@@ -18,6 +18,7 @@ rust_library(
         "src/existence_cache_store.rs",
         "src/fast_slow_store.rs",
         "src/filesystem_store.rs",
+        "src/gcs_store.rs",
         "src/grpc_store.rs",
         "src/lib.rs",
         "src/memory_store.rs",

diff --git a/nativelink-store/Cargo.toml b/nativelink-store/Cargo.toml
@@ -57,6 +57,10 @@ tokio-util = { version = "0.7.13" }
 tonic = { version = "0.12.3", features = ["transport", "tls"], default-features = false }
 tracing = { version = "0.1.41", default-features = false }
 uuid = { version = "1.11.0", default-features = false, features = ["v4", "serde"] }
+reqwest = { version = "0.12.9", features = ["json", "gzip", "stream"]}
+google-cloud-storage = "0.23.0"
+percent-encoding = "2.3.1"
+reqwest-middleware = "0.4.0"
 
 [dev-dependencies]
 nativelink-macro = { path = "../nativelink-macro" }

diff --git a/nativelink-store/src/default_store_factory.rs b/nativelink-store/src/default_store_factory.rs
@@ -29,6 +29,7 @@ use crate::dedup_store::DedupStore;
 use crate::existence_cache_store::ExistenceCacheStore;
 use crate::fast_slow_store::FastSlowStore;
 use crate::filesystem_store::FilesystemStore;
+use crate::gcs_store::GCSStore;
 use crate::grpc_store::GrpcStore;
 use crate::memory_store::MemoryStore;
 use crate::noop_store::NoopStore;
@@ -51,6 +52,7 @@ pub fn store_factory<'a>(
         let store: Arc<dyn StoreDriver> = match backend {
             StoreSpec::memory(spec) => MemoryStore::new(spec),
             StoreSpec::experimental_s3_store(spec) => S3Store::new(spec, SystemTime::now).await?,
+            StoreSpec::experimental_gcs_store(spec) => GCSStore::new(spec).await?,
             StoreSpec::redis_store(spec) => RedisStore::new(spec.clone())?,
             StoreSpec::verify(spec) => VerifyStore::new(
                 spec,