/build/source/nativelink-store/src/gcs_client/client.rs

Source
// 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 core::fmt::{Debug, Formatter};
use core::future::Future;
use core::time::Duration;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;

use async_trait::async_trait;
use google_cloud_auth::credentials::CredentialsFile;
use google_cloud_storage::client::{Client, ClientConfig};
use google_cloud_storage::http::Error as GcsError;
use google_cloud_storage::http::objects::Object;
use google_cloud_storage::http::objects::download::Range;
use google_cloud_storage::http::objects::get::GetObjectRequest;
use google_cloud_storage::http::objects::upload::{Media, UploadObjectRequest, UploadType};
use google_cloud_storage::http::resumable_upload_client::{ChunkSize, UploadStatus};
use nativelink_config::stores::ExperimentalGcsSpec;
use nativelink_error::{Code, Error, make_err};
use nativelink_util::buf_channel::DropCloserReadHalf;
use rand::Rng;
use tokio::fs;
use tokio::sync::Semaphore;
use tokio::time::sleep;

use crate::gcs_client::types::{
    CHUNK_SIZE, DEFAULT_CONCURRENT_UPLOADS, DEFAULT_CONTENT_TYPE, GcsObject,
    INITIAL_UPLOAD_RETRY_DELAY_MS, MAX_UPLOAD_RETRIES, MAX_UPLOAD_RETRY_DELAY_MS, ObjectPath,
    SIMPLE_UPLOAD_THRESHOLD, Timestamp,
};

/// A trait that defines the required GCS operations.
/// This abstraction allows for easier testing by mocking GCS responses.
#[async_trait]
pub trait GcsOperations: Send + Sync + Debug {
    /// Read metadata for a GCS object
    async fn read_object_metadata(&self, object: &ObjectPath) -> Result<Option<GcsObject>, Error>;

    /// Read the content of a GCS object, optionally with a range
    async fn read_object_content(
        &self,
        object_path: &ObjectPath,
        start: i64,
        end: Option<i64>,
    ) -> Result<Vec<u8>, Error>;

    /// Write object with simple upload (for smaller objects)
    async fn write_object(&self, object_path: &ObjectPath, content: Vec<u8>) -> Result<(), Error>;

    /// Start a resumable write operation and return the upload URL
    async fn start_resumable_write(&self, object_path: &ObjectPath) -> Result<String, Error>;

    /// Upload a chunk of data in a resumable upload session
    async fn upload_chunk(
        &self,
        upload_url: &str,
        object_path: &ObjectPath,
        data: Vec<u8>,
        offset: i64,
        end_offset: i64,
        is_final: bool,
    ) -> Result<(), Error>;

    /// Complete high-level operation to upload data from a reader
    async fn upload_from_reader(
        &self,
        object_path: &ObjectPath,
        reader: &mut DropCloserReadHalf,
        upload_id: &str,
        max_size: i64,
    ) -> Result<(), Error>;

    /// Check if an object exists
    async fn object_exists(&self, object_path: &ObjectPath) -> Result<bool, Error>;
}

/// Main client for interacting with Google Cloud Storage
pub struct GcsClient {
    client: Client,
    resumable_chunk_size: usize,
    semaphore: Arc<Semaphore>,
}

impl GcsClient {
    /// Create a new GCS client from the provided spec
    pub async fn new(spec: &ExperimentalGcsSpec) -> Result<Self, Error> {
        // Creating default config without authentication initially
        let mut client_config = ClientConfig::default();
        let mut auth_success = false;

        // Trying authentication with credentials file path from environment variable.
        // Google Cloud Auth expects the path to the credentials file to be specified
        // in the GOOGLE_APPLICATION_CREDENTIALS environment variable.
        // Check https://cloud.google.com/docs/authentication/application-default-credentials#GAC
        // for more details.
        if let Ok(creds_path) = std::env::var("GOOGLE_APPLICATION_CREDENTIALS") {
            let path_str = PathBuf::from(creds_path).to_string_lossy().to_string();

            tracing::info!("Attempting to load credentials from: {}", path_str);
            match CredentialsFile::new_from_file(path_str.clone()).await {
                Ok(creds_file) => {
                    match ClientConfig::default().with_credentials(creds_file).await {
                        Ok(config) => {
                            tracing::info!("Successfully authenticated with credentials file");
                            client_config = config;
                            auth_success = true;
                        }
                        Err(e) => {
                            tracing::warn!("Failed to configure client with credentials: {}", e);
                        }
                    }
                }
                Err(e) => {
                    tracing::warn!("Failed to load credentials from file: {}", e);
                }
            }
        }

        // Trying JSON credentials in environment variable if first method failed
        if !auth_success {
            if let Ok(creds_json) = std::env::var("GOOGLE_APPLICATION_CREDENTIALS_JSON") {
                tracing::info!("Attempting to use credentials from environment variable");

                // Creating temporary file
                let temp_dir = std::env::temp_dir();
                let temp_path = temp_dir.join(format!("gcs-creds-{}.json", uuid::Uuid::new_v4()));
                let temp_path_str = temp_path.to_string_lossy().to_string();

                // Writing JSON to temporary file
                if let Err(e) = fs::write(&temp_path, creds_json).await {
                    tracing::warn!("Failed to write credentials to temp file: {}", e);
                } else {
                    // Load credentials from temporary file
                    match CredentialsFile::new_from_file(temp_path_str).await {
                        Ok(creds_file) => {
                            match ClientConfig::default().with_credentials(creds_file).await {
                                Ok(config) => {
                                    tracing::info!(
                                        "Successfully authenticated with JSON credentials"
                                    );
                                    client_config = config;
                                    auth_success = true;
                                }
                                Err(e) => {
                                    tracing::warn!(
                                        "Failed to configure client with JSON credentials: {}",
                                        e
                                    );
                                }
                            }
                        }
                        Err(e) => {
                            tracing::warn!("Failed to load credentials from JSON: {}", e);
                        }
                    }

                    // Clean up temporary file
                    drop(fs::remove_file(temp_path).await);
                }
            }
        }

        if !auth_success {
            match ClientConfig::default().with_auth().await {
                Ok(config) => {
                    tracing::info!(
                        "Successfully authenticated with application default credentials"
                    );
                    client_config = config;
                }
                Err(e) => {
                    tracing::warn!("Failed to use application default credentials: {}", e);
                    tracing::info!("Continuing with unauthenticated access");
                }
            }
        }

        // Creating client with the configured authentication
        let client = Client::new(client_config);
        let resumable_chunk_size = spec.resumable_chunk_size.unwrap_or(CHUNK_SIZE);

        // Get max connections from config
        let max_connections = spec
            .common
            .multipart_max_concurrent_uploads
            .unwrap_or(DEFAULT_CONCURRENT_UPLOADS);

        Ok(Self {
            client,
            resumable_chunk_size,
            semaphore: Arc::new(Semaphore::new(max_connections)),
        })
    }

    /// Generic method to execute operations with connection limiting
    async fn with_connection<F, Fut, T>(&self, operation: F) -> Result<T, Error>
    where
        F: FnOnce() -> Fut + Send,
        Fut: Future<Output = Result<T, Error>> + Send,
    {
        let permit =
            self.semaphore.acquire().await.map_err(|e| {
                make_err!(Code::Internal, "Failed to acquire connection permit: {}", e)
            })?;

        let result = operation().await;
        drop(permit);
        result
    }

    /// Convert GCS object to our internal representation
    fn convert_to_gcs_object(&self, obj: Object) -> GcsObject {
        let update_time = obj.updated.map(|dt| Timestamp {
            seconds: dt.unix_timestamp(),
            nanos: 0,
        });

        GcsObject {
            name: obj.name,
            bucket: obj.bucket,
            size: obj.size,
            content_type: obj
                .content_type
                .unwrap_or_else(|| DEFAULT_CONTENT_TYPE.to_string()),
            update_time,
        }
    }

    /// Handle error from GCS operations
    fn handle_gcs_error(&self, err: &GcsError) -> Error {
        let code = match &err {
            GcsError::Response(resp) => match resp.code {
                404 => Code::NotFound,
                401 | 403 => Code::PermissionDenied,
                408 | 429 => Code::ResourceExhausted,
                500..=599 => Code::Unavailable,
                _ => Code::Unknown,
            },
            _ => Code::Internal,
        };

        make_err!(code, "GCS operation failed: {}", err)
    }

    /// Reading data from reader and upload in a single operation
    async fn read_and_upload_all(
        &self,
        object_path: &ObjectPath,
        reader: &mut DropCloserReadHalf,
        max_size: i64,
    ) -> Result<(), Error> {
        let initial_capacity = core::cmp::min(max_size as usize, 10 * 1024 * 1024);
        let mut data = Vec::with_capacity(initial_capacity);
        let mut total_size: i64 = 0;

        while total_size < max_size {
            let to_read =
                core::cmp::min(self.resumable_chunk_size, (max_size - total_size) as usize);
            let chunk = reader.consume(Some(to_read)).await?;

            if chunk.is_empty() {
                break;
            }

            data.extend_from_slice(&chunk);
            total_size += chunk.len() as i64;
        }

        self.write_object(object_path, data).await
    }

    /// Implementing a resumable upload using the resumable upload API
    async fn try_resumable_upload(
        &self,
        object_path: &ObjectPath,
        reader: &mut DropCloserReadHalf,
        max_size: i64,
    ) -> Result<(), Error> {
        self.with_connection(|| async {
            let request = UploadObjectRequest {
                bucket: object_path.bucket.clone(),
                ..Default::default()
            };

            let mut metadata = HashMap::<String, String>::new();
            metadata.insert("name".to_string(), object_path.path.clone());

            // Use Multipart upload type with metadata
            let upload_type = UploadType::Multipart(Box::new(Object {
                name: object_path.path.clone(),
                content_type: Some(DEFAULT_CONTENT_TYPE.to_string()),
                metadata: Some(metadata),
                ..Default::default()
            }));

            // Prepare resumable upload
            let uploader = self
                .client
                .prepare_resumable_upload(&request, &upload_type)
                .await
                .map_err(|e| self.handle_gcs_error(&e))?;

            // Upload data in chunks
            let mut offset: u64 = 0;
            let mut total_uploaded: i64 = 0;

            while total_uploaded < max_size {
                let to_read = core::cmp::min(
                    self.resumable_chunk_size,
                    (max_size - total_uploaded) as usize,
                );
                let chunk = reader.consume(Some(to_read)).await?;

                if chunk.is_empty() {
                    break;
                }

                let chunk_size = chunk.len() as u64;
                total_uploaded += chunk.len() as i64;

                let is_final = total_uploaded >= max_size || chunk.len() < to_read;
                let total_size = if is_final {
                    Some(offset + chunk_size)
                } else {
                    Some(max_size as u64)
                };

                let chunk_def = ChunkSize::new(offset, offset + chunk_size - 1, total_size);

                // Upload chunk
                let status = uploader
                    .upload_multiple_chunk(chunk, &chunk_def)
                    .await
                    .map_err(|e| self.handle_gcs_error(&e))?;

                // Update offset for next chunk
                offset += chunk_size;

                if let UploadStatus::Ok(_) = status {
                    break;
                }
            }

            // If nothing was uploaded, finalizing with empty content
            if offset == 0 {
                let chunk_def = ChunkSize::new(0, 0, Some(0));
                uploader
                    .upload_multiple_chunk(Vec::new(), &chunk_def)
                    .await
                    .map_err(|e| self.handle_gcs_error(&e))?;
            }

            // Check if the object exists
            match self.read_object_metadata(object_path).await? {
                Some(_) => Ok(()),
                None => Err(make_err!(
                    Code::Internal,
                    "Upload completed but object not found"
                )),
            }
        })
        .await
    }
}

impl Debug for GcsClient {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("GcsClient")
            .field("resumable_chunk_size", &self.resumable_chunk_size)
            .field("max_connections", &self.semaphore.available_permits())
            .finish_non_exhaustive()
    }
}

#[async_trait]
impl GcsOperations for GcsClient {
    async fn read_object_metadata(
        &self,
        object_path: &ObjectPath,
    ) -> Result<Option<GcsObject>, Error> {
        self.with_connection(|| async {
            let request = GetObjectRequest {
                bucket: object_path.bucket.clone(),
                object: object_path.path.clone(),
                ..Default::default()
            };

            match self.client.get_object(&request).await {
                Ok(obj) => Ok(Some(self.convert_to_gcs_object(obj))),
                Err(err) => {
                    if let GcsError::Response(resp) = &err {
                        if resp.code == 404 {
                            return Ok(None);
                        }
                    }
                    Err(self.handle_gcs_error(&err))
                }
            }
        })
        .await
    }

    async fn read_object_content(
        &self,
        object_path: &ObjectPath,
        start: i64,
        end: Option<i64>,
    ) -> Result<Vec<u8>, Error> {
        self.with_connection(|| async {
            let request = GetObjectRequest {
                bucket: object_path.bucket.clone(),
                object: object_path.path.clone(),
                ..Default::default()
            };

            let range = if start > 0 || end.is_some() {
                Range(
                    if start > 0 { Some(start as u64) } else { None },
                    end.map(|e| e as u64),
                )
            } else {
                Range(None, None)
            };

            // Download the object
            let content = self
                .client
                .download_object(&request, &range)
                .await
                .map_err(|e| self.handle_gcs_error(&e))?;

            Ok(content)
        })
        .await
    }

    async fn write_object(&self, object_path: &ObjectPath, content: Vec<u8>) -> Result<(), Error> {
        self.with_connection(|| async {
            let request = UploadObjectRequest {
                bucket: object_path.bucket.clone(),
                ..Default::default()
            };

            let media = Media::new(object_path.path.clone());
            let upload_type = UploadType::Simple(media);

            self.client
                .upload_object(&request, content, &upload_type)
                .await
                .map_err(|e| self.handle_gcs_error(&e))?;

            Ok(())
        })
        .await
    }

    async fn start_resumable_write(&self, object_path: &ObjectPath) -> Result<String, Error> {
        self.with_connection(|| async {
            let request = UploadObjectRequest {
                bucket: object_path.bucket.clone(),
                ..Default::default()
            };

            let upload_type = UploadType::Multipart(Box::new(Object {
                name: object_path.path.clone(),
                content_type: Some(DEFAULT_CONTENT_TYPE.to_string()),
                ..Default::default()
            }));

            // Start resumable upload session
            let uploader = self
                .client
                .prepare_resumable_upload(&request, &upload_type)
                .await
                .map_err(|e| self.handle_gcs_error(&e))?;

            Ok(uploader.url().to_string())
        })
        .await
    }

    async fn upload_chunk(
        &self,
        upload_url: &str,
        _object_path: &ObjectPath,
        data: Vec<u8>,
        offset: i64,
        end_offset: i64,
        is_final: bool,
    ) -> Result<(), Error> {
        self.with_connection(|| async {
            let uploader = self.client.get_resumable_upload(upload_url.to_string());

            let total_size = if is_final {
                Some(end_offset as u64)
            } else {
                None
            };

            let chunk_def = ChunkSize::new(offset as u64, end_offset as u64 - 1, total_size);

            // Upload chunk
            uploader
                .upload_multiple_chunk(data, &chunk_def)
                .await
                .map_err(|e| self.handle_gcs_error(&e))?;

            Ok(())
        })
        .await
    }

    async fn upload_from_reader(
        &self,
        object_path: &ObjectPath,
        reader: &mut DropCloserReadHalf,
        _upload_id: &str,
        max_size: i64,
    ) -> Result<(), Error> {
        let mut retry_count = 0;
        let mut retry_delay = INITIAL_UPLOAD_RETRY_DELAY_MS;

        loop {
            let result = if max_size < SIMPLE_UPLOAD_THRESHOLD {
                self.read_and_upload_all(object_path, reader, max_size)
                    .await
            } else {
                self.try_resumable_upload(object_path, reader, max_size)
                    .await
            };

            match result {
                Ok(()) => return Ok(()),
                Err(e) => {
                    let is_retriable = matches!(
                        e.code,
                        Code::Unavailable | Code::ResourceExhausted | Code::DeadlineExceeded
                    ) || (e.code == Code::Internal
                        && e.to_string().contains("connection"));

                    if !is_retriable || retry_count >= MAX_UPLOAD_RETRIES {
                        return Err(e);
                    }

                    if let Err(reset_err) = reader.try_reset_stream() {
                        return Err(e.merge(reset_err));
                    }

                    sleep(Duration::from_millis(retry_delay)).await;
                    retry_delay = core::cmp::min(retry_delay * 2, MAX_UPLOAD_RETRY_DELAY_MS);

                    let mut rng = rand::rng();
                    let jitter_factor = rng.random::<f64>().mul_add(0.4, 0.8);
                    retry_delay = (retry_delay as f64 * jitter_factor) as u64;

                    retry_count += 1;
                }
            }
        }
    }

    async fn object_exists(&self, object_path: &ObjectPath) -> Result<bool, Error> {
        let metadata = self.read_object_metadata(object_path).await?;
        Ok(metadata.is_some())
    }
}

Coverage Report

Created: 2025-05-08 18:47

Line	Count	Source
1		// Copyright 2024 The NativeLink Authors. All rights reserved.
2		//
3		// Licensed under the Apache License, Version 2.0 (the "License");
4		// you may not use this file except in compliance with the License.
5		// You may obtain a copy of the License at
6		//
7		// http://www.apache.org/licenses/LICENSE-2.0
8		//
9		// Unless required by applicable law or agreed to in writing, software
10		// distributed under the License is distributed on an "AS IS" BASIS,
11		// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12		// See the License for the specific language governing permissions and
13		// limitations under the License.
14
15		use core::fmt::{Debug, Formatter};
16		use core::future::Future;
17		use core::time::Duration;
18		use std::collections::HashMap;
19		use std::path::PathBuf;
20		use std::sync::Arc;
21
22		use async_trait::async_trait;
23		use google_cloud_auth::credentials::CredentialsFile;
24		use google_cloud_storage::client::{Client, ClientConfig};
25		use google_cloud_storage::http::Error as GcsError;
26		use google_cloud_storage::http::objects::Object;
27		use google_cloud_storage::http::objects::download::Range;
28		use google_cloud_storage::http::objects::get::GetObjectRequest;
29		use google_cloud_storage::http::objects::upload::{Media, UploadObjectRequest, UploadType};
30		use google_cloud_storage::http::resumable_upload_client::{ChunkSize, UploadStatus};
31		use nativelink_config::stores::ExperimentalGcsSpec;
32		use nativelink_error::{Code, Error, make_err};
33		use nativelink_util::buf_channel::DropCloserReadHalf;
34		use rand::Rng;
35		use tokio::fs;
36		use tokio::sync::Semaphore;
37		use tokio::time::sleep;
38
39		use crate::gcs_client::types::{
40		CHUNK_SIZE, DEFAULT_CONCURRENT_UPLOADS, DEFAULT_CONTENT_TYPE, GcsObject,
41		INITIAL_UPLOAD_RETRY_DELAY_MS, MAX_UPLOAD_RETRIES, MAX_UPLOAD_RETRY_DELAY_MS, ObjectPath,
42		SIMPLE_UPLOAD_THRESHOLD, Timestamp,
43		};
44
45		/// A trait that defines the required GCS operations.
46		/// This abstraction allows for easier testing by mocking GCS responses.
47		#[async_trait]
48		pub trait GcsOperations: Send + Sync + Debug {
49		/// Read metadata for a GCS object
50		async fn read_object_metadata(&self, object: &ObjectPath) -> Result<Option<GcsObject>, Error>;
51
52		/// Read the content of a GCS object, optionally with a range
53		async fn read_object_content(
54		&self,
55		object_path: &ObjectPath,
56		start: i64,
57		end: Option<i64>,
58		) -> Result<Vec<u8>, Error>;
59
60		/// Write object with simple upload (for smaller objects)
61		async fn write_object(&self, object_path: &ObjectPath, content: Vec<u8>) -> Result<(), Error>;
62
63		/// Start a resumable write operation and return the upload URL
64		async fn start_resumable_write(&self, object_path: &ObjectPath) -> Result<String, Error>;
65
66		/// Upload a chunk of data in a resumable upload session
67		async fn upload_chunk(
68		&self,
69		upload_url: &str,
70		object_path: &ObjectPath,
71		data: Vec<u8>,
72		offset: i64,
73		end_offset: i64,
74		is_final: bool,
75		) -> Result<(), Error>;
76
77		/// Complete high-level operation to upload data from a reader
78		async fn upload_from_reader(
79		&self,
80		object_path: &ObjectPath,
81		reader: &mut DropCloserReadHalf,
82		upload_id: &str,
83		max_size: i64,
84		) -> Result<(), Error>;
85
86		/// Check if an object exists
87		async fn object_exists(&self, object_path: &ObjectPath) -> Result<bool, Error>;
88		}
89
90		/// Main client for interacting with Google Cloud Storage
91		pub struct GcsClient {
92		client: Client,
93		resumable_chunk_size: usize,
94		semaphore: Arc<Semaphore>,
95		}
96
97		impl GcsClient {
98		/// Create a new GCS client from the provided spec
99	0	pub async fn new(spec: &ExperimentalGcsSpec) -> Result<Self, Error> {
100	0	// Creating default config without authentication initially
101	0	let mut client_config = ClientConfig::default();
102	0	let mut auth_success = false;
103
104		// Trying authentication with credentials file path from environment variable.
105		// Google Cloud Auth expects the path to the credentials file to be specified
106		// in the GOOGLE_APPLICATION_CREDENTIALS environment variable.
107		// Check https://cloud.google.com/docs/authentication/application-default-credentials#GAC
108		// for more details.
109	0	if let Ok(creds_path) = std::env::var("GOOGLE_APPLICATION_CREDENTIALS") { Branch (109:16): [True: 0, False: 0] Branch (109:16): [Folded - Ignored]
110	0	let path_str = PathBuf::from(creds_path).to_string_lossy().to_string();
111	0
112	0	tracing::info!("Attempting to load credentials from: {}", path_str);
113	0	match CredentialsFile::new_from_file(path_str.clone()).await {
114	0	Ok(creds_file) => {
115	0	match ClientConfig::default().with_credentials(creds_file).await {
116	0	Ok(config) => {
117	0	tracing::info!("Successfully authenticated with credentials file");
118	0	client_config = config;
119	0	auth_success = true;
120		}
121	0	Err(e) => {
122	0	tracing::warn!("Failed to configure client with credentials: {}", e);
123		}
124		}
125		}
126	0	Err(e) => {
127	0	tracing::warn!("Failed to load credentials from file: {}", e);
128		}
129		}
130	0	}
131
132		// Trying JSON credentials in environment variable if first method failed
133	0	if !auth_success { Branch (133:12): [True: 0, False: 0] Branch (133:12): [Folded - Ignored]
134	0	if let Ok(creds_json) = std::env::var("GOOGLE_APPLICATION_CREDENTIALS_JSON") { Branch (134:20): [True: 0, False: 0] Branch (134:20): [Folded - Ignored]
135	0	tracing::info!("Attempting to use credentials from environment variable");
136
137		// Creating temporary file
138	0	let temp_dir = std::env::temp_dir();
139	0	let temp_path = temp_dir.join(format!("gcs-creds-{}.json", uuid::Uuid::new_v4()));
140	0	let temp_path_str = temp_path.to_string_lossy().to_string();
141
142		// Writing JSON to temporary file
143	0	if let Err(e) = fs::write(&temp_path, creds_json).await { Branch (143:24): [True: 0, False: 0] Branch (143:24): [Folded - Ignored]
144	0	tracing::warn!("Failed to write credentials to temp file: {}", e);
145		} else {
146		// Load credentials from temporary file
147	0	match CredentialsFile::new_from_file(temp_path_str).await {
148	0	Ok(creds_file) => {
149	0	match ClientConfig::default().with_credentials(creds_file).await {
150	0	Ok(config) => {
151	0	tracing::info!(
152	0	"Successfully authenticated with JSON credentials"
153		);
154	0	client_config = config;
155	0	auth_success = true;
156		}
157	0	Err(e) => {
158	0	tracing::warn!(
159	0	"Failed to configure client with JSON credentials: {}",
160		e
161		);
162		}
163		}
164		}
165	0	Err(e) => {
166	0	tracing::warn!("Failed to load credentials from JSON: {}", e);
167		}
168		}
169
170		// Clean up temporary file
171	0	drop(fs::remove_file(temp_path).await);
172		}
173	0	}
174	0	}
175
176	0	if !auth_success { Branch (176:12): [True: 0, False: 0] Branch (176:12): [Folded - Ignored]
177	0	match ClientConfig::default().with_auth().await {
178	0	Ok(config) => {
179	0	tracing::info!(
180	0	"Successfully authenticated with application default credentials"
181		);
182	0	client_config = config;
183		}
184	0	Err(e) => {
185	0	tracing::warn!("Failed to use application default credentials: {}", e);
186	0	tracing::info!("Continuing with unauthenticated access");
187		}
188		}
189	0	}
190
191		// Creating client with the configured authentication
192	0	let client = Client::new(client_config);
193	0	let resumable_chunk_size = spec.resumable_chunk_size.unwrap_or(CHUNK_SIZE);
194	0
195	0	// Get max connections from config
196	0	let max_connections = spec
197	0	.common
198	0	.multipart_max_concurrent_uploads
199	0	.unwrap_or(DEFAULT_CONCURRENT_UPLOADS);
200	0
201	0	Ok(Self {
202	0	client,
203	0	resumable_chunk_size,
204	0	semaphore: Arc::new(Semaphore::new(max_connections)),
205	0	})
206	0	}
207
208		/// Generic method to execute operations with connection limiting
209	0	async fn with_connection<F, Fut, T>(&self, operation: F) -> Result<T, Error>
210	0	where
211	0	F: FnOnce() -> Fut + Send,
212	0	Fut: Future<Output = Result<T, Error>> + Send,
213	0	{
214	0	let permit =
215	0	self.semaphore.acquire().await.map_err(\|e\| {
216	0	make_err!(Code::Internal, "Failed to acquire connection permit: {}", e)
217	0	})?;
218
219	0	let result = operation().await;
220	0	drop(permit);
221	0	result
222	0	}
223
224		/// Convert GCS object to our internal representation
225	0	fn convert_to_gcs_object(&self, obj: Object) -> GcsObject {
226	0	let update_time = obj.updated.map(\|dt\| Timestamp {
227	0	seconds: dt.unix_timestamp(),
228	0	nanos: 0,
229	0	});
230
231		GcsObject {
232	0	name: obj.name,
233	0	bucket: obj.bucket,
234	0	size: obj.size,
235	0	content_type: obj
236	0	.content_type
237	0	.unwrap_or_else(\|\| DEFAULT_CONTENT_TYPE.to_string()),
238	0	update_time,
239	0	}
240	0	}
241
242		/// Handle error from GCS operations
243	0	fn handle_gcs_error(&self, err: &GcsError) -> Error {
244	0	let code = match &err {
245	0	GcsError::Response(resp) => match resp.code {
246	0	404 => Code::NotFound,
247	0	401 \| 403 => Code::PermissionDenied,
248	0	408 \| 429 => Code::ResourceExhausted,
249	0	500..=599 => Code::Unavailable,
250	0	_ => Code::Unknown,
251		},
252	0	_ => Code::Internal,
253		};
254
255	0	make_err!(code, "GCS operation failed: {}", err)
256	0	}
257
258		/// Reading data from reader and upload in a single operation
259	0	async fn read_and_upload_all(
260	0	&self,
261	0	object_path: &ObjectPath,
262	0	reader: &mut DropCloserReadHalf,
263	0	max_size: i64,
264	0	) -> Result<(), Error> {
265	0	let initial_capacity = core::cmp::min(max_size as usize, 10 * 1024 * 1024);
266	0	let mut data = Vec::with_capacity(initial_capacity);
267	0	let mut total_size: i64 = 0;
268
269	0	while total_size < max_size { Branch (269:15): [True: 0, False: 0] Branch (269:15): [Folded - Ignored]
270	0	let to_read =
271	0	core::cmp::min(self.resumable_chunk_size, (max_size - total_size) as usize);
272	0	let chunk = reader.consume(Some(to_read)).await?;
273
274	0	if chunk.is_empty() { Branch (274:16): [True: 0, False: 0] Branch (274:16): [Folded - Ignored]
275	0	break;
276	0	}
277	0
278	0	data.extend_from_slice(&chunk);
279	0	total_size += chunk.len() as i64;
280		}
281
282	0	self.write_object(object_path, data).await
283	0	}
284
285		/// Implementing a resumable upload using the resumable upload API
286	0	async fn try_resumable_upload(
287	0	&self,
288	0	object_path: &ObjectPath,
289	0	reader: &mut DropCloserReadHalf,
290	0	max_size: i64,
291	0	) -> Result<(), Error> {
292	0	self.with_connection(\|\| async {
293	0	let request = UploadObjectRequest {
294	0	bucket: object_path.bucket.clone(),
295	0	..Default::default()
296	0	};
297	0
298	0	let mut metadata = HashMap::<String, String>::new();
299	0	metadata.insert("name".to_string(), object_path.path.clone());
300	0
301	0	// Use Multipart upload type with metadata
302	0	let upload_type = UploadType::Multipart(Box::new(Object {
303	0	name: object_path.path.clone(),
304	0	content_type: Some(DEFAULT_CONTENT_TYPE.to_string()),
305	0	metadata: Some(metadata),
306	0	..Default::default()
307	0	}));
308
309		// Prepare resumable upload
310	0	let uploader = self
311	0	.client
312	0	.prepare_resumable_upload(&request, &upload_type)
313	0	.await
314	0	.map_err(\|e\| self.handle_gcs_error(&e))?;
315
316		// Upload data in chunks
317	0	let mut offset: u64 = 0;
318	0	let mut total_uploaded: i64 = 0;
319
320	0	while total_uploaded < max_size { Branch (320:19): [True: 0, False: 0] Branch (320:19): [Folded - Ignored]
321	0	let to_read = core::cmp::min(
322	0	self.resumable_chunk_size,
323	0	(max_size - total_uploaded) as usize,
324		);
325	0	let chunk = reader.consume(Some(to_read)).await?;
326
327	0	if chunk.is_empty() { Branch (327:20): [True: 0, False: 0] Branch (327:20): [Folded - Ignored]
328	0	break;
329	0	}
330	0
331	0	let chunk_size = chunk.len() as u64;
332	0	total_uploaded += chunk.len() as i64;
333
334	0	let is_final = total_uploaded >= max_size \|\| chunk.len() < to_read; Branch (334:32): [True: 0, False: 0] Branch (334:32): [Folded - Ignored]
335	0	let total_size = if is_final { Branch (335:37): [True: 0, False: 0] Branch (335:37): [Folded - Ignored]
336	0	Some(offset + chunk_size)
337		} else {
338	0	Some(max_size as u64)
339		};
340
341	0	let chunk_def = ChunkSize::new(offset, offset + chunk_size - 1, total_size);
342
343		// Upload chunk
344	0	let status = uploader
345	0	.upload_multiple_chunk(chunk, &chunk_def)
346	0	.await
347	0	.map_err(\|e\| self.handle_gcs_error(&e))?;
348
349		// Update offset for next chunk
350	0	offset += chunk_size;
351	0
352	0	if let UploadStatus::Ok(_) = status { Branch (352:24): [True: 0, False: 0] Branch (352:24): [Folded - Ignored]
353	0	break;
354	0	}
355		}
356
357		// If nothing was uploaded, finalizing with empty content
358	0	if offset == 0 { Branch (358:16): [True: 0, False: 0] Branch (358:16): [Folded - Ignored]
359	0	let chunk_def = ChunkSize::new(0, 0, Some(0));
360	0	uploader
361	0	.upload_multiple_chunk(Vec::new(), &chunk_def)
362	0	.await
363	0	.map_err(\|e\| self.handle_gcs_error(&e))?;
364	0	}
365
366		// Check if the object exists
367	0	match self.read_object_metadata(object_path).await? {
368	0	Some(_) => Ok(()),
369	0	None => Err(make_err!(
370	0	Code::Internal,
371	0	"Upload completed but object not found"
372	0	)),
373		}
374	0	})
375	0	.await
376	0	}
377		}
378
379		impl Debug for GcsClient {
380	0	fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
381	0	f.debug_struct("GcsClient")
382	0	.field("resumable_chunk_size", &self.resumable_chunk_size)
383	0	.field("max_connections", &self.semaphore.available_permits())
384	0	.finish_non_exhaustive()
385	0	}
386		}
387
388		#[async_trait]
389		impl GcsOperations for GcsClient {
390		async fn read_object_metadata(
391		&self,
392		object_path: &ObjectPath,
393	0	) -> Result<Option<GcsObject>, Error> {
394	0	self.with_connection(\|\| async {
395	0	let request = GetObjectRequest {
396	0	bucket: object_path.bucket.clone(),
397	0	object: object_path.path.clone(),
398	0	..Default::default()
399	0	};
400	0
401	0	match self.client.get_object(&request).await {
402	0	Ok(obj) => Ok(Some(self.convert_to_gcs_object(obj))),
403	0	Err(err) => {
404	0	if let GcsError::Response(resp) = &err { Branch (404:28): [True: 0, False: 0] Branch (404:28): [Folded - Ignored]
405	0	if resp.code == 404 { Branch (405:28): [True: 0, False: 0] Branch (405:28): [Folded - Ignored]
406	0	return Ok(None);
407	0	}
408	0	}
409	0	Err(self.handle_gcs_error(&err))
410		}
411		}
412	0	})
413	0	.await
414	0	}
415
416		async fn read_object_content(
417		&self,
418		object_path: &ObjectPath,
419		start: i64,
420		end: Option<i64>,
421	0	) -> Result<Vec<u8>, Error> {
422	0	self.with_connection(\|\| async {
423	0	let request = GetObjectRequest {
424	0	bucket: object_path.bucket.clone(),
425	0	object: object_path.path.clone(),
426	0	..Default::default()
427	0	};
428
429	0	let range = if start > 0 \|\| end.is_some() { Branch (429:28): [True: 0, False: 0] Branch (429:41): [True: 0, False: 0] Branch (429:28): [Folded - Ignored] Branch (429:41): [Folded - Ignored]
430		Range(
431	0	if start > 0 { Some(start as u64) } else { None }, Branch (431:24): [True: 0, False: 0] Branch (431:24): [Folded - Ignored]
432	0	end.map(\|e\| e as u64),
433		)
434		} else {
435	0	Range(None, None)
436		};
437
438		// Download the object
439	0	let content = self
440	0	.client
441	0	.download_object(&request, &range)
442	0	.await
443	0	.map_err(\|e\| self.handle_gcs_error(&e))?;
444
445	0	Ok(content)
446	0	})
447	0	.await
448	0	}
449
450	0	async fn write_object(&self, object_path: &ObjectPath, content: Vec<u8>) -> Result<(), Error> {
451	0	self.with_connection(\|\| async {
452	0	let request = UploadObjectRequest {
453	0	bucket: object_path.bucket.clone(),
454	0	..Default::default()
455	0	};
456	0
457	0	let media = Media::new(object_path.path.clone());
458	0	let upload_type = UploadType::Simple(media);
459	0
460	0	self.client
461	0	.upload_object(&request, content, &upload_type)
462	0	.await
463	0	.map_err(\|e\| self.handle_gcs_error(&e))?;
464
465	0	Ok(())
466	0	})
467	0	.await
468	0	}
469
470	0	async fn start_resumable_write(&self, object_path: &ObjectPath) -> Result<String, Error> {
471	0	self.with_connection(\|\| async {
472	0	let request = UploadObjectRequest {
473	0	bucket: object_path.bucket.clone(),
474	0	..Default::default()
475	0	};
476	0
477	0	let upload_type = UploadType::Multipart(Box::new(Object {
478	0	name: object_path.path.clone(),
479	0	content_type: Some(DEFAULT_CONTENT_TYPE.to_string()),
480	0	..Default::default()
481	0	}));
482
483		// Start resumable upload session
484	0	let uploader = self
485	0	.client
486	0	.prepare_resumable_upload(&request, &upload_type)
487	0	.await
488	0	.map_err(\|e\| self.handle_gcs_error(&e))?;
489
490	0	Ok(uploader.url().to_string())
491	0	})
492	0	.await
493	0	}
494
495		async fn upload_chunk(
496		&self,
497		upload_url: &str,
498		_object_path: &ObjectPath,
499		data: Vec<u8>,
500		offset: i64,
501		end_offset: i64,
502		is_final: bool,
503	0	) -> Result<(), Error> {
504	0	self.with_connection(\|\| async {
505	0	let uploader = self.client.get_resumable_upload(upload_url.to_string());
506
507	0	let total_size = if is_final { Branch (507:33): [True: 0, False: 0] Branch (507:33): [Folded - Ignored]
508	0	Some(end_offset as u64)
509		} else {
510	0	None
511		};
512
513	0	let chunk_def = ChunkSize::new(offset as u64, end_offset as u64 - 1, total_size);
514	0
515	0	// Upload chunk
516	0	uploader
517	0	.upload_multiple_chunk(data, &chunk_def)
518	0	.await
519	0	.map_err(\|e\| self.handle_gcs_error(&e))?;
520
521	0	Ok(())
522	0	})
523	0	.await
524	0	}
525
526		async fn upload_from_reader(
527		&self,
528		object_path: &ObjectPath,
529		reader: &mut DropCloserReadHalf,
530		_upload_id: &str,
531		max_size: i64,
532	0	) -> Result<(), Error> {
533	0	let mut retry_count = 0;
534	0	let mut retry_delay = INITIAL_UPLOAD_RETRY_DELAY_MS;
535
536		loop {
537	0	let result = if max_size < SIMPLE_UPLOAD_THRESHOLD { Branch (537:29): [True: 0, False: 0] Branch (537:29): [Folded - Ignored]
538	0	self.read_and_upload_all(object_path, reader, max_size)
539	0	.await
540		} else {
541	0	self.try_resumable_upload(object_path, reader, max_size)
542	0	.await
543		};
544
545	0	match result {
546	0	Ok(()) => return Ok(()),
547	0	Err(e) => {
548	0	let is_retriable = matches!(
549	0	e.code,
550		Code::Unavailable \| Code::ResourceExhausted \| Code::DeadlineExceeded
551	0	) \|\| (e.code == Code::Internal Branch (551:27): [True: 0, False: 0] Branch (551:27): [Folded - Ignored]
552	0	&& e.to_string().contains("connection"));
553
554	0	if !is_retriable \|\| retry_count >= MAX_UPLOAD_RETRIES { Branch (554:24): [True: 0, False: 0] Branch (554:41): [True: 0, False: 0] Branch (554:24): [Folded - Ignored] Branch (554:41): [Folded - Ignored]
555	0	return Err(e);
556	0	}
557
558	0	if let Err(reset_err) = reader.try_reset_stream() { Branch (558:28): [True: 0, False: 0] Branch (558:28): [Folded - Ignored]
559	0	return Err(e.merge(reset_err));
560	0	}
561	0
562	0	sleep(Duration::from_millis(retry_delay)).await;
563	0	retry_delay = core::cmp::min(retry_delay * 2, MAX_UPLOAD_RETRY_DELAY_MS);
564	0
565	0	let mut rng = rand::rng();
566	0	let jitter_factor = rng.random::<f64>().mul_add(0.4, 0.8);
567	0	retry_delay = (retry_delay as f64 * jitter_factor) as u64;
568	0
569	0	retry_count += 1;
570		}
571		}
572		}
573	0	}
574
575	0	async fn object_exists(&self, object_path: &ObjectPath) -> Result<bool, Error> {
576	0	let metadata = self.read_object_metadata(object_path).await?;
577	0	Ok(metadata.is_some())
578	0	}
579		}