mz_persist/
gcs.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Implementation of [Blob] backed by Google Cloud Storage.
11
12use std::fmt::{Debug, Formatter};
13use std::future::Future;
14use std::sync::Arc;
15use std::time::Duration;
16
17use anyhow::{Context, anyhow};
18use async_trait::async_trait;
19use bytes::Bytes;
20use google_cloud_gax::error::rpc::Code;
21use google_cloud_gax::options::RequestOptionsBuilder;
22use google_cloud_gax::paginator::ItemPaginator;
23use google_cloud_storage::client::{Storage, StorageControl};
24use uuid::Uuid;
25
26use mz_ore::bytes::SegmentedBytes;
27use mz_ore::metrics::MetricsRegistry;
28
29use crate::cfg::BlobKnobs;
30use crate::location::{Blob, BlobMetadata, Determinate, ExternalError};
31use crate::metrics::S3BlobMetrics;
32
33/// Configuration for opening a [GcsBlob]
34#[derive(Clone, Debug)]
35pub struct GcsBlobConfig {
36    metrics: S3BlobMetrics,
37    knobs: Arc<dyn BlobKnobs>,
38    storage_client: Storage,
39    storage_control_client: StorageControl,
40    bucket: String,
41    prefix: String,
42}
43
44impl GcsBlobConfig {
45    const EXTERNAL_TESTS_GCS_BUCKET: &'static str = "MZ_PERSIST_EXTERNAL_STORAGE_TEST_GCS_BUCKET";
46
47    /// Returns a new [GcsBlobConfig] for use in production.
48    ///
49    /// Stores objects in the given bucket prepended with the (possibly empty)
50    /// prefix. GCS credentials must be available in the process or environment.
51    pub async fn new(
52        bucket: String,
53        prefix: String,
54        knobs: Box<dyn BlobKnobs>,
55        metrics: S3BlobMetrics,
56    ) -> Result<Self, ExternalError> {
57        let storage_client = Storage::builder()
58            .build()
59            .await
60            .map_err(|e| anyhow!("GCS Storage builder error: {}", e))?;
61        let storage_control_client = StorageControl::builder()
62            .build()
63            .await
64            .map_err(|e| anyhow!("GCS StorageControl builder error: {}", e))?;
65
66        Ok(GcsBlobConfig {
67            metrics,
68            knobs: Arc::from(knobs),
69            storage_client,
70            storage_control_client,
71            bucket,
72            prefix,
73        })
74    }
75
76    /// Returns a new [GcsBlobConfig] for use in tests.
77    pub async fn new_for_test() -> Result<Option<Self>, ExternalError> {
78        struct TestBlobKnobs;
79        impl Debug for TestBlobKnobs {
80            fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
81                f.debug_struct("TestBlobKnobs").finish_non_exhaustive()
82            }
83        }
84        impl BlobKnobs for TestBlobKnobs {
85            fn operation_timeout(&self) -> Duration {
86                Duration::from_secs(30)
87            }
88
89            fn operation_attempt_timeout(&self) -> Duration {
90                Duration::from_secs(10)
91            }
92
93            fn connect_timeout(&self) -> Duration {
94                Duration::from_secs(5)
95            }
96
97            fn read_timeout(&self) -> Duration {
98                Duration::from_secs(5)
99            }
100
101            fn is_cc_active(&self) -> bool {
102                false
103            }
104        }
105
106        let bucket = match std::env::var(Self::EXTERNAL_TESTS_GCS_BUCKET) {
107            Ok(bucket) => bucket,
108            Err(_) => {
109                if mz_ore::env::is_var_truthy("CI") {
110                    panic!("CI is supposed to run this test but something has gone wrong!");
111                }
112                return Ok(None);
113            }
114        };
115        let prefix = Uuid::new_v4().to_string();
116        let metrics = S3BlobMetrics::new(&MetricsRegistry::new());
117        Ok(Some(
118            GcsBlobConfig::new(bucket, prefix, Box::new(TestBlobKnobs), metrics).await?,
119        ))
120    }
121}
122
123/// Implementation of [Blob] backed by GCS
124#[derive(Debug)]
125pub struct GcsBlob {
126    metrics: S3BlobMetrics,
127    knobs: Arc<dyn BlobKnobs>,
128    storage_client: Storage,
129    storage_control_client: StorageControl,
130    bucket: String,
131    prefix: String,
132}
133
134impl GcsBlob {
135    /// Opens the given location for non-exclusive read-write access.
136    pub async fn open(config: GcsBlobConfig) -> Result<Self, ExternalError> {
137        let ret = GcsBlob {
138            metrics: config.metrics,
139            knobs: config.knobs,
140            storage_client: config.storage_client,
141            storage_control_client: config.storage_control_client,
142            bucket: config.bucket,
143            prefix: config.prefix,
144        };
145        // Connect before returning success. Listing fails if the bucket is
146        // missing or we lack access to it, while reading a missing key would
147        // succeed against a missing bucket too.
148        ret.with_timeout("open", async {
149            ret.metrics.list_objects.inc();
150            ret.storage_control_client
151                .list_objects()
152                .set_parent(ret.formatted_bucket_path())
153                .set_prefix(format!("{}/", ret.prefix))
154                .set_page_size(1)
155                .with_attempt_timeout(ret.knobs.operation_attempt_timeout())
156                .send()
157                .await
158                .map_err(|e| ret.error("ListObjects", e))
159        })
160        .await?;
161        Ok(ret)
162    }
163
164    fn get_path(&self, key: &str) -> String {
165        format!("{}/{}", self.prefix, key)
166    }
167
168    /// Returns the bucket path formatted for gRPC API calls.
169    /// The gRPC API requires bucket names in the format `projects/_/buckets/{bucket}`.
170    fn formatted_bucket_path(&self) -> String {
171        format!("projects/_/buckets/{}", self.bucket)
172    }
173
174    /// Fails `op` if it takes longer than the operation timeout, read anew on
175    /// each call.
176    ///
177    /// NOTE: The SDK has no per-attempt timeout for `read_object` and
178    /// `write_object` (google-cloud-storage 1.20.0 leaves it as a TODO in
179    /// `perform_upload/unbuffered.rs`), nor a way to set connect or read
180    /// timeouts, so this is the only bound on a stalled data transfer.
181    async fn with_timeout<T>(
182        &self,
183        op: &str,
184        fut: impl Future<Output = Result<T, ExternalError>>,
185    ) -> Result<T, ExternalError> {
186        let timeout = self.knobs.operation_timeout();
187        match tokio::time::timeout(timeout, fut).await {
188            Ok(res) => res,
189            Err(_) => {
190                self.metrics.operation_timeouts.inc();
191                Err(ExternalError::from(anyhow!(
192                    "GCS {op} timed out after {timeout:?}"
193                )))
194            }
195        }
196    }
197
198    /// Converts an SDK error from the call `op` and counts it.
199    fn error(&self, op: &str, e: google_cloud_storage::Error) -> ExternalError {
200        let code = if let Some(status) = e.status() {
201            status.code.name().to_string()
202        } else if let Some(code) = e.http_status_code() {
203            format!("HTTP{code}")
204        } else if e.is_timeout() {
205            "Timeout".to_string()
206        } else if e.is_connect() {
207            "Connect".to_string()
208        } else if e.is_authentication() {
209            "Authentication".to_string()
210        } else if e.is_io() || e.is_transport() {
211            "Transport".to_string()
212        } else {
213            "Unknown".to_string()
214        };
215        self.metrics
216            .error_counts
217            .with_label_values(&[op, &code])
218            .inc();
219        ExternalError::from(anyhow!("GCS {op} error: {e}"))
220    }
221
222    async fn get_inner(&self, key: &str) -> Result<Option<SegmentedBytes>, ExternalError> {
223        let path = self.get_path(key);
224
225        self.metrics.get_part.inc();
226        let mut stream = match self
227            .storage_client
228            .read_object(&self.formatted_bucket_path(), path)
229            .send()
230            .await
231        {
232            Ok(stream) => stream,
233            Err(e) if is_not_found(&e) => return Ok(None),
234            Err(e) => return Err(self.error("ReadObject", e)),
235        };
236        let size = usize::try_from(stream.object().size).context("GCS sent negative size value")?;
237        let mut buf: Vec<u8> = Vec::with_capacity(size);
238        while let Some(data) = stream.next().await {
239            let data = data.map_err(|e| self.error("ReadObject", e))?;
240            buf.extend_from_slice(&data[..])
241        }
242        Ok(Some(SegmentedBytes::from(buf)))
243    }
244
245    async fn list_keys_and_metadata_inner(
246        &self,
247        key_prefix: &str,
248        f: &mut (dyn FnMut(BlobMetadata) + Send + Sync),
249    ) -> Result<(), ExternalError> {
250        let blob_key_prefix = self.get_path(key_prefix);
251        let strippable_root_prefix = format!("{}/", self.prefix);
252
253        let mut stream = self
254            .storage_control_client
255            .list_objects()
256            .set_parent(&self.formatted_bucket_path())
257            .set_prefix(blob_key_prefix)
258            .with_attempt_timeout(self.knobs.operation_attempt_timeout())
259            .by_item();
260        self.metrics.list_objects.inc();
261        while let Some(response) = stream.next().await {
262            let obj = response.map_err(|e| self.error("ListObjects", e))?;
263            if let Some(key) = obj.name.strip_prefix(&strippable_root_prefix) {
264                let size_in_bytes =
265                    u64::try_from(obj.size).context("GCS sent negative size value")?;
266                f(BlobMetadata { key, size_in_bytes })
267            } else {
268                return Err(ExternalError::from(anyhow!(
269                    "found key with invalid prefix: {}",
270                    obj.name
271                )));
272            }
273        }
274        Ok(())
275    }
276
277    async fn set_inner(&self, key: &str, value: Bytes) -> Result<(), ExternalError> {
278        let path = self.get_path(key);
279        self.metrics.set_single.inc();
280        self.storage_client
281            .write_object(&self.formatted_bucket_path(), path, value)
282            .send_unbuffered()
283            .await
284            .map_err(|e| self.error("WriteObject", e))?;
285        Ok(())
286    }
287
288    async fn delete_inner(&self, key: &str) -> Result<Option<usize>, ExternalError> {
289        let path = self.get_path(key);
290
291        // GCS deletes don't report the object size, so fetch it first. Making
292        // the delete conditional on the fetched generation keeps the reported
293        // size accurate.
294        self.metrics.delete_head.inc();
295        let obj = match self
296            .storage_control_client
297            .get_object()
298            .set_bucket(&self.formatted_bucket_path())
299            .set_object(&path)
300            .with_attempt_timeout(self.knobs.operation_attempt_timeout())
301            .send()
302            .await
303        {
304            Ok(obj) => obj,
305            Err(e) if is_not_found(&e) => return Ok(None),
306            Err(e) => return Err(self.error("GetObject", e)),
307        };
308        let size = usize::try_from(obj.size).context("GCS sent negative size value")?;
309        self.metrics.delete_object.inc();
310        match self
311            .storage_control_client
312            .delete_object()
313            .set_bucket(&self.formatted_bucket_path())
314            .set_object(path)
315            // NOTE: `set_generation` would permanently delete that generation,
316            // even in a versioned bucket, leaving nothing for `restore`.
317            .set_if_generation_match(obj.generation)
318            .with_attempt_timeout(self.knobs.operation_attempt_timeout())
319            .send()
320            .await
321        {
322            Ok(()) => Ok(Some(size)),
323            // Raced with a concurrent delete.
324            Err(e) if is_not_found(&e) => Ok(None),
325            Err(e) => Err(self.error("DeleteObject", e)),
326        }
327    }
328
329    async fn restore_inner(&self, key: &str) -> Result<(), ExternalError> {
330        let path = self.get_path(key);
331
332        match self
333            .storage_control_client
334            .get_object()
335            .set_bucket(&self.formatted_bucket_path())
336            .set_object(&path)
337            .with_attempt_timeout(self.knobs.operation_attempt_timeout())
338            .send()
339            .await
340        {
341            Ok(_) => return Ok(()),
342            Err(e) if is_not_found(&e) => {}
343            Err(e) => return Err(self.error("GetObject", e)),
344        }
345
346        // Restoring relies on object versioning: a delete leaves the object's
347        // data behind as a noncurrent generation, which we copy back to a live
348        // object. Soft-deleted objects are not considered.
349        let mut stream = self
350            .storage_control_client
351            .list_objects()
352            .set_parent(&self.formatted_bucket_path())
353            .set_prefix(&path)
354            .set_versions(true)
355            .with_attempt_timeout(self.knobs.operation_attempt_timeout())
356            .by_item();
357        let mut latest_generation = None;
358        while let Some(obj) = stream.next().await {
359            let obj = obj.map_err(|e| self.error("ListObjects", e))?;
360            // The prefix also matches longer keys.
361            if obj.name == path && latest_generation < Some(obj.generation) {
362                latest_generation = Some(obj.generation);
363            }
364        }
365        let Some(generation) = latest_generation else {
366            return Err(Determinate::new(anyhow!(
367                "unable to restore {key} in GCS: no valid version exists"
368            ))
369            .into());
370        };
371
372        // A rewrite may take several calls, each resuming from the previous
373        // call's token.
374        let mut rewrite_token = String::new();
375        loop {
376            let response = match self
377                .storage_control_client
378                .rewrite_object()
379                .set_source_bucket(self.formatted_bucket_path())
380                .set_source_object(&path)
381                .set_source_generation(generation)
382                .set_destination_bucket(self.formatted_bucket_path())
383                .set_destination_name(&path)
384                // Only restore if no live object exists.
385                .set_if_generation_match(0)
386                .set_rewrite_token(rewrite_token)
387                .with_attempt_timeout(self.knobs.operation_attempt_timeout())
388                .send()
389                .await
390            {
391                Ok(response) => response,
392                // A live object appeared concurrently.
393                Err(e) if is_failed_precondition(&e) => return Ok(()),
394                Err(e) => return Err(self.error("RewriteObject", e)),
395            };
396            if response.done {
397                return Ok(());
398            }
399            rewrite_token = response.rewrite_token;
400        }
401    }
402}
403
404#[async_trait]
405impl Blob for GcsBlob {
406    async fn get(&self, key: &str) -> Result<Option<SegmentedBytes>, ExternalError> {
407        self.with_timeout("get", self.get_inner(key)).await
408    }
409
410    async fn list_keys_and_metadata(
411        &self,
412        key_prefix: &str,
413        f: &mut (dyn FnMut(BlobMetadata) + Send + Sync),
414    ) -> Result<(), ExternalError> {
415        self.with_timeout("list", self.list_keys_and_metadata_inner(key_prefix, f))
416            .await
417    }
418
419    async fn set(&self, key: &str, value: Bytes) -> Result<(), ExternalError> {
420        self.with_timeout("set", self.set_inner(key, value)).await
421    }
422
423    async fn delete(&self, key: &str) -> Result<Option<usize>, ExternalError> {
424        self.with_timeout("delete", self.delete_inner(key)).await
425    }
426
427    async fn restore(&self, key: &str) -> Result<(), ExternalError> {
428        self.with_timeout("restore", self.restore_inner(key)).await
429    }
430}
431
432fn is_not_found(e: &google_cloud_storage::Error) -> bool {
433    e.status().is_some_and(|s| s.code == Code::NotFound) || e.http_status_code() == Some(404)
434}
435
436fn is_failed_precondition(e: &google_cloud_storage::Error) -> bool {
437    e.status()
438        .is_some_and(|s| s.code == Code::FailedPrecondition)
439        || e.http_status_code() == Some(412)
440}
441
442#[cfg(test)]
443mod tests {
444    use tracing::info;
445
446    use super::*;
447    use crate::location::tests::blob_impl_test;
448
449    #[cfg_attr(miri, ignore)] // error: unsupported operation: can't call foreign function `TLS_method` on OS `linux`
450    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
451    async fn gcs_blob() -> Result<(), ExternalError> {
452        let config = match GcsBlobConfig::new_for_test().await? {
453            Some(client) => client,
454            None => {
455                info!(
456                    "{} env not set: skipping test that uses external service",
457                    GcsBlobConfig::EXTERNAL_TESTS_GCS_BUCKET
458                );
459                return Ok(());
460            }
461        };
462        blob_impl_test(move |_path| {
463            let config = config.clone();
464            async move { GcsBlob::open(config).await }
465        })
466        .await
467    }
468
469    /// `blob_impl_test` accepts a backend that cannot restore, but `GcsBlob`
470    /// must restore from a versioned bucket, which the test bucket is.
471    #[cfg_attr(miri, ignore)] // error: unsupported operation: can't call foreign function `TLS_method` on OS `linux`
472    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
473    async fn gcs_blob_restore() -> Result<(), ExternalError> {
474        let Some(config) = GcsBlobConfig::new_for_test().await? else {
475            info!(
476                "{} env not set: skipping test that uses external service",
477                GcsBlobConfig::EXTERNAL_TESTS_GCS_BUCKET
478            );
479            return Ok(());
480        };
481        let blob = GcsBlob::open(config).await?;
482        blob.set("k", Bytes::from("data")).await?;
483        assert_eq!(blob.delete("k").await?, Some("data".len()));
484        assert_eq!(blob.get("k").await?, None);
485        blob.restore("k").await?;
486        assert_eq!(
487            blob.get("k").await?.map(|b| b.into_contiguous()),
488            Some(b"data".to_vec())
489        );
490        blob.delete("k").await?;
491        Ok(())
492    }
493
494    #[cfg_attr(miri, ignore)] // error: unsupported operation: can't call foreign function `TLS_method` on OS `linux`
495    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
496    async fn gcs_blob_open_missing_bucket() -> Result<(), ExternalError> {
497        let Some(mut config) = GcsBlobConfig::new_for_test().await? else {
498            info!(
499                "{} env not set: skipping test that uses external service",
500                GcsBlobConfig::EXTERNAL_TESTS_GCS_BUCKET
501            );
502            return Ok(());
503        };
504        config.bucket = format!("mz-missing-{}", Uuid::new_v4());
505        let err = GcsBlob::open(config)
506            .await
507            .expect_err("opening a missing bucket succeeded");
508        assert!(err.to_string().contains("ListObjects"), "{err}");
509        Ok(())
510    }
511}
mz_persist/gcs.rs

mz_persist/
gcs.rs
