From 0c58407e4374bd769fd59d41e0a8b173f6638886 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 | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/storage/src/storage/bidi_write.rs b/src/storage/src/storage/bidi_write.rs index daea72d79b..414e3b57a3 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -33,7 +33,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, @@ -50,7 +49,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( @@ -144,14 +142,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() From a08e1ec0afe3efa228a0106e50d9d6f8dfe28694 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 414e3b57a3..c93364e8a8 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -23,6 +23,7 @@ mod retry_redirect; #[allow(dead_code)] pub(crate) mod state; 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 392297e82cac9b1268181053f52733d8de349a56 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Thu, 9 Jul 2026 08:49:44 +0000 Subject: [PATCH 3/3] Incorporate feedback --- src/storage/src/storage/bidi_write.rs | 1 + src/storage/src/storage/bidi_write/worker.rs | 197 +++++++++++++------ 2 files changed, 140 insertions(+), 58 deletions(-) diff --git a/src/storage/src/storage/bidi_write.rs b/src/storage/src/storage/bidi_write.rs index c93364e8a8..286463381b 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -23,6 +23,7 @@ mod retry_redirect; #[allow(dead_code)] pub(crate) mod state; pub(crate) mod stub; +#[allow(dead_code)] mod worker; use crate::google::storage::v2::{BidiWriteObjectRequest, BidiWriteObjectResponse}; diff --git a/src/storage/src/storage/bidi_write/worker.rs b/src/storage/src/storage/bidi_write/worker.rs index d8b7880ed2..2b02d5a95e 100644 --- a/src/storage/src/storage/bidi_write/worker.rs +++ b/src/storage/src/storage/bidi_write/worker.rs @@ -19,23 +19,20 @@ 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>; +type LoopResult = std::result::Result; -/// The intent sent from the user to the background worker. +/// The intent sent from the foreground task 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>, @@ -47,6 +44,8 @@ pub struct Worker { _connector: super::connector::Connector, pending_flushes: std::collections::VecDeque>>, + /// Tracks if the client intends to complete the upload, by sending a Finalize intent. + finalized: bool, } impl Worker { @@ -54,6 +53,7 @@ impl Worker { Self { _connector: connector, pending_flushes: std::collections::VecDeque::new(), + finalized: false, } } } @@ -76,12 +76,13 @@ where match self.handle_response(m) { // Successful end of stream, return without error. None => break None, - // An unrecoverable in the stream or its data, return + // An unrecoverable error 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)) => {}, + // TODO(#5716): Update when implementing reconnect logic. // The stream reconnected successfully, update the local // variables and continue. Some(Ok(Some(connection))) => { @@ -91,21 +92,21 @@ where }, 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); + Some(intent) => { + let request = match intent { + UploadIntent::Append(req) => req, + UploadIntent::Flush(req, sender) => { + self.pending_flushes.push_back(sender); + req + } + UploadIntent::Finalize(req, sender) => { + self.pending_flushes.push_back(sender); + self.finalized = true; + req + } + }; 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))); + break Some(crate::Error::io(e)); } } None => { @@ -116,7 +117,7 @@ where } }; - if let Some(e) = error.clone() { + if let Some(e) = error { for sender in self.pending_flushes.drain(..) { let _ = sender.send(Err(crate::Error::io(e.to_string()))); } @@ -133,17 +134,19 @@ where 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", - )))); + // If the stream is unexpectedly closed by the server before the client + // intends to finalize the upload, treat it as an error to prevent silent + // failures on subsequent client writes. + if !self.pending_flushes.is_empty() || !self.finalized { + return Some(Err(crate::Error::io("stream closed unexpectedly"))); } return None; } - Err(e) => return Some(Err(Arc::new(crate::Error::io(e)))), + Err(e) => return Some(Err(crate::Error::io(e))), }; self.handle_response_success(response); - // TODO(#5716): Update when implementing reconnect logic. + + // TODO(#5716): Implement reconnect logic. Some(Ok(None)) } @@ -170,9 +173,15 @@ mod tests { 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); + type TestWorkerContext = ( + tokio::task::JoinHandle>, + mpsc::Sender, + mpsc::Receiver, + mpsc::Sender>, + ); + + fn spawn_test_worker() -> TestWorkerContext { + 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); @@ -184,25 +193,30 @@ mod tests { let worker = Worker::new(connector); let handle = tokio::spawn(worker.run(connection, rx)); - drop(response_tx); + (handle, tx, request_rx, response_tx) + } + + #[tokio::test] + async fn run_append() -> anyhow::Result<()> { + let (handle, tx, mut request_rx, _response_tx) = spawn_test_worker(); + + let append_request = BidiWriteObjectRequest { + write_offset: 10, + ..Default::default() + }; + tx.send(UploadIntent::Append(append_request)).await?; + + let stream_req = request_rx.recv().await.unwrap(); + assert_eq!(stream_req.write_offset, 10); + 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)); + async fn run_flush() -> anyhow::Result<()> { + let (handle, tx, mut request_rx, response_tx) = spawn_test_worker(); let (flush_tx, flush_rx) = oneshot::channel(); let flush_request = BidiWriteObjectRequest { @@ -212,41 +226,108 @@ mod tests { 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??; + drop(response_tx); 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); + async fn run_finalize() -> anyhow::Result<()> { + let (handle, tx, mut request_rx, response_tx) = spawn_test_worker(); - let mut mock = MockTestClient::new(); - mock.expect_start().never(); + let (finalize_tx, finalize_rx) = oneshot::channel(); + let finalize_request = BidiWriteObjectRequest { + finish_write: true, + ..Default::default() + }; + tx.send(UploadIntent::Finalize( + finalize_request.clone(), + finalize_tx, + )) + .await?; - let connector = mock_connector(mock); - let worker = Worker::new(connector); + let stream_req = request_rx.recv().await.unwrap(); + assert!(stream_req.finish_write); + + let server_resp = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(100)), + ..Default::default() + }; + response_tx.send(Ok(server_resp.clone())).await?; + + let received_resp = finalize_rx.await??; + assert_eq!(received_resp.write_status, server_resp.write_status); + + drop(response_tx); + handle.await??; + Ok(()) + } + + #[tokio::test] + async fn run_stop_on_closed_requests() -> anyhow::Result<()> { + let (handle, tx, _request_rx, _response_tx) = spawn_test_worker(); drop(tx); - worker.run(connection, rx).await?; + handle.await??; + Ok(()) + } + + #[tokio::test] + async fn run_server_closes_unexpectedly() -> anyhow::Result<()> { + let (handle, _tx, _request_rx, response_tx) = spawn_test_worker(); + + // Close the stream from the server side unexpectedly. + drop(response_tx); + + let result = handle.await?; + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "the transport reports an error: stream closed unexpectedly" + ); + + Ok(()) + } + + #[tokio::test] + async fn run_stream_error_during_flush() -> anyhow::Result<()> { + let (handle, tx, mut request_rx, response_tx) = spawn_test_worker(); + + 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?; + + let stream_req = request_rx.recv().await.unwrap(); + assert!(stream_req.flush); + + // Before the server responds, the stream unexpectedly closes. + drop(response_tx); + + let received_resp = flush_rx.await?; + assert!(received_resp.is_err()); + assert_eq!( + received_resp.unwrap_err().to_string(), + "the transport reports an error: the transport reports an error: stream closed unexpectedly" + ); + + let result = handle.await?; + assert!(result.is_err()); Ok(()) } }