diff --git a/src/storage/examples/src/lib.rs b/src/storage/examples/src/lib.rs index 4a304f3b58..ff1651ad94 100644 --- a/src/storage/examples/src/lib.rs +++ b/src/storage/examples/src/lib.rs @@ -28,7 +28,7 @@ use google_cloud_gax::{ }; use google_cloud_storage::client::{Storage, StorageControl}; use google_cloud_storage::model::bucket::{ - ObjectRetention, + CustomPlacementConfig, ObjectRetention, iam_config::UniformBucketLevelAccess, {HierarchicalNamespace, IamConfig}, }; @@ -714,6 +714,38 @@ pub async fn create_test_hns_bucket() -> anyhow::Result<(StorageControl, Bucket) Ok((client, create)) } +pub async fn create_test_rapid_bucket() -> anyhow::Result<(StorageControl, Bucket)> { + let project_id = std::env::var("GOOGLE_CLOUD_PROJECT")?; + let client = client_for_create_bucket().await?; + cleanup_stale_buckets(&client, &project_id).await?; + + let bucket_id = crate::random_bucket_id(); + + let create = client + .create_bucket() + .set_parent("projects/_") + .set_bucket_id(bucket_id) + .set_bucket( + Bucket::new() + .set_project(format!("projects/{project_id}")) + .set_location("us-central1") + .set_custom_placement_config( + CustomPlacementConfig::new().set_data_locations(["us-central1-a".to_string()]), + ) + .set_storage_class("RAPID") + .set_labels([("integration-test", "true")]) + .set_hierarchical_namespace(HierarchicalNamespace::new().set_enabled(true)) + .set_iam_config(IamConfig::new().set_uniform_bucket_level_access( + UniformBucketLevelAccess::new().set_enabled(true), + )), + ) + .with_idempotency(true) + .send() + .await?; + println!("create_test_rapid_bucket(): {create:?}"); + Ok((client, create)) +} + async fn client_for_create_bucket() -> anyhow::Result { let client = StorageControl::builder() .with_tracing() diff --git a/src/storage/src/lib.rs b/src/storage/src/lib.rs index b2f955aee9..7e83502c40 100644 --- a/src/storage/src/lib.rs +++ b/src/storage/src/lib.rs @@ -97,9 +97,15 @@ pub mod builder { pub mod storage { //! Request builders for [Storage][crate::client::Storage]. pub use crate::storage::client::ClientBuilder; + #[cfg(google_cloud_unstable_storage_bidi)] + #[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] + pub use crate::storage::open_appendable_object::OpenAppendableObject; pub use crate::storage::open_object::OpenObject; pub use crate::storage::post_policy::{PostPolicyV4Builder, PostPolicyV4Result}; pub use crate::storage::read_object::ReadObject; + #[cfg(google_cloud_unstable_storage_bidi)] + #[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] + pub use crate::storage::reopen_appendable_object::ReopenAppendableObject; pub use crate::storage::signed_url::SignedUrlBuilder; pub use crate::storage::write_object::WriteObject; } diff --git a/src/storage/src/model_ext.rs b/src/storage/src/model_ext.rs index 0b96eafb8e..894b91d0f8 100644 --- a/src/storage/src/model_ext.rs +++ b/src/storage/src/model_ext.rs @@ -22,6 +22,18 @@ use sha2::{Digest, Sha256}; mod open_object_request; pub use open_object_request::OpenObjectRequest; +#[cfg(google_cloud_unstable_storage_bidi)] +mod open_appendable_object_request; +#[cfg(google_cloud_unstable_storage_bidi)] +#[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] +pub use open_appendable_object_request::OpenAppendableObjectRequest; + +#[cfg(google_cloud_unstable_storage_bidi)] +mod reopen_appendable_object_request; +#[cfg(google_cloud_unstable_storage_bidi)] +#[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] +pub use reopen_appendable_object_request::ReopenAppendableObjectRequest; + /// ObjectHighlights contains select metadata from a [crate::model::Object]. #[derive(Clone, Debug, Default, PartialEq)] #[non_exhaustive] @@ -316,61 +328,6 @@ pub struct WriteObjectRequest { pub params: Option, } -#[cfg(google_cloud_unstable_storage_bidi)] -/// Represents the parameters of a request to open a new object for exclusive appends. -/// -/// Consumers of the `google-cloud-storage` crate rarely have a need to use this type directly, the most common exception is when mocking of the `Storage` client. -#[derive(Debug, PartialEq)] -#[non_exhaustive] -pub struct OpenAppendableObjectRequest { - /// The object attributes and pre-conditions for the open operation. - pub spec: crate::model::WriteObjectSpec, - /// Additional request parameters. - pub params: Option, -} - -#[cfg(google_cloud_unstable_storage_bidi)] -/// Represents the parameters of a request to reopen an existing object for appends. -/// -/// Consumers of the `google-cloud-storage` crate rarely have a need to use this type directly, the most common exception is when mocking of the `Storage` client. -#[derive(Debug, PartialEq)] -#[non_exhaustive] -pub struct ReopenAppendableObjectRequest { - /// The bucket containing the target object. - pub bucket: String, - /// The target object name. - pub object: String, - /// The target object generation to append to. - pub generation: i64, - /// If set, return an error if the current metageneration does not match the value. - pub if_metageneration_match: Option, - /// If set, return an error if the current metageneration matches the value. - pub if_metageneration_not_match: Option, - /// A routing token from a previous operation. - pub routing_token: Option, - /// A write handle from a previous operation. - pub write_handle: Option, - /// Additional request parameters. - pub params: Option, -} - -#[cfg(google_cloud_unstable_storage_bidi)] -impl From for crate::google::storage::v2::AppendObjectSpec { - fn from(value: ReopenAppendableObjectRequest) -> Self { - Self { - bucket: value.bucket, - object: value.object, - generation: value.generation, - if_metageneration_match: value.if_metageneration_match, - if_metageneration_not_match: value.if_metageneration_not_match, - routing_token: value.routing_token, - write_handle: value - .write_handle - .map(|h| crate::google::storage::v2::BidiWriteHandle { handle: h }), - } - } -} - #[cfg(test)] pub(crate) mod tests { use super::*; @@ -490,44 +447,4 @@ pub(crate) mod tests { assert_eq!(key_aes_256.to_string(), key_base64); Ok(()) } - - #[cfg(google_cloud_unstable_storage_bidi)] - #[test] - fn test_open_appendable_object_request() { - let req = OpenAppendableObjectRequest { - spec: crate::model::WriteObjectSpec::default(), - params: None, - }; - assert_eq!(req.spec.resource, None); - assert_eq!(req.params, None); - } - - #[cfg(google_cloud_unstable_storage_bidi)] - #[test] - fn test_reopen_appendable_object_request_from() { - let req = ReopenAppendableObjectRequest { - bucket: "my-bucket".into(), - object: "my-object".into(), - generation: 42, - if_metageneration_match: Some(1), - if_metageneration_not_match: Some(2), - routing_token: Some("token".into()), - write_handle: Some(bytes::Bytes::from("handle")), - params: None, - }; - - let spec = crate::google::storage::v2::AppendObjectSpec::from(req); - assert_eq!(spec.bucket, "my-bucket"); - assert_eq!(spec.object, "my-object"); - assert_eq!(spec.generation, 42); - assert_eq!(spec.if_metageneration_match, Some(1)); - assert_eq!(spec.if_metageneration_not_match, Some(2)); - assert_eq!(spec.routing_token, Some("token".into())); - assert_eq!( - spec.write_handle, - Some(crate::google::storage::v2::BidiWriteHandle { - handle: bytes::Bytes::from("handle") - }) - ); - } } diff --git a/src/storage/src/model_ext/open_appendable_object_request.rs b/src/storage/src/model_ext/open_appendable_object_request.rs new file mode 100644 index 0000000000..a402c89c33 --- /dev/null +++ b/src/storage/src/model_ext/open_appendable_object_request.rs @@ -0,0 +1,41 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] +/// Represents the parameters of a request to open a new object for exclusive appends. +/// +/// Consumers of the `google-cloud-storage` crate rarely have a need to use this type directly, the most common exception is when mocking of the `Storage` client. +#[derive(Clone, Debug, Default, PartialEq)] +#[non_exhaustive] +pub struct OpenAppendableObjectRequest { + /// The object attributes and pre-conditions for the open operation. + pub spec: crate::model::WriteObjectSpec, + /// Additional request parameters. + pub params: Option, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_open_appendable_object_request() { + let req = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default(), + params: None, + }; + assert_eq!(req.spec.resource, None); + assert_eq!(req.params, None); + } +} diff --git a/src/storage/src/model_ext/reopen_appendable_object_request.rs b/src/storage/src/model_ext/reopen_appendable_object_request.rs new file mode 100644 index 0000000000..fc4a154b73 --- /dev/null +++ b/src/storage/src/model_ext/reopen_appendable_object_request.rs @@ -0,0 +1,87 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] +/// Represents the parameters of a request to reopen an existing object for appends. +/// +/// Consumers of the `google-cloud-storage` crate rarely have a need to use this type directly, the most common exception is when mocking of the `Storage` client. +#[derive(Clone, Debug, Default, PartialEq)] +#[non_exhaustive] +pub struct ReopenAppendableObjectRequest { + /// The bucket containing the target object. + pub bucket: String, + /// The target object name. + pub object: String, + /// The target object generation to append to. + pub generation: i64, + /// If set, return an error if the current metageneration does not match the value. + pub if_metageneration_match: Option, + /// If set, return an error if the current metageneration matches the value. + pub if_metageneration_not_match: Option, + /// A routing token from a previous operation. + pub routing_token: Option, + /// A write handle from a previous operation. + pub write_handle: Option, + /// Additional request parameters. + pub params: Option, +} + +impl From for crate::google::storage::v2::AppendObjectSpec { + fn from(value: ReopenAppendableObjectRequest) -> Self { + Self { + bucket: value.bucket, + object: value.object, + generation: value.generation, + if_metageneration_match: value.if_metageneration_match, + if_metageneration_not_match: value.if_metageneration_not_match, + routing_token: value.routing_token, + write_handle: value + .write_handle + .map(|h| crate::google::storage::v2::BidiWriteHandle { handle: h }), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_reopen_appendable_object_request_from() { + let req = ReopenAppendableObjectRequest { + bucket: "my-bucket".into(), + object: "my-object".into(), + generation: 42, + if_metageneration_match: Some(1), + if_metageneration_not_match: Some(2), + routing_token: Some("token".into()), + write_handle: Some(bytes::Bytes::from("handle")), + params: None, + }; + + let spec = crate::google::storage::v2::AppendObjectSpec::from(req); + assert_eq!(spec.bucket, "my-bucket"); + assert_eq!(spec.object, "my-object"); + assert_eq!(spec.generation, 42); + assert_eq!(spec.if_metageneration_match, Some(1)); + assert_eq!(spec.if_metageneration_not_match, Some(2)); + assert_eq!(spec.routing_token, Some("token".into())); + assert_eq!( + spec.write_handle, + Some(crate::google::storage::v2::BidiWriteHandle { + handle: bytes::Bytes::from("handle") + }) + ); + } +} diff --git a/src/storage/src/storage.rs b/src/storage/src/storage.rs index 3db6a0e0ed..65e23cf40b 100644 --- a/src/storage/src/storage.rs +++ b/src/storage/src/storage.rs @@ -18,10 +18,14 @@ pub(crate) mod bidi_write; pub(crate) mod checksum; pub(crate) mod client; pub(crate) mod common_options; +#[cfg(google_cloud_unstable_storage_bidi)] +pub(crate) mod open_appendable_object; pub(crate) mod open_object; pub(crate) mod perform_upload; pub(crate) mod post_policy; pub(crate) mod read_object; +#[cfg(google_cloud_unstable_storage_bidi)] +pub(crate) mod reopen_appendable_object; pub mod request_options; pub(crate) mod signed_url; pub mod streaming_source; diff --git a/src/storage/src/storage/bidi_write.rs b/src/storage/src/storage/bidi_write.rs index b6576d33be..152048342d 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -14,9 +14,14 @@ //! Internal traits and types for Appendable Object Write (Bidi Write). +#![allow(dead_code)] + +pub(crate) mod connector; mod redirect; mod retry_redirect; pub(crate) mod stub; +pub(crate) mod transport; +mod worker; use crate::google::storage::v2::{BidiWriteObjectRequest, BidiWriteObjectResponse}; use crate::request_options::RequestOptions; @@ -27,7 +32,6 @@ use tokio::sync::mpsc::Receiver; /// A trait to mock `Streaming` in the unit tests. /// /// This is not a public trait, we only need this for our own testing. -#[allow(dead_code)] pub(crate) trait TonicStreaming: std::fmt::Debug + Send + 'static { fn next_message( &mut self, @@ -44,7 +48,6 @@ impl TonicStreaming for Streaming { /// A trait to mock `gaxi::grpc::Client` in the unit tests. /// /// This is not a public trait, we only need this for our own testing. -#[allow(dead_code)] pub(crate) trait Client: std::fmt::Debug + Send + 'static { type Stream: Sized; fn start( @@ -82,6 +85,9 @@ impl Client for gaxi::grpc::Client { } } +#[cfg(test)] +mod mocks; + #[cfg(test)] pub(crate) mod tests { use crate::Error; @@ -135,14 +141,12 @@ pub(crate) mod tests { ) } - #[allow(dead_code)] pub(crate) fn test_options() -> RequestOptions { let mut options = RequestOptions::new(); options.backoff_policy = Arc::new(test_backoff()); options } - #[allow(dead_code)] fn test_backoff() -> impl google_cloud_gax::backoff_policy::BackoffPolicy { use std::time::Duration; google_cloud_gax::exponential_backoff::ExponentialBackoffBuilder::new() diff --git a/src/storage/src/storage/bidi_write/connector.rs b/src/storage/src/storage/bidi_write/connector.rs new file mode 100644 index 0000000000..35421ebf74 --- /dev/null +++ b/src/storage/src/storage/bidi_write/connector.rs @@ -0,0 +1,745 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// TODO(#5716): Lift to shared bidi module + +use super::redirect::handle_redirect; +use super::retry_redirect::RetryRedirect; +use super::{Client, TonicStreaming}; +use crate::google::storage::v2::{ + AppendObjectSpec, BidiWriteObjectRequest, BidiWriteObjectResponse, CommonObjectRequestParams, + WriteObjectSpec, bidi_write_object_request::FirstMessage, + bidi_write_object_response::WriteStatus, +}; +use crate::request_options::RequestOptions; +use crate::storage::info::X_GOOG_API_CLIENT_HEADER; +use crate::{Error, Result}; +use gaxi::grpc::Client as GrpcClient; +use gaxi::grpc::tonic::{Extensions, GrpcMethod, Streaming}; +use gaxi::prost::ToProto; +use google_cloud_gax::error::binding::{ + BindingError, PathMismatch, SubstitutionFail, SubstitutionMismatch, +}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::sync::mpsc::Sender; +/// Represents a bidirectional streaming connection. +/// Contains the transmission channel for requests and the receiving stream for responses. +#[derive(Debug)] +pub struct Connection> { + pub tx: Sender, + pub rx: S, +} + +impl Connection { + pub fn new(tx: Sender, rx: S) -> Self { + Self { tx, rx } + } +} + +/// Represents the state of the initial request in the stream. +#[derive(Clone, Debug)] +pub enum AppendObjectSpecState { + Write(Box, Option), + Append(AppendObjectSpec), +} + +/// Establishes and handles the initial handshake for bidi streaming writes. +/// +/// Connecting and reconnecting bidirectional streaming writes requires: +/// - Constructing the initial `WriteObjectSpec` or `AppendObjectSpec`. +/// - Following routing token redirects correctly. +/// - Passing back the established `Connection` to the async worker. +/// +/// # Parameters +/// - `T`: a type implementing the [Client] trait, this is used in tests. +#[derive(Clone, Debug)] +pub struct Connector { + spec: Arc>, + options: RequestOptions, + client: T, + params: Option, +} + +impl Connector +where + T: Client + Clone + Send + 'static, + ::Stream: TonicStreaming, +{ + pub fn new(options: RequestOptions, client: T) -> Self { + Self { + spec: Arc::new(Mutex::new(AppendObjectSpecState::Write( + Box::default(), + None, + ))), + options, + client, + params: None, + } + } + + pub async fn connect_open( + &mut self, + req: crate::model_ext::OpenAppendableObjectRequest, + ) -> Result<(BidiWriteObjectResponse, Connection)> { + let resource = match req.spec.resource { + Some(r) => { + let object: crate::google::storage::v2::Object = gaxi::prost::ToProto::to_proto(r) + .map_err(|e: gaxi::prost::ConvertError| Error::io(e.to_string()))?; + Some(object) + } + None => None, + }; + let spec = WriteObjectSpec { + resource, + predefined_acl: req.spec.predefined_acl, + if_generation_match: req.spec.if_generation_match, + if_generation_not_match: req.spec.if_generation_not_match, + if_metageneration_match: req.spec.if_metageneration_match, + if_metageneration_not_match: req.spec.if_metageneration_not_match, + object_size: req.spec.object_size, + appendable: req.spec.appendable, + }; + self.params = req + .params + .map(|p| p.to_proto().map_err(|e| Error::io(e.to_string()))) + .transpose()?; + *self.spec.lock().expect("never poisoned") = + AppendObjectSpecState::Write(Box::new(spec), None); + self.connect_attempt_loop().await + } + + pub async fn connect_reopen( + &mut self, + req: crate::model_ext::ReopenAppendableObjectRequest, + ) -> Result<(BidiWriteObjectResponse, Connection)> { + let spec = AppendObjectSpec { + bucket: req.bucket, + object: req.object, + generation: req.generation, + routing_token: req.routing_token, + if_metageneration_match: req.if_metageneration_match, + if_metageneration_not_match: req.if_metageneration_not_match, + write_handle: req + .write_handle + .map(|handle| crate::google::storage::v2::BidiWriteHandle { handle }), + }; + self.params = req + .params + .map(|p| p.to_proto().map_err(|e| Error::io(e.to_string()))) + .transpose()?; + *self.spec.lock().expect("never poisoned") = AppendObjectSpecState::Append(spec); + self.connect_attempt_loop().await + } + + async fn connect_attempt_loop( + &mut self, + ) -> Result<(BidiWriteObjectResponse, Connection)> { + let throttler = self.options.retry_throttler.clone(); + let retry = Arc::new(RetryRedirect::new(self.options.retry_policy.clone())); + let backoff = self.options.backoff_policy.clone(); + let client = self.client.clone(); + let options = self.options.clone(); + let spec = self.spec.clone(); + let params = self.params.clone(); + let sleep = async |backoff| tokio::time::sleep(backoff).await; + let default_timeout = self.options.bidi_attempt_timeout; + + let inner = async move |d: Option| { + let attempt_timeout = std::cmp::min(default_timeout, d.unwrap_or(default_timeout)); + let attempt = + Self::connect_attempt(client.clone(), spec.clone(), &options, params.clone()); + match tokio::time::timeout(attempt_timeout, attempt).await { + Ok(r) => r, + Err(e) => Err(Error::timeout(e)), + } + }; + google_cloud_gax::retry_loop_internal::retry_loop( + inner, sleep, true, throttler, retry, backoff, + ) + .await + } + + async fn connect_attempt( + client: T, + spec: Arc>, + options: &RequestOptions, + params: Option, + ) -> Result<(BidiWriteObjectResponse, Connection)> { + let (first_message, routing_token) = match &*spec.lock().expect("never poisoned") { + AppendObjectSpecState::Write(spec, rt) => { + (FirstMessage::WriteObjectSpec(*spec.clone()), rt.clone()) + } + AppendObjectSpecState::Append(spec) => ( + FirstMessage::AppendObjectSpec(spec.clone()), + spec.routing_token.clone(), + ), + }; + + let state_lookup = matches!(first_message, FirstMessage::AppendObjectSpec(_)); + + let request = BidiWriteObjectRequest { + first_message: Some(first_message), + common_object_request_params: params, + state_lookup, + ..BidiWriteObjectRequest::default() + }; + + let bucket_name = request + .first_message + .as_ref() + .and_then(|m| match m { + FirstMessage::WriteObjectSpec(s) => s.resource.as_ref().map(|r| r.bucket.as_str()), + FirstMessage::AppendObjectSpec(s) => Some(s.bucket.as_str()), + _ => None, + }) + .unwrap_or_default(); + + if bucket_name + .strip_prefix("projects/_/buckets/") + .is_none_or(|x| x.is_empty()) + { + let problem = SubstitutionFail::MismatchExpecting( + bucket_name.to_string(), + "projects/_/buckets/*", + ); + let mismatch = SubstitutionMismatch { + field_name: "bucket", + problem, + }; + let mismatch = PathMismatch { + subs: vec![mismatch], + }; + let mismatch = BindingError { + paths: vec![mismatch], + }; + + return Err(crate::Error::binding(mismatch)); + } + + let mut x_goog_request_params = format!("bucket={bucket_name}"); + if let Some(token) = routing_token { + x_goog_request_params.push_str("&routing_token="); + x_goog_request_params.push_str(&token); + } + + let (tx, rx) = tokio::sync::mpsc::channel::(100); + tx.send(request).await.map_err(Error::io)?; + + let extensions = { + let mut e = Extensions::new(); + e.insert(GrpcMethod::new( + "google.storage.v2.Storage", + "BidiWriteObject", + )); + e + }; + let path = + http::uri::PathAndQuery::from_static("/google.storage.v2.Storage/BidiWriteObject"); + + let response = client + .start( + extensions, + path, + rx, + options, + &X_GOOG_API_CLIENT_HEADER, + &x_goog_request_params, + ) + .await?; + + let response = match response { + Ok(r) => r, + Err(status) => return Err(handle_redirect(spec, status)), + }; + + let (_metadata, mut stream, _) = response.into_parts(); + match stream.next_message().await { + Ok(Some(m)) => { + let mut guard = spec.lock().expect("never poisoned"); + let mut new_generation = None; + if let Some(WriteStatus::Resource(resource)) = &m.write_status { + new_generation = Some(resource.generation); + } + + let new_state = match &*guard { + AppendObjectSpecState::Write(s, _) => { + AppendObjectSpecState::Append(AppendObjectSpec { + bucket: s + .resource + .as_ref() + .map(|r| r.bucket.clone()) + .unwrap_or_default(), + object: s + .resource + .as_ref() + .map(|r| r.name.clone()) + .unwrap_or_default(), + generation: new_generation.unwrap_or(0), + write_handle: m.write_handle.clone(), + ..Default::default() + }) + } + AppendObjectSpecState::Append(s) => { + AppendObjectSpecState::Append(AppendObjectSpec { + generation: new_generation.unwrap_or(s.generation), + write_handle: m.write_handle.clone().or_else(|| s.write_handle.clone()), + ..s.clone() + }) + } + }; + *guard = new_state; + + Ok((m, Connection::new(tx, stream))) + } + Ok(None) => Err(Error::io("bidi_write_object stream closed before start")), + Err(status) => Err(handle_redirect(spec, status)), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::model_ext::OpenAppendableObjectRequest; + use crate::storage::request_options::RequestOptions; + use anyhow::Result; + use gaxi::grpc::Client as GrpcClient; + use google_cloud_auth::credentials::{Credentials, anonymous::Builder as Anonymous}; + use google_cloud_gax::retry_policy::NeverRetry; + use static_assertions::assert_impl_all; + use std::error::Error as _; + use std::sync::Arc; + + fn test_credentials() -> Credentials { + Anonymous::new().build() + } + + fn test_options() -> RequestOptions { + RequestOptions::new() + } + + #[test] + fn assertions() { + assert_impl_all!(Connector: Clone, std::fmt::Debug, Send, Sync); + } + + #[tokio::test] + async fn bad_endpoint() -> Result<()> { + fn need_send(_val: &T) {} + + let mut config = gaxi::options::ClientConfig::default(); + config.cred = Some(test_credentials()); + let client = GrpcClient::new(config, "http://127.0.0.1:1").await?; + + let mut options = test_options(); + options.retry_policy = Arc::new(NeverRetry); + let mut connector = Connector::new(options, client); + + let mut req = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default(), + params: None, + }; + req.spec = crate::model::WriteObjectSpec { + resource: Some(crate::model::Object { + bucket: "projects/_/buckets/test-only-bucket".into(), + name: "test-only-object".into(), + ..Default::default() + }), + ..Default::default() + }; + + let start = connector.connect_open(req); + need_send(&start); + + let err = start.await.unwrap_err(); + assert!(err.is_connect(), "{err:?}"); + let source = err.source().unwrap().to_string(); + assert!(source.contains("127.0.0.1:1"), "{source}"); + + Ok(()) + } + + #[tokio::test] + #[test_case::test_case("")] + #[test_case::test_case("my-bucket")] + async fn binding_error(bucket_name: &str) -> Result<()> { + use super::super::mocks::{MockTestClient, SharedMockClient}; + let mut mock = MockTestClient::new(); + // Binding errors are detected before a request is sent. + mock.expect_start().never(); + let client = SharedMockClient::new(mock); + + let mut connector = Connector::new(test_options(), client); + + let mut req = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default(), + params: None, + }; + req.spec = crate::model::WriteObjectSpec { + resource: Some(crate::model::Object { + bucket: bucket_name.into(), + name: "object".into(), + ..Default::default() + }), + ..Default::default() + }; + + let err = connector.connect_open(req).await.unwrap_err(); + assert!(err.is_binding(), "{err:?}"); + use google_cloud_gax::error::binding::{BindingError, SubstitutionFail}; + let source = err.source().and_then(|e| e.downcast_ref::()); + assert!(matches!(source, Some(BindingError { .. })), "{err:?}"); + // Extract all the field names that did not match, and expect a single name: + let mismatch = source + .iter() + .flat_map(|f| f.paths.iter()) + .flat_map(|f| f.subs.iter()) + .map(|f| f.field_name) + .collect::>(); + assert_eq!(mismatch, vec!["bucket"], "{err:?}"); + + // Extract all the problems: + let mismatch = source + .iter() + .flat_map(|f| f.paths.iter()) + .flat_map(|f| f.subs.iter()) + .map(|f| &f.problem) + .collect::>(); + assert!( + matches!( + mismatch.first(), + Some(SubstitutionFail::MismatchExpecting(n, p)) if n == bucket_name && *p == "projects/_/buckets/*" + ), + "{err:?}" + ); + Ok(()) + } + + #[tokio::test] + async fn start_error() -> Result<()> { + use super::super::mocks::{MockTestClient, SharedMockClient}; + use super::super::tests::permanent_error; + use gaxi::grpc::tonic::GrpcMethod; + use std::sync::Mutex; + + let receivers = Arc::new(Mutex::new(Vec::new())); + let save = receivers.clone(); + let mut mock = MockTestClient::new(); + mock.expect_start() + .return_once(move |extensions, path, rx, _options, header, params| { + // Verify all the parameters. We should have a couple of tests + // that do this, but should avoid doing so in every test. + assert!( + matches!(extensions.get::(), Some(m) if m.service() == "google.storage.v2.Storage" && m.method() == "BidiWriteObject") + ); + assert_eq!(path.path(), "/google.storage.v2.Storage/BidiWriteObject"); + assert_eq!(header, *crate::storage::info::X_GOOG_API_CLIENT_HEADER); + assert_eq!(params, "bucket=projects/_/buckets/test-bucket"); + save.lock().expect("never poisoned").push(rx); + Err(permanent_error()) + }); + let client = SharedMockClient::new(mock); + + let mut connector = Connector::new(test_options(), client); + + let mut req = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default(), + params: None, + }; + req.spec = crate::model::WriteObjectSpec { + resource: Some(crate::model::Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + ..Default::default() + }), + ..Default::default() + }; + + let err = connector.connect_open(req).await.unwrap_err(); + assert!(err.status().is_some(), "{err:?}"); + + let mut rx = { + let mut guard = receivers.lock().expect("never poisoned"); + let rx = guard.pop().expect("at least one receiver"); + assert!(guard.is_empty(), "{receivers:?}"); + rx + }; + + let first = rx.recv().await.expect("non-empty request"); + let spec = match first.first_message.as_ref().unwrap() { + crate::google::storage::v2::bidi_write_object_request::FirstMessage::WriteObjectSpec(s) => s, + _ => panic!("Expected WriteObjectSpec"), + }; + assert_eq!( + spec.resource.as_ref().map(|s| s.bucket.as_str()), + Some("projects/_/buckets/test-bucket") + ); + assert_eq!( + spec.resource.as_ref().map(|s| s.name.as_str()), + Some("test-object") + ); + + Ok(()) + } + + #[tokio::test] + async fn start_error_with_routing() -> Result<()> { + use super::super::mocks::{MockTestClient, SharedMockClient}; + use super::super::tests::permanent_error; + use gaxi::grpc::tonic::GrpcMethod; + use std::sync::Mutex; + + let receivers = Arc::new(Mutex::new(Vec::new())); + let save = receivers.clone(); + let mut mock = MockTestClient::new(); + mock.expect_start() + .return_once(move |extensions, path, rx, _options, header, params| { + // Verify all the parameters. We should have a couple of tests + // that do this, but should avoid doing so in every test. + assert!( + matches!( + extensions.get::(), + Some(m) if m.service() == "google.storage.v2.Storage" && m.method() == "BidiWriteObject" + ) + ); + assert_eq!(path.path(), "/google.storage.v2.Storage/BidiWriteObject"); + assert_eq!(header, *crate::storage::info::X_GOOG_API_CLIENT_HEADER); + let mut split = params.split('&').collect::>(); + split.sort(); + assert_eq!(split, vec!["bucket=projects/_/buckets/test-bucket", "routing_token=test-routing-token"]); + save.lock().expect("never poisoned").push(rx); + + Err(permanent_error()) + }); + let client = SharedMockClient::new(mock); + + let req = crate::model_ext::ReopenAppendableObjectRequest { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + generation: 345678, + routing_token: Some("test-routing-token".to_string()), + if_metageneration_match: None, + if_metageneration_not_match: None, + params: None, + write_handle: Some(bytes::Bytes::from_static(b"test-write-handle")), + }; + + let mut connector = Connector::new(test_options(), client); + let err = connector.connect_reopen(req).await.unwrap_err(); + assert!(err.status().is_some(), "{err:?}"); + + let mut rx = { + let mut guard = receivers.lock().expect("never poisoned"); + let rx = guard.pop().expect("at least one receiver"); + assert!(guard.is_empty(), "{guard:?}"); + rx + }; + let first = rx.recv().await.expect("non-empty request"); + let spec = match first.first_message.as_ref().unwrap() { + crate::google::storage::v2::bidi_write_object_request::FirstMessage::AppendObjectSpec(s) => s, + _ => panic!("Expected AppendObjectSpec"), + }; + assert_eq!(spec.bucket.as_str(), "projects/_/buckets/test-bucket"); + assert_eq!(spec.object.as_str(), "test-object"); + assert_eq!(spec.generation, 345678); + assert_eq!( + spec.write_handle.as_ref().map(|h| h.handle.clone()), + Some(bytes::Bytes::from_static(b"test-write-handle")) + ); + Ok(()) + } + + #[tokio::test] + async fn start_redirect_then_error() -> Result<()> { + use super::super::mocks::{MockTestClient, SharedMockClient}; + use super::super::tests::{permanent_error, redirect_status}; + use std::sync::Mutex; + + let mut seq = mockall::Sequence::new(); + let mut mock = MockTestClient::new(); + let receivers = Arc::new(Mutex::new(Vec::new())); + let save = receivers.clone(); + mock.expect_start() + .times(1) + .in_sequence(&mut seq) + .returning(move |_, _, rx, _, _, _| { + save.lock().expect("never poisoned").push(rx); + Ok(Err(redirect_status("r1"))) + }); + let save = receivers.clone(); + mock.expect_start() + .times(1) + .in_sequence(&mut seq) + .returning(move |_, _, rx, _, _, _| { + save.lock().expect("never poisoned").push(rx); + Err(permanent_error()) + }); + let client = SharedMockClient::new(mock); + + let mut connector = Connector::new(test_options(), client); + + let mut req = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default(), + params: None, + }; + req.spec = crate::model::WriteObjectSpec { + resource: Some(crate::model::Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + ..Default::default() + }), + ..Default::default() + }; + + let err = connector.connect_open(req).await.unwrap_err(); + assert_eq!(err.status(), permanent_error().status(), "{err:?}"); + + let got = connector.spec.lock().expect("never poisoned").clone(); + match got { + AppendObjectSpecState::Write(_, _) => panic!("Should be Append"), + AppendObjectSpecState::Append(got) => { + assert_eq!(got.routing_token.as_deref(), Some("r1")); + } + } + + let mut rx = receivers + .lock() + .expect("never poisoned") + .pop() + .expect("at least two receiver"); + // We pop the receivers, so this is the second receiver. This receiver should include an spec with the redirect options. + let got = rx.recv().await.expect("at least one request sent"); + let want = crate::google::storage::v2::AppendObjectSpec { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + routing_token: Some("r1".to_string()), + ..crate::google::storage::v2::AppendObjectSpec::default() + }; + let spec = match got.first_message.unwrap() { + crate::google::storage::v2::bidi_write_object_request::FirstMessage::AppendObjectSpec(s) => s, + _ => panic!("Expected AppendObjectSpec"), + }; + assert_eq!(spec.bucket, want.bucket); + assert_eq!(spec.object, want.object); + assert_eq!(spec.routing_token, want.routing_token); + + let mut rx = receivers + .lock() + .expect("never poisoned") + .pop() + .expect("at least two receiver"); + // We pop the receivers, so this is the second receiver. This receiver should include an spec with the redirect options. + let got = rx.recv().await.expect("at least one request sent"); + let want = crate::google::storage::v2::WriteObjectSpec { + resource: Some(crate::google::storage::v2::Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + ..Default::default() + }), + ..crate::google::storage::v2::WriteObjectSpec::default() + }; + let spec = match got.first_message.unwrap() { + crate::google::storage::v2::bidi_write_object_request::FirstMessage::WriteObjectSpec(s) => s, + _ => panic!("Expected WriteObjectSpec"), + }; + assert_eq!( + spec.resource.as_ref().unwrap().bucket, + want.resource.as_ref().unwrap().bucket + ); + assert_eq!( + spec.resource.as_ref().unwrap().name, + want.resource.as_ref().unwrap().name + ); + + Ok(()) + } + + #[tokio::test] + async fn start_immediately_closed() -> Result<()> { + use super::super::mocks::{MockTestClient, SharedMockClient}; + use gaxi::grpc::tonic::Response as TonicResponse; + use gaxi::grpc::tonic::Result as TonicResult; + + let (tx1, rx1) = tokio::sync::mpsc::channel::>(5); + let stream1 = TonicResponse::from(rx1); + drop(tx1); + let (tx2, rx2) = tokio::sync::mpsc::channel::>(5); + let stream2 = TonicResponse::from(rx2); + + let mut seq = mockall::Sequence::new(); + let mut mock = MockTestClient::new(); + mock.expect_start() + .times(1) + .in_sequence(&mut seq) + .return_once(move |_, _, _, _, _, _| Ok(Ok(stream1))); + mock.expect_start() + .times(1) + .in_sequence(&mut seq) + .return_once(move |_, _, _, _, _, _| Ok(Ok(stream2))); + let client = SharedMockClient::new(mock); + + let mut connector = Connector::new(test_options(), client); + + let initial = BidiWriteObjectResponse { + write_status: Some( + crate::google::storage::v2::bidi_write_object_response::WriteStatus::Resource( + crate::google::storage::v2::Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + generation: 123456, + ..crate::google::storage::v2::Object::default() + }, + ), + ), + write_handle: Some(crate::google::storage::v2::BidiWriteHandle { + handle: bytes::Bytes::from_static(b"test-handle-open"), + }), + ..BidiWriteObjectResponse::default() + }; + tx2.send(Ok(initial.clone())).await?; + + let mut req = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default(), + params: None, + }; + req.spec = crate::model::WriteObjectSpec { + resource: Some(crate::model::Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + ..Default::default() + }), + ..Default::default() + }; + + let (response, _connection) = connector.connect_open(req).await?; + assert_eq!(response, initial); + + let guard = connector.spec.lock().expect("never poisoned"); + if let AppendObjectSpecState::Append(s) = &*guard { + assert!(s.routing_token.is_none(), "{s:?}"); + assert_eq!(s.generation, 123456, "{s:?}"); + assert_eq!( + s.write_handle.as_ref().map(|h| h.handle.clone()), + Some(bytes::Bytes::from_static(b"test-handle-open")) + ); + } else { + panic!("Expected AppendObjectSpecState::Append"); + } + drop(tx2); + + Ok(()) + } +} diff --git a/src/storage/src/storage/bidi_write/mocks.rs b/src/storage/src/storage/bidi_write/mocks.rs new file mode 100644 index 0000000000..315603b4c8 --- /dev/null +++ b/src/storage/src/storage/bidi_write/mocks.rs @@ -0,0 +1,84 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// TODO(#5716): Lift to shared bidi module. + +use super::connector::Connector; +use super::tests::test_options; +use super::{Client, Receiver, RequestOptions, TonicStreaming}; +use crate::google::storage::v2::{BidiWriteObjectRequest, BidiWriteObjectResponse}; +use gaxi::grpc::tonic::{Extensions, Response as TonicResponse, Result as TonicResult}; +use std::sync::Arc; + +// mockall mocks are not `Clone` and we need a thing that can be cloned. +// The solution is to wrap the mock in a thing that implements the right +// trait. +#[derive(Clone, Debug)] +pub struct SharedMockClient(pub(crate) Arc); + +impl SharedMockClient { + pub fn new(mock: MockTestClient) -> Self { + Self(Arc::new(mock)) + } +} + +impl Client for SharedMockClient { + type Stream = MockStream; + + async fn start( + &self, + extensions: Extensions, + path: http::uri::PathAndQuery, + rx: Receiver, + options: &RequestOptions, + api_client_header: &'static str, + request_params: &str, + ) -> crate::Result>> { + self.0.start( + extensions, + path, + rx, + options, + api_client_header, + request_params, + ) + } +} + +impl TonicStreaming for Receiver> { + async fn next_message(&mut self) -> TonicResult> { + self.recv().await.transpose() + } +} + +#[mockall::automock] +pub trait TestClient: std::fmt::Debug { + fn start( + &self, + extensions: Extensions, + path: http::uri::PathAndQuery, + rx: Receiver, + options: &RequestOptions, + api_client_header: &'static str, + request_params: &str, + ) -> crate::Result>>; +} + +pub type MockStream = Receiver>; + +pub fn mock_connector(mock: MockTestClient) -> Connector { + let client = SharedMockClient::new(mock); + + Connector::new(test_options(), client.clone()) +} diff --git a/src/storage/src/storage/bidi_write/redirect.rs b/src/storage/src/storage/bidi_write/redirect.rs index e17a4a745a..7ea054398d 100644 --- a/src/storage/src/storage/bidi_write/redirect.rs +++ b/src/storage/src/storage/bidi_write/redirect.rs @@ -14,6 +14,7 @@ // TODO(#5716): Lift to shared bidi module +use super::connector::AppendObjectSpecState; use crate::Error; use crate::google::rpc::Status as RpcStatus; use crate::google::storage::v2::{AppendObjectSpec, BidiWriteObjectRedirectedError}; @@ -24,8 +25,7 @@ use google_cloud_gax::error::rpc::Code; use prost::Message; use std::sync::{Arc, Mutex}; -#[allow(dead_code)] -pub fn handle_redirect(spec: Arc>, status: Status) -> crate::Error { +pub fn handle_redirect(state: Arc>, status: Status) -> crate::Error { let Ok(rpc_status) = RpcStatus::decode(status.details()) else { return to_gax_error(status); }; @@ -34,12 +34,47 @@ pub fn handle_redirect(spec: Arc>, status: Status) -> cr .into_iter() .find_map(|d| d.to_msg::().ok()) { - let mut guard = spec.lock().expect("never poisoned"); - guard.routing_token = redirect.routing_token; - guard.write_handle = redirect.write_handle; - if let Some(generation) = redirect.generation { - guard.generation = generation; - } + let mut guard = state.lock().expect("never poisoned"); + + let new_state = match &*guard { + AppendObjectSpecState::Write(spec, _) => { + if let Some(generation) = redirect.generation { + let bucket = spec + .resource + .as_ref() + .map(|r| r.bucket.clone()) + .unwrap_or_default(); + let object = spec + .resource + .as_ref() + .map(|r| r.name.clone()) + .unwrap_or_default(); + let new_spec = AppendObjectSpec { + bucket, + object, + generation, + if_metageneration_match: spec.if_metageneration_match, + if_metageneration_not_match: spec.if_metageneration_not_match, + routing_token: redirect.routing_token, + write_handle: redirect.write_handle, + }; + AppendObjectSpecState::Append(new_spec) + } else { + AppendObjectSpecState::Write(spec.clone(), redirect.routing_token) + } + } + AppendObjectSpecState::Append(spec) => { + let mut new_spec = spec.clone(); + new_spec.routing_token = redirect.routing_token; + new_spec.write_handle = redirect.write_handle; + if let Some(g) = redirect.generation { + new_spec.generation = g; + } + AppendObjectSpecState::Append(new_spec) + } + }; + + *guard = new_state; } to_gax_error(status) } @@ -106,16 +141,65 @@ mod tests { write_handle: Some(BidiWriteHandle { handle: bytes::Bytes::from_static(b"initial-handle"), }), + if_metageneration_match: Some(10), + if_metageneration_not_match: Some(20), + }; + let state = Arc::new(Mutex::new(AppendObjectSpecState::Append(spec))); + + let got = handle_redirect(state.clone(), status); + assert!(got.status().is_some(), "{got:?}"); + let guard = state.lock().expect("not poisoned"); + if let AppendObjectSpecState::Append(ref spec) = *guard { + assert_eq!(spec.routing_token.as_deref(), routing); + assert_eq!(spec.write_handle, write_handle); + assert_eq!(spec.generation, 42); + assert_eq!(spec.if_metageneration_match, Some(10)); + assert_eq!(spec.if_metageneration_not_match, Some(20)); + } else { + panic!("Expected AppendObjectSpecState::Append"); + } + } + + #[test] + fn preserves_preconditions_on_write() { + use crate::google::storage::v2::WriteObjectSpec; + let redirect = BidiWriteObjectRedirectedError { + routing_token: Some("new-routing".to_string()), + write_handle: Some(BidiWriteHandle { + handle: bytes::Bytes::from_static(b"new-handle"), + }), + generation: Some(42), + }; + let redirect = prost_types::Any::from_msg(&redirect).unwrap(); + let status = RpcStatus { + code: Code::Aborted as i32, + message: "test-only".to_string(), + details: vec![redirect], + }; + let details = bytes::Bytes::from_owner(status.encode_to_vec()); + let status = Status::with_details(Code::Aborted, "test-only", details); + + let write_spec = WriteObjectSpec { + if_metageneration_match: Some(11), + if_metageneration_not_match: Some(22), ..Default::default() }; - let spec = Arc::new(Mutex::new(spec)); + let state = Arc::new(Mutex::new(AppendObjectSpecState::Write( + Box::new(write_spec), + None, + ))); - let got = handle_redirect(spec.clone(), status); + let got = handle_redirect(state.clone(), status); assert!(got.status().is_some(), "{got:?}"); - let guard = spec.lock().expect("not poisoned"); - assert_eq!(guard.routing_token.as_deref(), routing); - assert_eq!(guard.write_handle, write_handle); - assert_eq!(guard.generation, 42); + let guard = state.lock().expect("not poisoned"); + if let AppendObjectSpecState::Append(ref spec) = *guard { + assert_eq!(spec.routing_token.as_deref(), Some("new-routing")); + assert_eq!(spec.generation, 42); + assert_eq!(spec.if_metageneration_match, Some(11)); + assert_eq!(spec.if_metageneration_not_match, Some(22)); + } else { + panic!("Expected AppendObjectSpecState::Append"); + } } #[test] @@ -138,14 +222,18 @@ mod tests { write_handle: Some(initial_handle.clone()), ..Default::default() }; - let spec = Arc::new(Mutex::new(spec)); + let state = Arc::new(Mutex::new(AppendObjectSpecState::Append(spec))); - let got = handle_redirect(spec.clone(), status); + let got = handle_redirect(state.clone(), status); assert!(got.status().is_some(), "{got:?}"); - let guard = spec.lock().expect("not poisoned"); - assert_eq!(guard.routing_token.as_deref(), Some("initial-token")); - assert_eq!(guard.write_handle, Some(initial_handle)); - assert_eq!(guard.generation, 1); + let guard = state.lock().expect("not poisoned"); + if let AppendObjectSpecState::Append(ref spec) = *guard { + assert_eq!(spec.routing_token.as_deref(), Some("initial-token")); + assert_eq!(spec.write_handle, Some(initial_handle)); + assert_eq!(spec.generation, 1); + } else { + panic!("Expected AppendObjectSpecState::Append"); + } } #[test_case(permanent_error(), false)] diff --git a/src/storage/src/storage/bidi_write/retry_redirect.rs b/src/storage/src/storage/bidi_write/retry_redirect.rs index 993292ca02..b43c307c21 100644 --- a/src/storage/src/storage/bidi_write/retry_redirect.rs +++ b/src/storage/src/storage/bidi_write/retry_redirect.rs @@ -33,13 +33,11 @@ use std::sync::Arc; /// the decorator has no effect. If the error is "permanent", but happens to be /// a redirect, then it is treated as retryable. #[derive(Clone, Debug)] -#[allow(dead_code)] pub struct RetryRedirect { inner: T, } impl RetryRedirect { - #[allow(dead_code)] pub fn new(inner: T) -> Self { Self { inner } } diff --git a/src/storage/src/storage/bidi_write/transport.rs b/src/storage/src/storage/bidi_write/transport.rs new file mode 100644 index 0000000000..c479c2a9d6 --- /dev/null +++ b/src/storage/src/storage/bidi_write/transport.rs @@ -0,0 +1,438 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// TODO(#5716): Lift to shared bidi module + +use super::worker::UploadIntent; +use crate::google::storage::v2::ObjectChecksums; +use crate::google::storage::v2::{ + BidiWriteObjectRequest, ChecksummedData, bidi_write_object_request::Data, + bidi_write_object_response::WriteStatus, +}; +use crate::stub::AppendableObjectWriter; +use bytes::Bytes; +use gaxi::prost::FromProto; +use tokio::sync::mpsc::Sender; +use tokio::sync::oneshot; + +/// The high-level transport adapter for bidirectional streaming writes. +/// Translates user calls into gRPC requests and delegates them to the background worker. +#[derive(Debug)] +pub struct AppendableObjectWriterTransport { + tx: Sender, + generation: i64, + persisted_size: i64, + write_offset: i64, + running_crc32c: u32, +} + +impl AppendableObjectWriterTransport { + pub async fn new_open( + mut connector: super::connector::Connector, + req: crate::model_ext::OpenAppendableObjectRequest, + ) -> crate::Result + where + T: super::Client + Clone + Sync + Send + 'static, + ::Stream: super::TonicStreaming + Send + Sync, + { + let (initial, connection) = connector.connect_open(req).await?; + Self::start_worker(connector, initial, connection, 0) + } + + pub async fn new_reopen( + mut connector: super::connector::Connector, + req: crate::model_ext::ReopenAppendableObjectRequest, + ) -> crate::Result + where + T: super::Client + Clone + Sync + Send + 'static, + ::Stream: super::TonicStreaming + Send + Sync, + { + let generation = req.generation; + let (initial, connection) = connector.connect_reopen(req).await?; + Self::start_worker(connector, initial, connection, generation) + } + + fn start_worker( + connector: super::connector::Connector, + initial: crate::google::storage::v2::BidiWriteObjectResponse, + connection: super::connector::Connection<::Stream>, + mut generation: i64, + ) -> crate::Result + where + T: super::Client + Clone + Sync + Send + 'static, + ::Stream: super::TonicStreaming + Send + Sync, + { + let mut persisted_size = 0; + if let Some(WriteStatus::Resource(r)) = initial.write_status.as_ref() { + generation = r.generation; + persisted_size = r.size; + } else if let Some(WriteStatus::PersistedSize(s)) = initial.write_status.as_ref() { + persisted_size = *s; + } + + let mut running_crc32c = 0; + if let Some(crc) = initial + .persisted_data_checksums + .as_ref() + .and_then(|c| c.crc32c) + { + running_crc32c = crc; + } + + let (tx, rx) = tokio::sync::mpsc::channel(100); + let worker = super::worker::Worker::new(connector); + tokio::spawn(worker.run(connection, rx)); + + Ok(Self { + tx, + generation, + persisted_size, + write_offset: persisted_size, + running_crc32c, + }) + } +} + +impl AppendableObjectWriter for AppendableObjectWriterTransport { + async fn append(&mut self, chunk: Bytes) -> crate::Result<()> { + let length = chunk.len() as i64; + let crc32c = crc32c::crc32c(&chunk); + + let new_crc32c = crc32c::crc32c_append(self.running_crc32c, &chunk); + + let request = BidiWriteObjectRequest { + write_offset: self.write_offset, + data: Some(Data::ChecksummedData(ChecksummedData { + content: chunk, + crc32c: Some(crc32c), + })), + ..BidiWriteObjectRequest::default() + }; + + self.tx + .send(UploadIntent::Append(request)) + .await + .map_err(|e| crate::Error::io(e.to_string()))?; + + self.write_offset += length; + self.running_crc32c = new_crc32c; + + Ok(()) + } + + async fn flush(&mut self) -> crate::Result { + let (sender, receiver) = oneshot::channel(); + let request = BidiWriteObjectRequest { + flush: true, + state_lookup: true, + write_offset: self.write_offset, + ..BidiWriteObjectRequest::default() + }; + + self.tx + .send(UploadIntent::Flush(request, sender)) + .await + .map_err(|e| crate::Error::io(e.to_string()))?; + + let response = receiver + .await + .map_err(|e| crate::Error::io(e.to_string()))??; + let size = match response.write_status { + Some(WriteStatus::PersistedSize(s)) => s, + Some(WriteStatus::Resource(r)) => r.size, + None => return Err(crate::Error::io("flush response missing write_status")), + }; + self.persisted_size = size; + Ok(size) + } + + async fn finalize(self) -> crate::Result { + let (sender, receiver) = oneshot::channel(); + let request = BidiWriteObjectRequest { + finish_write: true, + flush: true, + write_offset: self.write_offset, + object_checksums: Some(ObjectChecksums { + crc32c: Some(self.running_crc32c), + md5_hash: vec![].into(), + }), + ..BidiWriteObjectRequest::default() + }; + + self.tx + .send(UploadIntent::Finalize(request, sender)) + .await + .map_err(|e| crate::Error::io(e.to_string()))?; + + let response = receiver + .await + .map_err(|e| crate::Error::io(e.to_string()))??; + + let resource = match response.write_status { + Some(WriteStatus::Resource(r)) => r, + _ => return Err(crate::Error::io("finalize did not return a resource")), + }; + + let object = FromProto::cnv(resource) + .map_err(|_| crate::Error::deser("converting resource to object"))?; + + Ok(object) + } + + async fn close(mut self) -> crate::Result { + self.flush().await + } + + fn generation(&self) -> i64 { + self.generation + } + + fn persisted_size(&self) -> i64 { + self.persisted_size + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::google::storage::v2::{ + BidiWriteObjectResponse, Object, bidi_write_object_response::WriteStatus, + }; + use tokio::sync::mpsc; + + #[tokio::test] + async fn success() -> anyhow::Result<()> { + let (tx, mut rx) = mpsc::channel(1); + let transport = AppendableObjectWriterTransport { + tx, + write_offset: 0, + running_crc32c: 0, + generation: 123456, + persisted_size: 0, + }; + + let handle = tokio::spawn(async move { + let mut transport = transport; + transport.append(bytes::Bytes::from("hello")).await.unwrap(); + transport.flush().await.unwrap(); + transport.finalize().await.unwrap(); + }); + + // Verify append + let intent = rx.recv().await.unwrap(); + if let UploadIntent::Append(req) = intent { + assert_eq!(req.write_offset, 0); + assert!(!req.finish_write); + } else { + panic!("expected Append"); + } + + // Verify flush + let intent = rx.recv().await.unwrap(); + if let UploadIntent::Flush(req, sender) = intent { + assert!(req.flush); + let resp = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(5)), + ..Default::default() + }; + sender.send(Ok(resp)).unwrap(); + } else { + panic!("expected Flush"); + } + + // Verify finalize + let intent = rx.recv().await.unwrap(); + if let UploadIntent::Finalize(req, sender) = intent { + assert!(req.finish_write); + let object = Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + size: 5, + generation: 123456, + ..Default::default() + }; + let resp = BidiWriteObjectResponse { + write_status: Some(WriteStatus::Resource(object)), + ..Default::default() + }; + sender.send(Ok(resp)).unwrap(); + } else { + panic!("expected Finalize"); + } + + handle.await?; + Ok(()) + } + + #[tokio::test] + async fn connect_error_open() -> anyhow::Result<()> { + use super::super::mocks::{MockTestClient, mock_connector}; + use super::super::tests::permanent_error; + use crate::model_ext::OpenAppendableObjectRequest; + + let mut mock = MockTestClient::new(); + mock.expect_start() + .return_once(move |_, _, _, _, _, _| Err(permanent_error())); + let connector = mock_connector(mock); + let mut req = OpenAppendableObjectRequest { + spec: Default::default(), + params: None, + }; + req.spec.resource = Some( + crate::model::Object::default() + .set_bucket("projects/_/buckets/test-bucket") + .set_name("test-object"), + ); + + let err = AppendableObjectWriterTransport::new_open(connector, req) + .await + .unwrap_err(); + assert_eq!(err.status(), permanent_error().status(), "{err:?}"); + Ok(()) + } + + #[tokio::test] + async fn connect_error_reopen() -> anyhow::Result<()> { + use super::super::mocks::{MockTestClient, mock_connector}; + use super::super::tests::permanent_error; + use crate::model_ext::ReopenAppendableObjectRequest; + + let mut mock = MockTestClient::new(); + mock.expect_start() + .return_once(move |_, _, _, _, _, _| Err(permanent_error())); + let connector = mock_connector(mock); + let req = ReopenAppendableObjectRequest { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + generation: 123, + if_metageneration_match: None, + if_metageneration_not_match: None, + routing_token: None, + write_handle: None, + params: None, + }; + + let err = AppendableObjectWriterTransport::new_reopen(connector, req) + .await + .unwrap_err(); + assert_eq!(err.status(), permanent_error().status(), "{err:?}"); + Ok(()) + } + + #[tokio::test] + async fn write_error() -> anyhow::Result<()> { + let (tx, rx) = mpsc::channel(1); + let mut transport = AppendableObjectWriterTransport { + tx, + write_offset: 0, + running_crc32c: 0, + generation: 123456, + persisted_size: 0, + }; + + // Simulate an early stream closure / dropped rx (like a worker dying) + drop(rx); + let err = transport + .append(bytes::Bytes::from("hello")) + .await + .unwrap_err(); + assert!(err.is_io(), "{err:?}"); + + // Assert that state was NOT modified due to the error + assert_eq!(transport.write_offset, 0); + assert_eq!(transport.running_crc32c, 0); + + Ok(()) + } + + #[tokio::test] + async fn finalize_error() -> anyhow::Result<()> { + let (tx, mut rx) = mpsc::channel(1); + let transport = AppendableObjectWriterTransport { + tx, + write_offset: 0, + running_crc32c: 0, + generation: 123456, + persisted_size: 0, + }; + + let handle = tokio::spawn(async move { + let transport = transport; + transport.finalize().await + }); + + // Verify finalize intent + let intent = rx.recv().await.unwrap(); + if let UploadIntent::Finalize(_, sender) = intent { + // Respond with an invalid WriteStatus (not Resource) + let resp = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(5)), + ..Default::default() + }; + sender.send(Ok(resp)).unwrap(); + } else { + panic!("expected Finalize"); + } + + let err = handle.await?.unwrap_err(); + assert!(err.is_io(), "{err:?}"); + + Ok(()) + } + + #[tokio::test] + async fn new_reopen_initial_persisted_size() -> anyhow::Result<()> { + use super::super::mocks::{MockTestClient, mock_connector}; + use crate::model_ext::ReopenAppendableObjectRequest; + use gaxi::grpc::tonic::Response as TonicResponse; + use gaxi::grpc::tonic::Result as TonicResult; + + let (tx1, rx1) = tokio::sync::mpsc::channel::>(5); + let stream1 = TonicResponse::from(rx1); + + let mut mock = MockTestClient::new(); + mock.expect_start() + .return_once(move |_, _, _, _, _, _| Ok(Ok(stream1))); + let connector = mock_connector(mock); + + let req = ReopenAppendableObjectRequest { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + generation: 123456, + if_metageneration_match: None, + if_metageneration_not_match: None, + routing_token: None, + write_handle: None, + params: None, + }; + + let initial_response = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(1024)), + persisted_data_checksums: Some(crate::google::storage::v2::ObjectChecksums { + crc32c: Some(9999), + md5_hash: vec![].into(), + }), + ..Default::default() + }; + tx1.send(Ok(initial_response)).await?; + + let transport = AppendableObjectWriterTransport::new_reopen(connector, req).await?; + + assert_eq!(transport.generation(), 123456); + assert_eq!(transport.persisted_size(), 1024); + assert_eq!(transport.write_offset, 1024); + assert_eq!(transport.running_crc32c, 9999); + Ok(()) + } +} diff --git a/src/storage/src/storage/bidi_write/worker.rs b/src/storage/src/storage/bidi_write/worker.rs new file mode 100644 index 0000000000..d8b7880ed2 --- /dev/null +++ b/src/storage/src/storage/bidi_write/worker.rs @@ -0,0 +1,252 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// TODO(#5716): Lift to shared bidi module + +use super::connector::Connection; +use super::{Client, TonicStreaming}; +use crate::error::WriteError; +use crate::google::storage::v2::{BidiWriteObjectRequest, BidiWriteObjectResponse}; +use gaxi::grpc::tonic::Result as TonicResult; +use std::sync::Arc; +use tokio::sync::mpsc::Receiver; +use tokio::sync::oneshot; + +type WriteResult = std::result::Result; +type LoopResult = std::result::Result>; + +/// The intent sent from the user to the background worker. +pub enum UploadIntent { + /// Append a chunk of data. + Append(BidiWriteObjectRequest), + /// Flush the data to the server. + Flush( + BidiWriteObjectRequest, + oneshot::Sender>, + ), + /// Finalize the upload. + Finalize( + BidiWriteObjectRequest, + oneshot::Sender>, + ), +} + +/// The background worker that manages the live gRPC stream. +pub struct Worker { + _connector: super::connector::Connector, + pending_flushes: + std::collections::VecDeque>>, +} + +impl Worker { + pub fn new(connector: super::connector::Connector) -> Self { + Self { + _connector: connector, + pending_flushes: std::collections::VecDeque::new(), + } + } +} + +impl Worker +where + C: Client + Clone + 'static, + ::Stream: TonicStreaming, +{ + pub async fn run( + mut self, + connection: Connection, + mut requests: Receiver, + ) -> LoopResult<()> { + let (mut rx, mut tx) = (connection.rx, connection.tx); + + let error = loop { + tokio::select! { + m = rx.next_message() => { + match self.handle_response(m) { + // Successful end of stream, return without error. + None => break None, + // An unrecoverable in the stream or its data, return + // the error. + Some(Err(e)) => break Some(e), + // New message on the stream handled successfully, + // continue. + Some(Ok(None)) => {}, + // The stream reconnected successfully, update the local + // variables and continue. + Some(Ok(Some(connection))) => { + (rx, tx) = (connection.rx, connection.tx); + } + } + }, + intent = requests.recv() => { + match intent { + Some(UploadIntent::Append(request)) => { + if let Err(e) = tx.send(request).await { + break Some(Arc::new(crate::Error::io(e))); + } + } + Some(UploadIntent::Flush(request, sender)) => { + self.pending_flushes.push_back(sender); + if let Err(e) = tx.send(request).await { + break Some(Arc::new(crate::Error::io(e))); + } + } + Some(UploadIntent::Finalize(request, sender)) => { + self.pending_flushes.push_back(sender); + if let Err(e) = tx.send(request).await { + break Some(Arc::new(crate::Error::io(e))); + } + } + None => { + break None; + } + } + } + } + }; + + if let Some(e) = error.clone() { + for sender in self.pending_flushes.drain(..) { + let _ = sender.send(Err(crate::Error::io(e.to_string()))); + } + return Err(e); + } + + Ok(()) + } + + pub fn handle_response( + &mut self, + message: TonicResult>, + ) -> Option>>> { + let response = match message { + Ok(Some(msg)) => msg, + Ok(None) => { + if !self.pending_flushes.is_empty() { + return Some(Err(Arc::new(crate::Error::io( + "stream closed unexpectedly", + )))); + } + return None; + } + Err(e) => return Some(Err(Arc::new(crate::Error::io(e)))), + }; + self.handle_response_success(response); + // TODO(#5716): Update when implementing reconnect logic. + Some(Ok(None)) + } + + pub fn handle_response_success(&mut self, response: BidiWriteObjectResponse) { + if let Some(sender) = self.pending_flushes.pop_front() { + let _ = sender.send(Ok(response)); + } else { + // Log unprompted server responses. + tracing::debug!( + "Received unprompted BidiWriteObjectResponse from server: {:?}", + response + ); + } + } +} + +#[cfg(test)] +mod tests { + use super::super::mocks::{MockTestClient, mock_connector}; + use super::*; + use crate::google::storage::v2::{ + BidiWriteObjectRequest, BidiWriteObjectResponse, bidi_write_object_response::WriteStatus, + }; + use tokio::sync::mpsc; + use tokio::sync::oneshot; + + #[tokio::test] + async fn run_immediately_closed() -> anyhow::Result<()> { + let (request_tx, _request_rx) = mpsc::channel(1); + let (response_tx, response_rx) = mpsc::channel(10); + let (tx, rx) = mpsc::channel(1); + let connection = Connection::new(request_tx, response_rx); + + let mut mock = MockTestClient::new(); + mock.expect_start().never(); + + let connector = mock_connector(mock); + let worker = Worker::new(connector); + let handle = tokio::spawn(worker.run(connection, rx)); + + drop(response_tx); + drop(tx); + handle.await??; + Ok(()) + } + + #[tokio::test] + async fn run_flush_response() -> anyhow::Result<()> { + let (request_tx, mut request_rx) = mpsc::channel(1); + let (response_tx, response_rx) = mpsc::channel(10); + let (tx, rx) = mpsc::channel(1); + let connection = Connection::new(request_tx, response_rx); + + let mut mock = MockTestClient::new(); + mock.expect_start().never(); + + let connector = mock_connector(mock); + let worker = Worker::new(connector); + let handle = tokio::spawn(worker.run(connection, rx)); + + let (flush_tx, flush_rx) = oneshot::channel(); + let flush_request = BidiWriteObjectRequest { + flush: true, + ..Default::default() + }; + tx.send(UploadIntent::Flush(flush_request.clone(), flush_tx)) + .await?; + + // The worker should send the request to the stream. + let stream_req = request_rx.recv().await.unwrap(); + assert!(stream_req.flush); + + // The server responds + let server_resp = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(100)), + ..Default::default() + }; + response_tx.send(Ok(server_resp.clone())).await?; + + // The flush sender should get the response + let received_resp = flush_rx.await??; + assert_eq!(received_resp.write_status, server_resp.write_status); + + drop(response_tx); + drop(tx); + handle.await??; + Ok(()) + } + + #[tokio::test] + async fn run_stop_on_closed_requests() -> anyhow::Result<()> { + let (request_tx, _request_rx) = mpsc::channel(1); + let (_response_tx, response_rx) = mpsc::channel(10); + let (tx, rx) = mpsc::channel(1); + let connection = Connection::new(request_tx, response_rx); + + let mut mock = MockTestClient::new(); + mock.expect_start().never(); + + let connector = mock_connector(mock); + let worker = Worker::new(connector); + drop(tx); + worker.run(connection, rx).await?; + Ok(()) + } +} diff --git a/src/storage/src/storage/client.rs b/src/storage/src/storage/client.rs index 9b9c1b8a44..259975f885 100644 --- a/src/storage/src/storage/client.rs +++ b/src/storage/src/storage/client.rs @@ -13,7 +13,11 @@ // limitations under the License. use super::request_options::RequestOptions; +#[cfg(google_cloud_unstable_storage_bidi)] +use crate::builder::storage::OpenAppendableObject; use crate::builder::storage::ReadObject; +#[cfg(google_cloud_unstable_storage_bidi)] +use crate::builder::storage::ReopenAppendableObject; use crate::builder::storage::WriteObject; use crate::read_resume_policy::ReadResumePolicy; use crate::storage::bidi::OpenObject; @@ -275,6 +279,88 @@ where { OpenObject::new(self.stub.clone(), bucket, object, self.options.clone()) } + + #[cfg(google_cloud_unstable_storage_bidi)] + #[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] + /// Opens a new appendable object for exclusive appends. + /// + /// This method allows you to incrementally append data to a single object. + /// + /// # Example + /// ``` + /// # use google_cloud_storage::client::Storage; + /// # use bytes::Bytes; + /// # async fn sample(client: &Storage) -> anyhow::Result<()> { + /// let mut writer = client + /// .open_appendable_object("projects/_/buckets/my-bucket", "my-object") + /// .send() + /// .await?; + /// + /// writer.append(Bytes::from_static(b"hello ")).await?; + /// writer.append(Bytes::from_static(b"world")).await?; + /// writer.flush().await?; + /// + /// let object = writer.finalize().await?; + /// println!("Finalized object size: {:?}", object.size); + /// # Ok(()) } + /// ``` + /// + /// # Parameters + /// * `bucket` - the bucket name containing the object. In + /// `projects/_/buckets/{bucket_id}` format. + /// * `object` - the object name. + #[cfg(google_cloud_unstable_storage_bidi)] + #[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] + pub fn open_appendable_object(&self, bucket: B, object: O) -> OpenAppendableObject + where + B: Into, + O: Into, + { + OpenAppendableObject::new(self.stub.clone(), bucket, object, self.options.clone()) + } + + #[cfg(google_cloud_unstable_storage_bidi)] + #[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] + /// Reopens an existing appendable object for exclusive appends. + /// + /// # Example + /// ``` + /// # use google_cloud_storage::client::Storage; + /// # use bytes::Bytes; + /// # async fn sample(client: &Storage) -> anyhow::Result<()> { + /// let mut writer = client + /// .reopen_appendable_object("projects/_/buckets/my-bucket", "my-object", 12345) + /// .send() + /// .await?; + /// + /// writer.append(Bytes::from_static(b"more bytes")).await?; + /// writer.flush().await?; + /// # Ok(()) } + /// ``` + /// + /// # Parameters + /// * `bucket` - the bucket name containing the object. In + /// `projects/_/buckets/{bucket_id}` format. + /// * `object` - the object name. + /// * `generation` - the object generation. + pub fn reopen_appendable_object( + &self, + bucket: B, + object: O, + generation: i64, + ) -> ReopenAppendableObject + where + B: Into, + O: Into, + { + ReopenAppendableObject::new( + self.stub.clone(), + bucket, + object, + generation, + self.options.clone(), + ) + } } impl Storage { diff --git a/src/storage/src/storage/open_appendable_object.rs b/src/storage/src/storage/open_appendable_object.rs new file mode 100644 index 0000000000..0cadc86c22 --- /dev/null +++ b/src/storage/src/storage/open_appendable_object.rs @@ -0,0 +1,368 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::Result; +use crate::appendable_object_writer::AppendableObjectWriter; +use crate::model_ext::{KeyAes256, OpenAppendableObjectRequest}; +use crate::request_options::RequestOptions; +use std::sync::Arc; +use std::time::Duration; + +#[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] +/// A request builder for [Storage::open_appendable_object][crate::client::Storage::open_appendable_object]. +/// +/// # Example +/// ``` +/// use google_cloud_storage::client::Storage; +/// # use google_cloud_storage::builder::storage::OpenAppendableObject; +/// async fn sample(client: &Storage) -> anyhow::Result<()> { +/// let builder: OpenAppendableObject = client +/// .open_appendable_object("projects/_/buckets/my-bucket", "my-object"); +/// let mut writer = builder +/// .send() +/// .await?; +/// writer.append(bytes::Bytes::from("hello")).await?; +/// writer.finalize().await?; +/// Ok(()) +/// } +/// ``` +#[derive(Clone, Debug)] +pub struct OpenAppendableObject { + stub: Arc, + request: OpenAppendableObjectRequest, + options: RequestOptions, +} + +impl OpenAppendableObject +where + S: crate::storage::stub::Storage + 'static, +{ + /// Sends the request, returning an appendable object writer. + /// + /// Example: + /// ```ignore + /// # use google_cloud_storage::{model_ext::KeyAes256, client::Storage}; + /// # async fn sample(client: &Storage) -> anyhow::Result<()> { + /// let mut writer = client + /// .open_appendable_object("projects/_/buckets/my-bucket", "my-object") + /// .send() + /// .await?; + /// writer.append(bytes::Bytes::from("world")).await?; + /// writer.finalize().await?; + /// # Ok(()) } + /// ``` + pub async fn send(self) -> Result { + self.stub + .open_appendable_object(self.request, self.options) + .await + } +} + +impl OpenAppendableObject { + pub(crate) fn new( + stub: std::sync::Arc, + bucket: B, + object: O, + options: RequestOptions, + ) -> Self + where + B: Into, + O: Into, + { + let request = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default() + .set_appendable(true) + .set_resource( + crate::model::Object::new() + .set_bucket(bucket.into()) + .set_name(object.into()), + ), + params: None, + }; + Self { + request, + options, + stub, + } + } + + /// Makes the operation conditional on whether the object's live generation + /// matches the given value. + pub fn set_if_generation_match(mut self, v: i64) -> Self { + self.request.spec.if_generation_match = Some(v); + self + } + + /// Makes the operation conditional on whether the object's live generation + /// does not match the given value. If no live object exists, the precondition + /// fails. + pub fn set_if_generation_not_match(mut self, v: i64) -> Self { + self.request.spec.if_generation_not_match = Some(v); + self + } + + /// Makes the operation conditional on whether the object's current + /// metageneration matches the given value. + pub fn set_if_metageneration_match(mut self, v: i64) -> Self { + self.request.spec.if_metageneration_match = Some(v); + self + } + + /// Makes the operation conditional on whether the object's current + /// metageneration does not match the given value. + pub fn set_if_metageneration_not_match(mut self, v: i64) -> Self { + self.request.spec.if_metageneration_not_match = Some(v); + self + } + + /// The encryption key used with the Customer-Supplied Encryption Keys + /// feature. In raw bytes format (not base64-encoded). + pub fn set_key(mut self, v: KeyAes256) -> Self { + self.request.params = Some(crate::model::CommonObjectRequestParams::from(v)); + self + } + + /// The retry policy used for this request. + pub fn with_retry_policy>( + mut self, + v: V, + ) -> Self { + self.options.retry_policy = v.into().into(); + self + } + + /// The backoff policy used for this request. + pub fn with_backoff_policy>( + mut self, + v: V, + ) -> Self { + self.options.backoff_policy = v.into().into(); + self + } + + /// The retry throttler used for this request. + pub fn with_retry_throttler>( + mut self, + v: V, + ) -> Self { + self.options.retry_throttler = v.into().into(); + self + } + + /// Configure per-attempt timeout. + pub fn with_attempt_timeout(mut self, v: Duration) -> Self { + self.options.set_bidi_attempt_timeout(v); + self + } + + /// Sets the `User-Agent` header for this request. + pub fn with_user_agent(mut self, user_agent: impl Into) -> Self { + self.options.user_agent = Some(user_agent.into()); + self + } + + /// Sets the project that will be billed for this request. + pub fn with_quota_project(mut self, project: impl Into) -> Self { + self.options.set_quota_project(project); + self + } + + fn mut_resource(&mut self) -> &mut crate::model::Object { + if self.request.spec.resource.is_none() { + self.request.spec.resource = Some(crate::model::Object::new()); + } + self.request.spec.resource.as_mut().unwrap() + } + + /// Sets the ACL for the new object. + pub fn set_acl(mut self, v: I) -> Self + where + I: IntoIterator, + V: Into, + { + self.mut_resource().acl = v.into_iter().map(|a| a.into()).collect(); + self + } + + /// Sets the [cache control] for the new object. + pub fn set_cache_control>(mut self, v: V) -> Self { + self.mut_resource().cache_control = v.into(); + self + } + + /// Sets the [content disposition] for the new object. + pub fn set_content_disposition>(mut self, v: V) -> Self { + self.mut_resource().content_disposition = v.into(); + self + } + + /// Sets the [content encoding] for the object data. + pub fn set_content_encoding>(mut self, v: V) -> Self { + self.mut_resource().content_encoding = v.into(); + self + } + + /// Sets the [content language] for the new object. + pub fn set_content_language>(mut self, v: V) -> Self { + self.mut_resource().content_language = v.into(); + self + } + + /// Sets the [content type] for the new object. + pub fn set_content_type>(mut self, v: V) -> Self { + self.mut_resource().content_type = v.into(); + self + } + + /// Sets the [custom time] for the new object. + pub fn set_custom_time>(mut self, v: V) -> Self { + self.mut_resource().custom_time = Some(v.into()); + self + } + + /// Sets the [event based hold] flag for the new object. + pub fn set_event_based_hold>(mut self, v: V) -> Self { + self.mut_resource().event_based_hold = Some(v.into()); + self + } + + /// Sets the [custom metadata] for the new object. + pub fn set_metadata(mut self, i: I) -> Self + where + I: IntoIterator, + K: Into, + V: Into, + { + self.mut_resource().metadata = i.into_iter().map(|(k, v)| (k.into(), v.into())).collect(); + self + } + + /// Sets the [retention configuration] for the new object. + pub fn set_retention(mut self, v: V) -> Self + where + V: Into, + { + self.mut_resource().retention = Some(v.into()); + self + } + + /// Sets the [storage class] for the new object. + pub fn set_storage_class(mut self, v: V) -> Self + where + V: Into, + { + self.mut_resource().storage_class = v.into(); + self + } + + /// Sets the [temporary hold] flag for the new object. + pub fn set_temporary_hold>(mut self, v: V) -> Self { + self.mut_resource().temporary_hold = v.into(); + self + } + + /// Sets the resource name of the [Customer-managed encryption key] for this + /// object. + pub fn set_kms_key(mut self, v: V) -> Self + where + V: Into, + { + self.mut_resource().kms_key = v.into(); + self + } + + /// Configure this object to use one of the [predefined ACLs]. + pub fn set_predefined_acl(mut self, v: V) -> Self + where + V: Into, + { + self.request.spec.predefined_acl = v.into(); + self + } + + /// Sets the object custom contexts. + pub fn set_contexts(mut self, v: V) -> Self + where + V: Into, + { + self.mut_resource().contexts = Some(v.into()); + self + } +} + +#[cfg(test)] +#[cfg(google_cloud_unstable_storage_bidi)] +mod tests { + use super::*; + use crate::client::Storage; + use crate::request_options::RequestOptions; + use anyhow::Result; + use google_cloud_auth::credentials::anonymous::Builder as Anonymous; + use static_assertions::assert_impl_all; + use std::sync::Arc; + + #[derive(Debug)] + struct StorageStub; + impl crate::stub::Storage for StorageStub {} + + const BUCKET_NAME: &str = "projects/_/buckets/test-bucket"; + const OBJECT_NAME: &str = "test-object"; + + #[tokio::test] + async fn traits() -> Result<()> { + assert_impl_all!(OpenAppendableObject: Clone, std::fmt::Debug, Send, Sync); + let client = Storage::builder() + .with_credentials(Anonymous::new().build()) + .build() + .await?; + + fn need_send(_val: &T) {} + fn need_static(_val: &T) {} + + let open = client.open_appendable_object(BUCKET_NAME, OBJECT_NAME); + need_static(&open); + let fut = client + .open_appendable_object(BUCKET_NAME, OBJECT_NAME) + .send(); + need_send(&fut); + need_static(&fut); + Ok(()) + } + + #[tokio::test] + async fn attributes() -> Result<()> { + let options = RequestOptions::new(); + let builder = OpenAppendableObject::new( + Arc::new(StorageStub), + BUCKET_NAME.to_string(), + OBJECT_NAME.to_string(), + options, + ) + .set_if_generation_match(234) + .set_if_generation_not_match(345) + .set_if_metageneration_match(456) + .set_if_metageneration_not_match(567) + .set_content_type("text/plain"); + + assert_eq!(builder.request.spec.if_generation_match, Some(234)); + assert_eq!(builder.request.spec.if_generation_not_match, Some(345)); + assert_eq!(builder.request.spec.if_metageneration_match, Some(456)); + assert_eq!(builder.request.spec.if_metageneration_not_match, Some(567)); + assert_eq!( + builder.request.spec.resource.as_ref().unwrap().content_type, + "text/plain" + ); + Ok(()) + } +} diff --git a/src/storage/src/storage/reopen_appendable_object.rs b/src/storage/src/storage/reopen_appendable_object.rs new file mode 100644 index 0000000000..b79b671e61 --- /dev/null +++ b/src/storage/src/storage/reopen_appendable_object.rs @@ -0,0 +1,243 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::Result; +use crate::appendable_object_writer::AppendableObjectWriter; +use crate::model_ext::{KeyAes256, ReopenAppendableObjectRequest}; +use crate::request_options::RequestOptions; +use std::sync::Arc; +use std::time::Duration; + +#[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] +/// A request builder for [Storage::reopen_appendable_object][crate::client::Storage::reopen_appendable_object]. +/// +/// # Example +/// ``` +/// use google_cloud_storage::client::Storage; +/// # use google_cloud_storage::builder::storage::ReopenAppendableObject; +/// async fn sample(client: &Storage) -> anyhow::Result<()> { +/// let builder: ReopenAppendableObject = client +/// .reopen_appendable_object("projects/_/buckets/my-bucket", "my-object", 123456); +/// let mut writer = builder +/// .send() +/// .await?; +/// writer.append(bytes::Bytes::from("hello")).await?; +/// writer.finalize().await?; +/// Ok(()) +/// } +/// ``` +#[derive(Clone, Debug)] +pub struct ReopenAppendableObject { + stub: Arc, + request: ReopenAppendableObjectRequest, + options: RequestOptions, +} + +impl ReopenAppendableObject +where + S: crate::storage::stub::Storage + 'static, +{ + /// Sends the request, returning an appendable object writer. + /// + /// Example: + /// ```ignore + /// # use google_cloud_storage::{model_ext::KeyAes256, client::Storage}; + /// # async fn sample(client: &Storage) -> anyhow::Result<()> { + /// let mut writer = client + /// .reopen_appendable_object("projects/_/buckets/my-bucket", "my-object", 123456) + /// .send() + /// .await?; + /// writer.append(bytes::Bytes::from("world")).await?; + /// writer.finalize().await?; + /// # Ok(()) } + /// ``` + pub async fn send(self) -> Result { + self.stub + .reopen_appendable_object(self.request, self.options) + .await + } +} + +impl ReopenAppendableObject { + pub(crate) fn new( + stub: std::sync::Arc, + bucket: B, + object: O, + generation: i64, + options: RequestOptions, + ) -> Self + where + B: Into, + O: Into, + { + let request = ReopenAppendableObjectRequest { + bucket: bucket.into(), + object: object.into(), + generation, + if_metageneration_match: None, + if_metageneration_not_match: None, + routing_token: None, + write_handle: None, + params: None, + }; + Self { + request, + options, + stub, + } + } + + /// Makes the operation conditional on whether the object's current + /// metageneration matches the given value. + pub fn set_if_metageneration_match(mut self, v: i64) -> Self { + self.request.if_metageneration_match = Some(v); + self + } + + /// Makes the operation conditional on whether the object's current + /// metageneration does not match the given value. + pub fn set_if_metageneration_not_match(mut self, v: i64) -> Self { + self.request.if_metageneration_not_match = Some(v); + self + } + + /// Explicitly provide the routing token from the previous operation. + pub fn set_routing_token(mut self, token: impl Into) -> Self { + self.request.routing_token = Some(token.into()); + self + } + + /// Explicitly provide the write handle from the previous operation. + pub fn set_write_handle(mut self, handle: bytes::Bytes) -> Self { + self.request.write_handle = Some(handle); + self + } + + /// The encryption key used with the Customer-Supplied Encryption Keys + /// feature. In raw bytes format (not base64-encoded). + pub fn set_key(mut self, v: KeyAes256) -> Self { + self.request.params = Some(crate::model::CommonObjectRequestParams::from(v)); + self + } + + /// The retry policy used for this request. + pub fn with_retry_policy>( + mut self, + v: V, + ) -> Self { + self.options.retry_policy = v.into().into(); + self + } + + /// The backoff policy used for this request. + pub fn with_backoff_policy>( + mut self, + v: V, + ) -> Self { + self.options.backoff_policy = v.into().into(); + self + } + + /// The retry throttler used for this request. + pub fn with_retry_throttler>( + mut self, + v: V, + ) -> Self { + self.options.retry_throttler = v.into().into(); + self + } + + /// Configure per-attempt timeout. + pub fn with_attempt_timeout(mut self, v: Duration) -> Self { + self.options.set_bidi_attempt_timeout(v); + self + } + + /// Sets the `User-Agent` header for this request. + pub fn with_user_agent(mut self, user_agent: impl Into) -> Self { + self.options.user_agent = Some(user_agent.into()); + self + } + + /// Sets the project that will be billed for this request. + pub fn with_quota_project(mut self, project: impl Into) -> Self { + self.options.set_quota_project(project); + self + } +} + +#[cfg(test)] +#[cfg(google_cloud_unstable_storage_bidi)] +mod tests { + use super::*; + use crate::client::Storage; + use crate::request_options::RequestOptions; + use anyhow::Result; + use google_cloud_auth::credentials::anonymous::Builder as Anonymous; + use static_assertions::assert_impl_all; + use std::sync::Arc; + + #[derive(Debug)] + struct StorageStub; + impl crate::stub::Storage for StorageStub {} + + const BUCKET_NAME: &str = "projects/_/buckets/test-bucket"; + const OBJECT_NAME: &str = "test-object"; + + #[tokio::test] + async fn traits() -> Result<()> { + assert_impl_all!(ReopenAppendableObject: Clone, std::fmt::Debug, Send, Sync); + let client = Storage::builder() + .with_credentials(Anonymous::new().build()) + .build() + .await?; + + fn need_send(_val: &T) {} + fn need_static(_val: &T) {} + + let open = client.reopen_appendable_object(BUCKET_NAME, OBJECT_NAME, 123456); + need_static(&open); + let fut = client + .reopen_appendable_object(BUCKET_NAME, OBJECT_NAME, 123456) + .send(); + need_send(&fut); + need_static(&fut); + Ok(()) + } + + #[tokio::test] + async fn attributes() -> Result<()> { + let options = RequestOptions::new(); + let builder = ReopenAppendableObject::new( + Arc::new(StorageStub), + BUCKET_NAME.to_string(), + OBJECT_NAME.to_string(), + 123456, + options, + ) + .set_if_metageneration_match(456) + .set_if_metageneration_not_match(567) + .set_routing_token("token") + .set_write_handle(bytes::Bytes::from("handle")); + + assert_eq!(builder.request.if_metageneration_match, Some(456)); + assert_eq!(builder.request.if_metageneration_not_match, Some(567)); + assert_eq!(builder.request.routing_token.as_deref(), Some("token")); + assert_eq!( + builder.request.write_handle, + Some(bytes::Bytes::from("handle")) + ); + Ok(()) + } +} diff --git a/src/storage/src/storage/stub.rs b/src/storage/src/storage/stub.rs index e1957991ca..873a49769f 100644 --- a/src/storage/src/storage/stub.rs +++ b/src/storage/src/storage/stub.rs @@ -83,6 +83,30 @@ pub trait Storage: std::fmt::Debug + Send + Sync { { unimplemented_stub::<(Descriptor, Vec)>() } + + #[cfg(google_cloud_unstable_storage_bidi)] + /// Implements [crate::client::Storage::open_appendable_object]. + fn open_appendable_object( + &self, + _request: crate::model_ext::OpenAppendableObjectRequest, + _options: RequestOptions, + ) -> impl std::future::Future< + Output = Result, + > + Send { + unimplemented_stub::() + } + + #[cfg(google_cloud_unstable_storage_bidi)] + /// Implements [crate::client::Storage::reopen_appendable_object]. + fn reopen_appendable_object( + &self, + _request: crate::model_ext::ReopenAppendableObjectRequest, + _options: RequestOptions, + ) -> impl std::future::Future< + Output = Result, + > + Send { + unimplemented_stub::() + } } /// Defines the trait used to implement [crate::object_descriptor::ObjectDescriptor]. diff --git a/src/storage/src/storage/tracing.rs b/src/storage/src/storage/tracing.rs index 00c6ca448d..2822d80ac1 100644 --- a/src/storage/src/storage/tracing.rs +++ b/src/storage/src/storage/tracing.rs @@ -106,6 +106,130 @@ impl ObjectDescriptorStub for TracingObjectDescriptor { + inner: T, +} + +#[cfg(google_cloud_unstable_storage_bidi)] +impl TracingAppendableObjectWriter { + pub(crate) fn new(inner: T) -> Self { + Self { inner } + } +} + +#[cfg(google_cloud_unstable_storage_bidi)] +impl AppendableObjectWriterStub + for TracingAppendableObjectWriter> +{ + #[tracing::instrument( + name = "append", + err, + skip(self, chunk), + fields( + otel.kind = "Internal", + rpc.system.name = "grpc", + gcp.client.service = INSTRUMENTATION.service_name, + gcp.client.version = INSTRUMENTATION.client_version, + gcp.client.repo = GCP_CLIENT_REPO_GOOGLEAPIS, + gcp.client.artifact = INSTRUMENTATION.client_artifact, + gcp.schema.url = SCHEMA_URL_VALUE, + append.chunk_size = chunk.len(), + append.generation = self.inner.generation(), + ) + )] + async fn append(&mut self, chunk: bytes::Bytes) -> crate::Result<()> { + self.inner.append(chunk).await + } + + #[tracing::instrument( + name = "flush", + err, + skip(self), + fields( + otel.kind = "Internal", + rpc.system.name = "grpc", + gcp.client.service = INSTRUMENTATION.service_name, + gcp.client.version = INSTRUMENTATION.client_version, + gcp.client.repo = GCP_CLIENT_REPO_GOOGLEAPIS, + gcp.client.artifact = INSTRUMENTATION.client_artifact, + gcp.schema.url = SCHEMA_URL_VALUE, + flush.generation = self.inner.generation(), + flush.persisted_size = tracing::field::Empty, + ) + )] + async fn flush(&mut self) -> crate::Result { + let res = self.inner.flush().await; + if res.is_ok() { + tracing::Span::current().record("flush.persisted_size", self.inner.persisted_size()); + } + res + } + + #[tracing::instrument( + name = "finalize", + err, + skip(self), + fields( + otel.kind = "Internal", + rpc.system.name = "grpc", + gcp.client.service = INSTRUMENTATION.service_name, + gcp.client.version = INSTRUMENTATION.client_version, + gcp.client.repo = GCP_CLIENT_REPO_GOOGLEAPIS, + gcp.client.artifact = INSTRUMENTATION.client_artifact, + gcp.schema.url = SCHEMA_URL_VALUE, + finalize.generation = self.inner.generation(), + finalize.persisted_size = tracing::field::Empty, + ) + )] + async fn finalize(self) -> crate::Result { + let res = self.inner.finalize().await; + if let Ok(obj) = &res { + tracing::Span::current().record("finalize.persisted_size", obj.size); + } + res + } + + #[tracing::instrument( + name = "close", + err, + skip(self), + fields( + otel.kind = "Internal", + rpc.system.name = "grpc", + gcp.client.service = INSTRUMENTATION.service_name, + gcp.client.version = INSTRUMENTATION.client_version, + gcp.client.repo = GCP_CLIENT_REPO_GOOGLEAPIS, + gcp.client.artifact = INSTRUMENTATION.client_artifact, + gcp.schema.url = SCHEMA_URL_VALUE, + close.generation = self.inner.generation(), + close.persisted_size = tracing::field::Empty, + ) + )] + async fn close(self) -> crate::Result { + let res = self.inner.close().await; + if let Ok(size) = &res { + tracing::Span::current().record("close.persisted_size", size); + } + res + } + + fn generation(&self) -> i64 { + self.inner.generation() + } + + fn persisted_size(&self) -> i64 { + self.inner.persisted_size() + } +} + #[cfg(test)] mod tests { use super::*; @@ -139,4 +263,63 @@ mod tests { assert_eq!(descriptor.headers(), headers); let _reader = descriptor.read_range(ReadRange::segment(100, 200)).await; } + + #[cfg(google_cloud_unstable_storage_bidi)] + mockall::mock! { + #[derive(Debug)] + Writer {} + impl crate::stub::AppendableObjectWriter for Writer { + async fn append(&mut self, chunk: bytes::Bytes) -> crate::Result<()>; + async fn flush(&mut self) -> crate::Result; + async fn finalize(self) -> crate::Result; + async fn close(self) -> crate::Result; + fn generation(&self) -> i64; + fn persisted_size(&self) -> i64; + } + } + + #[cfg(google_cloud_unstable_storage_bidi)] + #[tokio::test] + async fn appendable_writer_forwards_calls() { + use crate::appendable_object_writer::AppendableObjectWriter; + use bytes::Bytes; + + let mut mock = MockWriter::new(); + mock.expect_generation().returning(|| 123); + mock.expect_persisted_size().returning(|| 456); + mock.expect_append().returning(|_| Ok(())); + mock.expect_flush().returning(|| Ok(789)); + + let writer = AppendableObjectWriter::new(mock); + let mut writer = + AppendableObjectWriter::new(TracingAppendableObjectWriter::new(writer.into_parts())); + + writer.append(Bytes::from("test")).await.unwrap(); + assert_eq!(writer.flush().await.unwrap(), 789); + assert_eq!(writer.generation(), 123); + assert_eq!(writer.persisted_size(), 456); + + let mut mock = MockWriter::new(); + mock.expect_generation().returning(|| 123); + mock.expect_persisted_size().returning(|| 456); + mock.expect_finalize().returning(|| { + Ok(crate::model::Object { + size: 999, + ..Default::default() + }) + }); + let writer = AppendableObjectWriter::new(mock); + let writer = + AppendableObjectWriter::new(TracingAppendableObjectWriter::new(writer.into_parts())); + assert_eq!(writer.finalize().await.unwrap().size, 999); + + let mut mock = MockWriter::new(); + mock.expect_generation().returning(|| 123); + mock.expect_persisted_size().returning(|| 456); + mock.expect_close().returning(|| Ok(100)); + let writer = AppendableObjectWriter::new(mock); + let writer = + AppendableObjectWriter::new(TracingAppendableObjectWriter::new(writer.into_parts())); + assert_eq!(writer.close().await.unwrap(), 100); + } } diff --git a/src/storage/src/storage/transport.rs b/src/storage/src/storage/transport.rs index 7a70735b57..c799cd0149 100644 --- a/src/storage/src/storage/transport.rs +++ b/src/storage/src/storage/transport.rs @@ -23,6 +23,14 @@ use crate::storage::perform_upload::PerformUpload; use crate::storage::read_object::Reader; use crate::storage::request_options::RequestOptions; use crate::storage::streaming_source::{Seek, StreamingSource}; +#[cfg(google_cloud_unstable_storage_bidi)] +use crate::{ + appendable_object_writer::AppendableObjectWriter, + model_ext::{OpenAppendableObjectRequest, ReopenAppendableObjectRequest}, + storage::bidi_write::{ + connector::Connector as BidiWriteConnector, transport::AppendableObjectWriterTransport, + }, +}; use crate::{ model_ext::OpenObjectRequest, object_descriptor::ObjectDescriptor, storage::bidi::connector::Connector, storage::bidi::transport::ObjectDescriptorTransport, @@ -262,6 +270,96 @@ impl Storage { .collect::>(); Ok((descriptor, readers)) } + + #[cfg(google_cloud_unstable_storage_bidi)] + async fn open_appendable_object_plain( + &self, + request: OpenAppendableObjectRequest, + options: RequestOptions, + ) -> Result { + let connector = BidiWriteConnector::new(options, self.inner.grpc.clone()); + let transport = AppendableObjectWriterTransport::new_open(connector, request).await?; + Ok(AppendableObjectWriter::new(transport)) + } + + #[cfg(google_cloud_unstable_storage_bidi)] + #[tracing::instrument(name = "open_appendable_object", level = tracing::Level::DEBUG, ret, err(Debug))] + async fn open_appendable_object_tracing( + &self, + request: OpenAppendableObjectRequest, + options: RequestOptions, + ) -> Result { + let resource_name = format!( + "//storage.googleapis.com/{}", + request + .spec + .resource + .as_ref() + .map(|r| r.bucket.as_str()) + .unwrap_or_default() + ); + let (_span, pending) = gaxi::client_request_signals!( + metric: self.metric.clone(), + info: *INSTRUMENTATION, + method: "client::Storage::open_appendable_object", + async { + if let Some(recorder) = RequestRecorder::current() { + recorder.on_client_request( + ClientRequestAttributes::default() + .set_rpc_method("google.storage.v2.Storage/BidiWriteObject") + .set_url_template("/upload/storage/v1/b/{bucket}/o") + .set_resource_name(resource_name), + ); + } + self.open_appendable_object_plain(request, options).await + } + ); + let writer = pending.await?; + Ok(AppendableObjectWriter::new( + super::tracing::TracingAppendableObjectWriter::new(writer.into_parts()), + )) + } + + #[cfg(google_cloud_unstable_storage_bidi)] + async fn reopen_appendable_object_plain( + &self, + request: ReopenAppendableObjectRequest, + options: RequestOptions, + ) -> Result { + let connector = BidiWriteConnector::new(options, self.inner.grpc.clone()); + let transport = AppendableObjectWriterTransport::new_reopen(connector, request).await?; + Ok(AppendableObjectWriter::new(transport)) + } + + #[cfg(google_cloud_unstable_storage_bidi)] + #[tracing::instrument(name = "reopen_appendable_object", level = tracing::Level::DEBUG, ret, err(Debug))] + async fn reopen_appendable_object_tracing( + &self, + request: ReopenAppendableObjectRequest, + options: RequestOptions, + ) -> Result { + let resource_name = format!("//storage.googleapis.com/{}", request.bucket); + let (_span, pending) = gaxi::client_request_signals!( + metric: self.metric.clone(), + info: *INSTRUMENTATION, + method: "client::Storage::reopen_appendable_object", + async { + if let Some(recorder) = RequestRecorder::current() { + recorder.on_client_request( + ClientRequestAttributes::default() + .set_rpc_method("google.storage.v2.Storage/BidiWriteObject") + .set_url_template("/upload/storage/v1/b/{bucket}/o") + .set_resource_name(resource_name), + ); + } + self.reopen_appendable_object_plain(request, options).await + } + ); + let writer = pending.await?; + Ok(AppendableObjectWriter::new( + super::tracing::TracingAppendableObjectWriter::new(writer.into_parts()), + )) + } } impl super::stub::Storage for Storage { @@ -325,6 +423,32 @@ impl super::stub::Storage for Storage { } self.open_object_plain(request, options).await } + + #[cfg(google_cloud_unstable_storage_bidi)] + async fn open_appendable_object( + &self, + request: OpenAppendableObjectRequest, + options: RequestOptions, + ) -> Result { + if self.tracing { + return self.open_appendable_object_tracing(request, options).await; + } + self.open_appendable_object_plain(request, options).await + } + + #[cfg(google_cloud_unstable_storage_bidi)] + async fn reopen_appendable_object( + &self, + request: ReopenAppendableObjectRequest, + options: RequestOptions, + ) -> Result { + if self.tracing { + return self + .reopen_appendable_object_tracing(request, options) + .await; + } + self.reopen_appendable_object_plain(request, options).await + } } #[cfg(test)] @@ -657,6 +781,250 @@ mod tests { Ok(()) } + #[cfg(google_cloud_unstable_storage_bidi)] + #[tokio::test] + async fn open_appendable_object() -> anyhow::Result<()> { + use gaxi::grpc::tonic::Status as TonicStatus; + use google_cloud_gax::error::rpc::Code; + use storage_grpc_mock::{MockStorage, start}; + + let guard = TestLayer::initialize(); + let mut mock = MockStorage::new(); + mock.expect_bidi_write_object() + .return_once(|_| Err(TonicStatus::not_found("not here"))); + let (endpoint, _server) = start("0.0.0.0:0", mock).await?; + + let client = crate::client::Storage::builder() + .with_credentials(Anonymous::new().build()) + .with_endpoint(endpoint.clone()) + .with_tracing() + .build() + .await?; + let response = client + .open_appendable_object("projects/_/buckets/test-bucket", "test-object") + .send() + .await; + assert!( + matches!(response, Err(ref e) if e.status().is_some_and(|s| s.code == Code::NotFound)), + "{response:?}" + ); + let captured = TestLayer::capture(&guard); + check_debug_log(&captured, "open_appendable_object"); + + client_request_span(&captured, "open_appendable_object", "NOT_FOUND", "grpc"); + + Ok(()) + } + + /// Models a complete lifecycle: `open` -> `append` -> `flush` -> `close`. + /// This tests all spans and attributes without redundant test cases for + /// identical transport logic. + #[cfg(google_cloud_unstable_storage_bidi)] + #[tokio::test] + async fn open_appendable_object_success() -> anyhow::Result<()> { + use bytes::Bytes; + use gaxi::grpc::tonic::{Response as TonicResponse, Result as TonicResult}; + use storage_grpc_mock::google::storage::v2::{ + BidiWriteObjectResponse, Object, bidi_write_object_response::WriteStatus, + }; + use storage_grpc_mock::{MockStorage, start}; + const BUCKET_NAME: &str = "projects/_/buckets/test-bucket"; + const OBJECT_NAME: &str = "test-object"; + const BIND_ADDRESS: &str = "0.0.0.0:0"; + + let guard = TestLayer::initialize(); + + let (tx, rx) = tokio::sync::mpsc::channel::>(10); + + let initial_response = BidiWriteObjectResponse { + write_status: Some(WriteStatus::Resource(Object { + bucket: BUCKET_NAME.to_string(), + name: OBJECT_NAME.to_string(), + generation: 98765, + ..Default::default() + })), + ..BidiWriteObjectResponse::default() + }; + + let flush_response = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(5)), + ..BidiWriteObjectResponse::default() + }; + + // The first response is the initial handshake/metadata response expected by + // `connector.rs` immediately upon opening the stream, before any data is sent. + tx.send(Ok(initial_response)).await?; + + let mut mock = MockStorage::new(); + mock.expect_bidi_write_object().return_once(move |req| { + let mut stream = req.into_inner(); + tokio::spawn(async move { + while let Some(Ok(msg)) = stream.recv().await { + if msg.flush { + // The second response is sent ONLY when the client explicitly requests a flush. + // `writer.append()` does not wait for a response, but `writer.flush()` does. + let _ = tx.send(Ok(flush_response.clone())).await; + } + } + }); + Ok(TonicResponse::new(rx)) + }); + let (endpoint, _server) = start(BIND_ADDRESS, mock).await?; + + let client = crate::client::Storage::builder() + .with_credentials(Anonymous::new().build()) + .with_endpoint(endpoint.clone()) + .with_tracing() + .build() + .await?; + let mut writer = client + .open_appendable_object(BUCKET_NAME, OBJECT_NAME) + .send() + .await?; + + writer + .append(Bytes::from_static(b"hello")) + .await + .expect("append failed"); + writer.flush().await.expect("flush failed"); + writer.close().await.expect("close failed"); + + let captured = TestLayer::capture(&guard); + let _span = captured + .iter() + .find(|s| s.name == "client_request") + .unwrap_or_else(|| panic!("missing `client_request` span in capture: {captured:#?}")); + + check_bidi_write_span_attributes(&captured, "append", 98765, Some(5)); + check_bidi_write_span_attributes(&captured, "flush", 98765, Some(5)); + check_bidi_write_span_attributes(&captured, "close", 98765, Some(5)); + + Ok(()) + } + + #[cfg(google_cloud_unstable_storage_bidi)] + #[tokio::test] + async fn reopen_appendable_object() -> anyhow::Result<()> { + use gaxi::grpc::tonic::Status as TonicStatus; + use google_cloud_gax::error::rpc::Code; + use storage_grpc_mock::{MockStorage, start}; + + let guard = TestLayer::initialize(); + let mut mock = MockStorage::new(); + mock.expect_bidi_write_object() + .return_once(|_| Err(TonicStatus::not_found("not here"))); + let (endpoint, _server) = start("0.0.0.0:0", mock).await?; + + let client = crate::client::Storage::builder() + .with_credentials(Anonymous::new().build()) + .with_endpoint(endpoint.clone()) + .with_tracing() + .build() + .await?; + let response = client + .reopen_appendable_object("projects/_/buckets/test-bucket", "test-object", 12345) + .send() + .await; + assert!( + matches!(response, Err(ref e) if e.status().is_some_and(|s| s.code == Code::NotFound)), + "{response:?}" + ); + let captured = TestLayer::capture(&guard); + check_debug_log(&captured, "reopen_appendable_object"); + + client_request_span(&captured, "reopen_appendable_object", "NOT_FOUND", "grpc"); + + Ok(()) + } + + /// Models a complete lifecycle: `reopen` -> `append` -> `flush` -> `finalize`. + /// This complements the `open_appendable_object_success` test case to to test + /// the remaining terminal operation (`finalize`) along with its spans, + /// avoiding redundant test cases for identical transport logic. + #[cfg(google_cloud_unstable_storage_bidi)] + #[tokio::test] + async fn reopen_appendable_object_success() -> anyhow::Result<()> { + use bytes::Bytes; + use gaxi::grpc::tonic::{Response as TonicResponse, Result as TonicResult}; + use storage_grpc_mock::google::storage::v2::{ + BidiWriteObjectResponse, bidi_write_object_response::WriteStatus, + }; + use storage_grpc_mock::{MockStorage, start}; + const BUCKET_NAME: &str = "projects/_/buckets/test-bucket"; + const OBJECT_NAME: &str = "test-object"; + const BIND_ADDRESS: &str = "0.0.0.0:0"; + + let guard = TestLayer::initialize(); + + let (tx, rx) = tokio::sync::mpsc::channel::>(10); + let response = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(100)), + ..BidiWriteObjectResponse::default() + }; + + // The first response is the initial handshake/metadata response expected by + // `connector.rs` immediately upon opening the stream, before any data is sent. + tx.send(Ok(response.clone())).await?; + + let finalize_response = BidiWriteObjectResponse { + write_status: Some(WriteStatus::Resource( + storage_grpc_mock::google::storage::v2::Object { + size: 105, + ..Default::default() + }, + )), + ..BidiWriteObjectResponse::default() + }; + + let mut mock = MockStorage::new(); + mock.expect_bidi_write_object().return_once(move |req| { + let mut stream = req.into_inner(); + tokio::spawn(async move { + while let Some(Ok(msg)) = stream.recv().await { + if msg.finish_write { + let _ = tx.send(Ok(finalize_response.clone())).await; + } else if msg.flush { + // The second response is sent ONLY when the client explicitly requests a flush. + // `writer.append()` does not wait for a response, but `writer.flush()` does. + let _ = tx.send(Ok(response.clone())).await; + } + } + }); + Ok(TonicResponse::from(rx)) + }); + let (endpoint, _server) = start(BIND_ADDRESS, mock).await?; + + let client = crate::client::Storage::builder() + .with_credentials(Anonymous::new().build()) + .with_endpoint(endpoint.clone()) + .with_tracing() + .build() + .await?; + let mut writer = client + .reopen_appendable_object(BUCKET_NAME, OBJECT_NAME, 12345) + .send() + .await?; + + writer + .append(Bytes::from_static(b"hello")) + .await + .expect("append failed"); + writer.flush().await.expect("flush failed"); + writer.finalize().await.expect("finalize failed"); + + let captured = TestLayer::capture(&guard); + let _span = captured + .iter() + .find(|s| s.name == "client_request") + .unwrap_or_else(|| panic!("missing `client_request` span in capture: {captured:#?}")); + + check_bidi_write_span_attributes(&captured, "append", 12345, Some(5)); + check_bidi_write_span_attributes(&captured, "flush", 12345, Some(100)); + check_bidi_write_span_attributes(&captured, "finalize", 12345, Some(105)); + + Ok(()) + } + #[track_caller] fn check_debug_log(captured: &Vec, method: &'static str) { let span = captured @@ -723,4 +1091,62 @@ mod tests { "mismatch = {mismatch:?}\ngot = {got:?}\nwant = {want:?}" ); } + + #[allow(dead_code)] + #[track_caller] + fn check_bidi_write_span_attributes( + captured: &[CapturedSpan], + method: &str, + expected_generation: i64, + expected_size: Option, + ) { + let span = captured + .iter() + .find(|s| { + s.name == method && s.attributes.contains_key(&format!("{method}.generation")) + }) + .unwrap_or_else(|| panic!("missing `{method}` span in capture: {captured:#?}")); + + let expected_strs = [ + ("otel.kind", "Internal"), + ("rpc.system.name", "grpc"), + ("gcp.client.service", "storage"), + ("gcp.client.version", env!("CARGO_PKG_VERSION")), + ("gcp.client.repo", "googleapis/google-cloud-rust"), + ("gcp.client.artifact", "google-cloud-storage"), + ("gcp.schema.url", "https://opentelemetry.io/schemas/1.39.0"), + ]; + + for (k, v) in expected_strs { + assert_eq!( + span.attributes.get(k), + Some(&AttributeValue::String(v.to_string().into())), + "mismatched attribute `{k}` in `{method}` span" + ); + } + + let gen_k = format!("{method}.generation"); + assert_eq!( + span.attributes.get(gen_k.as_str()), + Some(&AttributeValue::Int64(expected_generation)), + "mismatched `{gen_k}`" + ); + + if let Some(s) = expected_size { + if method == "append" { + assert_eq!( + span.attributes.get("append.chunk_size"), + Some(&AttributeValue::UInt64(s as u64)), + "mismatched `append.chunk_size`" + ); + } else { + let size_k = format!("{method}.persisted_size"); + assert_eq!( + span.attributes.get(size_k.as_str()), + Some(&AttributeValue::Int64(s)), + "mismatched `{size_k}`" + ); + } + } + } } diff --git a/src/storage/tests/mocking.rs b/src/storage/tests/mocking.rs index 529a83917f..9af3401db2 100644 --- a/src/storage/tests/mocking.rs +++ b/src/storage/tests/mocking.rs @@ -54,6 +54,18 @@ mod tests { _request: OpenObjectRequest, _options: RequestOptions, ) -> Result<(ObjectDescriptor, Vec)>; + #[cfg(google_cloud_unstable_storage_bidi)] + async fn open_appendable_object( + &self, + _request: gcs::model_ext::OpenAppendableObjectRequest, + _options: RequestOptions, + ) -> Result; + #[cfg(google_cloud_unstable_storage_bidi)] + async fn reopen_appendable_object( + &self, + _request: gcs::model_ext::ReopenAppendableObjectRequest, + _options: RequestOptions, + ) -> Result; } } @@ -143,6 +155,40 @@ mod tests { Ok(()) } + #[tokio::test] + #[cfg(google_cloud_unstable_storage_bidi)] + async fn mock_open_appendable_object_fail() { + let status = Status::default().set_code(Code::Aborted); + let want = status.clone(); + let mut mock = MockStorage::new(); + mock.expect_open_appendable_object() + .return_once(move |_, _| Err(Error::service(status))); + let client = gcs::client::Storage::from_stub(mock); + let err = client + .open_appendable_object("projects/_/buckets/my-bucket", "my-object") + .send() + .await + .unwrap_err(); + assert_eq!(err.status(), Some(&want), "{err:?}"); + } + + #[tokio::test] + #[cfg(google_cloud_unstable_storage_bidi)] + async fn mock_reopen_appendable_object_fail() { + let status = Status::default().set_code(Code::Aborted); + let want = status.clone(); + let mut mock = MockStorage::new(); + mock.expect_reopen_appendable_object() + .return_once(move |_, _| Err(Error::service(status))); + let client = gcs::client::Storage::from_stub(mock); + let err = client + .reopen_appendable_object("projects/_/buckets/my-bucket", "my-object", 1234) + .send() + .await + .unwrap_err(); + assert_eq!(err.status(), Some(&want), "{err:?}"); + } + #[tokio::test] async fn mock_write_object_buffered() { let mut mock = MockStorage::new(); diff --git a/tests/storage/src/bidi_write.rs b/tests/storage/src/bidi_write.rs new file mode 100644 index 0000000000..ab41dd3f06 --- /dev/null +++ b/tests/storage/src/bidi_write.rs @@ -0,0 +1,180 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use bytes::Bytes; +use google_cloud_storage::client::Storage; +use std::time::SystemTime; + +pub async fn run(bucket_name: &str) -> anyhow::Result<()> { + let client = Storage::builder().build().await?; + test_bidi_write_single_block(&client, bucket_name).await?; + test_bidi_write_chunked_appends(&client, bucket_name).await?; + test_bidi_write_resume_append(&client, bucket_name).await?; + test_bidi_write_drop_stream(&client, bucket_name).await?; + Ok(()) +} + +async fn check_object_contents( + client: &Storage, + bucket: &str, + object: &str, + expected: &[u8], +) -> anyhow::Result<()> { + let mut resp = client.read_object(bucket, object).send().await?; + let mut buf = Vec::new(); + while let Some(chunk) = resp.next().await { + buf.extend_from_slice(&chunk?); + } + assert_eq!(buf, expected); + Ok(()) +} + +async fn test_bidi_write_single_block(client: &Storage, bucket: &str) -> anyhow::Result<()> { + let now = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let object_name = format!("test_bidi_write_single_block_{now}"); + + let mut writer = client + .open_appendable_object(bucket, &object_name) + .send() + .await?; + + writer.append(Bytes::from("hello world")).await?; + let object_metadata = writer.finalize().await?; + + assert_eq!(object_metadata.size, 11); + assert_eq!(object_metadata.name, object_name); + check_object_contents(client, bucket, &object_name, b"hello world").await?; + + Ok(()) +} + +async fn test_bidi_write_chunked_appends(client: &Storage, bucket: &str) -> anyhow::Result<()> { + let now = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let object_name = format!("test_bidi_write_chunked_appends_{now}"); + + let mut writer = client + .open_appendable_object(bucket, &object_name) + .send() + .await?; + + writer.append(Bytes::from("chunk1")).await?; + writer.append(Bytes::from("chunk2")).await?; + writer.flush().await?; + writer.append(Bytes::from("chunk3")).await?; + let object_metadata = writer.finalize().await?; + + assert_eq!(object_metadata.size, 18); + assert_eq!(object_metadata.name, object_name); + check_object_contents(client, bucket, &object_name, b"chunk1chunk2chunk3").await?; + + Ok(()) +} + +async fn test_bidi_write_resume_append(client: &Storage, bucket: &str) -> anyhow::Result<()> { + let now = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let object_name = format!("test_bidi_write_resume_append_{now}"); + + let mut writer = client + .open_appendable_object(bucket, &object_name) + .send() + .await?; + + writer.append(Bytes::from("hello ")).await?; + let generation = writer.generation(); + let persisted_size = writer.close().await?; + assert_eq!(persisted_size, 6); + + // Reopen. + let mut writer2 = client + .reopen_appendable_object(bucket, &object_name, generation) + .send() + .await?; + + assert_eq!(writer2.persisted_size(), 6); + writer2.append(Bytes::from("world")).await?; + let object_metadata = writer2.finalize().await?; + + assert_eq!(object_metadata.size, 11); + assert_eq!(object_metadata.name, object_name); + check_object_contents(client, bucket, &object_name, b"hello world").await?; + + Ok(()) +} + +async fn test_bidi_write_drop_stream(client: &Storage, bucket: &str) -> anyhow::Result<()> { + // Test reopen and finalize after a dropped connection. We'll append some + // data and drop the connection without flushing or finalizing. Since we + // didn't flush the data, it may or may not be persisted. + let now = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let object_name = format!("test_bidi_write_drop_stream_{now}"); + + let mut writer = client + .open_appendable_object(bucket, &object_name) + .send() + .await?; + + writer.append(Bytes::from("hello ")).await?; + writer.flush().await?; + let generation = writer.generation(); + writer + .append(Bytes::from("data that won't be flushed")) + .await?; + + // Explicitly drop the writer without flush/finalize/close. + drop(writer); + + // Give the server a moment to recognize the dropped connection. + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + + // Reopen using the generation. + let mut writer2 = client + .reopen_appendable_object(bucket, &object_name, generation) + .send() + .await?; + + // We only flushed "hello " (6 bytes). The data may or may not have + // been persisted, so the persisted size must be at least 6. + assert!(writer2.persisted_size() >= 6); + + writer2.append(Bytes::from("world")).await?; + let object_metadata = writer2.finalize().await?; + + assert!(object_metadata.size >= 11); + assert_eq!(object_metadata.name, object_name); + + // We can't strictly check exact contents because the data may or may not + // have been persisted, but we know it starts with "hello " and ends with + // "world". + let mut resp = client.read_object(bucket, &object_name).send().await?; + let mut buf = Vec::new(); + while let Some(chunk) = resp.next().await { + buf.extend_from_slice(&chunk?); + } + assert!(buf.starts_with(b"hello ")); + assert!(buf.ends_with(b"world")); + + Ok(()) +} diff --git a/tests/storage/src/lib.rs b/tests/storage/src/lib.rs index 03e17ccea0..cb9396f9e1 100644 --- a/tests/storage/src/lib.rs +++ b/tests/storage/src/lib.rs @@ -13,6 +13,8 @@ // limitations under the License. pub mod bidi_read; +#[cfg(google_cloud_unstable_storage_bidi)] +pub mod bidi_write; pub mod read_object; pub mod write_object; @@ -36,7 +38,9 @@ use google_cloud_test_utils::resource_names::random_bucket_id; use google_cloud_test_utils::runtime_config::{project_id, test_service_account}; use google_cloud_wkt::FieldMask; use std::time::Duration; -pub use storage_samples::{cleanup_stale_buckets, create_test_bucket, create_test_hns_bucket}; +pub use storage_samples::{ + cleanup_stale_buckets, create_test_bucket, create_test_hns_bucket, create_test_rapid_bucket, +}; pub async fn objects(builder: StorageBuilder, bucket_name: &str, prefix: &str) -> Result<()> { let client = builder.build().await?; diff --git a/tests/storage/tests/driver.rs b/tests/storage/tests/driver.rs index 67514f5027..4450356693 100644 --- a/tests/storage/tests/driver.rs +++ b/tests/storage/tests/driver.rs @@ -115,7 +115,7 @@ mod storage { #[tokio::test(flavor = "multi_thread")] async fn run_storage_bidi() -> anyhow::Result<()> { let _guard = enable_tracing(); - let (control, bucket) = integration_tests_storage::create_test_hns_bucket() + let (control, bucket) = integration_tests_storage::create_test_bucket() .await .inspect_err(anydump)?; let result = integration_tests_storage::bidi_read::run(&bucket.name) @@ -128,4 +128,22 @@ mod storage { .inspect_err(anydump); result } + + #[tokio::test(flavor = "multi_thread")] + #[cfg(google_cloud_unstable_storage_bidi)] + async fn run_storage_bidi_write() -> anyhow::Result<()> { + let _guard = enable_tracing(); + let (control, bucket) = integration_tests_storage::create_test_rapid_bucket() + .await + .inspect_err(anydump)?; + let result = integration_tests_storage::bidi_write::run(&bucket.name) + .await + .inspect_err(anydump); + let _ = + storage_samples::cleanup_bucket(control, bucket.name.clone(), bucket.project.clone()) + .await + .inspect_err(|e| tracing::error!("error cleaning up bucket {}: {e:?}", bucket.name)) + .inspect_err(anydump); + result + } }