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

Source
// Copyright 2024 The NativeLink Authors. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// 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::sync::Arc;

use bytes::Bytes;
use futures::Stream;
use gcloud_auth::credentials::CredentialsFile;
use gcloud_storage::client::{Client, ClientConfig};
use gcloud_storage::http::Error as GcsError;
use gcloud_storage::http::objects::Object;
use gcloud_storage::http::objects::download::Range;
use gcloud_storage::http::objects::get::GetObjectRequest;
use gcloud_storage::http::objects::upload::{Media, UploadObjectRequest, UploadType};
use gcloud_storage::http::resumable_upload_client::{ChunkSize, UploadStatus};
use nativelink_config::stores::ExperimentalGcsSpec;
use nativelink_error::{Code, Error, ResultExt, make_err};
use nativelink_util::buf_channel::DropCloserReadHalf;
use rand::Rng;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use tokio::time::sleep;

use crate::common_s3_utils::install_default_rustls_crypto_provider;
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.
pub trait GcsOperations: Send + Sync + Debug {
    /// Read metadata for a GCS object
    fn read_object_metadata(
        &self,
        object: &ObjectPath,
    ) -> impl Future<Output = Result<Option<GcsObject>, Error>> + Send;

    /// Read the content of a GCS object, optionally with a range
    fn read_object_content(
        &self,
        object_path: &ObjectPath,
        start: u64,
        end: Option<u64>,
    ) -> impl Future<
        Output = Result<Box<dyn Stream<Item = Result<Bytes, Error>> + Send + Unpin>, Error>,
    > + Send;

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

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

    /// Upload a chunk of data in a resumable upload session
    fn upload_chunk(
        &self,
        upload_url: &str,
        object_path: &ObjectPath,
        data: Bytes,
        offset: u64,
        end_offset: u64,
        total_size: Option<u64>,
    ) -> impl Future<Output = Result<(), Error>> + Send;

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

    /// Check if an object exists
    fn object_exists(
        &self,
        object_path: &ObjectPath,
    ) -> impl Future<Output = Result<bool, Error>> + Send;
}

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

impl GcsClient {
    fn create_client_config(spec: &ExperimentalGcsSpec) -> Result<ClientConfig, Error> {
        install_default_rustls_crypto_provider();

        let mut client_config = ClientConfig::default();
        let connect_timeout = if spec.connection_timeout_s > 0 {
            Duration::from_secs(spec.connection_timeout_s)
        } else {
            Duration::from_secs(3)
        };
        let read_timeout = if spec.read_timeout_s > 0 {
            Duration::from_secs(spec.read_timeout_s)
        } else {
            Duration::from_secs(3)
        };
        let client = reqwest::ClientBuilder::new()
            .connect_timeout(connect_timeout)
            .read_timeout(read_timeout)
            .build()
            .map_err(|e| {0
                Error::from_std_err(Code::Internal, &e).append("Unable to create GCS client")
            })?;
        let mid_client = reqwest_middleware::ClientBuilder::new(client).build();
        client_config.http = Some(mid_client);
        Ok(client_config)
    }

    /// Create a new GCS client from the provided spec
    pub async fn new(spec: &ExperimentalGcsSpec) -> Result<Self, Error> {
        // Attempt to get the authentication from a file with the environment
        // variable GOOGLE_APPLICATION_CREDENTIALS or directly from the
        // environment in variable GOOGLE_APPLICATION_CREDENTIALS_JSON.  If that
        // fails, attempt to get authentication from the environment.
        let maybe_client_config = match CredentialsFile::new().await {
            Ok(credentials) => {
                Self::create_client_config(spec)?
                    .with_credentials(credentials)
                    .await
            }
            Err(_) => Self::create_client_config(spec)?.with_auth().await,
        }
        .map_err(|e| {
            Error::from_std_err(Code::Internal, &e)
                .append("Failed to create client config with credentials")
        });

        // If authentication is required then error, otherwise use anonymous.
        let client_config = if spec.authentication_required {
            maybe_client_config.err_tip(|| "Authentication required and none found.")?
        } else {
            maybe_client_config
                .or_else(|_| Self::create_client_config(spec).map(ClientConfig::anonymous))?
        };

        // 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)),
        })
    }

    /// Create a mock GCS client for testing
    pub fn new_mock(spec: &ExperimentalGcsSpec, endpoint: String) -> Result<Self, Error> {
        let mut client_config = ClientConfig::anonymous(Self::create_client_config(spec)?0);
        client_config.storage_endpoint = endpoint;

        let client = Client::new(client_config);
        let resumable_chunk_size = spec.resumable_chunk_size.unwrap_or(CHUNK_SIZE);
        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| {0
            Error::from_std_err(Code::Internal, &e).append("Failed to acquire connection permit")
        })?;

        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(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,
            },
            GcsError::HttpClient(resp) => match resp.status() {
                Some(http::StatusCode::NOT_FOUND) => Code::NotFound2,
                Some(http::StatusCode::UNAUTHORIZED | http::StatusCode::FORBIDDEN16) => {
                    Code::PermissionDenied
                }
                Some(http::StatusCode::REQUEST_TIMEOUT | http::StatusCode::TOO_MANY_REQUESTS12) => {
                    Code::ResourceExhausted
                }
                Some(code8) if code.is_server_error()8 => Code::Unavailable8,
                _ => Code::Unknown,
            },
            _ => Code::Internal,
        };

        Error::from_std_err(code, &err).append("GCS operation failed")
    }

    /// 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: u64,
    ) -> Result<(), Error> {
        let initial_capacity = core::cmp::min(
            usize::try_from(max_size).unwrap_or(usize::MAX),
            10 * 1024 * 1024,
        );
        let mut data = Vec::with_capacity(initial_capacity);
        let max_size = usize::try_from(max_size).unwrap_or(usize::MAX);
        let mut total_size = 0usize;

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

            if chunk.is_empty() {
                break;
            }

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

        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: u64,
    ) -> 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 max_size = usize::try_from(max_size).unwrap_or(usize::MAX);
            let mut total_uploaded = 0usize;

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

                if chunk.is_empty() {
                    break;
                }

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

                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()
    }
}

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(resp0) = &err
                        && resp.code == 404
                    {
                        return Ok(None);
                    }
                    Err(Self::handle_gcs_error(&err))
                }
            }
        })
        .await
    }

    async fn read_object_content(
        &self,
        object_path: &ObjectPath,
        start: u64,
        end: Option<u64>,
    ) -> Result<Box<dyn Stream<Item = Result<Bytes, Error>> + Send + Unpin>, Error> {
        type StreamItem = Result<Bytes, gcloud_storage::http::Error>;
        struct ReadStream<T: Stream<Item = StreamItem> + Send + Unpin> {
            stream: T,
            permit: Option<OwnedSemaphorePermit>,
        }

        impl<T: Stream<Item = StreamItem> + Send + Unpin> Stream for ReadStream<T> {
            type Item = Result<Bytes, Error>;

            fn poll_next(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
            ) -> core::task::Poll<Option<Self::Item>> {
                match std::pin::pin!(&mut self.stream).poll_next(cx) {
                    core::task::Poll::Ready(Some(Ok(bytes))) => {
                        core::task::Poll::Ready(Some(Ok(bytes)))
                    }
                    core::task::Poll::Ready(Some(Err(err))) => {
                        self.permit.take();
                        core::task::Poll::Ready(Some(Err(GcsClient::handle_gcs_error(&err))))
                    }
                    core::task::Poll::Ready(None) => {
                        self.permit.take();
                        core::task::Poll::Ready(None)
                    }
                    core::task::Poll::Pending => core::task::Poll::Pending,
                }
            }

            fn size_hint(&self) -> (usize, Option<usize>) {
                self.stream.size_hint()
            }
        }

        let permit = self.semaphore.clone().acquire_owned().await.map_err(|e| {
            Error::from_std_err(Code::Internal, &e).append("Failed to acquire connection permit")
        })?;
        let request = GetObjectRequest {
            bucket: object_path.bucket.clone(),
            object: object_path.path.clone(),
            ..Default::default()
        };

        let start = (start > 0).then_some(start);
        let range = Range(start, end);

        // Download the object
        let stream = self
            .client
            .download_streamed_object(&request, &range)
            .await
            .map_err(|e| Self::handle_gcs_error(&e))?;

        Ok(Box::new(ReadStream {
            stream,
            permit: Some(permit),
        }))
    }

    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: Bytes,
        offset: u64,
        end_offset: u64,
        total_size: Option<u64>,
    ) -> Result<(), Error> {
        self.with_connection(|| async {
            let uploader = self.client.get_resumable_upload(upload_url.to_string());

            let last_byte = if end_offset == 0 { 0 } else { end_offset - 1 };
            let chunk_def = ChunkSize::new(offset, last_byte, 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: u64,
    ) -> 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 = Duration::from_millis(retry_delay)
                        .mul_f64(jitter_factor)
                        .as_millis()
                        .try_into()
                        .unwrap_or(u64::MAX);

                    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: 2026-07-21 15:28

Line	Count	Source
1		// Copyright 2024 The NativeLink Authors. All rights reserved.
2		//
3		// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (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		// See LICENSE file for details
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::sync::Arc;
20
21		use bytes::Bytes;
22		use futures::Stream;
23		use gcloud_auth::credentials::CredentialsFile;
24		use gcloud_storage::client::{Client, ClientConfig};
25		use gcloud_storage::http::Error as GcsError;
26		use gcloud_storage::http::objects::Object;
27		use gcloud_storage::http::objects::download::Range;
28		use gcloud_storage::http::objects::get::GetObjectRequest;
29		use gcloud_storage::http::objects::upload::{Media, UploadObjectRequest, UploadType};
30		use gcloud_storage::http::resumable_upload_client::{ChunkSize, UploadStatus};
31		use nativelink_config::stores::ExperimentalGcsSpec;
32		use nativelink_error::{Code, Error, ResultExt, make_err};
33		use nativelink_util::buf_channel::DropCloserReadHalf;
34		use rand::Rng;
35		use tokio::sync::{OwnedSemaphorePermit, Semaphore};
36		use tokio::time::sleep;
37
38		use crate::common_s3_utils::install_default_rustls_crypto_provider;
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		pub trait GcsOperations: Send + Sync + Debug {
48		/// Read metadata for a GCS object
49		fn read_object_metadata(
50		&self,
51		object: &ObjectPath,
52		) -> impl Future<Output = Result<Option<GcsObject>, Error>> + Send;
53
54		/// Read the content of a GCS object, optionally with a range
55		fn read_object_content(
56		&self,
57		object_path: &ObjectPath,
58		start: u64,
59		end: Option<u64>,
60		) -> impl Future<
61		Output = Result<Box<dyn Stream<Item = Result<Bytes, Error>> + Send + Unpin>, Error>,
62		> + Send;
63
64		/// Write object with simple upload (for smaller objects)
65		fn write_object(
66		&self,
67		object_path: &ObjectPath,
68		content: Vec<u8>,
69		) -> impl Future<Output = Result<(), Error>> + Send;
70
71		/// Start a resumable write operation and return the upload URL
72		fn start_resumable_write(
73		&self,
74		object_path: &ObjectPath,
75		) -> impl Future<Output = Result<String, Error>> + Send;
76
77		/// Upload a chunk of data in a resumable upload session
78		fn upload_chunk(
79		&self,
80		upload_url: &str,
81		object_path: &ObjectPath,
82		data: Bytes,
83		offset: u64,
84		end_offset: u64,
85		total_size: Option<u64>,
86		) -> impl Future<Output = Result<(), Error>> + Send;
87
88		/// Complete high-level operation to upload data from a reader
89		fn upload_from_reader(
90		&self,
91		object_path: &ObjectPath,
92		reader: &mut DropCloserReadHalf,
93		upload_id: &str,
94		max_size: u64,
95		) -> impl Future<Output = Result<(), Error>>;
96
97		/// Check if an object exists
98		fn object_exists(
99		&self,
100		object_path: &ObjectPath,
101		) -> impl Future<Output = Result<bool, Error>> + Send;
102		}
103
104		/// Main client for interacting with Google Cloud Storage
105		pub struct GcsClient {
106		client: Client,
107		resumable_chunk_size: usize,
108		semaphore: Arc<Semaphore>,
109		}
110
111		impl GcsClient {
112	21	fn create_client_config(spec: &ExperimentalGcsSpec) -> Result<ClientConfig, Error> {
113	21	install_default_rustls_crypto_provider();
114
115	21	let mut client_config = ClientConfig::default();
116	21	let connect_timeout = if spec.connection_timeout_s > 0 {
117	0	Duration::from_secs(spec.connection_timeout_s)
118		} else {
119	21	Duration::from_secs(3)
120		};
121	21	let read_timeout = if spec.read_timeout_s > 0 {
122	0	Duration::from_secs(spec.read_timeout_s)
123		} else {
124	21	Duration::from_secs(3)
125		};
126	21	let client = reqwest::ClientBuilder::new()
127	21	.connect_timeout(connect_timeout)
128	21	.read_timeout(read_timeout)
129	21	.build()
130	21	.map_err(\|e\| {0
131	0	Error::from_std_err(Code::Internal, &e).append("Unable to create GCS client")
132	0	})?;
133	21	let mid_client = reqwest_middleware::ClientBuilder::new(client).build();
134	21	client_config.http = Some(mid_client);
135	21	Ok(client_config)
136	21	}
137
138		/// Create a new GCS client from the provided spec
139	0	pub async fn new(spec: &ExperimentalGcsSpec) -> Result<Self, Error> {
140		// Attempt to get the authentication from a file with the environment
141		// variable GOOGLE_APPLICATION_CREDENTIALS or directly from the
142		// environment in variable GOOGLE_APPLICATION_CREDENTIALS_JSON. If that
143		// fails, attempt to get authentication from the environment.
144	0	let maybe_client_config = match CredentialsFile::new().await {
145	0	Ok(credentials) => {
146	0	Self::create_client_config(spec)?
147	0	.with_credentials(credentials)
148	0	.await
149		}
150	0	Err(_) => Self::create_client_config(spec)?.with_auth().await,
151		}
152	0	.map_err(\|e\| {
153	0	Error::from_std_err(Code::Internal, &e)
154	0	.append("Failed to create client config with credentials")
155	0	});
156
157		// If authentication is required then error, otherwise use anonymous.
158	0	let client_config = if spec.authentication_required {
159	0	maybe_client_config.err_tip(\|\| "Authentication required and none found.")?
160		} else {
161	0	maybe_client_config
162	0	.or_else(\|_\| Self::create_client_config(spec).map(ClientConfig::anonymous))?
163		};
164
165		// Creating client with the configured authentication
166	0	let client = Client::new(client_config);
167	0	let resumable_chunk_size = spec.resumable_chunk_size.unwrap_or(CHUNK_SIZE);
168
169		// Get max connections from config
170	0	let max_connections = spec
171	0	.common
172	0	.multipart_max_concurrent_uploads
173	0	.unwrap_or(DEFAULT_CONCURRENT_UPLOADS);
174
175	0	Ok(Self {
176	0	client,
177	0	resumable_chunk_size,
178	0	semaphore: Arc::new(Semaphore::new(max_connections)),
179	0	})
180	0	}
181
182		/// Create a mock GCS client for testing
183	21	pub fn new_mock(spec: &ExperimentalGcsSpec, endpoint: String) -> Result<Self, Error> {
184	21	let mut client_config = ClientConfig::anonymous(Self::create_client_config(spec) ?0 );
185	21	client_config.storage_endpoint = endpoint;
186
187	21	let client = Client::new(client_config);
188	21	let resumable_chunk_size = spec.resumable_chunk_size.unwrap_or(CHUNK_SIZE);
189	21	let max_connections = spec
190	21	.common
191	21	.multipart_max_concurrent_uploads
192	21	.unwrap_or(DEFAULT_CONCURRENT_UPLOADS);
193
194	21	Ok(Self {
195	21	client,
196	21	resumable_chunk_size,
197	21	semaphore: Arc::new(Semaphore::new(max_connections)),
198	21	})
199	21	}
200
201		/// Generic method to execute operations with connection limiting
202	21	async fn with_connection<F, Fut, T>(&self, operation: F) -> Result<T, Error>
203	21	where
204	21	F: FnOnce() -> Fut + Send,
205	21	Fut: Future<Output = Result<T, Error>> + Send,
206	21	{
207	21	let permit = self.semaphore.acquire().await.map_err(\|e\| {0
208	0	Error::from_std_err(Code::Internal, &e).append("Failed to acquire connection permit")
209	0	})?;
210
211	21	let result = operation().await;
212	21	drop(permit);
213	21	result
214	21	}
215
216		/// Convert GCS object to our internal representation
217	0	fn convert_to_gcs_object(&self, obj: Object) -> GcsObject {
218	0	let update_time = obj.updated.map(\|dt\| Timestamp {
219	0	seconds: dt.unix_timestamp(),
220		nanos: 0,
221	0	});
222
223		GcsObject {
224	0	name: obj.name,
225	0	bucket: obj.bucket,
226	0	size: obj.size,
227	0	content_type: obj
228	0	.content_type
229	0	.unwrap_or_else(\|\| DEFAULT_CONTENT_TYPE.to_string()),
230	0	update_time,
231		}
232	0	}
233
234		/// Handle error from GCS operations
235	21	fn handle_gcs_error(err: &GcsError) -> Error {
236	21	let code = match &err {
237	0	GcsError::Response(resp) => match resp.code {
238	0	404 => Code::NotFound,
239	0	401 \| 403 => Code::PermissionDenied,
240	0	408 \| 429 => Code::ResourceExhausted,
241	0	500..=599 => Code::Unavailable,
242	0	_ => Code::Unknown,
243		},
244	21	GcsError::HttpClient(resp) => match resp.status() {
245	20	Some(http::StatusCode::NOT_FOUND) => Code::NotFound2 ,
246	18	Some(http::StatusCode::UNAUTHORIZED \| http::StatusCode::FORBIDDEN16 ) => {
247	4	Code::PermissionDenied
248		}
249	14	Some(http::StatusCode::REQUEST_TIMEOUT \| http::StatusCode::TOO_MANY_REQUESTS12 ) => {
250	4	Code::ResourceExhausted
251		}
252	10	Some( code8 ) if code.is_server_error( )8 => Code::Unavailable8 ,
253	3	_ => Code::Unknown,
254		},
255	0	_ => Code::Internal,
256		};
257
258	21	Error::from_std_err(code, &err).append("GCS operation failed")
259	21	}
260
261		/// Reading data from reader and upload in a single operation
262	0	async fn read_and_upload_all(
263	0	&self,
264	0	object_path: &ObjectPath,
265	0	reader: &mut DropCloserReadHalf,
266	0	max_size: u64,
267	0	) -> Result<(), Error> {
268	0	let initial_capacity = core::cmp::min(
269	0	usize::try_from(max_size).unwrap_or(usize::MAX),
270	0	10 * 1024 * 1024,
271		);
272	0	let mut data = Vec::with_capacity(initial_capacity);
273	0	let max_size = usize::try_from(max_size).unwrap_or(usize::MAX);
274	0	let mut total_size = 0usize;
275
276	0	while total_size < max_size {
277	0	let to_read = core::cmp::min(self.resumable_chunk_size, max_size - total_size);
278	0	let chunk = reader.consume(Some(to_read)).await?;
279
280	0	if chunk.is_empty() {
281	0	break;
282	0	}
283
284	0	data.extend_from_slice(&chunk);
285	0	total_size += chunk.len();
286		}
287
288	0	self.write_object(object_path, data).await
289	0	}
290
291		/// Implementing a resumable upload using the resumable upload API
292	0	async fn try_resumable_upload(
293	0	&self,
294	0	object_path: &ObjectPath,
295	0	reader: &mut DropCloserReadHalf,
296	0	max_size: u64,
297	0	) -> Result<(), Error> {
298	0	self.with_connection(\|\| async {
299	0	let request = UploadObjectRequest {
300	0	bucket: object_path.bucket.clone(),
301	0	..Default::default()
302	0	};
303
304	0	let mut metadata = HashMap::<String, String>::new();
305	0	metadata.insert("name".to_string(), object_path.path.clone());
306
307		// Use Multipart upload type with metadata
308	0	let upload_type = UploadType::Multipart(Box::new(Object {
309	0	name: object_path.path.clone(),
310	0	content_type: Some(DEFAULT_CONTENT_TYPE.to_string()),
311	0	metadata: Some(metadata),
312	0	..Default::default()
313	0	}));
314
315		// Prepare resumable upload
316	0	let uploader = self
317	0	.client
318	0	.prepare_resumable_upload(&request, &upload_type)
319	0	.await
320	0	.map_err(\|e\| Self::handle_gcs_error(&e))?;
321
322		// Upload data in chunks
323	0	let mut offset: u64 = 0;
324	0	let max_size = usize::try_from(max_size).unwrap_or(usize::MAX);
325	0	let mut total_uploaded = 0usize;
326
327	0	while total_uploaded < max_size {
328	0	let to_read = core::cmp::min(self.resumable_chunk_size, max_size - total_uploaded);
329	0	let chunk = reader.consume(Some(to_read)).await?;
330
331	0	if chunk.is_empty() {
332	0	break;
333	0	}
334
335	0	let chunk_size = chunk.len() as u64;
336	0	total_uploaded += chunk.len();
337
338	0	let is_final = total_uploaded >= max_size \|\| chunk.len() < to_read;
339	0	let total_size = if is_final {
340	0	Some(offset + chunk_size)
341		} else {
342	0	Some(max_size as u64)
343		};
344
345	0	let chunk_def = ChunkSize::new(offset, offset + chunk_size - 1, total_size);
346
347		// Upload chunk
348	0	let status = uploader
349	0	.upload_multiple_chunk(chunk, &chunk_def)
350	0	.await
351	0	.map_err(\|e\| Self::handle_gcs_error(&e))?;
352
353		// Update offset for next chunk
354	0	offset += chunk_size;
355
356	0	if let UploadStatus::Ok(_) = status {
357	0	break;
358	0	}
359		}
360
361		// If nothing was uploaded, finalizing with empty content
362	0	if offset == 0 {
363	0	let chunk_def = ChunkSize::new(0, 0, Some(0));
364	0	uploader
365	0	.upload_multiple_chunk(Vec::new(), &chunk_def)
366	0	.await
367	0	.map_err(\|e\| Self::handle_gcs_error(&e))?;
368	0	}
369
370		// Check if the object exists
371	0	match self.read_object_metadata(object_path).await? {
372	0	Some(_) => Ok(()),
373	0	None => Err(make_err!(
374	0	Code::Internal,
375	0	"Upload completed but object not found"
376	0	)),
377		}
378	0	})
379	0	.await
380	0	}
381		}
382
383		impl Debug for GcsClient {
384	0	fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
385	0	f.debug_struct("GcsClient")
386	0	.field("resumable_chunk_size", &self.resumable_chunk_size)
387	0	.field("max_connections", &self.semaphore.available_permits())
388	0	.finish_non_exhaustive()
389	0	}
390		}
391
392		impl GcsOperations for GcsClient {
393	21	async fn read_object_metadata(
394	21	&self,
395	21	object_path: &ObjectPath,
396	21	) -> Result<Option<GcsObject>, Error> {
397	21	self.with_connection(\|\| async {
398	21	let request = GetObjectRequest {
399	21	bucket: object_path.bucket.clone(),
400	21	object: object_path.path.clone(),
401	21	..Default::default()
402	21	};
403
404	21	match self.client.get_object(&request).await {
405	0	Ok(obj) => Ok(Some(self.convert_to_gcs_object(obj))),
406	21	Err(err) => {
407	21	if let GcsError::Response( resp0 ) = &err
408	0	&& resp.code == 404
409		{
410	0	return Ok(None);
411	21	}
412	21	Err(Self::handle_gcs_error(&err))
413		}
414		}
415	42	})
416	21	.await
417	21	}
418
419	0	async fn read_object_content(
420	0	&self,
421	0	object_path: &ObjectPath,
422	0	start: u64,
423	0	end: Option<u64>,
424	0	) -> Result<Box<dyn Stream<Item = Result<Bytes, Error>> + Send + Unpin>, Error> {
425		type StreamItem = Result<Bytes, gcloud_storage::http::Error>;
426		struct ReadStream<T: Stream<Item = StreamItem> + Send + Unpin> {
427		stream: T,
428		permit: Option<OwnedSemaphorePermit>,
429		}
430
431		impl<T: Stream<Item = StreamItem> + Send + Unpin> Stream for ReadStream<T> {
432		type Item = Result<Bytes, Error>;
433
434	0	fn poll_next(
435	0	mut self: core::pin::Pin<&mut Self>,
436	0	cx: &mut core::task::Context<'_>,
437	0	) -> core::task::Poll<Option<Self::Item>> {
438	0	match std::pin::pin!(&mut self.stream).poll_next(cx) {
439	0	core::task::Poll::Ready(Some(Ok(bytes))) => {
440	0	core::task::Poll::Ready(Some(Ok(bytes)))
441		}
442	0	core::task::Poll::Ready(Some(Err(err))) => {
443	0	self.permit.take();
444	0	core::task::Poll::Ready(Some(Err(GcsClient::handle_gcs_error(&err))))
445		}
446		core::task::Poll::Ready(None) => {
447	0	self.permit.take();
448	0	core::task::Poll::Ready(None)
449		}
450	0	core::task::Poll::Pending => core::task::Poll::Pending,
451		}
452	0	}
453
454	0	fn size_hint(&self) -> (usize, Option<usize>) {
455	0	self.stream.size_hint()
456	0	}
457		}
458
459	0	let permit = self.semaphore.clone().acquire_owned().await.map_err(\|e\| {
460	0	Error::from_std_err(Code::Internal, &e).append("Failed to acquire connection permit")
461	0	})?;
462	0	let request = GetObjectRequest {
463	0	bucket: object_path.bucket.clone(),
464	0	object: object_path.path.clone(),
465	0	..Default::default()
466	0	};
467
468	0	let start = (start > 0).then_some(start);
469	0	let range = Range(start, end);
470
471		// Download the object
472	0	let stream = self
473	0	.client
474	0	.download_streamed_object(&request, &range)
475	0	.await
476	0	.map_err(\|e\| Self::handle_gcs_error(&e))?;
477
478	0	Ok(Box::new(ReadStream {
479	0	stream,
480	0	permit: Some(permit),
481	0	}))
482	0	}
483
484	0	async fn write_object(&self, object_path: &ObjectPath, content: Vec<u8>) -> Result<(), Error> {
485	0	self.with_connection(\|\| async {
486	0	let request = UploadObjectRequest {
487	0	bucket: object_path.bucket.clone(),
488	0	..Default::default()
489	0	};
490
491	0	let media = Media::new(object_path.path.clone());
492	0	let upload_type = UploadType::Simple(media);
493
494	0	self.client
495	0	.upload_object(&request, content, &upload_type)
496	0	.await
497	0	.map_err(\|e\| Self::handle_gcs_error(&e))?;
498
499	0	Ok(())
500	0	})
501	0	.await
502	0	}
503
504	0	async fn start_resumable_write(&self, object_path: &ObjectPath) -> Result<String, Error> {
505	0	self.with_connection(\|\| async {
506	0	let request = UploadObjectRequest {
507	0	bucket: object_path.bucket.clone(),
508	0	..Default::default()
509	0	};
510
511	0	let upload_type = UploadType::Multipart(Box::new(Object {
512	0	name: object_path.path.clone(),
513	0	content_type: Some(DEFAULT_CONTENT_TYPE.to_string()),
514	0	..Default::default()
515	0	}));
516
517		// Start resumable upload session
518	0	let uploader = self
519	0	.client
520	0	.prepare_resumable_upload(&request, &upload_type)
521	0	.await
522	0	.map_err(\|e\| Self::handle_gcs_error(&e))?;
523
524	0	Ok(uploader.url().to_string())
525	0	})
526	0	.await
527	0	}
528
529	0	async fn upload_chunk(
530	0	&self,
531	0	upload_url: &str,
532	0	_object_path: &ObjectPath,
533	0	data: Bytes,
534	0	offset: u64,
535	0	end_offset: u64,
536	0	total_size: Option<u64>,
537	0	) -> Result<(), Error> {
538	0	self.with_connection(\|\| async {
539	0	let uploader = self.client.get_resumable_upload(upload_url.to_string());
540
541	0	let last_byte = if end_offset == 0 { 0 } else { end_offset - 1 };
542	0	let chunk_def = ChunkSize::new(offset, last_byte, total_size);
543
544		// Upload chunk
545	0	uploader
546	0	.upload_multiple_chunk(data, &chunk_def)
547	0	.await
548	0	.map_err(\|e\| Self::handle_gcs_error(&e))?;
549
550	0	Ok(())
551	0	})
552	0	.await
553	0	}
554
555	0	async fn upload_from_reader(
556	0	&self,
557	0	object_path: &ObjectPath,
558	0	reader: &mut DropCloserReadHalf,
559	0	_upload_id: &str,
560	0	max_size: u64,
561	0	) -> Result<(), Error> {
562	0	let mut retry_count = 0;
563	0	let mut retry_delay = INITIAL_UPLOAD_RETRY_DELAY_MS;
564
565		loop {
566	0	let result = if max_size < SIMPLE_UPLOAD_THRESHOLD {
567	0	self.read_and_upload_all(object_path, reader, max_size)
568	0	.await
569		} else {
570	0	self.try_resumable_upload(object_path, reader, max_size)
571	0	.await
572		};
573
574	0	match result {
575	0	Ok(()) => return Ok(()),
576	0	Err(e) => {
577	0	let is_retriable = matches!(
578	0	e.code,
579		Code::Unavailable \| Code::ResourceExhausted \| Code::DeadlineExceeded
580	0	) \|\| (e.code == Code::Internal
581	0	&& e.to_string().contains("connection"));
582
583	0	if !is_retriable \|\| retry_count >= MAX_UPLOAD_RETRIES {
584	0	return Err(e);
585	0	}
586
587	0	if let Err(reset_err) = reader.try_reset_stream() {
588	0	return Err(e.merge(reset_err));
589	0	}
590
591	0	sleep(Duration::from_millis(retry_delay)).await;
592	0	retry_delay = core::cmp::min(retry_delay * 2, MAX_UPLOAD_RETRY_DELAY_MS);
593
594	0	let mut rng = rand::rng();
595	0	let jitter_factor = rng.random::<f64>().mul_add(0.4, 0.8);
596	0	retry_delay = Duration::from_millis(retry_delay)
597	0	.mul_f64(jitter_factor)
598	0	.as_millis()
599	0	.try_into()
600	0	.unwrap_or(u64::MAX);
601
602	0	retry_count += 1;
603		}
604		}
605		}
606	0	}
607
608	0	async fn object_exists(&self, object_path: &ObjectPath) -> Result<bool, Error> {
609	0	let metadata = self.read_object_metadata(object_path).await?;
610	0	Ok(metadata.is_some())
611	0	}
612		}