feat: add support for GCS Store using tonic/grpc

TraceMachina · Dec 18, 2024 · 5e4393d · 5e4393d
1 parent 24adc11
commit 5e4393d
Show file tree

Hide file tree

Showing 9 changed files with 469 additions and 473 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -70,7 +70,6 @@ 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
@@ -730,21 +730,56 @@ pub struct EvictionPolicy {
 #[derive(Serialize, Deserialize, Debug, Default, Clone)]
 #[serde(deny_unknown_fields)]
 pub struct GCSSpec {
-    /// Region name for GCS objects.
+    /// GCS region or location. Example: US, US-CENTRAL1, EUROPE-WEST1.
     #[serde(default, deserialize_with = "convert_string_with_shellexpand")]
-    pub region: String,
-
-    /// Optional key prefix for GCS objects.
-    #[serde(default)]
-    pub key_prefix: Option<String>,
+    pub location: String,
 
     /// Bucket name to use as the backend.
     #[serde(default, deserialize_with = "convert_string_with_shellexpand")]
     pub bucket: String,
 
+    /// Optional prefix for object keys. If None, no prefix will be used.
+    #[serde(default)]
+    pub key_prefix: Option<String>,
+
     /// Retry configuration to use when a network request fails.
     #[serde(default)]
     pub retry: Retry,
+
+    /// Time in seconds after which an object is considered "expired."
+    /// Allows external tools to clean up unused objects.
+    /// Default: 0. Zero means never consider an object expired.
+    #[serde(default, deserialize_with = "convert_duration_with_shellexpand")]
+    pub consider_expired_after_s: u32,
+
+    /// Maximum buffer size to retain in case of a retryable error during upload.
+    /// Setting this to zero will disable upload buffering.
+    /// Default: 5MB.
+    pub max_retry_buffer_per_request: Option<usize>,
+
+    /// Enable resumable uploads for large objects.
+    /// Default: true.
+    #[serde(default)]
+    pub enable_resumable_uploads: bool,
+
+    /// The maximum size of chunks (in bytes) for resumable uploads.
+    /// Default: 8MB.
+    pub resumable_chunk_size: Option<usize>,
+
+    /// Allow unencrypted HTTP connections. Only use this for local testing.
+    /// Default: false
+    #[serde(default)]
+    pub insecure_allow_http: bool,
+
+    /// Disable http/2 connections and only use http/1.1.
+    /// Default: false
+    #[serde(default)]
+    pub disable_http2: bool,
+
+    /// Optional configuration for client authentication.
+    /// Example: Path to a service account JSON key file or environment-based authentication.
+    #[serde(default)]
+    pub auth_key_file: Option<String>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Default, Clone)]

diff --git a/nativelink-error/BUILD.bazel b/nativelink-error/BUILD.bazel
@@ -15,7 +15,6 @@ rust_library(
         "//nativelink-metric",
         "//nativelink-proto",
         "@crates//:fred",
-        "@crates//:google-cloud-storage",
         "@crates//:hex",
         "@crates//:prost",
         "@crates//:prost-types",

diff --git a/nativelink-error/Cargo.toml b/nativelink-error/Cargo.toml
@@ -19,4 +19,3 @@ prost-types = { version = "0.13.4", default-features = false }
 serde = { version = "1.0.216", default-features = false }
 tokio = { version = "1.42.0", features = ["fs", "rt-multi-thread", "signal", "io-util"], default-features = false }
 tonic = { version = "0.12.3", features = ["transport", "tls"], default-features = false }
-google-cloud-storage = "0.23.0"
diff --git a/nativelink-error/src/lib.rs b/nativelink-error/src/lib.rs
@@ -14,7 +14,6 @@
 
 use std::convert::Into;
 
-pub use google_cloud_storage::http::Error as GcsError;
 use nativelink_metric::{
     MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent,
 };
@@ -220,36 +219,6 @@ impl From<Code> for Error {
     }
 }
 
-impl From<GcsError> for Error {
-    fn from(err: GcsError) -> Self {
-        match err {
-            GcsError::Response(error_response) => {
-                make_err!(
-                    Code::Unavailable,
-                    "GCS Response Error: {:?}",
-                    error_response
-                )
-            }
-            GcsError::HttpClient(error) => {
-                make_err!(Code::Unavailable, "GCS HTTP Client Error: {:?}", error)
-            }
-            GcsError::HttpMiddleware(error) => {
-                make_err!(Code::Unavailable, "GCS HTTP Middleware Error: {:?}", error)
-            }
-            GcsError::TokenSource(error) => {
-                make_err!(Code::Unauthenticated, "GCS Token Source Error: {:?}", error)
-            }
-            GcsError::InvalidRangeHeader(header) => {
-                make_err!(
-                    Code::InvalidArgument,
-                    "GCS Invalid Range Header: {:?}",
-                    header
-                )
-            }
-        }
-    }
-}
-
 impl From<fred::error::Error> for Error {
     fn from(error: fred::error::Error) -> Self {
         use fred::error::ErrorKind::{

diff --git a/nativelink-store/BUILD.bazel b/nativelink-store/BUILD.bazel
@@ -53,22 +53,20 @@ rust_library(
         "@crates//:bytes",
         "@crates//:bytes-utils",
         "@crates//:const_format",
+        "@crates//:crc32c",
         "@crates//:filetime",
         "@crates//:fred",
         "@crates//:futures",
-        "@crates//:google-cloud-storage",
+        "@crates//:googleapis-tonic-google-storage-v2",
         "@crates//:hex",
         "@crates//:http-body",
         "@crates//:hyper-0.14.31",
         "@crates//:hyper-rustls",
         "@crates//:lz4_flex",
         "@crates//:parking_lot",
         "@crates//:patricia_tree",
-        "@crates//:percent-encoding",
         "@crates//:prost",
         "@crates//:rand",
-        "@crates//:reqwest",
-        "@crates//:reqwest-middleware",
         "@crates//:serde",
         "@crates//:tokio",
         "@crates//:tokio-stream",

diff --git a/nativelink-store/Cargo.toml b/nativelink-store/Cargo.toml
@@ -57,10 +57,8 @@ 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"
+googleapis-tonic-google-storage-v2 = "0.16.0"
+crc32c = "0.6.8"
 
 [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
@@ -52,7 +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::experimental_gcs_store(spec) => GCSStore::new(spec, SystemTime::now).await?,
             StoreSpec::redis_store(spec) => RedisStore::new(spec.clone())?,
             StoreSpec::verify(spec) => VerifyStore::new(
                 spec,