/build/source/nativelink-store/src/gcs_store.rs

Source
// Copyright 2024 Trace Machina, Inc. 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, MAX_UPLOAD_SIZE,
    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: u64,
    #[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(),
            )
            .with_client("gcs"),
            consider_expired_after_s: u64::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();
                            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!0(optimization, StoreOptimizations::LazyExistenceOnSync)
    }

    fn max_upload_size(&self) -> u64 {
        MAX_UPLOAD_SIZE
    }

    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)?.saturating_add(1)))
                .await?;
            let content_len = content.len() as u64;
            upload_size.validate_size(content_len)?;
            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() {
                upload_size.validate_size(offset)?;
                break;
            }
            let end = offset + chunk.len() as u64;
            if matches!(upload_size, UploadSizeInfo::MaxSize(_)) {
                upload_size.validate_size(end)?;
            }
            if let UploadSizeInfo::ExactSize(expected) = upload_size
                && (end >= expected || chunk.len() < self.max_chunk_size)
            {
                upload_size.validate_size(end)?;
                // Do not commit the final chunk until the producer confirms EOF.
                if !reader.recv().await?.is_empty() {
                    return Err(make_err!(
                        Code::InvalidArgument,
                        "Data after exact upload size"
                    ));
                }
            }
            // 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 stream4 = 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_chunk8) = stream.next().await {
                        match next_chunk {
                            Ok(bytes) => {
                                // HTTP/2 peers may emit zero-length DATA frames
                                // (e.g. a bare END_STREAM frame), which surface
                                // here as empty chunks. `send()` rejects empty
                                // buffers by contract, so skip them.
                                if bytes.is_empty() {
                                    continue;
                                }
                                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-11-13 23:29

Line	Count	Source
1		// Copyright 2024 Trace Machina, Inc. 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, MAX_UPLOAD_SIZE,
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: u64,
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	26	pub fn new_with_ops(
84	26	spec: &ExperimentalGcsSpec,
85	26	client: Arc<Client>,
86	26	now_fn: NowFn,
87	26	) -> 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	26	let max_connections = spec
92	26	.common
93	26	.multipart_max_concurrent_uploads
94	26	.unwrap_or(DEFAULT_CONCURRENT_UPLOADS);
95
96	26	let jitter_amt = spec.common.retry.jitter;
97	26	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	26	let max_chunk_size =
105	26	core::cmp::min(spec.resumable_chunk_size.unwrap_or(CHUNK_SIZE), CHUNK_SIZE);
106
107	26	let max_chunk_size = if max_chunk_size.is_multiple_of(CHUNK_MULTIPLE) {
108	0	max_chunk_size
109		} else {
110	26	((max_chunk_size + CHUNK_MULTIPLE / 2) / CHUNK_MULTIPLE) * CHUNK_MULTIPLE
111		};
112
113	26	let max_retry_buffer_size = spec
114	26	.common
115	26	.max_retry_buffer_per_request
116	26	.unwrap_or(DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST);
117
118		// The retry buffer should be at least as big as the chunk size.
119	26	let max_retry_buffer_size = if max_retry_buffer_size < max_chunk_size {
120	0	max_chunk_size
121		} else {
122	26	max_retry_buffer_size
123		};
124
125	26	Ok(Arc::new(Self {
126	26	client,
127	26	now_fn,
128	26	bucket: spec.bucket.clone(),
129	26	key_prefix: spec
130	26	.common
131	26	.key_prefix
132	26	.as_ref()
133	26	.unwrap_or(&String::new())
134	26	.clone(),
135	26	retrier: Retrier::new(
136	26	Arc::new(\|duration\| Box::pin2 ( sleep2 ( duration2 ))),
137	26	jitter_fn,
138	26	spec.common.retry.clone(),
139		)
140	26	.with_client("gcs"),
141	26	consider_expired_after_s: u64::from(spec.common.consider_expired_after_s),
142	26	max_retry_buffer_size,
143	26	max_chunk_size,
144	26	max_concurrent_uploads: max_connections,
145		}))
146	26	}
147
148	10	async fn has(self: Pin<&Self>, key: &StoreKey<'_>) -> Result<Option<u64>, Error> {
149	10	let object_path = self.make_object_path(key);
150	10	let client = &self.client;
151	10	let consider_expired_after_s = self.consider_expired_after_s;
152	10	let now_fn = &self.now_fn;
153
154	10	self.retrier
155	10	.retry(unfold(object_path, move \|object_path\| async move {
156	10	match client.read_object_metadata(&object_path).await.err_tip(\|\| {2
157	2	format!(
158		"Error while trying to read - bucket: {} path: {}",
159		object_path.bucket, object_path.path
160		)
161	2	}) {
162	5	Ok(Some(metadata)) => {
163	5	if consider_expired_after_s != 0
164	3	&& let Some(update_time) = &metadata.update_time
165		{
166	3	let now_s = now_fn().unix_timestamp();
167	3	if update_time.seconds + consider_expired_after_s <= now_s {
168	2	return Some((RetryResult::Ok(None), object_path));
169	1	}
170	2	}
171
172	3	if metadata.size >= 0 {
173	3	Some((RetryResult::Ok(Some(metadata.size as u64)), object_path))
174		} else {
175	0	Some((
176	0	RetryResult::Err(make_err!(
177	0	Code::InvalidArgument,
178	0	"Invalid metadata size in GCS: {}",
179	0	metadata.size
180	0	)),
181	0	object_path,
182	0	))
183		}
184		}
185	3	Ok(None) => Some((RetryResult::Ok(None), object_path)),
186	2	Err( e1 ) if e.code == Code::NotFoun d1 => {
187	1	Some((RetryResult::Ok(None), object_path))
188		}
189	1	Err(e) => Some((RetryResult::Retry(e), object_path)),
190		}
191	20	}))
192	10	.await
193	10	}
194
195	24	fn make_object_path(&self, key: &StoreKey) -> ObjectPath {
196	24	ObjectPath::new(
197	24	self.bucket.clone(),
198	24	&format!("{}{}", self.key_prefix, key.as_str()),
199		)
200	24	}
201		}
202
203		#[async_trait]
204		impl<I, Client, NowFn> StoreDriver for GcsStore<Client, NowFn>
205		where
206		I: InstantWrapper,
207		Client: GcsOperations + 'static,
208		NowFn: Fn() -> I + Send + Sync + Unpin + 'static,
209		{
210	0	async fn post_init(self: Arc<Self>) -> Result<(), Error> {
211		Ok(())
212	0	}
213
214		async fn has_with_results(
215		self: Pin<&Self>,
216		keys: &[StoreKey<'_>],
217		results: &mut [Option<u64>],
218	9	) -> Result<(), Error> {
219		keys.iter()
220		.zip(results.iter_mut())
221	11	.map(\|(key, result)\| async move {
222	11	if is_zero_digest(key.borrow()) {
223	1	*result = Some(0);
224	1	return Ok(());
225	10	}
226	10	*result = self.has(key).await ?1 ;
227	9	Ok(())
228	22	})
229		.collect::<FuturesUnordered<_>>()
230		.try_collect()
231		.await
232	9	}
233
234	4	fn optimized_for(&self, optimization: StoreOptimizations) -> bool {
235	4	matches!0 (optimization, StoreOptimizations::LazyExistenceOnSync)
236	4	}
237
238	4	fn max_upload_size(&self) -> u64 {
239	4	MAX_UPLOAD_SIZE
240	4	}
241
242		async fn update(
243		self: Pin<&Self>,
244		digest: StoreKey<'_>,
245		mut reader: DropCloserReadHalf,
246		upload_size: UploadSizeInfo,
247	8	) -> Result<u64, Error> {
248		if is_zero_digest(digest.borrow()) {
249	2	return reader.recv().await.and_then(\|should_be_empty\| {
250	2	if should_be_empty.is_empty() {
251	1	Ok(0)
252		} else {
253	1	Err(make_err!(Code::Internal, "Zero byte hash not empty"))
254		}
255	2	});
256		}
257
258		let object_path = self.make_object_path(&digest);
259
260		reader.set_max_recent_data_size(
261		u64::try_from(self.max_retry_buffer_size)
262		.err_tip(\|\| "Could not convert max_retry_buffer_size to u64")?,
263		);
264
265		// For small files with exact size, we'll use simple upload
266		if let UploadSizeInfo::ExactSize(size) = upload_size
267		&& size < MIN_MULTIPART_SIZE
268		{
269		let content = reader
270		.consume(Some(usize::try_from(size)?.saturating_add(1)))
271		.await?;
272		let content_len = content.len() as u64;
273		upload_size.validate_size(content_len)?;
274		let client = &self.client;
275
276		return self
277		.retrier
278	1	.retry(unfold(content, \|content\| async {
279	1	match client.write_object(&object_path, content.to_vec()).await {
280	1	Ok(()) => Some((RetryResult::Ok(content_len), content)),
281	0	Err(e) => Some((RetryResult::Retry(e), content)),
282		}
283	2	}))
284		.await;
285		}
286
287		// For larger files, we'll use resumable upload
288		// Stream and upload data in chunks
289		let mut offset = 0u64;
290		let mut total_size = if let UploadSizeInfo::ExactSize(size) = upload_size {
291		Some(size)
292		} else {
293		None
294		};
295		let mut upload_id: Option<String> = None;
296		let client = &self.client;
297
298		loop {
299		let chunk = reader.consume(Some(self.max_chunk_size)).await?;
300		if chunk.is_empty() {
301		upload_size.validate_size(offset)?;
302		break;
303		}
304		let end = offset + chunk.len() as u64;
305		if matches!(upload_size, UploadSizeInfo::MaxSize(_)) {
306		upload_size.validate_size(end)?;
307		}
308		if let UploadSizeInfo::ExactSize(expected) = upload_size
309		&& (end >= expected \|\| chunk.len() < self.max_chunk_size)
310		{
311		upload_size.validate_size(end)?;
312		// Do not commit the final chunk until the producer confirms EOF.
313		if !reader.recv().await?.is_empty() {
314		return Err(make_err!(
315		Code::InvalidArgument,
316		"Data after exact upload size"
317		));
318		}
319		}
320		// If a full chunk wasn't read, then this is the full length.
321		if chunk.len() < self.max_chunk_size {
322		total_size = Some(offset + chunk.len() as u64);
323		}
324
325		let upload_id_ref = if let Some(upload_id_ref) = &upload_id {
326		upload_id_ref
327		} else {
328		// Initiate the upload session on the first non-empty chunk.
329		upload_id = Some(
330		self.retrier
331	1	.retry(unfold((), \|()\| async {
332	1	match client.start_resumable_write(&object_path).await {
333	1	Ok(id) => Some((RetryResult::Ok(id), ())),
334	0	Err(e) => Some((
335	0	RetryResult::Retry(make_err!(
336	0	Code::Aborted,
337	0	"Failed to start resumable upload: {:?}",
338	0	e
339	0	)),
340	0	(),
341	0	)),
342		}
343	2	}))
344		.await?,
345		);
346		upload_id.as_deref().unwrap()
347		};
348
349		let current_offset = offset;
350		offset += chunk.len() as u64;
351
352		// Uploading the chunk with a retry
353		let object_path_ref = &object_path;
354		self.retrier
355	6	.retry(unfold(chunk, \|chunk\| async move {
356	6	match client
357	6	.upload_chunk(
358	6	upload_id_ref,
359	6	object_path_ref,
360	6	chunk.clone(),
361	6	current_offset,
362	6	offset,
363	6	total_size,
364	6	)
365	6	.await
366		{
367	6	Ok(()) => Some((RetryResult::Ok(()), chunk)),
368	0	Err(e) => Some((RetryResult::Retry(e), chunk)),
369		}
370	12	}))
371		.await?;
372		}
373
374		// Handle the case that the stream was of unknown length and
375		// happened to be an exact multiple of chunk size.
376		if let Some(upload_id_ref) = &upload_id {
377		if total_size.is_none() {
378		let object_path_ref = &object_path;
379		self.retrier
380	0	.retry(unfold((), \|()\| async move {
381	0	match client
382	0	.upload_chunk(
383	0	upload_id_ref,
384	0	object_path_ref,
385	0	Bytes::new(),
386	0	offset,
387	0	offset,
388	0	Some(offset),
389	0	)
390	0	.await
391		{
392	0	Ok(()) => Some((RetryResult::Ok(offset), ())),
393	0	Err(e) => Some((RetryResult::Retry(e), ())),
394		}
395	0	}))
396		.await?;
397		}
398		} else {
399		// Handle streamed empty file.
400		return self
401		.retrier
402	0	.retry(unfold((), \|()\| async {
403	0	match client.write_object(&object_path, Vec::new()).await {
404	0	Ok(()) => Some((RetryResult::Ok(0), ())),
405	0	Err(e) => Some((RetryResult::Retry(e), ())),
406		}
407	0	}))
408		.await;
409		}
410
411		// Verifying if the upload was successful
412		self.retrier
413	1	.retry(unfold((), \|()\| async {
414	1	match client.object_exists(&object_path).await {
415	1	Ok(true) => Some((RetryResult::Ok(()), ())),
416	0	Ok(false) => Some((
417	0	RetryResult::Retry(make_err!(
418	0	Code::Internal,
419	0	"Object not found after upload completion"
420	0	)),
421	0	(),
422	0	)),
423	0	Err(e) => Some((RetryResult::Retry(e), ())),
424		}
425	2	}))
426		.await?;
427
428		Ok(offset)
429	8	}
430
431		async fn get_part(
432		self: Pin<&Self>,
433		key: StoreKey<'_>,
434		writer: &mut DropCloserWriteHalf,
435		offset: u64,
436		length: Option<u64>,
437	9	) -> Result<(), Error> {
438		if is_zero_digest(key.borrow()) {
439		writer.send_eof()?;
440		return Ok(());
441		}
442
443		let object_path = self.make_object_path(&key);
444	1	let end_offset = length.map(\|len\| offset + len);
445		let client = &self.client;
446
447		let object_path_ref = &object_path;
448		self.retrier
449		.retry(unfold(
450		(offset, writer),
451	10	\|(mut offset, writer)\| async move {
452	10	let mut stream4 = match client
453	10	.read_object_content(object_path_ref, offset, end_offset)
454	10	.await
455		{
456	4	Ok(stream) => stream,
457		// NotFound is intentionally not special-cased here:
458		// the retrier doesn't retry NotFound by default, but
459		// emitting `Retry` lets `retry.retry_on_errors` opt
460		// reads into retrying read-after-write races where
461		// an object is still finalizing or being repopulated.
462	6	Err(e) => return Some((RetryResult::Retry(e), (offset, writer))),
463		};
464
465	12	while let Some( next_chunk8 ) = stream.next().await {
466	8	match next_chunk {
467	8	Ok(bytes) => {
468		// HTTP/2 peers may emit zero-length DATA frames
469		// (e.g. a bare END_STREAM frame), which surface
470		// here as empty chunks. `send()` rejects empty
471		// buffers by contract, so skip them.
472	8	if bytes.is_empty() {
473	3	continue;
474	5	}
475	5	offset += bytes.len() as u64;
476	5	if let Err( err0 ) = writer.send(bytes).await {
477	0	return Some((RetryResult::Err(err), (offset, writer)));
478	5	}
479		}
480	0	Err(err) => return Some((RetryResult::Retry(err), (offset, writer))),
481		}
482		}
483
484	4	if let Err( err0 ) = writer.send_eof() {
485	0	return Some((RetryResult::Err(err), (offset, writer)));
486	4	}
487
488	4	Some((RetryResult::Ok(()), (offset, writer)))
489	20	},
490		))
491		.await
492	9	}
493
494	6	fn inner_store(&self, _digest: Option<StoreKey>) -> &'_ dyn StoreDriver {
495	6	self
496	6	}
497
498	0	fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) {
499	0	self
500	0	}
501
502	0	fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
503	0	self
504	0	}
505
506	0	fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) {
507	0	registry.register_indicator(self);
508	0	}
509
510	0	fn register_remove_callback(self: Arc<Self>, _callback: RemoveCallback) -> Result<(), Error> {
511		// As we're backed by GCS, this store doesn't actually drop stuff
512		// so we can actually just ignore this
513	0	Ok(())
514	0	}
515		}
516
517		#[async_trait]
518		impl<I, Client, NowFn> HealthStatusIndicator for GcsStore<Client, NowFn>
519		where
520		I: InstantWrapper,
521		Client: GcsOperations + 'static,
522		NowFn: Fn() -> I + Send + Sync + Unpin + 'static,
523		{
524	2	fn get_name(&self) -> &'static str {
525	2	"GcsStore"
526	2	}
527
528		/// Lightweight probe: a single `object_exists` against a fixed
529		/// never-existing path. Shares no resources with production traffic
530		/// and stays well under the `HealthServer` per-indicator budget.
531	2	async fn check_health(&self, _namespace: Cow<'static, str>) -> HealthStatus {
532		const HEALTH_PROBE_TIMEOUT: Duration = Duration::from_secs(2);
533
534		let probe_path = ObjectPath::new(
535		self.bucket.clone(),
536		"__nativelink_health_probe__/does-not-exist",
537		);
538
539		let probe = self.client.object_exists(&probe_path);
540		match timeout(HEALTH_PROBE_TIMEOUT, probe).await {
541		Ok(Ok(_)) => HealthStatus::new_ok(self, "GcsStore::check_health: ok".into()),
542		Ok(Err(e)) => {
543		warn!(?e, "GcsStore::check_health: object_exists errored");
544		HealthStatus::new_failed(
545		self,
546		format!("GcsStore::check_health: object_exists errored: {e}").into(),
547		)
548		}
549		Err(_) => {
550		warn!(
551		timeout_secs = HEALTH_PROBE_TIMEOUT.as_secs(),
552		"GcsStore::check_health: probe timed out",
553		);
554		HealthStatus::Timeout {
555		struct_name: self.struct_name(),
556		}
557		}
558		}
559	2	}
560		}