/build/source/nativelink-store/src/gcs_store.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;
use core::pin::Pin;
use core::time::Duration;
use std::borrow::Cow;
use std::sync::Arc;

use async_trait::async_trait;
use bytes::Bytes;
use futures::stream::{FuturesUnordered, unfold};
use futures::{StreamExt, TryStreamExt};
use nativelink_config::stores::ExperimentalGcsSpec;
use nativelink_error::{Code, Error, ResultExt, make_err};
use nativelink_metric::MetricsComponent;
use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatus, HealthStatusIndicator};
use nativelink_util::instant_wrapper::InstantWrapper;
use nativelink_util::retry::{Retrier, RetryResult};
use nativelink_util::store_trait::{
    RemoveCallback, StoreDriver, StoreKey, StoreOptimizations, UploadSizeInfo,
};
use rand::Rng;
use tokio::time::{sleep, timeout};
use tracing::warn;

use crate::cas_utils::is_zero_digest;
use crate::gcs_client::client::{GcsClient, GcsOperations};
use crate::gcs_client::types::{
    CHUNK_SIZE, DEFAULT_CONCURRENT_UPLOADS, DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST,
    MIN_MULTIPART_SIZE, ObjectPath,
};

#[derive(MetricsComponent, Debug)]
pub struct GcsStore<Client: GcsOperations, NowFn> {
    client: Arc<Client>,
    now_fn: NowFn,
    #[metric(help = "The bucket name for the GCS store")]
    bucket: String,
    #[metric(help = "The key prefix for the GCS store")]
    key_prefix: String,
    retrier: Retrier,
    #[metric(help = "The number of seconds to consider an object expired")]
    consider_expired_after_s: i64,
    #[metric(help = "The number of bytes to buffer for retrying requests")]
    max_retry_buffer_size: usize,
    #[metric(help = "The size of chunks for resumable uploads")]
    max_chunk_size: usize,
    #[metric(help = "The number of concurrent uploads allowed")]
    max_concurrent_uploads: usize,
}

impl<I, NowFn> GcsStore<GcsClient, NowFn>
where
    I: InstantWrapper,
    NowFn: Fn() -> I + Send + Sync + Unpin + 'static,
{
    pub async fn new(spec: &ExperimentalGcsSpec, now_fn: NowFn) -> Result<Arc<Self>, Error> {
        let client = Arc::new(GcsClient::new(spec).await?);
        Self::new_with_ops(spec, client, now_fn)
    }
}

impl<I, Client, NowFn> GcsStore<Client, NowFn>
where
    I: InstantWrapper,
    Client: GcsOperations + Send + Sync,
    NowFn: Fn() -> I + Send + Sync + Unpin + 'static,
{
    // Primarily used for injecting a mock or real operations implementation
    pub fn new_with_ops(
        spec: &ExperimentalGcsSpec,
        client: Arc<Client>,
        now_fn: NowFn,
    ) -> Result<Arc<Self>, Error> {
        // Chunks must be a multiple of 256kb according to the documentation.
        const CHUNK_MULTIPLE: usize = 256 * 1024;

        let max_connections = spec
            .common
            .multipart_max_concurrent_uploads
            .unwrap_or(DEFAULT_CONCURRENT_UPLOADS);

        let jitter_amt = spec.common.retry.jitter;
        let jitter_fn = Arc::new(move |delay: Duration| {2
            if jitter_amt == 0.0 {
                return delay;
            }
            delay.mul_f32(jitter_amt.mul_add(rand::rng().random::<f32>() - 0.5, 1.))
        });

        let max_chunk_size =
            core::cmp::min(spec.resumable_chunk_size.unwrap_or(CHUNK_SIZE), CHUNK_SIZE);

        let max_chunk_size = if max_chunk_size.is_multiple_of(CHUNK_MULTIPLE) {
            max_chunk_size
        } else {
            ((max_chunk_size + CHUNK_MULTIPLE / 2) / CHUNK_MULTIPLE) * CHUNK_MULTIPLE
        };

        let max_retry_buffer_size = spec
            .common
            .max_retry_buffer_per_request
            .unwrap_or(DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST);

        // The retry buffer should be at least as big as the chunk size.
        let max_retry_buffer_size = if max_retry_buffer_size < max_chunk_size {
            max_chunk_size
        } else {
            max_retry_buffer_size
        };

        Ok(Arc::new(Self {
            client,
            now_fn,
            bucket: spec.bucket.clone(),
            key_prefix: spec
                .common
                .key_prefix
                .as_ref()
                .unwrap_or(&String::new())
                .clone(),
            retrier: Retrier::new(
                Arc::new(|duration| Box::pin2(sleep2(duration2))),
                jitter_fn,
                spec.common.retry.clone(),
            ),
            consider_expired_after_s: i64::from(spec.common.consider_expired_after_s),
            max_retry_buffer_size,
            max_chunk_size,
            max_concurrent_uploads: max_connections,
        }))
    }

    async fn has(self: Pin<&Self>, key: &StoreKey<'_>) -> Result<Option<u64>, Error> {
        let object_path = self.make_object_path(key);
        let client = &self.client;
        let consider_expired_after_s = self.consider_expired_after_s;
        let now_fn = &self.now_fn;

        self.retrier
            .retry(unfold(object_path, move |object_path| async move {
                match client.read_object_metadata(&object_path).await.err_tip(|| {2
                    format!(
                        "Error while trying to read - bucket: {} path: {}",
                        object_path.bucket, object_path.path
                    )
                }) {
                    Ok(Some(metadata)) => {
                        if consider_expired_after_s != 0
                            && let Some(update_time) = &metadata.update_time
                        {
                            let now_s = now_fn().unix_timestamp() as i64;
                            if update_time.seconds + consider_expired_after_s <= now_s {
                                return Some((RetryResult::Ok(None), object_path));
                            }
                        }

                        if metadata.size >= 0 {
                            Some((RetryResult::Ok(Some(metadata.size as u64)), object_path))
                        } else {
                            Some((
                                RetryResult::Err(make_err!(
                                    Code::InvalidArgument,
                                    "Invalid metadata size in GCS: {}",
                                    metadata.size
                                )),
                                object_path,
                            ))
                        }
                    }
                    Ok(None) => Some((RetryResult::Ok(None), object_path)),
                    Err(e1) if e.code == Code::NotFound1 => {
                        Some((RetryResult::Ok(None), object_path))
                    }
                    Err(e) => Some((RetryResult::Retry(e), object_path)),
                }
            }))
            .await
    }

    fn make_object_path(&self, key: &StoreKey) -> ObjectPath {
        ObjectPath::new(
            self.bucket.clone(),
            &format!("{}{}", self.key_prefix, key.as_str()),
        )
    }
}

#[async_trait]
impl<I, Client, NowFn> StoreDriver for GcsStore<Client, NowFn>
where
    I: InstantWrapper,
    Client: GcsOperations + 'static,
    NowFn: Fn() -> I + Send + Sync + Unpin + 'static,
{
    async fn post_init(self: Arc<Self>) -> Result<(), Error> {
        Ok(())
    }

    async fn has_with_results(
        self: Pin<&Self>,
        keys: &[StoreKey<'_>],
        results: &mut [Option<u64>],
    ) -> Result<(), Error> {
        keys.iter()
            .zip(results.iter_mut())
            .map(|(key, result)| async move {
                if is_zero_digest(key.borrow()) {
                    *result = Some(0);
                    return Ok(());
                }
                *result = self.has(key).await?1;
                Ok(())
            })
            .collect::<FuturesUnordered<_>>()
            .try_collect()
            .await
    }

    fn optimized_for(&self, optimization: StoreOptimizations) -> bool {
        matches!(optimization, StoreOptimizations::LazyExistenceOnSync)
    }

    async fn update(
        self: Pin<&Self>,
        digest: StoreKey<'_>,
        mut reader: DropCloserReadHalf,
        upload_size: UploadSizeInfo,
    ) -> Result<u64, Error> {
        if is_zero_digest(digest.borrow()) {
            return reader.recv().await.and_then(|should_be_empty| {
                if should_be_empty.is_empty() {
                    Ok(0)
                } else {
                    Err(make_err!(Code::Internal, "Zero byte hash not empty"))
                }
            });
        }

        let object_path = self.make_object_path(&digest);

        reader.set_max_recent_data_size(
            u64::try_from(self.max_retry_buffer_size)
                .err_tip(|| "Could not convert max_retry_buffer_size to u64")?,
        );

        // For small files with exact size, we'll use simple upload
        if let UploadSizeInfo::ExactSize(size) = upload_size
            && size < MIN_MULTIPART_SIZE
        {
            let content = reader.consume(Some(usize::try_from(size)?)).await?;
            let content_len = content.len() as u64;
            let client = &self.client;

            return self
                .retrier
                .retry(unfold(content, |content| async {
                    match client.write_object(&object_path, content.to_vec()).await {
                        Ok(()) => Some((RetryResult::Ok(content_len), content)),
                        Err(e) => Some((RetryResult::Retry(e), content)),
                    }
                }))
                .await;
        }

        // For larger files, we'll use resumable upload
        // Stream and upload data in chunks
        let mut offset = 0u64;
        let mut total_size = if let UploadSizeInfo::ExactSize(size) = upload_size {
            Some(size)
        } else {
            None
        };
        let mut upload_id: Option<String> = None;
        let client = &self.client;

        loop {
            let chunk = reader.consume(Some(self.max_chunk_size)).await?;
            if chunk.is_empty() {
                break;
            }
            // If a full chunk wasn't read, then this is the full length.
            if chunk.len() < self.max_chunk_size {
                total_size = Some(offset + chunk.len() as u64);
            }

            let upload_id_ref = if let Some(upload_id_ref) = &upload_id {
                upload_id_ref
            } else {
                // Initiate the upload session on the first non-empty chunk.
                upload_id = Some(
                    self.retrier
                        .retry(unfold((), |()| async {
                            match client.start_resumable_write(&object_path).await {
                                Ok(id) => Some((RetryResult::Ok(id), ())),
                                Err(e) => Some((
                                    RetryResult::Retry(make_err!(
                                        Code::Aborted,
                                        "Failed to start resumable upload: {:?}",
                                        e
                                    )),
                                    (),
                                )),
                            }
                        }))
                        .await?,
                );
                upload_id.as_deref().unwrap()
            };

            let current_offset = offset;
            offset += chunk.len() as u64;

            // Uploading the chunk with a retry
            let object_path_ref = &object_path;
            self.retrier
                .retry(unfold(chunk, |chunk| async move {
                    match client
                        .upload_chunk(
                            upload_id_ref,
                            object_path_ref,
                            chunk.clone(),
                            current_offset,
                            offset,
                            total_size,
                        )
                        .await
                    {
                        Ok(()) => Some((RetryResult::Ok(()), chunk)),
                        Err(e) => Some((RetryResult::Retry(e), chunk)),
                    }
                }))
                .await?;
        }

        // Handle the case that the stream was of unknown length and
        // happened to be an exact multiple of chunk size.
        if let Some(upload_id_ref) = &upload_id {
            if total_size.is_none() {
                let object_path_ref = &object_path;
                self.retrier
                    .retry(unfold((), |()| async move {
                        match client
                            .upload_chunk(
                                upload_id_ref,
                                object_path_ref,
                                Bytes::new(),
                                offset,
                                offset,
                                Some(offset),
                            )
                            .await
                        {
                            Ok(()) => Some((RetryResult::Ok(offset), ())),
                            Err(e) => Some((RetryResult::Retry(e), ())),
                        }
                    }))
                    .await?;
            }
        } else {
            // Handle streamed empty file.
            return self
                .retrier
                .retry(unfold((), |()| async {
                    match client.write_object(&object_path, Vec::new()).await {
                        Ok(()) => Some((RetryResult::Ok(0), ())),
                        Err(e) => Some((RetryResult::Retry(e), ())),
                    }
                }))
                .await;
        }

        // Verifying if the upload was successful
        self.retrier
            .retry(unfold((), |()| async {
                match client.object_exists(&object_path).await {
                    Ok(true) => Some((RetryResult::Ok(()), ())),
                    Ok(false) => Some((
                        RetryResult::Retry(make_err!(
                            Code::Internal,
                            "Object not found after upload completion"
                        )),
                        (),
                    )),
                    Err(e) => Some((RetryResult::Retry(e), ())),
                }
            }))
            .await?;

        Ok(offset)
    }

    async fn get_part(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        writer: &mut DropCloserWriteHalf,
        offset: u64,
        length: Option<u64>,
    ) -> Result<(), Error> {
        if is_zero_digest(key.borrow()) {
            writer.send_eof()?;
            return Ok(());
        }

        let object_path = self.make_object_path(&key);
        let end_offset = length.map(|len| offset + len);
        let client = &self.client;

        let object_path_ref = &object_path;
        self.retrier
            .retry(unfold(
                (offset, writer),
                |(mut offset, writer)| async move {
                    let mut stream2 = match client
                        .read_object_content(object_path_ref, offset, end_offset)
                        .await
                    {
                        Ok(stream) => stream,
                        // NotFound is intentionally not special-cased here:
                        // the retrier doesn't retry NotFound by default, but
                        // emitting `Retry` lets `retry.retry_on_errors` opt
                        // reads into retrying read-after-write races where
                        // an object is still finalizing or being repopulated.
                        Err(e) => return Some((RetryResult::Retry(e), (offset, writer))),
                    };

                    while let Some(next_chunk2) = stream.next().await {
                        match next_chunk {
                            Ok(bytes) => {
                                offset += bytes.len() as u64;
                                if let Err(err0) = writer.send(bytes).await {
                                    return Some((RetryResult::Err(err), (offset, writer)));
                                }
                            }
                            Err(err) => return Some((RetryResult::Retry(err), (offset, writer))),
                        }
                    }

                    if let Err(err0) = writer.send_eof() {
                        return Some((RetryResult::Err(err), (offset, writer)));
                    }

                    Some((RetryResult::Ok(()), (offset, writer)))
                },
            ))
            .await
    }

    fn inner_store(&self, _digest: Option<StoreKey>) -> &'_ dyn StoreDriver {
        self
    }

    fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) {
        self
    }

    fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
        self
    }

    fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) {
        registry.register_indicator(self);
    }

    fn register_remove_callback(self: Arc<Self>, _callback: RemoveCallback) -> Result<(), Error> {
        // As we're backed by GCS, this store doesn't actually drop stuff
        // so we can actually just ignore this
        Ok(())
    }
}

#[async_trait]
impl<I, Client, NowFn> HealthStatusIndicator for GcsStore<Client, NowFn>
where
    I: InstantWrapper,
    Client: GcsOperations + 'static,
    NowFn: Fn() -> I + Send + Sync + Unpin + 'static,
{
    fn get_name(&self) -> &'static str {
        "GcsStore"
    }

    /// Lightweight probe: a single `object_exists` against a fixed
    /// never-existing path. Shares no resources with production traffic
    /// and stays well under the `HealthServer` per-indicator budget.
    async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus {
        const HEALTH_PROBE_TIMEOUT: Duration = Duration::from_secs(2);

        let probe_path = ObjectPath::new(
            self.bucket.clone(),
            "__nativelink_health_probe__/does-not-exist",
        );

        let probe = self.client.object_exists(&probe_path);
        match timeout(HEALTH_PROBE_TIMEOUT, probe).await {
            Ok(Ok(_)) => HealthStatus::new_ok(self, "GcsStore::check_health: ok".into()),
            Ok(Err(e)) => {
                warn!(?e, "GcsStore::check_health: object_exists errored");
                HealthStatus::new_failed(
                    self,
                    format!("GcsStore::check_health: object_exists errored: {e}").into(),
                )
            }
            Err(_) => {
                warn!(
                    timeout_secs = HEALTH_PROBE_TIMEOUT.as_secs(),
                    "GcsStore::check_health: probe timed out",
                );
                HealthStatus::Timeout {
                    struct_name: self.struct_name(),
                }
            }
        }
    }
}

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;
16		use core::pin::Pin;
17		use core::time::Duration;
18		use std::borrow::Cow;
19		use std::sync::Arc;
20
21		use async_trait::async_trait;
22		use bytes::Bytes;
23		use futures::stream::{FuturesUnordered, unfold};
24		use futures::{StreamExt, TryStreamExt};
25		use nativelink_config::stores::ExperimentalGcsSpec;
26		use nativelink_error::{Code, Error, ResultExt, make_err};
27		use nativelink_metric::MetricsComponent;
28		use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
29		use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatus, HealthStatusIndicator};
30		use nativelink_util::instant_wrapper::InstantWrapper;
31		use nativelink_util::retry::{Retrier, RetryResult};
32		use nativelink_util::store_trait::{
33		RemoveCallback, StoreDriver, StoreKey, StoreOptimizations, UploadSizeInfo,
34		};
35		use rand::Rng;
36		use tokio::time::{sleep, timeout};
37		use tracing::warn;
38
39		use crate::cas_utils::is_zero_digest;
40		use crate::gcs_client::client::{GcsClient, GcsOperations};
41		use crate::gcs_client::types::{
42		CHUNK_SIZE, DEFAULT_CONCURRENT_UPLOADS, DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST,
43		MIN_MULTIPART_SIZE, ObjectPath,
44		};
45
46		#[derive(MetricsComponent, Debug)]
47		pub struct GcsStore<Client: GcsOperations, NowFn> {
48		client: Arc<Client>,
49		now_fn: NowFn,
50		#[metric(help = "The bucket name for the GCS store")]
51		bucket: String,
52		#[metric(help = "The key prefix for the GCS store")]
53		key_prefix: String,
54		retrier: Retrier,
55		#[metric(help = "The number of seconds to consider an object expired")]
56		consider_expired_after_s: i64,
57		#[metric(help = "The number of bytes to buffer for retrying requests")]
58		max_retry_buffer_size: usize,
59		#[metric(help = "The size of chunks for resumable uploads")]
60		max_chunk_size: usize,
61		#[metric(help = "The number of concurrent uploads allowed")]
62		max_concurrent_uploads: usize,
63		}
64
65		impl<I, NowFn> GcsStore<GcsClient, NowFn>
66		where
67		I: InstantWrapper,
68		NowFn: Fn() -> I + Send + Sync + Unpin + 'static,
69		{
70	0	pub async fn new(spec: &ExperimentalGcsSpec, now_fn: NowFn) -> Result<Arc<Self>, Error> {
71	0	let client = Arc::new(GcsClient::new(spec).await?);
72	0	Self::new_with_ops(spec, client, now_fn)
73	0	}
74		}
75
76		impl<I, Client, NowFn> GcsStore<Client, NowFn>
77		where
78		I: InstantWrapper,
79		Client: GcsOperations + Send + Sync,
80		NowFn: Fn() -> I + Send + Sync + Unpin + 'static,
81		{
82		// Primarily used for injecting a mock or real operations implementation
83	20	pub fn new_with_ops(
84	20	spec: &ExperimentalGcsSpec,
85	20	client: Arc<Client>,
86	20	now_fn: NowFn,
87	20	) -> Result<Arc<Self>, Error> {
88		// Chunks must be a multiple of 256kb according to the documentation.
89		const CHUNK_MULTIPLE: usize = 256 * 1024;
90
91	20	let max_connections = spec
92	20	.common
93	20	.multipart_max_concurrent_uploads
94	20	.unwrap_or(DEFAULT_CONCURRENT_UPLOADS);
95
96	20	let jitter_amt = spec.common.retry.jitter;
97	20	let jitter_fn = Arc::new(move \|delay: Duration\| {2
98	2	if jitter_amt == 0.0 {
99	2	return delay;
100	0	}
101	0	delay.mul_f32(jitter_amt.mul_add(rand::rng().random::<f32>() - 0.5, 1.))
102	2	});
103
104	20	let max_chunk_size =
105	20	core::cmp::min(spec.resumable_chunk_size.unwrap_or(CHUNK_SIZE), CHUNK_SIZE);
106
107	20	let max_chunk_size = if max_chunk_size.is_multiple_of(CHUNK_MULTIPLE) {
108	0	max_chunk_size
109		} else {
110	20	((max_chunk_size + CHUNK_MULTIPLE / 2) / CHUNK_MULTIPLE) * CHUNK_MULTIPLE
111		};
112
113	20	let max_retry_buffer_size = spec
114	20	.common
115	20	.max_retry_buffer_per_request
116	20	.unwrap_or(DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST);
117
118		// The retry buffer should be at least as big as the chunk size.
119	20	let max_retry_buffer_size = if max_retry_buffer_size < max_chunk_size {
120	0	max_chunk_size
121		} else {
122	20	max_retry_buffer_size
123		};
124
125	20	Ok(Arc::new(Self {
126	20	client,
127	20	now_fn,
128	20	bucket: spec.bucket.clone(),
129	20	key_prefix: spec
130	20	.common
131	20	.key_prefix
132	20	.as_ref()
133	20	.unwrap_or(&String::new())
134	20	.clone(),
135	20	retrier: Retrier::new(
136	20	Arc::new(\|duration\| Box::pin2 ( sleep2 ( duration2 ))),
137	20	jitter_fn,
138	20	spec.common.retry.clone(),
139		),
140	20	consider_expired_after_s: i64::from(spec.common.consider_expired_after_s),
141	20	max_retry_buffer_size,
142	20	max_chunk_size,
143	20	max_concurrent_uploads: max_connections,
144		}))
145	20	}
146
147	9	async fn has(self: Pin<&Self>, key: &StoreKey<'_>) -> Result<Option<u64>, Error> {
148	9	let object_path = self.make_object_path(key);
149	9	let client = &self.client;
150	9	let consider_expired_after_s = self.consider_expired_after_s;
151	9	let now_fn = &self.now_fn;
152
153	9	self.retrier
154	9	.retry(unfold(object_path, move \|object_path\| async move {
155	9	match client.read_object_metadata(&object_path).await.err_tip(\|\| {2
156	2	format!(
157		"Error while trying to read - bucket: {} path: {}",
158		object_path.bucket, object_path.path
159		)
160	2	}) {
161	4	Ok(Some(metadata)) => {
162	4	if consider_expired_after_s != 0
163	2	&& let Some(update_time) = &metadata.update_time
164		{
165	2	let now_s = now_fn().unix_timestamp() as i64;
166	2	if update_time.seconds + consider_expired_after_s <= now_s {
167	1	return Some((RetryResult::Ok(None), object_path));
168	1	}
169	2	}
170
171	3	if metadata.size >= 0 {
172	3	Some((RetryResult::Ok(Some(metadata.size as u64)), object_path))
173		} else {
174	0	Some((
175	0	RetryResult::Err(make_err!(
176	0	Code::InvalidArgument,
177	0	"Invalid metadata size in GCS: {}",
178	0	metadata.size
179	0	)),
180	0	object_path,
181	0	))
182		}
183		}
184	3	Ok(None) => Some((RetryResult::Ok(None), object_path)),
185	2	Err( e1 ) if e.code == Code::NotFoun d1 => {
186	1	Some((RetryResult::Ok(None), object_path))
187		}
188	1	Err(e) => Some((RetryResult::Retry(e), object_path)),
189		}
190	18	}))
191	9	.await
192	9	}
193
194	16	fn make_object_path(&self, key: &StoreKey) -> ObjectPath {
195	16	ObjectPath::new(
196	16	self.bucket.clone(),
197	16	&format!("{}{}", self.key_prefix, key.as_str()),
198		)
199	16	}
200		}
201
202		#[async_trait]
203		impl<I, Client, NowFn> StoreDriver for GcsStore<Client, NowFn>
204		where
205		I: InstantWrapper,
206		Client: GcsOperations + 'static,
207		NowFn: Fn() -> I + Send + Sync + Unpin + 'static,
208		{
209	0	async fn post_init(self: Arc<Self>) -> Result<(), Error> {
210		Ok(())
211	0	}
212
213		async fn has_with_results(
214		self: Pin<&Self>,
215		keys: &[StoreKey<'_>],
216		results: &mut [Option<u64>],
217	8	) -> Result<(), Error> {
218		keys.iter()
219		.zip(results.iter_mut())
220	10	.map(\|(key, result)\| async move {
221	10	if is_zero_digest(key.borrow()) {
222	1	*result = Some(0);
223	1	return Ok(());
224	9	}
225	9	*result = self.has(key).await ?1 ;
226	8	Ok(())
227	20	})
228		.collect::<FuturesUnordered<_>>()
229		.try_collect()
230		.await
231	8	}
232
233	0	fn optimized_for(&self, optimization: StoreOptimizations) -> bool {
234	0	matches!(optimization, StoreOptimizations::LazyExistenceOnSync)
235	0	}
236
237		async fn update(
238		self: Pin<&Self>,
239		digest: StoreKey<'_>,
240		mut reader: DropCloserReadHalf,
241		upload_size: UploadSizeInfo,
242	4	) -> Result<u64, Error> {
243		if is_zero_digest(digest.borrow()) {
244	2	return reader.recv().await.and_then(\|should_be_empty\| {
245	2	if should_be_empty.is_empty() {
246	1	Ok(0)
247		} else {
248	1	Err(make_err!(Code::Internal, "Zero byte hash not empty"))
249		}
250	2	});
251		}
252
253		let object_path = self.make_object_path(&digest);
254
255		reader.set_max_recent_data_size(
256		u64::try_from(self.max_retry_buffer_size)
257		.err_tip(\|\| "Could not convert max_retry_buffer_size to u64")?,
258		);
259
260		// For small files with exact size, we'll use simple upload
261		if let UploadSizeInfo::ExactSize(size) = upload_size
262		&& size < MIN_MULTIPART_SIZE
263		{
264		let content = reader.consume(Some(usize::try_from(size)?)).await?;
265		let content_len = content.len() as u64;
266		let client = &self.client;
267
268		return self
269		.retrier
270	1	.retry(unfold(content, \|content\| async {
271	1	match client.write_object(&object_path, content.to_vec()).await {
272	1	Ok(()) => Some((RetryResult::Ok(content_len), content)),
273	0	Err(e) => Some((RetryResult::Retry(e), content)),
274		}
275	2	}))
276		.await;
277		}
278
279		// For larger files, we'll use resumable upload
280		// Stream and upload data in chunks
281		let mut offset = 0u64;
282		let mut total_size = if let UploadSizeInfo::ExactSize(size) = upload_size {
283		Some(size)
284		} else {
285		None
286		};
287		let mut upload_id: Option<String> = None;
288		let client = &self.client;
289
290		loop {
291		let chunk = reader.consume(Some(self.max_chunk_size)).await?;
292		if chunk.is_empty() {
293		break;
294		}
295		// If a full chunk wasn't read, then this is the full length.
296		if chunk.len() < self.max_chunk_size {
297		total_size = Some(offset + chunk.len() as u64);
298		}
299
300		let upload_id_ref = if let Some(upload_id_ref) = &upload_id {
301		upload_id_ref
302		} else {
303		// Initiate the upload session on the first non-empty chunk.
304		upload_id = Some(
305		self.retrier
306	1	.retry(unfold((), \|()\| async {
307	1	match client.start_resumable_write(&object_path).await {
308	1	Ok(id) => Some((RetryResult::Ok(id), ())),
309	0	Err(e) => Some((
310	0	RetryResult::Retry(make_err!(
311	0	Code::Aborted,
312	0	"Failed to start resumable upload: {:?}",
313	0	e
314	0	)),
315	0	(),
316	0	)),
317		}
318	2	}))
319		.await?,
320		);
321		upload_id.as_deref().unwrap()
322		};
323
324		let current_offset = offset;
325		offset += chunk.len() as u64;
326
327		// Uploading the chunk with a retry
328		let object_path_ref = &object_path;
329		self.retrier
330	6	.retry(unfold(chunk, \|chunk\| async move {
331	6	match client
332	6	.upload_chunk(
333	6	upload_id_ref,
334	6	object_path_ref,
335	6	chunk.clone(),
336	6	current_offset,
337	6	offset,
338	6	total_size,
339	6	)
340	6	.await
341		{
342	6	Ok(()) => Some((RetryResult::Ok(()), chunk)),
343	0	Err(e) => Some((RetryResult::Retry(e), chunk)),
344		}
345	12	}))
346		.await?;
347		}
348
349		// Handle the case that the stream was of unknown length and
350		// happened to be an exact multiple of chunk size.
351		if let Some(upload_id_ref) = &upload_id {
352		if total_size.is_none() {
353		let object_path_ref = &object_path;
354		self.retrier
355	0	.retry(unfold((), \|()\| async move {
356	0	match client
357	0	.upload_chunk(
358	0	upload_id_ref,
359	0	object_path_ref,
360	0	Bytes::new(),
361	0	offset,
362	0	offset,
363	0	Some(offset),
364	0	)
365	0	.await
366		{
367	0	Ok(()) => Some((RetryResult::Ok(offset), ())),
368	0	Err(e) => Some((RetryResult::Retry(e), ())),
369		}
370	0	}))
371		.await?;
372		}
373		} else {
374		// Handle streamed empty file.
375		return self
376		.retrier
377	0	.retry(unfold((), \|()\| async {
378	0	match client.write_object(&object_path, Vec::new()).await {
379	0	Ok(()) => Some((RetryResult::Ok(0), ())),
380	0	Err(e) => Some((RetryResult::Retry(e), ())),
381		}
382	0	}))
383		.await;
384		}
385
386		// Verifying if the upload was successful
387		self.retrier
388	1	.retry(unfold((), \|()\| async {
389	1	match client.object_exists(&object_path).await {
390	1	Ok(true) => Some((RetryResult::Ok(()), ())),
391	0	Ok(false) => Some((
392	0	RetryResult::Retry(make_err!(
393	0	Code::Internal,
394	0	"Object not found after upload completion"
395	0	)),
396	0	(),
397	0	)),
398	0	Err(e) => Some((RetryResult::Retry(e), ())),
399		}
400	2	}))
401		.await?;
402
403		Ok(offset)
404	4	}
405
406		async fn get_part(
407		self: Pin<&Self>,
408		key: StoreKey<'_>,
409		writer: &mut DropCloserWriteHalf,
410		offset: u64,
411		length: Option<u64>,
412	6	) -> Result<(), Error> {
413		if is_zero_digest(key.borrow()) {
414		writer.send_eof()?;
415		return Ok(());
416		}
417
418		let object_path = self.make_object_path(&key);
419	1	let end_offset = length.map(\|len\| offset + len);
420		let client = &self.client;
421
422		let object_path_ref = &object_path;
423		self.retrier
424		.retry(unfold(
425		(offset, writer),
426	7	\|(mut offset, writer)\| async move {
427	7	let mut stream2 = match client
428	7	.read_object_content(object_path_ref, offset, end_offset)
429	7	.await
430		{
431	2	Ok(stream) => stream,
432		// NotFound is intentionally not special-cased here:
433		// the retrier doesn't retry NotFound by default, but
434		// emitting `Retry` lets `retry.retry_on_errors` opt
435		// reads into retrying read-after-write races where
436		// an object is still finalizing or being repopulated.
437	5	Err(e) => return Some((RetryResult::Retry(e), (offset, writer))),
438		};
439
440	4	while let Some( next_chunk2 ) = stream.next().await {
441	2	match next_chunk {
442	2	Ok(bytes) => {
443	2	offset += bytes.len() as u64;
444	2	if let Err( err0 ) = writer.send(bytes).await {
445	0	return Some((RetryResult::Err(err), (offset, writer)));
446	2	}
447		}
448	0	Err(err) => return Some((RetryResult::Retry(err), (offset, writer))),
449		}
450		}
451
452	2	if let Err( err0 ) = writer.send_eof() {
453	0	return Some((RetryResult::Err(err), (offset, writer)));
454	2	}
455
456	2	Some((RetryResult::Ok(()), (offset, writer)))
457	14	},
458		))
459		.await
460	6	}
461
462	0	fn inner_store(&self, _digest: Option<StoreKey>) -> &'_ dyn StoreDriver {
463	0	self
464	0	}
465
466	0	fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) {
467	0	self
468	0	}
469
470	0	fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
471	0	self
472	0	}
473
474	0	fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) {
475	0	registry.register_indicator(self);
476	0	}
477
478	0	fn register_remove_callback(self: Arc<Self>, _callback: RemoveCallback) -> Result<(), Error> {
479		// As we're backed by GCS, this store doesn't actually drop stuff
480		// so we can actually just ignore this
481	0	Ok(())
482	0	}
483		}
484
485		#[async_trait]
486		impl<I, Client, NowFn> HealthStatusIndicator for GcsStore<Client, NowFn>
487		where
488		I: InstantWrapper,
489		Client: GcsOperations + 'static,
490		NowFn: Fn() -> I + Send + Sync + Unpin + 'static,
491		{
492	0	fn get_name(&self) -> &'static str {
493	0	"GcsStore"
494	0	}
495
496		/// Lightweight probe: a single `object_exists` against a fixed
497		/// never-existing path. Shares no resources with production traffic
498		/// and stays well under the `HealthServer` per-indicator budget.
499	2	async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus {
500		const HEALTH_PROBE_TIMEOUT: Duration = Duration::from_secs(2);
501
502		let probe_path = ObjectPath::new(
503		self.bucket.clone(),
504		"__nativelink_health_probe__/does-not-exist",
505		);
506
507		let probe = self.client.object_exists(&probe_path);
508		match timeout(HEALTH_PROBE_TIMEOUT, probe).await {
509		Ok(Ok(_)) => HealthStatus::new_ok(self, "GcsStore::check_health: ok".into()),
510		Ok(Err(e)) => {
511		warn!(?e, "GcsStore::check_health: object_exists errored");
512		HealthStatus::new_failed(
513		self,
514		format!("GcsStore::check_health: object_exists errored: {e}").into(),
515		)
516		}
517		Err(_) => {
518		warn!(
519		timeout_secs = HEALTH_PROBE_TIMEOUT.as_secs(),
520		"GcsStore::check_health: probe timed out",
521		);
522		HealthStatus::Timeout {
523		struct_name: self.struct_name(),
524		}
525		}
526		}
527	2	}
528		}