From e86f327de889a8823834e8bb06a162aaeb7da058 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Wed, 8 Jul 2026 01:47:34 +0000 Subject: [PATCH 1/3] feat(storage): implement appendable upload connector --- src/storage/src/storage/bidi_write.rs | 10 +- .../src/storage/bidi_write/connector.rs | 745 ++++++++++++++++++ src/storage/src/storage/bidi_write/mocks.rs | 84 ++ .../src/storage/bidi_write/redirect.rs | 128 ++- .../src/storage/bidi_write/retry_redirect.rs | 2 - 5 files changed, 943 insertions(+), 26 deletions(-) create mode 100644 src/storage/src/storage/bidi_write/connector.rs create mode 100644 src/storage/src/storage/bidi_write/mocks.rs diff --git a/src/storage/src/storage/bidi_write.rs b/src/storage/src/storage/bidi_write.rs index b6576d33be..ee197bfcaf 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -14,6 +14,9 @@ //! Internal traits and types for Appendable Object Write (Bidi Write). +#![allow(dead_code)] + +pub(crate) mod connector; mod redirect; mod retry_redirect; pub(crate) mod stub; @@ -27,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, @@ -44,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( @@ -82,6 +83,9 @@ impl Client for gaxi::grpc::Client { } } +#[cfg(test)] +mod mocks; + #[cfg(test)] pub(crate) mod tests { use crate::Error; @@ -135,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 new file mode 100644 index 0000000000..35421ebf74 --- /dev/null +++ b/src/storage/src/storage/bidi_write/connector.rs @@ -0,0 +1,745 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// TODO(#5716): Lift to shared bidi module + +use super::redirect::handle_redirect; +use super::retry_redirect::RetryRedirect; +use super::{Client, TonicStreaming}; +use crate::google::storage::v2::{ + AppendObjectSpec, BidiWriteObjectRequest, BidiWriteObjectResponse, CommonObjectRequestParams, + WriteObjectSpec, bidi_write_object_request::FirstMessage, + bidi_write_object_response::WriteStatus, +}; +use crate::request_options::RequestOptions; +use crate::storage::info::X_GOOG_API_CLIENT_HEADER; +use crate::{Error, Result}; +use gaxi::grpc::Client as GrpcClient; +use gaxi::grpc::tonic::{Extensions, GrpcMethod, Streaming}; +use gaxi::prost::ToProto; +use google_cloud_gax::error::binding::{ + BindingError, PathMismatch, SubstitutionFail, SubstitutionMismatch, +}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::sync::mpsc::Sender; +/// Represents a bidirectional streaming connection. +/// Contains the transmission channel for requests and the receiving stream for responses. +#[derive(Debug)] +pub struct Connection> { + pub tx: Sender, + pub rx: S, +} + +impl Connection { + pub fn new(tx: Sender, rx: S) -> Self { + Self { tx, rx } + } +} + +/// Represents the state of the initial request in the stream. +#[derive(Clone, Debug)] +pub enum AppendObjectSpecState { + Write(Box, Option), + Append(AppendObjectSpec), +} + +/// Establishes and handles the initial handshake for bidi streaming writes. +/// +/// Connecting and reconnecting bidirectional streaming writes requires: +/// - Constructing the initial `WriteObjectSpec` or `AppendObjectSpec`. +/// - Following routing token redirects correctly. +/// - Passing back the established `Connection` to the async worker. +/// +/// # Parameters +/// - `T`: a type implementing the [Client] trait, this is used in tests. +#[derive(Clone, Debug)] +pub struct Connector { + spec: Arc>, + options: RequestOptions, + client: T, + params: Option, +} + +impl Connector +where + T: Client + Clone + Send + 'static, + ::Stream: TonicStreaming, +{ + pub fn new(options: RequestOptions, client: T) -> Self { + Self { + spec: Arc::new(Mutex::new(AppendObjectSpecState::Write( + Box::default(), + None, + ))), + options, + client, + params: None, + } + } + + pub async fn connect_open( + &mut self, + req: crate::model_ext::OpenAppendableObjectRequest, + ) -> Result<(BidiWriteObjectResponse, Connection)> { + let resource = match req.spec.resource { + Some(r) => { + let object: crate::google::storage::v2::Object = gaxi::prost::ToProto::to_proto(r) + .map_err(|e: gaxi::prost::ConvertError| Error::io(e.to_string()))?; + Some(object) + } + None => None, + }; + let spec = WriteObjectSpec { + resource, + predefined_acl: req.spec.predefined_acl, + if_generation_match: req.spec.if_generation_match, + if_generation_not_match: req.spec.if_generation_not_match, + if_metageneration_match: req.spec.if_metageneration_match, + if_metageneration_not_match: req.spec.if_metageneration_not_match, + object_size: req.spec.object_size, + appendable: req.spec.appendable, + }; + self.params = req + .params + .map(|p| p.to_proto().map_err(|e| Error::io(e.to_string()))) + .transpose()?; + *self.spec.lock().expect("never poisoned") = + AppendObjectSpecState::Write(Box::new(spec), None); + self.connect_attempt_loop().await + } + + pub async fn connect_reopen( + &mut self, + req: crate::model_ext::ReopenAppendableObjectRequest, + ) -> Result<(BidiWriteObjectResponse, Connection)> { + let spec = AppendObjectSpec { + bucket: req.bucket, + object: req.object, + generation: req.generation, + routing_token: req.routing_token, + if_metageneration_match: req.if_metageneration_match, + if_metageneration_not_match: req.if_metageneration_not_match, + write_handle: req + .write_handle + .map(|handle| crate::google::storage::v2::BidiWriteHandle { handle }), + }; + self.params = req + .params + .map(|p| p.to_proto().map_err(|e| Error::io(e.to_string()))) + .transpose()?; + *self.spec.lock().expect("never poisoned") = AppendObjectSpecState::Append(spec); + self.connect_attempt_loop().await + } + + async fn connect_attempt_loop( + &mut self, + ) -> Result<(BidiWriteObjectResponse, Connection)> { + let throttler = self.options.retry_throttler.clone(); + let retry = Arc::new(RetryRedirect::new(self.options.retry_policy.clone())); + let backoff = self.options.backoff_policy.clone(); + let client = self.client.clone(); + let options = self.options.clone(); + let spec = self.spec.clone(); + let params = self.params.clone(); + let sleep = async |backoff| tokio::time::sleep(backoff).await; + let default_timeout = self.options.bidi_attempt_timeout; + + let inner = async move |d: Option| { + let attempt_timeout = std::cmp::min(default_timeout, d.unwrap_or(default_timeout)); + let attempt = + Self::connect_attempt(client.clone(), spec.clone(), &options, params.clone()); + match tokio::time::timeout(attempt_timeout, attempt).await { + Ok(r) => r, + Err(e) => Err(Error::timeout(e)), + } + }; + google_cloud_gax::retry_loop_internal::retry_loop( + inner, sleep, true, throttler, retry, backoff, + ) + .await + } + + async fn connect_attempt( + client: T, + spec: Arc>, + options: &RequestOptions, + params: Option, + ) -> Result<(BidiWriteObjectResponse, Connection)> { + let (first_message, routing_token) = match &*spec.lock().expect("never poisoned") { + AppendObjectSpecState::Write(spec, rt) => { + (FirstMessage::WriteObjectSpec(*spec.clone()), rt.clone()) + } + AppendObjectSpecState::Append(spec) => ( + FirstMessage::AppendObjectSpec(spec.clone()), + spec.routing_token.clone(), + ), + }; + + let state_lookup = matches!(first_message, FirstMessage::AppendObjectSpec(_)); + + let request = BidiWriteObjectRequest { + first_message: Some(first_message), + common_object_request_params: params, + state_lookup, + ..BidiWriteObjectRequest::default() + }; + + let bucket_name = request + .first_message + .as_ref() + .and_then(|m| match m { + FirstMessage::WriteObjectSpec(s) => s.resource.as_ref().map(|r| r.bucket.as_str()), + FirstMessage::AppendObjectSpec(s) => Some(s.bucket.as_str()), + _ => None, + }) + .unwrap_or_default(); + + if bucket_name + .strip_prefix("projects/_/buckets/") + .is_none_or(|x| x.is_empty()) + { + let problem = SubstitutionFail::MismatchExpecting( + bucket_name.to_string(), + "projects/_/buckets/*", + ); + let mismatch = SubstitutionMismatch { + field_name: "bucket", + problem, + }; + let mismatch = PathMismatch { + subs: vec![mismatch], + }; + let mismatch = BindingError { + paths: vec![mismatch], + }; + + return Err(crate::Error::binding(mismatch)); + } + + let mut x_goog_request_params = format!("bucket={bucket_name}"); + if let Some(token) = routing_token { + x_goog_request_params.push_str("&routing_token="); + x_goog_request_params.push_str(&token); + } + + let (tx, rx) = tokio::sync::mpsc::channel::(100); + tx.send(request).await.map_err(Error::io)?; + + let extensions = { + let mut e = Extensions::new(); + e.insert(GrpcMethod::new( + "google.storage.v2.Storage", + "BidiWriteObject", + )); + e + }; + let path = + http::uri::PathAndQuery::from_static("/google.storage.v2.Storage/BidiWriteObject"); + + let response = client + .start( + extensions, + path, + rx, + options, + &X_GOOG_API_CLIENT_HEADER, + &x_goog_request_params, + ) + .await?; + + let response = match response { + Ok(r) => r, + Err(status) => return Err(handle_redirect(spec, status)), + }; + + let (_metadata, mut stream, _) = response.into_parts(); + match stream.next_message().await { + Ok(Some(m)) => { + let mut guard = spec.lock().expect("never poisoned"); + let mut new_generation = None; + if let Some(WriteStatus::Resource(resource)) = &m.write_status { + new_generation = Some(resource.generation); + } + + let new_state = match &*guard { + AppendObjectSpecState::Write(s, _) => { + AppendObjectSpecState::Append(AppendObjectSpec { + bucket: s + .resource + .as_ref() + .map(|r| r.bucket.clone()) + .unwrap_or_default(), + object: s + .resource + .as_ref() + .map(|r| r.name.clone()) + .unwrap_or_default(), + generation: new_generation.unwrap_or(0), + write_handle: m.write_handle.clone(), + ..Default::default() + }) + } + AppendObjectSpecState::Append(s) => { + AppendObjectSpecState::Append(AppendObjectSpec { + generation: new_generation.unwrap_or(s.generation), + write_handle: m.write_handle.clone().or_else(|| s.write_handle.clone()), + ..s.clone() + }) + } + }; + *guard = new_state; + + Ok((m, Connection::new(tx, stream))) + } + Ok(None) => Err(Error::io("bidi_write_object stream closed before start")), + Err(status) => Err(handle_redirect(spec, status)), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::model_ext::OpenAppendableObjectRequest; + use crate::storage::request_options::RequestOptions; + use anyhow::Result; + use gaxi::grpc::Client as GrpcClient; + use google_cloud_auth::credentials::{Credentials, anonymous::Builder as Anonymous}; + use google_cloud_gax::retry_policy::NeverRetry; + use static_assertions::assert_impl_all; + use std::error::Error as _; + use std::sync::Arc; + + fn test_credentials() -> Credentials { + Anonymous::new().build() + } + + fn test_options() -> RequestOptions { + RequestOptions::new() + } + + #[test] + fn assertions() { + assert_impl_all!(Connector: Clone, std::fmt::Debug, Send, Sync); + } + + #[tokio::test] + async fn bad_endpoint() -> Result<()> { + fn need_send(_val: &T) {} + + let mut config = gaxi::options::ClientConfig::default(); + config.cred = Some(test_credentials()); + let client = GrpcClient::new(config, "http://127.0.0.1:1").await?; + + let mut options = test_options(); + options.retry_policy = Arc::new(NeverRetry); + let mut connector = Connector::new(options, client); + + let mut req = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default(), + params: None, + }; + req.spec = crate::model::WriteObjectSpec { + resource: Some(crate::model::Object { + bucket: "projects/_/buckets/test-only-bucket".into(), + name: "test-only-object".into(), + ..Default::default() + }), + ..Default::default() + }; + + let start = connector.connect_open(req); + need_send(&start); + + let err = start.await.unwrap_err(); + assert!(err.is_connect(), "{err:?}"); + let source = err.source().unwrap().to_string(); + assert!(source.contains("127.0.0.1:1"), "{source}"); + + Ok(()) + } + + #[tokio::test] + #[test_case::test_case("")] + #[test_case::test_case("my-bucket")] + async fn binding_error(bucket_name: &str) -> Result<()> { + use super::super::mocks::{MockTestClient, SharedMockClient}; + let mut mock = MockTestClient::new(); + // Binding errors are detected before a request is sent. + mock.expect_start().never(); + let client = SharedMockClient::new(mock); + + let mut connector = Connector::new(test_options(), client); + + let mut req = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default(), + params: None, + }; + req.spec = crate::model::WriteObjectSpec { + resource: Some(crate::model::Object { + bucket: bucket_name.into(), + name: "object".into(), + ..Default::default() + }), + ..Default::default() + }; + + let err = connector.connect_open(req).await.unwrap_err(); + assert!(err.is_binding(), "{err:?}"); + use google_cloud_gax::error::binding::{BindingError, SubstitutionFail}; + let source = err.source().and_then(|e| e.downcast_ref::()); + assert!(matches!(source, Some(BindingError { .. })), "{err:?}"); + // Extract all the field names that did not match, and expect a single name: + let mismatch = source + .iter() + .flat_map(|f| f.paths.iter()) + .flat_map(|f| f.subs.iter()) + .map(|f| f.field_name) + .collect::>(); + assert_eq!(mismatch, vec!["bucket"], "{err:?}"); + + // Extract all the problems: + let mismatch = source + .iter() + .flat_map(|f| f.paths.iter()) + .flat_map(|f| f.subs.iter()) + .map(|f| &f.problem) + .collect::>(); + assert!( + matches!( + mismatch.first(), + Some(SubstitutionFail::MismatchExpecting(n, p)) if n == bucket_name && *p == "projects/_/buckets/*" + ), + "{err:?}" + ); + Ok(()) + } + + #[tokio::test] + async fn start_error() -> Result<()> { + use super::super::mocks::{MockTestClient, SharedMockClient}; + use super::super::tests::permanent_error; + use gaxi::grpc::tonic::GrpcMethod; + use std::sync::Mutex; + + let receivers = Arc::new(Mutex::new(Vec::new())); + let save = receivers.clone(); + let mut mock = MockTestClient::new(); + mock.expect_start() + .return_once(move |extensions, path, rx, _options, header, params| { + // Verify all the parameters. We should have a couple of tests + // that do this, but should avoid doing so in every test. + assert!( + matches!(extensions.get::(), Some(m) if m.service() == "google.storage.v2.Storage" && m.method() == "BidiWriteObject") + ); + assert_eq!(path.path(), "/google.storage.v2.Storage/BidiWriteObject"); + assert_eq!(header, *crate::storage::info::X_GOOG_API_CLIENT_HEADER); + assert_eq!(params, "bucket=projects/_/buckets/test-bucket"); + save.lock().expect("never poisoned").push(rx); + Err(permanent_error()) + }); + let client = SharedMockClient::new(mock); + + let mut connector = Connector::new(test_options(), client); + + let mut req = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default(), + params: None, + }; + req.spec = crate::model::WriteObjectSpec { + resource: Some(crate::model::Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + ..Default::default() + }), + ..Default::default() + }; + + let err = connector.connect_open(req).await.unwrap_err(); + assert!(err.status().is_some(), "{err:?}"); + + let mut rx = { + let mut guard = receivers.lock().expect("never poisoned"); + let rx = guard.pop().expect("at least one receiver"); + assert!(guard.is_empty(), "{receivers:?}"); + rx + }; + + let first = rx.recv().await.expect("non-empty request"); + let spec = match first.first_message.as_ref().unwrap() { + crate::google::storage::v2::bidi_write_object_request::FirstMessage::WriteObjectSpec(s) => s, + _ => panic!("Expected WriteObjectSpec"), + }; + assert_eq!( + spec.resource.as_ref().map(|s| s.bucket.as_str()), + Some("projects/_/buckets/test-bucket") + ); + assert_eq!( + spec.resource.as_ref().map(|s| s.name.as_str()), + Some("test-object") + ); + + Ok(()) + } + + #[tokio::test] + async fn start_error_with_routing() -> Result<()> { + use super::super::mocks::{MockTestClient, SharedMockClient}; + use super::super::tests::permanent_error; + use gaxi::grpc::tonic::GrpcMethod; + use std::sync::Mutex; + + let receivers = Arc::new(Mutex::new(Vec::new())); + let save = receivers.clone(); + let mut mock = MockTestClient::new(); + mock.expect_start() + .return_once(move |extensions, path, rx, _options, header, params| { + // Verify all the parameters. We should have a couple of tests + // that do this, but should avoid doing so in every test. + assert!( + matches!( + extensions.get::(), + Some(m) if m.service() == "google.storage.v2.Storage" && m.method() == "BidiWriteObject" + ) + ); + assert_eq!(path.path(), "/google.storage.v2.Storage/BidiWriteObject"); + assert_eq!(header, *crate::storage::info::X_GOOG_API_CLIENT_HEADER); + let mut split = params.split('&').collect::>(); + split.sort(); + assert_eq!(split, vec!["bucket=projects/_/buckets/test-bucket", "routing_token=test-routing-token"]); + save.lock().expect("never poisoned").push(rx); + + Err(permanent_error()) + }); + let client = SharedMockClient::new(mock); + + let req = crate::model_ext::ReopenAppendableObjectRequest { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + generation: 345678, + routing_token: Some("test-routing-token".to_string()), + if_metageneration_match: None, + if_metageneration_not_match: None, + params: None, + write_handle: Some(bytes::Bytes::from_static(b"test-write-handle")), + }; + + let mut connector = Connector::new(test_options(), client); + let err = connector.connect_reopen(req).await.unwrap_err(); + assert!(err.status().is_some(), "{err:?}"); + + let mut rx = { + let mut guard = receivers.lock().expect("never poisoned"); + let rx = guard.pop().expect("at least one receiver"); + assert!(guard.is_empty(), "{guard:?}"); + rx + }; + let first = rx.recv().await.expect("non-empty request"); + let spec = match first.first_message.as_ref().unwrap() { + crate::google::storage::v2::bidi_write_object_request::FirstMessage::AppendObjectSpec(s) => s, + _ => panic!("Expected AppendObjectSpec"), + }; + assert_eq!(spec.bucket.as_str(), "projects/_/buckets/test-bucket"); + assert_eq!(spec.object.as_str(), "test-object"); + assert_eq!(spec.generation, 345678); + assert_eq!( + spec.write_handle.as_ref().map(|h| h.handle.clone()), + Some(bytes::Bytes::from_static(b"test-write-handle")) + ); + Ok(()) + } + + #[tokio::test] + async fn start_redirect_then_error() -> Result<()> { + use super::super::mocks::{MockTestClient, SharedMockClient}; + use super::super::tests::{permanent_error, redirect_status}; + use std::sync::Mutex; + + let mut seq = mockall::Sequence::new(); + let mut mock = MockTestClient::new(); + let receivers = Arc::new(Mutex::new(Vec::new())); + let save = receivers.clone(); + mock.expect_start() + .times(1) + .in_sequence(&mut seq) + .returning(move |_, _, rx, _, _, _| { + save.lock().expect("never poisoned").push(rx); + Ok(Err(redirect_status("r1"))) + }); + let save = receivers.clone(); + mock.expect_start() + .times(1) + .in_sequence(&mut seq) + .returning(move |_, _, rx, _, _, _| { + save.lock().expect("never poisoned").push(rx); + Err(permanent_error()) + }); + let client = SharedMockClient::new(mock); + + let mut connector = Connector::new(test_options(), client); + + let mut req = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default(), + params: None, + }; + req.spec = crate::model::WriteObjectSpec { + resource: Some(crate::model::Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + ..Default::default() + }), + ..Default::default() + }; + + let err = connector.connect_open(req).await.unwrap_err(); + assert_eq!(err.status(), permanent_error().status(), "{err:?}"); + + let got = connector.spec.lock().expect("never poisoned").clone(); + match got { + AppendObjectSpecState::Write(_, _) => panic!("Should be Append"), + AppendObjectSpecState::Append(got) => { + assert_eq!(got.routing_token.as_deref(), Some("r1")); + } + } + + let mut rx = receivers + .lock() + .expect("never poisoned") + .pop() + .expect("at least two receiver"); + // We pop the receivers, so this is the second receiver. This receiver should include an spec with the redirect options. + let got = rx.recv().await.expect("at least one request sent"); + let want = crate::google::storage::v2::AppendObjectSpec { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + routing_token: Some("r1".to_string()), + ..crate::google::storage::v2::AppendObjectSpec::default() + }; + let spec = match got.first_message.unwrap() { + crate::google::storage::v2::bidi_write_object_request::FirstMessage::AppendObjectSpec(s) => s, + _ => panic!("Expected AppendObjectSpec"), + }; + assert_eq!(spec.bucket, want.bucket); + assert_eq!(spec.object, want.object); + assert_eq!(spec.routing_token, want.routing_token); + + let mut rx = receivers + .lock() + .expect("never poisoned") + .pop() + .expect("at least two receiver"); + // We pop the receivers, so this is the second receiver. This receiver should include an spec with the redirect options. + let got = rx.recv().await.expect("at least one request sent"); + let want = crate::google::storage::v2::WriteObjectSpec { + resource: Some(crate::google::storage::v2::Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + ..Default::default() + }), + ..crate::google::storage::v2::WriteObjectSpec::default() + }; + let spec = match got.first_message.unwrap() { + crate::google::storage::v2::bidi_write_object_request::FirstMessage::WriteObjectSpec(s) => s, + _ => panic!("Expected WriteObjectSpec"), + }; + assert_eq!( + spec.resource.as_ref().unwrap().bucket, + want.resource.as_ref().unwrap().bucket + ); + assert_eq!( + spec.resource.as_ref().unwrap().name, + want.resource.as_ref().unwrap().name + ); + + Ok(()) + } + + #[tokio::test] + async fn start_immediately_closed() -> Result<()> { + use super::super::mocks::{MockTestClient, SharedMockClient}; + use gaxi::grpc::tonic::Response as TonicResponse; + use gaxi::grpc::tonic::Result as TonicResult; + + let (tx1, rx1) = tokio::sync::mpsc::channel::>(5); + let stream1 = TonicResponse::from(rx1); + drop(tx1); + let (tx2, rx2) = tokio::sync::mpsc::channel::>(5); + let stream2 = TonicResponse::from(rx2); + + let mut seq = mockall::Sequence::new(); + let mut mock = MockTestClient::new(); + mock.expect_start() + .times(1) + .in_sequence(&mut seq) + .return_once(move |_, _, _, _, _, _| Ok(Ok(stream1))); + mock.expect_start() + .times(1) + .in_sequence(&mut seq) + .return_once(move |_, _, _, _, _, _| Ok(Ok(stream2))); + let client = SharedMockClient::new(mock); + + let mut connector = Connector::new(test_options(), client); + + let initial = BidiWriteObjectResponse { + write_status: Some( + crate::google::storage::v2::bidi_write_object_response::WriteStatus::Resource( + crate::google::storage::v2::Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + generation: 123456, + ..crate::google::storage::v2::Object::default() + }, + ), + ), + write_handle: Some(crate::google::storage::v2::BidiWriteHandle { + handle: bytes::Bytes::from_static(b"test-handle-open"), + }), + ..BidiWriteObjectResponse::default() + }; + tx2.send(Ok(initial.clone())).await?; + + let mut req = OpenAppendableObjectRequest { + spec: crate::model::WriteObjectSpec::default(), + params: None, + }; + req.spec = crate::model::WriteObjectSpec { + resource: Some(crate::model::Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + ..Default::default() + }), + ..Default::default() + }; + + let (response, _connection) = connector.connect_open(req).await?; + assert_eq!(response, initial); + + let guard = connector.spec.lock().expect("never poisoned"); + if let AppendObjectSpecState::Append(s) = &*guard { + assert!(s.routing_token.is_none(), "{s:?}"); + assert_eq!(s.generation, 123456, "{s:?}"); + assert_eq!( + s.write_handle.as_ref().map(|h| h.handle.clone()), + Some(bytes::Bytes::from_static(b"test-handle-open")) + ); + } else { + panic!("Expected AppendObjectSpecState::Append"); + } + drop(tx2); + + Ok(()) + } +} diff --git a/src/storage/src/storage/bidi_write/mocks.rs b/src/storage/src/storage/bidi_write/mocks.rs new file mode 100644 index 0000000000..315603b4c8 --- /dev/null +++ b/src/storage/src/storage/bidi_write/mocks.rs @@ -0,0 +1,84 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// TODO(#5716): Lift to shared bidi module. + +use super::connector::Connector; +use super::tests::test_options; +use super::{Client, Receiver, RequestOptions, TonicStreaming}; +use crate::google::storage::v2::{BidiWriteObjectRequest, BidiWriteObjectResponse}; +use gaxi::grpc::tonic::{Extensions, Response as TonicResponse, Result as TonicResult}; +use std::sync::Arc; + +// mockall mocks are not `Clone` and we need a thing that can be cloned. +// The solution is to wrap the mock in a thing that implements the right +// trait. +#[derive(Clone, Debug)] +pub struct SharedMockClient(pub(crate) Arc); + +impl SharedMockClient { + pub fn new(mock: MockTestClient) -> Self { + Self(Arc::new(mock)) + } +} + +impl Client for SharedMockClient { + type Stream = MockStream; + + async fn start( + &self, + extensions: Extensions, + path: http::uri::PathAndQuery, + rx: Receiver, + options: &RequestOptions, + api_client_header: &'static str, + request_params: &str, + ) -> crate::Result>> { + self.0.start( + extensions, + path, + rx, + options, + api_client_header, + request_params, + ) + } +} + +impl TonicStreaming for Receiver> { + async fn next_message(&mut self) -> TonicResult> { + self.recv().await.transpose() + } +} + +#[mockall::automock] +pub trait TestClient: std::fmt::Debug { + fn start( + &self, + extensions: Extensions, + path: http::uri::PathAndQuery, + rx: Receiver, + options: &RequestOptions, + api_client_header: &'static str, + request_params: &str, + ) -> crate::Result>>; +} + +pub type MockStream = Receiver>; + +pub fn mock_connector(mock: MockTestClient) -> Connector { + let client = SharedMockClient::new(mock); + + Connector::new(test_options(), client.clone()) +} diff --git a/src/storage/src/storage/bidi_write/redirect.rs b/src/storage/src/storage/bidi_write/redirect.rs index e17a4a745a..7ea054398d 100644 --- a/src/storage/src/storage/bidi_write/redirect.rs +++ b/src/storage/src/storage/bidi_write/redirect.rs @@ -14,6 +14,7 @@ // TODO(#5716): Lift to shared bidi module +use super::connector::AppendObjectSpecState; use crate::Error; use crate::google::rpc::Status as RpcStatus; use crate::google::storage::v2::{AppendObjectSpec, BidiWriteObjectRedirectedError}; @@ -24,8 +25,7 @@ use google_cloud_gax::error::rpc::Code; use prost::Message; use std::sync::{Arc, Mutex}; -#[allow(dead_code)] -pub fn handle_redirect(spec: Arc>, status: Status) -> crate::Error { +pub fn handle_redirect(state: Arc>, status: Status) -> crate::Error { let Ok(rpc_status) = RpcStatus::decode(status.details()) else { return to_gax_error(status); }; @@ -34,12 +34,47 @@ pub fn handle_redirect(spec: Arc>, status: Status) -> cr .into_iter() .find_map(|d| d.to_msg::().ok()) { - let mut guard = spec.lock().expect("never poisoned"); - guard.routing_token = redirect.routing_token; - guard.write_handle = redirect.write_handle; - if let Some(generation) = redirect.generation { - guard.generation = generation; - } + let mut guard = state.lock().expect("never poisoned"); + + let new_state = match &*guard { + AppendObjectSpecState::Write(spec, _) => { + if let Some(generation) = redirect.generation { + let bucket = spec + .resource + .as_ref() + .map(|r| r.bucket.clone()) + .unwrap_or_default(); + let object = spec + .resource + .as_ref() + .map(|r| r.name.clone()) + .unwrap_or_default(); + let new_spec = AppendObjectSpec { + bucket, + object, + generation, + if_metageneration_match: spec.if_metageneration_match, + if_metageneration_not_match: spec.if_metageneration_not_match, + routing_token: redirect.routing_token, + write_handle: redirect.write_handle, + }; + AppendObjectSpecState::Append(new_spec) + } else { + AppendObjectSpecState::Write(spec.clone(), redirect.routing_token) + } + } + AppendObjectSpecState::Append(spec) => { + let mut new_spec = spec.clone(); + new_spec.routing_token = redirect.routing_token; + new_spec.write_handle = redirect.write_handle; + if let Some(g) = redirect.generation { + new_spec.generation = g; + } + AppendObjectSpecState::Append(new_spec) + } + }; + + *guard = new_state; } to_gax_error(status) } @@ -106,16 +141,65 @@ mod tests { write_handle: Some(BidiWriteHandle { handle: bytes::Bytes::from_static(b"initial-handle"), }), + if_metageneration_match: Some(10), + if_metageneration_not_match: Some(20), + }; + let state = Arc::new(Mutex::new(AppendObjectSpecState::Append(spec))); + + let got = handle_redirect(state.clone(), status); + assert!(got.status().is_some(), "{got:?}"); + let guard = state.lock().expect("not poisoned"); + if let AppendObjectSpecState::Append(ref spec) = *guard { + assert_eq!(spec.routing_token.as_deref(), routing); + assert_eq!(spec.write_handle, write_handle); + assert_eq!(spec.generation, 42); + assert_eq!(spec.if_metageneration_match, Some(10)); + assert_eq!(spec.if_metageneration_not_match, Some(20)); + } else { + panic!("Expected AppendObjectSpecState::Append"); + } + } + + #[test] + fn preserves_preconditions_on_write() { + use crate::google::storage::v2::WriteObjectSpec; + let redirect = BidiWriteObjectRedirectedError { + routing_token: Some("new-routing".to_string()), + write_handle: Some(BidiWriteHandle { + handle: bytes::Bytes::from_static(b"new-handle"), + }), + generation: Some(42), + }; + let redirect = prost_types::Any::from_msg(&redirect).unwrap(); + let status = RpcStatus { + code: Code::Aborted as i32, + message: "test-only".to_string(), + details: vec![redirect], + }; + let details = bytes::Bytes::from_owner(status.encode_to_vec()); + let status = Status::with_details(Code::Aborted, "test-only", details); + + let write_spec = WriteObjectSpec { + if_metageneration_match: Some(11), + if_metageneration_not_match: Some(22), ..Default::default() }; - let spec = Arc::new(Mutex::new(spec)); + let state = Arc::new(Mutex::new(AppendObjectSpecState::Write( + Box::new(write_spec), + None, + ))); - let got = handle_redirect(spec.clone(), status); + let got = handle_redirect(state.clone(), status); assert!(got.status().is_some(), "{got:?}"); - let guard = spec.lock().expect("not poisoned"); - assert_eq!(guard.routing_token.as_deref(), routing); - assert_eq!(guard.write_handle, write_handle); - assert_eq!(guard.generation, 42); + let guard = state.lock().expect("not poisoned"); + if let AppendObjectSpecState::Append(ref spec) = *guard { + assert_eq!(spec.routing_token.as_deref(), Some("new-routing")); + assert_eq!(spec.generation, 42); + assert_eq!(spec.if_metageneration_match, Some(11)); + assert_eq!(spec.if_metageneration_not_match, Some(22)); + } else { + panic!("Expected AppendObjectSpecState::Append"); + } } #[test] @@ -138,14 +222,18 @@ mod tests { write_handle: Some(initial_handle.clone()), ..Default::default() }; - let spec = Arc::new(Mutex::new(spec)); + let state = Arc::new(Mutex::new(AppendObjectSpecState::Append(spec))); - let got = handle_redirect(spec.clone(), status); + let got = handle_redirect(state.clone(), status); assert!(got.status().is_some(), "{got:?}"); - let guard = spec.lock().expect("not poisoned"); - assert_eq!(guard.routing_token.as_deref(), Some("initial-token")); - assert_eq!(guard.write_handle, Some(initial_handle)); - assert_eq!(guard.generation, 1); + let guard = state.lock().expect("not poisoned"); + if let AppendObjectSpecState::Append(ref spec) = *guard { + assert_eq!(spec.routing_token.as_deref(), Some("initial-token")); + assert_eq!(spec.write_handle, Some(initial_handle)); + assert_eq!(spec.generation, 1); + } else { + panic!("Expected AppendObjectSpecState::Append"); + } } #[test_case(permanent_error(), false)] diff --git a/src/storage/src/storage/bidi_write/retry_redirect.rs b/src/storage/src/storage/bidi_write/retry_redirect.rs index 993292ca02..b43c307c21 100644 --- a/src/storage/src/storage/bidi_write/retry_redirect.rs +++ b/src/storage/src/storage/bidi_write/retry_redirect.rs @@ -33,13 +33,11 @@ use std::sync::Arc; /// the decorator has no effect. If the error is "permanent", but happens to be /// a redirect, then it is treated as retryable. #[derive(Clone, Debug)] -#[allow(dead_code)] pub struct RetryRedirect { inner: T, } impl RetryRedirect { - #[allow(dead_code)] pub fn new(inner: T) -> Self { Self { inner } } From a763c07c6a4471335ab1bae5d283f9fcd5f58197 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Wed, 8 Jul 2026 01:48:26 +0000 Subject: [PATCH 2/3] feat(storage): implement appendable upload background worker --- src/storage/src/storage/bidi_write.rs | 1 + src/storage/src/storage/bidi_write/worker.rs | 252 +++++++++++++++++++ 2 files changed, 253 insertions(+) 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 ee197bfcaf..8535580874 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -20,6 +20,7 @@ pub(crate) mod connector; mod redirect; mod retry_redirect; pub(crate) mod stub; +mod worker; use crate::google::storage::v2::{BidiWriteObjectRequest, BidiWriteObjectResponse}; use crate::request_options::RequestOptions; diff --git a/src/storage/src/storage/bidi_write/worker.rs b/src/storage/src/storage/bidi_write/worker.rs new file mode 100644 index 0000000000..d8b7880ed2 --- /dev/null +++ b/src/storage/src/storage/bidi_write/worker.rs @@ -0,0 +1,252 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// TODO(#5716): Lift to shared bidi module + +use super::connector::Connection; +use super::{Client, TonicStreaming}; +use crate::error::WriteError; +use crate::google::storage::v2::{BidiWriteObjectRequest, BidiWriteObjectResponse}; +use gaxi::grpc::tonic::Result as TonicResult; +use std::sync::Arc; +use tokio::sync::mpsc::Receiver; +use tokio::sync::oneshot; + +type WriteResult = std::result::Result; +type LoopResult = std::result::Result>; + +/// The intent sent from the user to the background worker. +pub enum UploadIntent { + /// Append a chunk of data. + Append(BidiWriteObjectRequest), + /// Flush the data to the server. + Flush( + BidiWriteObjectRequest, + oneshot::Sender>, + ), + /// Finalize the upload. + Finalize( + BidiWriteObjectRequest, + oneshot::Sender>, + ), +} + +/// The background worker that manages the live gRPC stream. +pub struct Worker { + _connector: super::connector::Connector, + pending_flushes: + std::collections::VecDeque>>, +} + +impl Worker { + pub fn new(connector: super::connector::Connector) -> Self { + Self { + _connector: connector, + pending_flushes: std::collections::VecDeque::new(), + } + } +} + +impl Worker +where + C: Client + Clone + 'static, + ::Stream: TonicStreaming, +{ + pub async fn run( + mut self, + connection: Connection, + mut requests: Receiver, + ) -> LoopResult<()> { + let (mut rx, mut tx) = (connection.rx, connection.tx); + + let error = loop { + tokio::select! { + m = rx.next_message() => { + match self.handle_response(m) { + // Successful end of stream, return without error. + None => break None, + // An unrecoverable in the stream or its data, return + // the error. + Some(Err(e)) => break Some(e), + // New message on the stream handled successfully, + // continue. + Some(Ok(None)) => {}, + // The stream reconnected successfully, update the local + // variables and continue. + Some(Ok(Some(connection))) => { + (rx, tx) = (connection.rx, connection.tx); + } + } + }, + intent = requests.recv() => { + match intent { + Some(UploadIntent::Append(request)) => { + if let Err(e) = tx.send(request).await { + break Some(Arc::new(crate::Error::io(e))); + } + } + Some(UploadIntent::Flush(request, sender)) => { + self.pending_flushes.push_back(sender); + if let Err(e) = tx.send(request).await { + break Some(Arc::new(crate::Error::io(e))); + } + } + Some(UploadIntent::Finalize(request, sender)) => { + self.pending_flushes.push_back(sender); + if let Err(e) = tx.send(request).await { + break Some(Arc::new(crate::Error::io(e))); + } + } + None => { + break None; + } + } + } + } + }; + + if let Some(e) = error.clone() { + for sender in self.pending_flushes.drain(..) { + let _ = sender.send(Err(crate::Error::io(e.to_string()))); + } + return Err(e); + } + + Ok(()) + } + + pub fn handle_response( + &mut self, + message: TonicResult>, + ) -> Option>>> { + let response = match message { + Ok(Some(msg)) => msg, + Ok(None) => { + if !self.pending_flushes.is_empty() { + return Some(Err(Arc::new(crate::Error::io( + "stream closed unexpectedly", + )))); + } + return None; + } + Err(e) => return Some(Err(Arc::new(crate::Error::io(e)))), + }; + self.handle_response_success(response); + // TODO(#5716): Update when implementing reconnect logic. + Some(Ok(None)) + } + + pub fn handle_response_success(&mut self, response: BidiWriteObjectResponse) { + if let Some(sender) = self.pending_flushes.pop_front() { + let _ = sender.send(Ok(response)); + } else { + // Log unprompted server responses. + tracing::debug!( + "Received unprompted BidiWriteObjectResponse from server: {:?}", + response + ); + } + } +} + +#[cfg(test)] +mod tests { + use super::super::mocks::{MockTestClient, mock_connector}; + use super::*; + use crate::google::storage::v2::{ + BidiWriteObjectRequest, BidiWriteObjectResponse, bidi_write_object_response::WriteStatus, + }; + use tokio::sync::mpsc; + use tokio::sync::oneshot; + + #[tokio::test] + async fn run_immediately_closed() -> anyhow::Result<()> { + let (request_tx, _request_rx) = mpsc::channel(1); + let (response_tx, response_rx) = mpsc::channel(10); + let (tx, rx) = mpsc::channel(1); + let connection = Connection::new(request_tx, response_rx); + + let mut mock = MockTestClient::new(); + mock.expect_start().never(); + + let connector = mock_connector(mock); + let worker = Worker::new(connector); + let handle = tokio::spawn(worker.run(connection, rx)); + + drop(response_tx); + drop(tx); + handle.await??; + Ok(()) + } + + #[tokio::test] + async fn run_flush_response() -> anyhow::Result<()> { + let (request_tx, mut request_rx) = mpsc::channel(1); + let (response_tx, response_rx) = mpsc::channel(10); + let (tx, rx) = mpsc::channel(1); + let connection = Connection::new(request_tx, response_rx); + + let mut mock = MockTestClient::new(); + mock.expect_start().never(); + + let connector = mock_connector(mock); + let worker = Worker::new(connector); + let handle = tokio::spawn(worker.run(connection, rx)); + + let (flush_tx, flush_rx) = oneshot::channel(); + let flush_request = BidiWriteObjectRequest { + flush: true, + ..Default::default() + }; + tx.send(UploadIntent::Flush(flush_request.clone(), flush_tx)) + .await?; + + // The worker should send the request to the stream. + let stream_req = request_rx.recv().await.unwrap(); + assert!(stream_req.flush); + + // The server responds + let server_resp = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(100)), + ..Default::default() + }; + response_tx.send(Ok(server_resp.clone())).await?; + + // The flush sender should get the response + let received_resp = flush_rx.await??; + assert_eq!(received_resp.write_status, server_resp.write_status); + + drop(response_tx); + drop(tx); + handle.await??; + Ok(()) + } + + #[tokio::test] + async fn run_stop_on_closed_requests() -> anyhow::Result<()> { + let (request_tx, _request_rx) = mpsc::channel(1); + let (_response_tx, response_rx) = mpsc::channel(10); + let (tx, rx) = mpsc::channel(1); + let connection = Connection::new(request_tx, response_rx); + + let mut mock = MockTestClient::new(); + mock.expect_start().never(); + + let connector = mock_connector(mock); + let worker = Worker::new(connector); + drop(tx); + worker.run(connection, rx).await?; + Ok(()) + } +} From 31a22b9f0b4605201203fe6de03f9949b9c6356c Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Wed, 8 Jul 2026 01:49:38 +0000 Subject: [PATCH 3/3] feat(storage): implement appendable upload transport --- src/storage/src/storage/bidi_write.rs | 1 + .../src/storage/bidi_write/transport.rs | 438 ++++++++++++++++++ 2 files changed, 439 insertions(+) create mode 100644 src/storage/src/storage/bidi_write/transport.rs diff --git a/src/storage/src/storage/bidi_write.rs b/src/storage/src/storage/bidi_write.rs index 8535580874..152048342d 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -20,6 +20,7 @@ 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}; diff --git a/src/storage/src/storage/bidi_write/transport.rs b/src/storage/src/storage/bidi_write/transport.rs new file mode 100644 index 0000000000..c479c2a9d6 --- /dev/null +++ b/src/storage/src/storage/bidi_write/transport.rs @@ -0,0 +1,438 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// TODO(#5716): Lift to shared bidi module + +use super::worker::UploadIntent; +use crate::google::storage::v2::ObjectChecksums; +use crate::google::storage::v2::{ + BidiWriteObjectRequest, ChecksummedData, bidi_write_object_request::Data, + bidi_write_object_response::WriteStatus, +}; +use crate::stub::AppendableObjectWriter; +use bytes::Bytes; +use gaxi::prost::FromProto; +use tokio::sync::mpsc::Sender; +use tokio::sync::oneshot; + +/// The high-level transport adapter for bidirectional streaming writes. +/// Translates user calls into gRPC requests and delegates them to the background worker. +#[derive(Debug)] +pub struct AppendableObjectWriterTransport { + tx: Sender, + generation: i64, + persisted_size: i64, + write_offset: i64, + running_crc32c: u32, +} + +impl AppendableObjectWriterTransport { + pub async fn new_open( + mut connector: super::connector::Connector, + req: crate::model_ext::OpenAppendableObjectRequest, + ) -> crate::Result + where + T: super::Client + Clone + Sync + Send + 'static, + ::Stream: super::TonicStreaming + Send + Sync, + { + let (initial, connection) = connector.connect_open(req).await?; + Self::start_worker(connector, initial, connection, 0) + } + + pub async fn new_reopen( + mut connector: super::connector::Connector, + req: crate::model_ext::ReopenAppendableObjectRequest, + ) -> crate::Result + where + T: super::Client + Clone + Sync + Send + 'static, + ::Stream: super::TonicStreaming + Send + Sync, + { + let generation = req.generation; + let (initial, connection) = connector.connect_reopen(req).await?; + Self::start_worker(connector, initial, connection, generation) + } + + fn start_worker( + connector: super::connector::Connector, + initial: crate::google::storage::v2::BidiWriteObjectResponse, + connection: super::connector::Connection<::Stream>, + mut generation: i64, + ) -> crate::Result + where + T: super::Client + Clone + Sync + Send + 'static, + ::Stream: super::TonicStreaming + Send + Sync, + { + let mut persisted_size = 0; + if let Some(WriteStatus::Resource(r)) = initial.write_status.as_ref() { + generation = r.generation; + persisted_size = r.size; + } else if let Some(WriteStatus::PersistedSize(s)) = initial.write_status.as_ref() { + persisted_size = *s; + } + + let mut running_crc32c = 0; + if let Some(crc) = initial + .persisted_data_checksums + .as_ref() + .and_then(|c| c.crc32c) + { + running_crc32c = crc; + } + + let (tx, rx) = tokio::sync::mpsc::channel(100); + let worker = super::worker::Worker::new(connector); + tokio::spawn(worker.run(connection, rx)); + + Ok(Self { + tx, + generation, + persisted_size, + write_offset: persisted_size, + running_crc32c, + }) + } +} + +impl AppendableObjectWriter for AppendableObjectWriterTransport { + async fn append(&mut self, chunk: Bytes) -> crate::Result<()> { + let length = chunk.len() as i64; + let crc32c = crc32c::crc32c(&chunk); + + let new_crc32c = crc32c::crc32c_append(self.running_crc32c, &chunk); + + let request = BidiWriteObjectRequest { + write_offset: self.write_offset, + data: Some(Data::ChecksummedData(ChecksummedData { + content: chunk, + crc32c: Some(crc32c), + })), + ..BidiWriteObjectRequest::default() + }; + + self.tx + .send(UploadIntent::Append(request)) + .await + .map_err(|e| crate::Error::io(e.to_string()))?; + + self.write_offset += length; + self.running_crc32c = new_crc32c; + + Ok(()) + } + + async fn flush(&mut self) -> crate::Result { + let (sender, receiver) = oneshot::channel(); + let request = BidiWriteObjectRequest { + flush: true, + state_lookup: true, + write_offset: self.write_offset, + ..BidiWriteObjectRequest::default() + }; + + self.tx + .send(UploadIntent::Flush(request, sender)) + .await + .map_err(|e| crate::Error::io(e.to_string()))?; + + let response = receiver + .await + .map_err(|e| crate::Error::io(e.to_string()))??; + let size = match response.write_status { + Some(WriteStatus::PersistedSize(s)) => s, + Some(WriteStatus::Resource(r)) => r.size, + None => return Err(crate::Error::io("flush response missing write_status")), + }; + self.persisted_size = size; + Ok(size) + } + + async fn finalize(self) -> crate::Result { + let (sender, receiver) = oneshot::channel(); + let request = BidiWriteObjectRequest { + finish_write: true, + flush: true, + write_offset: self.write_offset, + object_checksums: Some(ObjectChecksums { + crc32c: Some(self.running_crc32c), + md5_hash: vec![].into(), + }), + ..BidiWriteObjectRequest::default() + }; + + self.tx + .send(UploadIntent::Finalize(request, sender)) + .await + .map_err(|e| crate::Error::io(e.to_string()))?; + + let response = receiver + .await + .map_err(|e| crate::Error::io(e.to_string()))??; + + let resource = match response.write_status { + Some(WriteStatus::Resource(r)) => r, + _ => return Err(crate::Error::io("finalize did not return a resource")), + }; + + let object = FromProto::cnv(resource) + .map_err(|_| crate::Error::deser("converting resource to object"))?; + + Ok(object) + } + + async fn close(mut self) -> crate::Result { + self.flush().await + } + + fn generation(&self) -> i64 { + self.generation + } + + fn persisted_size(&self) -> i64 { + self.persisted_size + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::google::storage::v2::{ + BidiWriteObjectResponse, Object, bidi_write_object_response::WriteStatus, + }; + use tokio::sync::mpsc; + + #[tokio::test] + async fn success() -> anyhow::Result<()> { + let (tx, mut rx) = mpsc::channel(1); + let transport = AppendableObjectWriterTransport { + tx, + write_offset: 0, + running_crc32c: 0, + generation: 123456, + persisted_size: 0, + }; + + let handle = tokio::spawn(async move { + let mut transport = transport; + transport.append(bytes::Bytes::from("hello")).await.unwrap(); + transport.flush().await.unwrap(); + transport.finalize().await.unwrap(); + }); + + // Verify append + let intent = rx.recv().await.unwrap(); + if let UploadIntent::Append(req) = intent { + assert_eq!(req.write_offset, 0); + assert!(!req.finish_write); + } else { + panic!("expected Append"); + } + + // Verify flush + let intent = rx.recv().await.unwrap(); + if let UploadIntent::Flush(req, sender) = intent { + assert!(req.flush); + let resp = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(5)), + ..Default::default() + }; + sender.send(Ok(resp)).unwrap(); + } else { + panic!("expected Flush"); + } + + // Verify finalize + let intent = rx.recv().await.unwrap(); + if let UploadIntent::Finalize(req, sender) = intent { + assert!(req.finish_write); + let object = Object { + bucket: "projects/_/buckets/test-bucket".into(), + name: "test-object".into(), + size: 5, + generation: 123456, + ..Default::default() + }; + let resp = BidiWriteObjectResponse { + write_status: Some(WriteStatus::Resource(object)), + ..Default::default() + }; + sender.send(Ok(resp)).unwrap(); + } else { + panic!("expected Finalize"); + } + + handle.await?; + Ok(()) + } + + #[tokio::test] + async fn connect_error_open() -> anyhow::Result<()> { + use super::super::mocks::{MockTestClient, mock_connector}; + use super::super::tests::permanent_error; + use crate::model_ext::OpenAppendableObjectRequest; + + let mut mock = MockTestClient::new(); + mock.expect_start() + .return_once(move |_, _, _, _, _, _| Err(permanent_error())); + let connector = mock_connector(mock); + let mut req = OpenAppendableObjectRequest { + spec: Default::default(), + params: None, + }; + req.spec.resource = Some( + crate::model::Object::default() + .set_bucket("projects/_/buckets/test-bucket") + .set_name("test-object"), + ); + + let err = AppendableObjectWriterTransport::new_open(connector, req) + .await + .unwrap_err(); + assert_eq!(err.status(), permanent_error().status(), "{err:?}"); + Ok(()) + } + + #[tokio::test] + async fn connect_error_reopen() -> anyhow::Result<()> { + use super::super::mocks::{MockTestClient, mock_connector}; + use super::super::tests::permanent_error; + use crate::model_ext::ReopenAppendableObjectRequest; + + let mut mock = MockTestClient::new(); + mock.expect_start() + .return_once(move |_, _, _, _, _, _| Err(permanent_error())); + let connector = mock_connector(mock); + let req = ReopenAppendableObjectRequest { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + generation: 123, + if_metageneration_match: None, + if_metageneration_not_match: None, + routing_token: None, + write_handle: None, + params: None, + }; + + let err = AppendableObjectWriterTransport::new_reopen(connector, req) + .await + .unwrap_err(); + assert_eq!(err.status(), permanent_error().status(), "{err:?}"); + Ok(()) + } + + #[tokio::test] + async fn write_error() -> anyhow::Result<()> { + let (tx, rx) = mpsc::channel(1); + let mut transport = AppendableObjectWriterTransport { + tx, + write_offset: 0, + running_crc32c: 0, + generation: 123456, + persisted_size: 0, + }; + + // Simulate an early stream closure / dropped rx (like a worker dying) + drop(rx); + let err = transport + .append(bytes::Bytes::from("hello")) + .await + .unwrap_err(); + assert!(err.is_io(), "{err:?}"); + + // Assert that state was NOT modified due to the error + assert_eq!(transport.write_offset, 0); + assert_eq!(transport.running_crc32c, 0); + + Ok(()) + } + + #[tokio::test] + async fn finalize_error() -> anyhow::Result<()> { + let (tx, mut rx) = mpsc::channel(1); + let transport = AppendableObjectWriterTransport { + tx, + write_offset: 0, + running_crc32c: 0, + generation: 123456, + persisted_size: 0, + }; + + let handle = tokio::spawn(async move { + let transport = transport; + transport.finalize().await + }); + + // Verify finalize intent + let intent = rx.recv().await.unwrap(); + if let UploadIntent::Finalize(_, sender) = intent { + // Respond with an invalid WriteStatus (not Resource) + let resp = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(5)), + ..Default::default() + }; + sender.send(Ok(resp)).unwrap(); + } else { + panic!("expected Finalize"); + } + + let err = handle.await?.unwrap_err(); + assert!(err.is_io(), "{err:?}"); + + Ok(()) + } + + #[tokio::test] + async fn new_reopen_initial_persisted_size() -> anyhow::Result<()> { + use super::super::mocks::{MockTestClient, mock_connector}; + use crate::model_ext::ReopenAppendableObjectRequest; + use gaxi::grpc::tonic::Response as TonicResponse; + use gaxi::grpc::tonic::Result as TonicResult; + + let (tx1, rx1) = tokio::sync::mpsc::channel::>(5); + let stream1 = TonicResponse::from(rx1); + + let mut mock = MockTestClient::new(); + mock.expect_start() + .return_once(move |_, _, _, _, _, _| Ok(Ok(stream1))); + let connector = mock_connector(mock); + + let req = ReopenAppendableObjectRequest { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + generation: 123456, + if_metageneration_match: None, + if_metageneration_not_match: None, + routing_token: None, + write_handle: None, + params: None, + }; + + let initial_response = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(1024)), + persisted_data_checksums: Some(crate::google::storage::v2::ObjectChecksums { + crc32c: Some(9999), + md5_hash: vec![].into(), + }), + ..Default::default() + }; + tx1.send(Ok(initial_response)).await?; + + let transport = AppendableObjectWriterTransport::new_reopen(connector, req).await?; + + assert_eq!(transport.generation(), 123456); + assert_eq!(transport.persisted_size(), 1024); + assert_eq!(transport.write_offset, 1024); + assert_eq!(transport.running_crc32c, 9999); + Ok(()) + } +}