From 07b7af794b75bd84e2e0da63f6320ae412e7c1ff Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Thu, 25 Jun 2026 07:19:32 +0000 Subject: [PATCH 01/13] feat(storage): implement transport, connector, worker --- src/storage/src/storage/bidi_write.rs | 6 + .../src/storage/bidi_write/connector.rs | 717 ++++++++++++++++++ src/storage/src/storage/bidi_write/mocks.rs | 83 ++ .../src/storage/bidi_write/redirect.rs | 76 +- .../src/storage/bidi_write/retry_redirect.rs | 2 - .../src/storage/bidi_write/transport.rs | 372 +++++++++ src/storage/src/storage/bidi_write/worker.rs | 240 ++++++ 7 files changed, 1475 insertions(+), 21 deletions(-) create mode 100644 src/storage/src/storage/bidi_write/connector.rs create mode 100644 src/storage/src/storage/bidi_write/mocks.rs create mode 100644 src/storage/src/storage/bidi_write/transport.rs create mode 100644 src/storage/src/storage/bidi_write/worker.rs diff --git a/src/storage/src/storage/bidi_write.rs b/src/storage/src/storage/bidi_write.rs index b6576d33be..59d9682fe3 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -14,9 +14,12 @@ //! Internal traits and types for Appendable Object Write (Bidi Write). +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; @@ -82,6 +85,9 @@ impl Client for gaxi::grpc::Client { } } +#[cfg(test)] +mod mocks; + #[cfg(test)] pub(crate) mod tests { use crate::Error; 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..0c500bd022 --- /dev/null +++ b/src/storage/src/storage/bidi_write/connector.rs @@ -0,0 +1,717 @@ +// Copyright 2025 Google LLC +#![allow(dead_code)] +// +// 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, 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 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), + 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 { + pub(crate) spec: Arc>, + options: RequestOptions, + client: T, +} + +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()))), + options, + client, + } + } + + 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 = crate::google::storage::v2::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.spec.lock().expect("never poisoned") = AppendObjectSpecState::Write(Box::new(spec)); + 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, + ..Default::default() + }; + *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 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); + 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, + ) -> Result<(BidiWriteObjectResponse, Connection)> { + let first_message = match &*spec.lock().expect("never poisoned") { + AppendObjectSpecState::Write(spec) => FirstMessage::WriteObjectSpec(*spec.clone()), + AppendObjectSpecState::Append(spec) => FirstMessage::AppendObjectSpec(spec.clone()), + }; + + let request = BidiWriteObjectRequest { + first_message: Some(first_message), + ..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()) + { + use google_cloud_gax::error::binding::*; + 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 x_goog_request_params = request + .first_message + .as_ref() + .and_then(|m| match m { + FirstMessage::WriteObjectSpec(_) => None, // WriteObjectSpec doesn't have routing_token + FirstMessage::AppendObjectSpec(s) => s.routing_token.as_deref(), + _ => None, + }) + .into_iter() + .fold(format!("bucket={bucket_name}"), |s, token| { + s + &format!("&routing_token={token}") + }); + + let (tx, rx) = tokio::sync::mpsc::channel::(100); + tx.send(request.clone()).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 = 0; + if let Some(WriteStatus::Resource(resource)) = &m.write_status { + new_generation = 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, + write_handle: m.write_handle.clone(), + ..Default::default() + }) + } + AppendObjectSpecState::Append(s) => { + AppendObjectSpecState::Append(AppendObjectSpec { + generation: new_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: None, + }; + + 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); + 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..c86bbb11bf --- /dev/null +++ b/src/storage/src/storage/bidi_write/mocks.rs @@ -0,0 +1,83 @@ +// 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; +use tokio::sync::mpsc::Sender; + +// 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 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..b7b8cc8e73 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,42 @@ 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 (bucket, object) = match &*guard { + AppendObjectSpecState::Write(spec) => { + 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(); + (bucket, object) + } + AppendObjectSpecState::Append(spec) => (spec.bucket.clone(), spec.object.clone()), + }; + + let mut new_spec = AppendObjectSpec { + bucket, + object, + generation: redirect.generation.unwrap_or(0), + routing_token: redirect.routing_token, + write_handle: redirect.write_handle, + ..Default::default() + }; + + if let AppendObjectSpecState::Append(old_spec) = &*guard + && redirect.generation.is_none() + { + new_spec.generation = old_spec.generation; } + + *guard = AppendObjectSpecState::Append(new_spec); + break; } to_gax_error(status) } @@ -108,14 +138,18 @@ mod tests { }), ..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(), 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(), routing); + assert_eq!(spec.write_handle, write_handle); + assert_eq!(spec.generation, 42); + } else { + panic!("Expected AppendObjectSpecState::Append"); + } } #[test] @@ -138,14 +172,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..92c6034d52 --- /dev/null +++ b/src/storage/src/storage/bidi_write/transport.rs @@ -0,0 +1,372 @@ +// Copyright 2025 Google LLC +#![allow(dead_code)] +// +// 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::{ + BidiWriteObjectRequest, ChecksummedData, bidi_write_object_request::Data, + bidi_write_object_response::WriteStatus, +}; +use crate::stub::AppendableObjectWriter; +use bytes::Bytes; +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, + crc32c_persisted: 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) + } + + 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 (initial, connection) = connector.connect_reopen(req).await?; + Self::start_worker(connector, initial, connection) + } + + fn start_worker( + connector: super::connector::Connector, + initial: crate::google::storage::v2::BidiWriteObjectResponse, + connection: super::connector::Connection<::Stream>, + ) -> crate::Result + where + T: super::Client + Clone + Sync + Send + 'static, + ::Stream: super::TonicStreaming + Send + Sync, + { + let mut generation = 0; + if let Some(WriteStatus::Resource(r)) = initial.write_status.as_ref() { + generation = r.generation; + } + + 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: 0, + write_offset: 0, + crc32c_persisted: 0, + }) + } +} + +impl AppendableObjectWriter for AppendableObjectWriterTransport { + async fn append(&mut self, chunk: Bytes) -> crate::Result<()> { + let length = chunk.len() as i64; + let crc32c = crc32c::crc32c(&chunk); + + self.crc32c_persisted = crc32c::crc32c_append(self.crc32c_persisted, &chunk); + + let request = BidiWriteObjectRequest { + write_offset: self.write_offset, + data: Some(Data::ChecksummedData(ChecksummedData { + content: chunk, + crc32c: Some(crc32c), + })), + ..BidiWriteObjectRequest::default() + }; + + self.write_offset += length; + self.tx + .send(UploadIntent::Append(request)) + .await + .map_err(|e| crate::Error::io(e.to_string()))?; + + Ok(()) + } + + async fn flush(&mut self) -> crate::Result { + let (sender, receiver) = oneshot::channel(); + let request = BidiWriteObjectRequest { + flush: 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 => 0, + }; + self.persisted_size = size; + Ok(size) + } + + async fn finalize(self) -> crate::Result { + use crate::google::storage::v2::ObjectChecksums; + use gaxi::prost::FromProto; + + 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.crc32c_persisted), + 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, + crc32c_persisted: 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_eq!(req.finish_write, false); + } else { + panic!("expected Append"); + } + + // Verify flush + let intent = rx.recv().await.unwrap(); + if let UploadIntent::Flush(req, sender) = intent { + assert_eq!(req.flush, true); + 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_eq!(req.finish_write, true); + 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, + crc32c_persisted: 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:?}"); + + Ok(()) + } + + #[tokio::test] + async fn finalize_error() -> anyhow::Result<()> { + let (tx, mut rx) = mpsc::channel(1); + let transport = AppendableObjectWriterTransport { + tx, + write_offset: 0, + crc32c_persisted: 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(()) + } +} 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..31d144e0af --- /dev/null +++ b/src/storage/src/storage/bidi_write/worker.rs @@ -0,0 +1,240 @@ +// Copyright 2025 Google LLC +#![allow(dead_code)] +// +// 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) => 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)); + } + } +} + +#[cfg(test)] +mod tests { + use super::super::mocks::{MockTestClient, mock_connector, mock_stream}; + 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) = mock_stream(); + 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) = mock_stream(); + 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_eq!(stream_req.flush, true); + + // 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) = mock_stream(); + 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(()) + } +} From 005537a50aaf70a6788a05243af0c27adb7e859f Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Fri, 26 Jun 2026 09:43:34 +0000 Subject: [PATCH 02/13] Incorporate feedback --- .../src/storage/bidi_write/connector.rs | 32 +++++--- .../src/storage/bidi_write/transport.rs | 80 +++++++++++++++---- 2 files changed, 88 insertions(+), 24 deletions(-) diff --git a/src/storage/src/storage/bidi_write/connector.rs b/src/storage/src/storage/bidi_write/connector.rs index 0c500bd022..79bfb2d8c9 100644 --- a/src/storage/src/storage/bidi_write/connector.rs +++ b/src/storage/src/storage/bidi_write/connector.rs @@ -19,14 +19,19 @@ use super::redirect::handle_redirect; use super::retry_redirect::RetryRedirect; use super::{Client, TonicStreaming}; use crate::google::storage::v2::{ - AppendObjectSpec, BidiWriteObjectRequest, BidiWriteObjectResponse, WriteObjectSpec, - bidi_write_object_request::FirstMessage, bidi_write_object_response::WriteStatus, + 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; @@ -62,9 +67,10 @@ pub enum AppendObjectSpecState { /// - `T`: a type implementing the [Client] trait, this is used in tests. #[derive(Clone, Debug)] pub struct Connector { - pub(crate) spec: Arc>, + spec: Arc>, options: RequestOptions, client: T, + params: Option, } impl Connector @@ -77,6 +83,7 @@ where spec: Arc::new(Mutex::new(AppendObjectSpecState::Write(Box::default()))), options, client, + params: None, } } @@ -92,7 +99,8 @@ where } None => None, }; - let spec = crate::google::storage::v2::WriteObjectSpec { + self.params = req.params.map(|p| p.to_proto().expect("valid params")); + let spec = WriteObjectSpec { resource, predefined_acl: req.spec.predefined_acl, if_generation_match: req.spec.if_generation_match, @@ -119,6 +127,7 @@ where if_metageneration_not_match: req.if_metageneration_not_match, ..Default::default() }; + self.params = req.params.map(|p| p.to_proto().expect("valid params")); *self.spec.lock().expect("never poisoned") = AppendObjectSpecState::Append(spec); self.connect_attempt_loop().await } @@ -132,12 +141,14 @@ where 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); + 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)), @@ -153,6 +164,7 @@ where client: T, spec: Arc>, options: &RequestOptions, + params: Option, ) -> Result<(BidiWriteObjectResponse, Connection)> { let first_message = match &*spec.lock().expect("never poisoned") { AppendObjectSpecState::Write(spec) => FirstMessage::WriteObjectSpec(*spec.clone()), @@ -161,6 +173,7 @@ where let request = BidiWriteObjectRequest { first_message: Some(first_message), + common_object_request_params: params, ..BidiWriteObjectRequest::default() }; @@ -178,7 +191,6 @@ where .strip_prefix("projects/_/buckets/") .is_none_or(|x| x.is_empty()) { - use google_cloud_gax::error::binding::*; let problem = SubstitutionFail::MismatchExpecting( bucket_name.to_string(), "projects/_/buckets/*", @@ -244,9 +256,9 @@ where match stream.next_message().await { Ok(Some(m)) => { let mut guard = spec.lock().expect("never poisoned"); - let mut new_generation = 0; + let mut new_generation = None; if let Some(WriteStatus::Resource(resource)) = &m.write_status { - new_generation = resource.generation; + new_generation = Some(resource.generation); } let new_state = match &*guard { @@ -262,14 +274,14 @@ where .as_ref() .map(|r| r.name.clone()) .unwrap_or_default(), - generation: new_generation, + generation: new_generation.unwrap_or(0), write_handle: m.write_handle.clone(), ..Default::default() }) } AppendObjectSpecState::Append(s) => { AppendObjectSpecState::Append(AppendObjectSpec { - generation: new_generation, + generation: new_generation.unwrap_or(s.generation), write_handle: m.write_handle.clone().or_else(|| s.write_handle.clone()), ..s.clone() }) diff --git a/src/storage/src/storage/bidi_write/transport.rs b/src/storage/src/storage/bidi_write/transport.rs index 92c6034d52..8d9dc38902 100644 --- a/src/storage/src/storage/bidi_write/transport.rs +++ b/src/storage/src/storage/bidi_write/transport.rs @@ -16,12 +16,14 @@ // 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; @@ -46,7 +48,7 @@ impl AppendableObjectWriterTransport { ::Stream: super::TonicStreaming + Send + Sync, { let (initial, connection) = connector.connect_open(req).await?; - Self::start_worker(connector, initial, connection) + Self::start_worker(connector, initial, connection, 0) } pub async fn new_reopen( @@ -57,22 +59,27 @@ impl AppendableObjectWriterTransport { 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) + 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 generation = 0; + 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 (tx, rx) = tokio::sync::mpsc::channel(100); @@ -82,8 +89,8 @@ impl AppendableObjectWriterTransport { Ok(Self { tx, generation, - persisted_size: 0, - write_offset: 0, + persisted_size, + write_offset: persisted_size, crc32c_persisted: 0, }) } @@ -94,7 +101,7 @@ impl AppendableObjectWriter for AppendableObjectWriterTransport { let length = chunk.len() as i64; let crc32c = crc32c::crc32c(&chunk); - self.crc32c_persisted = crc32c::crc32c_append(self.crc32c_persisted, &chunk); + let new_crc32c = crc32c::crc32c_append(self.crc32c_persisted, &chunk); let request = BidiWriteObjectRequest { write_offset: self.write_offset, @@ -105,12 +112,14 @@ impl AppendableObjectWriter for AppendableObjectWriterTransport { ..BidiWriteObjectRequest::default() }; - self.write_offset += length; self.tx .send(UploadIntent::Append(request)) .await .map_err(|e| crate::Error::io(e.to_string()))?; + self.write_offset += length; + self.crc32c_persisted = new_crc32c; + Ok(()) } @@ -133,16 +142,13 @@ impl AppendableObjectWriter for AppendableObjectWriterTransport { let size = match response.write_status { Some(WriteStatus::PersistedSize(s)) => s, Some(WriteStatus::Resource(r)) => r.size, - None => 0, + None => return Err(crate::Error::io("flush response missing write_status")), }; self.persisted_size = size; Ok(size) } async fn finalize(self) -> crate::Result { - use crate::google::storage::v2::ObjectChecksums; - use gaxi::prost::FromProto; - let (sender, receiver) = oneshot::channel(); let request = BidiWriteObjectRequest { finish_write: true, @@ -274,9 +280,11 @@ mod tests { spec: Default::default(), params: None, }; - req.spec.resource = Some(crate::model::Object::default() - .set_bucket("projects/_/buckets/test-bucket") - .set_name("test-object")); + 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 @@ -332,6 +340,10 @@ mod tests { .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.crc32c_persisted, 0); + Ok(()) } @@ -369,4 +381,44 @@ mod tests { 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)), + ..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); + Ok(()) + } } From a4ef212a59e1087bcff12818367fe17993caf8b2 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Sat, 27 Jun 2026 08:25:15 +0000 Subject: [PATCH 03/13] Incorporate feedback --- .../src/storage/bidi_write/connector.rs | 10 +++++++-- .../src/storage/bidi_write/transport.rs | 22 +++++++++++++++---- src/storage/src/storage/bidi_write/worker.rs | 8 ++++++- 3 files changed, 33 insertions(+), 7 deletions(-) diff --git a/src/storage/src/storage/bidi_write/connector.rs b/src/storage/src/storage/bidi_write/connector.rs index 79bfb2d8c9..5339a26b66 100644 --- a/src/storage/src/storage/bidi_write/connector.rs +++ b/src/storage/src/storage/bidi_write/connector.rs @@ -125,7 +125,9 @@ where routing_token: req.routing_token, if_metageneration_match: req.if_metageneration_match, if_metageneration_not_match: req.if_metageneration_not_match, - ..Default::default() + write_handle: req + .write_handle + .map(|handle| crate::google::storage::v2::BidiWriteHandle { handle }), }; self.params = req.params.map(|p| p.to_proto().expect("valid params")); *self.spec.lock().expect("never poisoned") = AppendObjectSpecState::Append(spec); @@ -522,7 +524,7 @@ mod tests { if_metageneration_match: None, if_metageneration_not_match: None, params: None, - write_handle: None, + write_handle: Some(bytes::Bytes::from_static(b"test-write-handle")), }; let mut connector = Connector::new(test_options(), client); @@ -543,6 +545,10 @@ mod tests { 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(()) } diff --git a/src/storage/src/storage/bidi_write/transport.rs b/src/storage/src/storage/bidi_write/transport.rs index 8d9dc38902..423179aeb9 100644 --- a/src/storage/src/storage/bidi_write/transport.rs +++ b/src/storage/src/storage/bidi_write/transport.rs @@ -82,6 +82,15 @@ impl AppendableObjectWriterTransport { persisted_size = *s; } + let mut crc32c_persisted = 0; + if let Some(crc) = initial + .persisted_data_checksums + .as_ref() + .and_then(|c| c.crc32c) + { + crc32c_persisted = crc; + } + let (tx, rx) = tokio::sync::mpsc::channel(100); let worker = super::worker::Worker::new(connector); tokio::spawn(worker.run(connection, rx)); @@ -91,7 +100,7 @@ impl AppendableObjectWriterTransport { generation, persisted_size, write_offset: persisted_size, - crc32c_persisted: 0, + crc32c_persisted, }) } } @@ -224,7 +233,7 @@ mod tests { let intent = rx.recv().await.unwrap(); if let UploadIntent::Append(req) = intent { assert_eq!(req.write_offset, 0); - assert_eq!(req.finish_write, false); + assert!(!req.finish_write); } else { panic!("expected Append"); } @@ -232,7 +241,7 @@ mod tests { // Verify flush let intent = rx.recv().await.unwrap(); if let UploadIntent::Flush(req, sender) = intent { - assert_eq!(req.flush, true); + assert!(req.flush); let resp = BidiWriteObjectResponse { write_status: Some(WriteStatus::PersistedSize(5)), ..Default::default() @@ -245,7 +254,7 @@ mod tests { // Verify finalize let intent = rx.recv().await.unwrap(); if let UploadIntent::Finalize(req, sender) = intent { - assert_eq!(req.finish_write, true); + assert!(req.finish_write); let object = Object { bucket: "projects/_/buckets/test-bucket".into(), name: "test-object".into(), @@ -410,6 +419,10 @@ mod tests { 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?; @@ -419,6 +432,7 @@ mod tests { assert_eq!(transport.generation(), 123456); assert_eq!(transport.persisted_size(), 1024); assert_eq!(transport.write_offset, 1024); + assert_eq!(transport.crc32c_persisted, 9999); Ok(()) } } diff --git a/src/storage/src/storage/bidi_write/worker.rs b/src/storage/src/storage/bidi_write/worker.rs index 31d144e0af..bd19f386c2 100644 --- a/src/storage/src/storage/bidi_write/worker.rs +++ b/src/storage/src/storage/bidi_write/worker.rs @@ -144,6 +144,12 @@ where 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 + ); } } } @@ -202,7 +208,7 @@ mod tests { // The worker should send the request to the stream. let stream_req = request_rx.recv().await.unwrap(); - assert_eq!(stream_req.flush, true); + assert!(stream_req.flush); // The server responds let server_resp = BidiWriteObjectResponse { From b529ee411f45546de90d508f9ec43c56549a63de Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Sun, 28 Jun 2026 09:24:14 +0000 Subject: [PATCH 04/13] Incorporate feedback --- .../src/storage/bidi_write/redirect.rs | 87 ++++++++++++++----- 1 file changed, 66 insertions(+), 21 deletions(-) diff --git a/src/storage/src/storage/bidi_write/redirect.rs b/src/storage/src/storage/bidi_write/redirect.rs index b7b8cc8e73..8e5812a892 100644 --- a/src/storage/src/storage/bidi_write/redirect.rs +++ b/src/storage/src/storage/bidi_write/redirect.rs @@ -36,7 +36,7 @@ pub fn handle_redirect(state: Arc>, status: Status) { let mut guard = state.lock().expect("never poisoned"); - let (bucket, object) = match &*guard { + let new_state = match &*guard { AppendObjectSpecState::Write(spec) => { let bucket = spec .resource @@ -48,28 +48,29 @@ pub fn handle_redirect(state: Arc>, status: Status) .as_ref() .map(|r| r.name.clone()) .unwrap_or_default(); - (bucket, object) + let new_spec = AppendObjectSpec { + bucket, + object, + generation: redirect.generation.unwrap_or(0), + 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) + } + 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) } - AppendObjectSpecState::Append(spec) => (spec.bucket.clone(), spec.object.clone()), - }; - - let mut new_spec = AppendObjectSpec { - bucket, - object, - generation: redirect.generation.unwrap_or(0), - routing_token: redirect.routing_token, - write_handle: redirect.write_handle, - ..Default::default() }; - if let AppendObjectSpecState::Append(old_spec) = &*guard - && redirect.generation.is_none() - { - new_spec.generation = old_spec.generation; - } - - *guard = AppendObjectSpecState::Append(new_spec); - break; + *guard = new_state; } to_gax_error(status) } @@ -136,7 +137,8 @@ mod tests { write_handle: Some(BidiWriteHandle { handle: bytes::Bytes::from_static(b"initial-handle"), }), - ..Default::default() + if_metageneration_match: Some(10), + if_metageneration_not_match: Some(20), }; let state = Arc::new(Mutex::new(AppendObjectSpecState::Append(spec))); @@ -147,6 +149,49 @@ mod tests { 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 state = Arc::new(Mutex::new(AppendObjectSpecState::Write(Box::new( + write_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(), 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"); } From 1a15a5af8084a29c4e28fd3ed74b3b003947f37e Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Mon, 29 Jun 2026 07:54:39 +0000 Subject: [PATCH 05/13] Incorporate feedback --- src/storage/src/appendable_object_writer.rs | 1 - src/storage/src/storage/bidi_write.rs | 4 --- .../src/storage/bidi_write/connector.rs | 34 +++++++++---------- .../src/storage/bidi_write/transport.rs | 3 +- src/storage/src/storage/bidi_write/worker.rs | 21 ++++++++---- 5 files changed, 32 insertions(+), 31 deletions(-) diff --git a/src/storage/src/appendable_object_writer.rs b/src/storage/src/appendable_object_writer.rs index 2969a892b3..2e5050d932 100644 --- a/src/storage/src/appendable_object_writer.rs +++ b/src/storage/src/appendable_object_writer.rs @@ -191,7 +191,6 @@ impl AppendableObjectWriter { } } - #[allow(dead_code)] pub(crate) fn into_parts(self) -> Box { self.inner } diff --git a/src/storage/src/storage/bidi_write.rs b/src/storage/src/storage/bidi_write.rs index 59d9682fe3..1d2d645432 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -30,7 +30,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, @@ -47,7 +46,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( @@ -141,14 +139,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 index 5339a26b66..3cf453398b 100644 --- a/src/storage/src/storage/bidi_write/connector.rs +++ b/src/storage/src/storage/bidi_write/connector.rs @@ -1,5 +1,4 @@ -// Copyright 2025 Google LLC -#![allow(dead_code)] +// 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. @@ -99,7 +98,6 @@ where } None => None, }; - self.params = req.params.map(|p| p.to_proto().expect("valid params")); let spec = WriteObjectSpec { resource, predefined_acl: req.spec.predefined_acl, @@ -110,6 +108,10 @@ where 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)); self.connect_attempt_loop().await } @@ -129,7 +131,10 @@ where .write_handle .map(|handle| crate::google::storage::v2::BidiWriteHandle { handle }), }; - self.params = req.params.map(|p| p.to_proto().expect("valid params")); + 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 } @@ -211,21 +216,16 @@ where return Err(crate::Error::binding(mismatch)); } - let x_goog_request_params = request - .first_message - .as_ref() - .and_then(|m| match m { - FirstMessage::WriteObjectSpec(_) => None, // WriteObjectSpec doesn't have routing_token - FirstMessage::AppendObjectSpec(s) => s.routing_token.as_deref(), - _ => None, - }) - .into_iter() - .fold(format!("bucket={bucket_name}"), |s, token| { - s + &format!("&routing_token={token}") - }); + let mut x_goog_request_params = format!("bucket={bucket_name}"); + if let Some(FirstMessage::AppendObjectSpec(s)) = &request.first_message { + if let Some(token) = &s.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.clone()).await.map_err(Error::io)?; + tx.send(request).await.map_err(Error::io)?; let extensions = { let mut e = Extensions::new(); diff --git a/src/storage/src/storage/bidi_write/transport.rs b/src/storage/src/storage/bidi_write/transport.rs index 423179aeb9..5a3bae31f3 100644 --- a/src/storage/src/storage/bidi_write/transport.rs +++ b/src/storage/src/storage/bidi_write/transport.rs @@ -1,5 +1,4 @@ -// Copyright 2025 Google LLC -#![allow(dead_code)] +// 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. diff --git a/src/storage/src/storage/bidi_write/worker.rs b/src/storage/src/storage/bidi_write/worker.rs index bd19f386c2..048906c1e2 100644 --- a/src/storage/src/storage/bidi_write/worker.rs +++ b/src/storage/src/storage/bidi_write/worker.rs @@ -1,5 +1,4 @@ -// Copyright 2025 Google LLC -#![allow(dead_code)] +// 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. @@ -104,6 +103,7 @@ where } } Some(UploadIntent::Finalize(request, sender)) => { + self.finalized = true; self.pending_flushes.push_back(sender); if let Err(e) = tx.send(request).await { break Some(Arc::new(crate::Error::io(e))); @@ -133,7 +133,14 @@ where ) -> Option>>> { let response = match message { Ok(Some(msg)) => msg, - Ok(None) => return None, + 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); @@ -156,7 +163,7 @@ where #[cfg(test)] mod tests { - use super::super::mocks::{MockTestClient, mock_connector, mock_stream}; + use super::super::mocks::{MockTestClient, mock_connector}; use super::*; use crate::google::storage::v2::{ BidiWriteObjectRequest, BidiWriteObjectResponse, bidi_write_object_response::WriteStatus, @@ -167,7 +174,7 @@ mod tests { #[tokio::test] async fn run_immediately_closed() -> anyhow::Result<()> { let (request_tx, _request_rx) = mpsc::channel(1); - let (response_tx, response_rx) = mock_stream(); + let (response_tx, response_rx) = mpsc::channel(10); let (tx, rx) = mpsc::channel(1); let connection = Connection::new(request_tx, response_rx); @@ -187,7 +194,7 @@ mod tests { #[tokio::test] async fn run_flush_response() -> anyhow::Result<()> { let (request_tx, mut request_rx) = mpsc::channel(1); - let (response_tx, response_rx) = mock_stream(); + let (response_tx, response_rx) = mpsc::channel(10); let (tx, rx) = mpsc::channel(1); let connection = Connection::new(request_tx, response_rx); @@ -230,7 +237,7 @@ mod tests { #[tokio::test] async fn run_stop_on_closed_requests() -> anyhow::Result<()> { let (request_tx, _request_rx) = mpsc::channel(1); - let (_response_tx, response_rx) = mock_stream(); + let (_response_tx, response_rx) = mpsc::channel(10); let (tx, rx) = mpsc::channel(1); let connection = Connection::new(request_tx, response_rx); From 80f331131fe06a6d5a75b50321aa308c17e96641 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Mon, 29 Jun 2026 08:19:50 +0000 Subject: [PATCH 06/13] Fix failing tests, format, lint, clippy, dead code --- src/storage/src/appendable_object_writer.rs | 1 + src/storage/src/storage/bidi_write.rs | 2 ++ .../src/storage/bidi_write/connector.rs | 10 ++++---- src/storage/src/storage/bidi_write/mocks.rs | 3 ++- .../src/storage/bidi_write/transport.rs | 24 +++++++++---------- src/storage/src/storage/bidi_write/worker.rs | 1 - 6 files changed, 22 insertions(+), 19 deletions(-) diff --git a/src/storage/src/appendable_object_writer.rs b/src/storage/src/appendable_object_writer.rs index 2e5050d932..2969a892b3 100644 --- a/src/storage/src/appendable_object_writer.rs +++ b/src/storage/src/appendable_object_writer.rs @@ -191,6 +191,7 @@ impl AppendableObjectWriter { } } + #[allow(dead_code)] pub(crate) fn into_parts(self) -> Box { self.inner } diff --git a/src/storage/src/storage/bidi_write.rs b/src/storage/src/storage/bidi_write.rs index 1d2d645432..152048342d 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -14,6 +14,8 @@ //! Internal traits and types for Appendable Object Write (Bidi Write). +#![allow(dead_code)] + pub(crate) mod connector; mod redirect; mod retry_redirect; diff --git a/src/storage/src/storage/bidi_write/connector.rs b/src/storage/src/storage/bidi_write/connector.rs index 3cf453398b..6f57bd3412 100644 --- a/src/storage/src/storage/bidi_write/connector.rs +++ b/src/storage/src/storage/bidi_write/connector.rs @@ -217,11 +217,11 @@ where } let mut x_goog_request_params = format!("bucket={bucket_name}"); - if let Some(FirstMessage::AppendObjectSpec(s)) = &request.first_message { - if let Some(token) = &s.routing_token { - x_goog_request_params.push_str("&routing_token="); - x_goog_request_params.push_str(token); - } + if let Some(FirstMessage::AppendObjectSpec(s)) = &request.first_message + && let Some(token) = &s.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); diff --git a/src/storage/src/storage/bidi_write/mocks.rs b/src/storage/src/storage/bidi_write/mocks.rs index c86bbb11bf..315603b4c8 100644 --- a/src/storage/src/storage/bidi_write/mocks.rs +++ b/src/storage/src/storage/bidi_write/mocks.rs @@ -20,7 +20,6 @@ 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; -use tokio::sync::mpsc::Sender; // 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 @@ -76,6 +75,8 @@ pub trait TestClient: std::fmt::Debug { ) -> crate::Result>>; } +pub type MockStream = Receiver>; + pub fn mock_connector(mock: MockTestClient) -> Connector { let client = SharedMockClient::new(mock); diff --git a/src/storage/src/storage/bidi_write/transport.rs b/src/storage/src/storage/bidi_write/transport.rs index 5a3bae31f3..17ea163ede 100644 --- a/src/storage/src/storage/bidi_write/transport.rs +++ b/src/storage/src/storage/bidi_write/transport.rs @@ -34,7 +34,7 @@ pub struct AppendableObjectWriterTransport { generation: i64, persisted_size: i64, write_offset: i64, - crc32c_persisted: u32, + running_crc32c: u32, } impl AppendableObjectWriterTransport { @@ -81,13 +81,13 @@ impl AppendableObjectWriterTransport { persisted_size = *s; } - let mut crc32c_persisted = 0; + let mut running_crc32c = 0; if let Some(crc) = initial .persisted_data_checksums .as_ref() .and_then(|c| c.crc32c) { - crc32c_persisted = crc; + running_crc32c = crc; } let (tx, rx) = tokio::sync::mpsc::channel(100); @@ -99,7 +99,7 @@ impl AppendableObjectWriterTransport { generation, persisted_size, write_offset: persisted_size, - crc32c_persisted, + running_crc32c, }) } } @@ -109,7 +109,7 @@ impl AppendableObjectWriter for AppendableObjectWriterTransport { let length = chunk.len() as i64; let crc32c = crc32c::crc32c(&chunk); - let new_crc32c = crc32c::crc32c_append(self.crc32c_persisted, &chunk); + let new_crc32c = crc32c::crc32c_append(self.running_crc32c, &chunk); let request = BidiWriteObjectRequest { write_offset: self.write_offset, @@ -126,7 +126,7 @@ impl AppendableObjectWriter for AppendableObjectWriterTransport { .map_err(|e| crate::Error::io(e.to_string()))?; self.write_offset += length; - self.crc32c_persisted = new_crc32c; + self.running_crc32c = new_crc32c; Ok(()) } @@ -163,7 +163,7 @@ impl AppendableObjectWriter for AppendableObjectWriterTransport { flush: true, write_offset: self.write_offset, object_checksums: Some(ObjectChecksums { - crc32c: Some(self.crc32c_persisted), + crc32c: Some(self.running_crc32c), md5_hash: vec![].into(), }), ..BidiWriteObjectRequest::default() @@ -216,7 +216,7 @@ mod tests { let transport = AppendableObjectWriterTransport { tx, write_offset: 0, - crc32c_persisted: 0, + running_crc32c: 0, generation: 123456, persisted_size: 0, }; @@ -335,7 +335,7 @@ mod tests { let mut transport = AppendableObjectWriterTransport { tx, write_offset: 0, - crc32c_persisted: 0, + running_crc32c: 0, generation: 123456, persisted_size: 0, }; @@ -350,7 +350,7 @@ mod tests { // Assert that state was NOT modified due to the error assert_eq!(transport.write_offset, 0); - assert_eq!(transport.crc32c_persisted, 0); + assert_eq!(transport.running_crc32c, 0); Ok(()) } @@ -361,7 +361,7 @@ mod tests { let transport = AppendableObjectWriterTransport { tx, write_offset: 0, - crc32c_persisted: 0, + running_crc32c: 0, generation: 123456, persisted_size: 0, }; @@ -431,7 +431,7 @@ mod tests { assert_eq!(transport.generation(), 123456); assert_eq!(transport.persisted_size(), 1024); assert_eq!(transport.write_offset, 1024); - assert_eq!(transport.crc32c_persisted, 9999); + 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 index 048906c1e2..d8b7880ed2 100644 --- a/src/storage/src/storage/bidi_write/worker.rs +++ b/src/storage/src/storage/bidi_write/worker.rs @@ -103,7 +103,6 @@ where } } Some(UploadIntent::Finalize(request, sender)) => { - self.finalized = true; self.pending_flushes.push_back(sender); if let Err(e) = tx.send(request).await { break Some(Arc::new(crate::Error::io(e))); From 9cca9869423490eb313ea68edf7517285e531f2e Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Mon, 29 Jun 2026 16:15:11 +0000 Subject: [PATCH 07/13] Fix bugs --- .../src/storage/bidi_write/connector.rs | 34 +++++++----- .../src/storage/bidi_write/redirect.rs | 53 ++++++++++--------- .../src/storage/bidi_write/transport.rs | 1 + 3 files changed, 52 insertions(+), 36 deletions(-) diff --git a/src/storage/src/storage/bidi_write/connector.rs b/src/storage/src/storage/bidi_write/connector.rs index 6f57bd3412..35421ebf74 100644 --- a/src/storage/src/storage/bidi_write/connector.rs +++ b/src/storage/src/storage/bidi_write/connector.rs @@ -51,7 +51,7 @@ impl Connection { /// Represents the state of the initial request in the stream. #[derive(Clone, Debug)] pub enum AppendObjectSpecState { - Write(Box), + Write(Box, Option), Append(AppendObjectSpec), } @@ -79,7 +79,10 @@ where { pub fn new(options: RequestOptions, client: T) -> Self { Self { - spec: Arc::new(Mutex::new(AppendObjectSpecState::Write(Box::default()))), + spec: Arc::new(Mutex::new(AppendObjectSpecState::Write( + Box::default(), + None, + ))), options, client, params: None, @@ -112,7 +115,8 @@ where .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)); + *self.spec.lock().expect("never poisoned") = + AppendObjectSpecState::Write(Box::new(spec), None); self.connect_attempt_loop().await } @@ -173,14 +177,22 @@ where options: &RequestOptions, params: Option, ) -> Result<(BidiWriteObjectResponse, Connection)> { - let first_message = match &*spec.lock().expect("never poisoned") { - AppendObjectSpecState::Write(spec) => FirstMessage::WriteObjectSpec(*spec.clone()), - AppendObjectSpecState::Append(spec) => FirstMessage::AppendObjectSpec(spec.clone()), + 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() }; @@ -217,11 +229,9 @@ where } let mut x_goog_request_params = format!("bucket={bucket_name}"); - if let Some(FirstMessage::AppendObjectSpec(s)) = &request.first_message - && let Some(token) = &s.routing_token - { + if let Some(token) = routing_token { x_goog_request_params.push_str("&routing_token="); - x_goog_request_params.push_str(token); + x_goog_request_params.push_str(&token); } let (tx, rx) = tokio::sync::mpsc::channel::(100); @@ -264,7 +274,7 @@ where } let new_state = match &*guard { - AppendObjectSpecState::Write(s) => { + AppendObjectSpecState::Write(s, _) => { AppendObjectSpecState::Append(AppendObjectSpec { bucket: s .resource @@ -599,7 +609,7 @@ mod tests { let got = connector.spec.lock().expect("never poisoned").clone(); match got { - AppendObjectSpecState::Write(_) => panic!("Should be Append"), + AppendObjectSpecState::Write(_, _) => panic!("Should be Append"), AppendObjectSpecState::Append(got) => { assert_eq!(got.routing_token.as_deref(), Some("r1")); } diff --git a/src/storage/src/storage/bidi_write/redirect.rs b/src/storage/src/storage/bidi_write/redirect.rs index 8e5812a892..7ea054398d 100644 --- a/src/storage/src/storage/bidi_write/redirect.rs +++ b/src/storage/src/storage/bidi_write/redirect.rs @@ -37,27 +37,31 @@ pub fn handle_redirect(state: Arc>, status: Status) let mut guard = state.lock().expect("never poisoned"); let new_state = match &*guard { - AppendObjectSpecState::Write(spec) => { - 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: redirect.generation.unwrap_or(0), - 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) + 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(); @@ -180,9 +184,10 @@ mod tests { if_metageneration_not_match: Some(22), ..Default::default() }; - let state = Arc::new(Mutex::new(AppendObjectSpecState::Write(Box::new( - write_spec, - )))); + let state = Arc::new(Mutex::new(AppendObjectSpecState::Write( + Box::new(write_spec), + None, + ))); let got = handle_redirect(state.clone(), status); assert!(got.status().is_some(), "{got:?}"); diff --git a/src/storage/src/storage/bidi_write/transport.rs b/src/storage/src/storage/bidi_write/transport.rs index 17ea163ede..c479c2a9d6 100644 --- a/src/storage/src/storage/bidi_write/transport.rs +++ b/src/storage/src/storage/bidi_write/transport.rs @@ -135,6 +135,7 @@ impl AppendableObjectWriter for AppendableObjectWriterTransport { let (sender, receiver) = oneshot::channel(); let request = BidiWriteObjectRequest { flush: true, + state_lookup: true, write_offset: self.write_offset, ..BidiWriteObjectRequest::default() }; From 060e389d0982c386850ecaf22f028b8bb27b56cd Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Fri, 26 Jun 2026 07:32:12 +0000 Subject: [PATCH 08/13] feat(storage): implement API, request builders, integration tests --- src/storage/src/lib.rs | 6 + src/storage/src/model_ext.rs | 107 +---- .../open_appendable_object_request.rs | 42 ++ .../reopen_appendable_object_request.rs | 89 +++++ src/storage/src/storage.rs | 4 + src/storage/src/storage/client.rs | 84 ++++ .../src/storage/open_appendable_object.rs | 372 ++++++++++++++++++ .../src/storage/reopen_appendable_object.rs | 246 ++++++++++++ src/storage/src/storage/stub.rs | 24 ++ src/storage/src/storage/transport.rs | 50 +++ src/storage/tests/mocking.rs | 46 +++ tests/storage/src/bidi_write.rs | 147 +++++++ tests/storage/src/lib.rs | 2 + tests/storage/tests/driver.rs | 7 + 14 files changed, 1131 insertions(+), 95 deletions(-) create mode 100644 src/storage/src/model_ext/open_appendable_object_request.rs create mode 100644 src/storage/src/model_ext/reopen_appendable_object_request.rs create mode 100644 src/storage/src/storage/open_appendable_object.rs create mode 100644 src/storage/src/storage/reopen_appendable_object.rs create mode 100644 tests/storage/src/bidi_write.rs 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..c085244355 --- /dev/null +++ b/src/storage/src/model_ext/open_appendable_object_request.rs @@ -0,0 +1,42 @@ +// Copyright 2025 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(google_cloud_unstable_storage_bidi)] +#[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..8dfbd3ec7d --- /dev/null +++ b/src/storage/src/model_ext/reopen_appendable_object_request.rs @@ -0,0 +1,89 @@ +// Copyright 2025 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(google_cloud_unstable_storage_bidi)] +#[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, +} + +#[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)] +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/client.rs b/src/storage/src/storage/client.rs index 9b9c1b8a44..7fa019e1df 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,86 @@ 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. + 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..30ab7a65e8 --- /dev/null +++ b/src/storage/src/storage/open_appendable_object.rs @@ -0,0 +1,372 @@ +// Copyright 2025 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(google_cloud_unstable_storage_bidi)] +#[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, +} + +#[cfg(google_cloud_unstable_storage_bidi)] +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 + } +} + +#[cfg(google_cloud_unstable_storage_bidi)] +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_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 { + if let Some(res) = &mut self.request.spec.resource { + *res = res.clone().set_generation(v); + } + 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..49d689403d --- /dev/null +++ b/src/storage/src/storage/reopen_appendable_object.rs @@ -0,0 +1,246 @@ +// Copyright 2025 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(google_cloud_unstable_storage_bidi)] +#[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, +} + +#[cfg(google_cloud_unstable_storage_bidi)] +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 + } +} + +#[cfg(google_cloud_unstable_storage_bidi)] +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/transport.rs b/src/storage/src/storage/transport.rs index 7a70735b57..1d5d5a3c13 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,28 @@ 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)] + 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)) + } } impl super::stub::Storage for Storage { @@ -325,6 +355,26 @@ 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 { + // Tracing is deferred to PR 5. + 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 { + // Tracing is deferred to PR 5. + self.reopen_appendable_object_plain(request, options).await + } } #[cfg(test)] 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..d7de21c1b4 --- /dev/null +++ b/tests/storage/src/bidi_write.rs @@ -0,0 +1,147 @@ +// Copyright 2025 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_poison_stream(&client, bucket_name).await?; + 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 obj = writer.finalize().await?; + + assert_eq!(obj.size, 11); + assert_eq!(obj.name, object_name); + + 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 obj = writer.finalize().await?; + + assert_eq!(obj.size, 18); + assert_eq!(obj.name, object_name); + + 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 obj = writer2.finalize().await?; + + assert_eq!(obj.size, 11); + assert_eq!(obj.name, object_name); + + Ok(()) +} + +async fn test_bidi_write_poison_stream(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_poison_stream_{now}"); + + let mut writer = client + .open_appendable_object(bucket, &object_name) + .send() + .await?; + + writer.append(Bytes::from("hello ")).await?; + writer.flush().await?; // Ensure we get the generation by flushing at least once + let generation = writer.generation(); + writer + .append(Bytes::from("poison data that won't be flushed")) + .await?; + + // Explicitly drop the writer without finalizing or closing, poisoning the stream. + drop(writer); + + // Give the server a moment to recognize the dropped connection (not strictly necessary but good practice) + 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 ", so persisted size might be 6 (though the server might have acked the rest, let's just assert it succeeds in reopening and we can append more). + writer2.append(Bytes::from("world")).await?; + let obj = writer2.finalize().await?; + + // Since we didn't flush the poison data, it may or may not be there. GCS makes no guarantees about unflushed data on disconnect. + // The main test is that we can successfully reopen and finalize after a dropped connection. + assert!(obj.size >= 11); + assert_eq!(obj.name, object_name); + + Ok(()) +} diff --git a/tests/storage/src/lib.rs b/tests/storage/src/lib.rs index 03e17ccea0..ac6242ec01 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; diff --git a/tests/storage/tests/driver.rs b/tests/storage/tests/driver.rs index 67514f5027..4d5433ab3d 100644 --- a/tests/storage/tests/driver.rs +++ b/tests/storage/tests/driver.rs @@ -121,6 +121,13 @@ mod storage { let result = integration_tests_storage::bidi_read::run(&bucket.name) .await .inspect_err(anydump); + #[cfg(google_cloud_unstable_storage_bidi)] + let result = match result { + Ok(()) => integration_tests_storage::bidi_write::run(&bucket.name) + .await + .inspect_err(anydump), + Err(e) => Err(e), + }; let _ = storage_samples::cleanup_bucket(control, bucket.name.clone(), bucket.project.clone()) .await From ae9609c0e447df61e2bebdb80a41040bea1321be Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Sat, 27 Jun 2026 09:13:17 +0000 Subject: [PATCH 09/13] Incorporate feedback --- .../src/storage/open_appendable_object.rs | 15 +++-- src/storage/src/storage/transport.rs | 62 +++++++++++++++++++ tests/storage/src/bidi_write.rs | 54 ++++++++++++---- tests/storage/tests/driver.rs | 2 +- 4 files changed, 112 insertions(+), 21 deletions(-) diff --git a/src/storage/src/storage/open_appendable_object.rs b/src/storage/src/storage/open_appendable_object.rs index 30ab7a65e8..72e9508d99 100644 --- a/src/storage/src/storage/open_appendable_object.rs +++ b/src/storage/src/storage/open_appendable_object.rs @@ -84,11 +84,13 @@ impl OpenAppendableObject { O: Into, { let request = OpenAppendableObjectRequest { - spec: crate::model::WriteObjectSpec::default().set_resource( - crate::model::Object::new() - .set_bucket(bucket.into()) - .set_name(object.into()), - ), + 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 { @@ -101,9 +103,6 @@ impl OpenAppendableObject { /// 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 { - if let Some(res) = &mut self.request.spec.resource { - *res = res.clone().set_generation(v); - } self.request.spec.if_generation_match = Some(v); self } diff --git a/src/storage/src/storage/transport.rs b/src/storage/src/storage/transport.rs index 1d5d5a3c13..cd2d69e54c 100644 --- a/src/storage/src/storage/transport.rs +++ b/src/storage/src/storage/transport.rs @@ -707,6 +707,68 @@ 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 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:?}" + ); + + // Tracing is deferred to PR 5, so we don't check `client_request_span` here yet. + 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 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:?}" + ); + + // Tracing is deferred to PR 5, so we don't check `client_request_span` here yet. + Ok(()) + } + #[track_caller] fn check_debug_log(captured: &Vec, method: &'static str) { let span = captured diff --git a/tests/storage/src/bidi_write.rs b/tests/storage/src/bidi_write.rs index d7de21c1b4..ebc0445f28 100644 --- a/tests/storage/src/bidi_write.rs +++ b/tests/storage/src/bidi_write.rs @@ -25,6 +25,21 @@ pub async fn run(bucket_name: &str) -> anyhow::Result<()> { 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) @@ -38,10 +53,11 @@ async fn test_bidi_write_single_block(client: &Storage, bucket: &str) -> anyhow: .await?; writer.append(Bytes::from("hello world")).await?; - let obj = writer.finalize().await?; + let object_metadata = writer.finalize().await?; - assert_eq!(obj.size, 11); - assert_eq!(obj.name, object_name); + 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(()) } @@ -62,10 +78,11 @@ async fn test_bidi_write_chunked_appends(client: &Storage, bucket: &str) -> anyh writer.append(Bytes::from("chunk2")).await?; writer.flush().await?; writer.append(Bytes::from("chunk3")).await?; - let obj = writer.finalize().await?; + let object_metadata = writer.finalize().await?; - assert_eq!(obj.size, 18); - assert_eq!(obj.name, object_name); + assert_eq!(object_metadata.size, 18); + assert_eq!(object_metadata.name, object_name); + check_object_contents(client, bucket, &object_name, b"chunk1chunk2chunk3").await?; Ok(()) } @@ -95,10 +112,11 @@ async fn test_bidi_write_resume_append(client: &Storage, bucket: &str) -> anyhow assert_eq!(writer2.persisted_size(), 6); writer2.append(Bytes::from("world")).await?; - let obj = writer2.finalize().await?; + let object_metadata = writer2.finalize().await?; - assert_eq!(obj.size, 11); - assert_eq!(obj.name, object_name); + 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(()) } @@ -134,14 +152,26 @@ async fn test_bidi_write_poison_stream(client: &Storage, bucket: &str) -> anyhow .send() .await?; + assert!(writer2.persisted_size() >= 6); + // We only flushed "hello ", so persisted size might be 6 (though the server might have acked the rest, let's just assert it succeeds in reopening and we can append more). writer2.append(Bytes::from("world")).await?; - let obj = writer2.finalize().await?; + let object_metadata = writer2.finalize().await?; // Since we didn't flush the poison data, it may or may not be there. GCS makes no guarantees about unflushed data on disconnect. // The main test is that we can successfully reopen and finalize after a dropped connection. - assert!(obj.size >= 11); - assert_eq!(obj.name, object_name); + assert!(object_metadata.size >= 11); + assert_eq!(object_metadata.name, object_name); + + // We can't strictly check exact contents because the poison 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/tests/driver.rs b/tests/storage/tests/driver.rs index 4d5433ab3d..f5c11c178c 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) From 037569691b94b06374f6a9831f73288c318da9bb Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Sun, 28 Jun 2026 10:10:49 +0000 Subject: [PATCH 10/13] Incorporate feedback --- tests/storage/src/bidi_write.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/storage/src/bidi_write.rs b/tests/storage/src/bidi_write.rs index ebc0445f28..b68dda9e96 100644 --- a/tests/storage/src/bidi_write.rs +++ b/tests/storage/src/bidi_write.rs @@ -143,7 +143,7 @@ async fn test_bidi_write_poison_stream(client: &Storage, bucket: &str) -> anyhow // Explicitly drop the writer without finalizing or closing, poisoning the stream. drop(writer); - // Give the server a moment to recognize the dropped connection (not strictly necessary but good practice) + // Give the server a moment to recognize the dropped connection. tokio::time::sleep(std::time::Duration::from_millis(500)).await; // Reopen using the generation From 9ceced853f1eb4cef25b21ec35f8cc9b480205e9 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Mon, 29 Jun 2026 08:28:55 +0000 Subject: [PATCH 11/13] Clean up code --- .../open_appendable_object_request.rs | 2 +- .../reopen_appendable_object_request.rs | 2 +- src/storage/src/storage/client.rs | 2 ++ .../src/storage/open_appendable_object.rs | 2 +- .../src/storage/reopen_appendable_object.rs | 2 +- tests/storage/src/bidi_write.rs | 30 ++++++++++--------- tests/storage/tests/driver.rs | 2 +- 7 files changed, 23 insertions(+), 19 deletions(-) diff --git a/src/storage/src/model_ext/open_appendable_object_request.rs b/src/storage/src/model_ext/open_appendable_object_request.rs index c085244355..ce8c98a510 100644 --- a/src/storage/src/model_ext/open_appendable_object_request.rs +++ b/src/storage/src/model_ext/open_appendable_object_request.rs @@ -1,4 +1,4 @@ -// Copyright 2025 Google LLC +// 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. diff --git a/src/storage/src/model_ext/reopen_appendable_object_request.rs b/src/storage/src/model_ext/reopen_appendable_object_request.rs index 8dfbd3ec7d..c51c796304 100644 --- a/src/storage/src/model_ext/reopen_appendable_object_request.rs +++ b/src/storage/src/model_ext/reopen_appendable_object_request.rs @@ -1,4 +1,4 @@ -// Copyright 2025 Google LLC +// 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. diff --git a/src/storage/src/storage/client.rs b/src/storage/src/storage/client.rs index 7fa019e1df..259975f885 100644 --- a/src/storage/src/storage/client.rs +++ b/src/storage/src/storage/client.rs @@ -309,6 +309,8 @@ where /// * `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, diff --git a/src/storage/src/storage/open_appendable_object.rs b/src/storage/src/storage/open_appendable_object.rs index 72e9508d99..6f017c8db3 100644 --- a/src/storage/src/storage/open_appendable_object.rs +++ b/src/storage/src/storage/open_appendable_object.rs @@ -1,4 +1,4 @@ -// Copyright 2025 Google LLC +// 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. diff --git a/src/storage/src/storage/reopen_appendable_object.rs b/src/storage/src/storage/reopen_appendable_object.rs index 49d689403d..c93405cbe5 100644 --- a/src/storage/src/storage/reopen_appendable_object.rs +++ b/src/storage/src/storage/reopen_appendable_object.rs @@ -1,4 +1,4 @@ -// Copyright 2025 Google LLC +// 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. diff --git a/tests/storage/src/bidi_write.rs b/tests/storage/src/bidi_write.rs index b68dda9e96..0e5f9077b2 100644 --- a/tests/storage/src/bidi_write.rs +++ b/tests/storage/src/bidi_write.rs @@ -1,4 +1,4 @@ -// Copyright 2025 Google LLC +// 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. @@ -21,7 +21,7 @@ pub async fn run(bucket_name: &str) -> anyhow::Result<()> { 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_poison_stream(&client, bucket_name).await?; + test_bidi_write_drop_stream(&client, bucket_name).await?; Ok(()) } @@ -104,7 +104,7 @@ async fn test_bidi_write_resume_append(client: &Storage, bucket: &str) -> anyhow let persisted_size = writer.close().await?; assert_eq!(persisted_size, 6); - // Reopen + // Reopen. let mut writer2 = client .reopen_appendable_object(bucket, &object_name, generation) .send() @@ -121,12 +121,15 @@ async fn test_bidi_write_resume_append(client: &Storage, bucket: &str) -> anyhow Ok(()) } -async fn test_bidi_write_poison_stream(client: &Storage, bucket: &str) -> anyhow::Result<()> { +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_poison_stream_{now}"); + let object_name = format!("test_bidi_write_drop_stream_{now}"); let mut writer = client .open_appendable_object(bucket, &object_name) @@ -134,37 +137,36 @@ async fn test_bidi_write_poison_stream(client: &Storage, bucket: &str) -> anyhow .await?; writer.append(Bytes::from("hello ")).await?; - writer.flush().await?; // Ensure we get the generation by flushing at least once let generation = writer.generation(); writer - .append(Bytes::from("poison data that won't be flushed")) + .append(Bytes::from("data that won't be flushed")) .await?; - // Explicitly drop the writer without finalizing or closing, poisoning the stream. + // 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 + // 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); - // We only flushed "hello ", so persisted size might be 6 (though the server might have acked the rest, let's just assert it succeeds in reopening and we can append more). writer2.append(Bytes::from("world")).await?; let object_metadata = writer2.finalize().await?; - // Since we didn't flush the poison data, it may or may not be there. GCS makes no guarantees about unflushed data on disconnect. - // The main test is that we can successfully reopen and finalize after a dropped connection. assert!(object_metadata.size >= 11); assert_eq!(object_metadata.name, object_name); - // We can't strictly check exact contents because the poison data may or may not have been persisted, - // but we know it starts with "hello " and ends with "world". + // 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 { diff --git a/tests/storage/tests/driver.rs b/tests/storage/tests/driver.rs index f5c11c178c..4d5433ab3d 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_bucket() + let (control, bucket) = integration_tests_storage::create_test_hns_bucket() .await .inspect_err(anydump)?; let result = integration_tests_storage::bidi_read::run(&bucket.name) From 5e066c190d13d6cb0ef4b65c786b2f130f885cc8 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Mon, 29 Jun 2026 16:05:09 +0000 Subject: [PATCH 12/13] Update integration test --- src/storage/examples/src/lib.rs | 34 ++++++++++++++++++- .../open_appendable_object_request.rs | 1 - .../reopen_appendable_object_request.rs | 2 -- .../src/storage/open_appendable_object.rs | 3 -- .../src/storage/reopen_appendable_object.rs | 3 -- tests/storage/src/bidi_write.rs | 1 + tests/storage/src/lib.rs | 4 ++- tests/storage/tests/driver.rs | 25 ++++++++++---- 8 files changed, 55 insertions(+), 18 deletions(-) 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/model_ext/open_appendable_object_request.rs b/src/storage/src/model_ext/open_appendable_object_request.rs index ce8c98a510..a402c89c33 100644 --- a/src/storage/src/model_ext/open_appendable_object_request.rs +++ b/src/storage/src/model_ext/open_appendable_object_request.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#[cfg(google_cloud_unstable_storage_bidi)] #[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] /// Represents the parameters of a request to open a new object for exclusive appends. /// diff --git a/src/storage/src/model_ext/reopen_appendable_object_request.rs b/src/storage/src/model_ext/reopen_appendable_object_request.rs index c51c796304..fc4a154b73 100644 --- a/src/storage/src/model_ext/reopen_appendable_object_request.rs +++ b/src/storage/src/model_ext/reopen_appendable_object_request.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#[cfg(google_cloud_unstable_storage_bidi)] #[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] /// Represents the parameters of a request to reopen an existing object for appends. /// @@ -38,7 +37,6 @@ pub struct ReopenAppendableObjectRequest { pub params: Option, } -#[cfg(google_cloud_unstable_storage_bidi)] impl From for crate::google::storage::v2::AppendObjectSpec { fn from(value: ReopenAppendableObjectRequest) -> Self { Self { diff --git a/src/storage/src/storage/open_appendable_object.rs b/src/storage/src/storage/open_appendable_object.rs index 6f017c8db3..0cadc86c22 100644 --- a/src/storage/src/storage/open_appendable_object.rs +++ b/src/storage/src/storage/open_appendable_object.rs @@ -19,7 +19,6 @@ use crate::request_options::RequestOptions; use std::sync::Arc; use std::time::Duration; -#[cfg(google_cloud_unstable_storage_bidi)] #[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] /// A request builder for [Storage::open_appendable_object][crate::client::Storage::open_appendable_object]. /// @@ -45,7 +44,6 @@ pub struct OpenAppendableObject { options: RequestOptions, } -#[cfg(google_cloud_unstable_storage_bidi)] impl OpenAppendableObject where S: crate::storage::stub::Storage + 'static, @@ -71,7 +69,6 @@ where } } -#[cfg(google_cloud_unstable_storage_bidi)] impl OpenAppendableObject { pub(crate) fn new( stub: std::sync::Arc, diff --git a/src/storage/src/storage/reopen_appendable_object.rs b/src/storage/src/storage/reopen_appendable_object.rs index c93405cbe5..b79b671e61 100644 --- a/src/storage/src/storage/reopen_appendable_object.rs +++ b/src/storage/src/storage/reopen_appendable_object.rs @@ -19,7 +19,6 @@ use crate::request_options::RequestOptions; use std::sync::Arc; use std::time::Duration; -#[cfg(google_cloud_unstable_storage_bidi)] #[cfg_attr(docsrs, doc(cfg(feature = "unstable-stream")))] /// A request builder for [Storage::reopen_appendable_object][crate::client::Storage::reopen_appendable_object]. /// @@ -45,7 +44,6 @@ pub struct ReopenAppendableObject { options: RequestOptions, } -#[cfg(google_cloud_unstable_storage_bidi)] impl ReopenAppendableObject where S: crate::storage::stub::Storage + 'static, @@ -71,7 +69,6 @@ where } } -#[cfg(google_cloud_unstable_storage_bidi)] impl ReopenAppendableObject { pub(crate) fn new( stub: std::sync::Arc, diff --git a/tests/storage/src/bidi_write.rs b/tests/storage/src/bidi_write.rs index 0e5f9077b2..ab41dd3f06 100644 --- a/tests/storage/src/bidi_write.rs +++ b/tests/storage/src/bidi_write.rs @@ -137,6 +137,7 @@ async fn test_bidi_write_drop_stream(client: &Storage, bucket: &str) -> anyhow:: .await?; writer.append(Bytes::from("hello ")).await?; + writer.flush().await?; let generation = writer.generation(); writer .append(Bytes::from("data that won't be flushed")) diff --git a/tests/storage/src/lib.rs b/tests/storage/src/lib.rs index ac6242ec01..cb9396f9e1 100644 --- a/tests/storage/src/lib.rs +++ b/tests/storage/src/lib.rs @@ -38,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 4d5433ab3d..4450356693 100644 --- a/tests/storage/tests/driver.rs +++ b/tests/storage/tests/driver.rs @@ -115,19 +115,30 @@ 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) .await .inspect_err(anydump); - #[cfg(google_cloud_unstable_storage_bidi)] - let result = match result { - Ok(()) => integration_tests_storage::bidi_write::run(&bucket.name) + let _ = + storage_samples::cleanup_bucket(control, bucket.name.clone(), bucket.project.clone()) .await - .inspect_err(anydump), - Err(e) => Err(e), - }; + .inspect_err(|e| tracing::error!("error cleaning up bucket {}: {e:?}", bucket.name)) + .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 From aed02a534f4fa08fab7e9f786e1f1b5043d0d68a Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Sat, 27 Jun 2026 10:12:34 +0000 Subject: [PATCH 13/13] feat(storage): add tracing and decorators --- src/storage/src/storage/tracing.rs | 183 +++++++++++++++ src/storage/src/storage/transport.rs | 322 ++++++++++++++++++++++++++- 2 files changed, 501 insertions(+), 4 deletions(-) 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 cd2d69e54c..c799cd0149 100644 --- a/src/storage/src/storage/transport.rs +++ b/src/storage/src/storage/transport.rs @@ -282,6 +282,44 @@ impl Storage { 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, @@ -292,6 +330,36 @@ impl Storage { 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 { @@ -362,7 +430,9 @@ impl super::stub::Storage for Storage { request: OpenAppendableObjectRequest, options: RequestOptions, ) -> Result { - // Tracing is deferred to PR 5. + if self.tracing { + return self.open_appendable_object_tracing(request, options).await; + } self.open_appendable_object_plain(request, options).await } @@ -372,7 +442,11 @@ impl super::stub::Storage for Storage { request: ReopenAppendableObjectRequest, options: RequestOptions, ) -> Result { - // Tracing is deferred to PR 5. + if self.tracing { + return self + .reopen_appendable_object_tracing(request, options) + .await; + } self.reopen_appendable_object_plain(request, options).await } } @@ -714,6 +788,7 @@ mod tests { 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"))); @@ -733,8 +808,97 @@ mod tests { 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)); - // Tracing is deferred to PR 5, so we don't check `client_request_span` here yet. Ok(()) } @@ -745,6 +909,7 @@ mod tests { 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"))); @@ -764,8 +929,99 @@ mod tests { 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)); - // Tracing is deferred to PR 5, so we don't check `client_request_span` here yet. Ok(()) } @@ -835,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}`" + ); + } + } + } }