From 4ccbe87c8e5460712173c15e6934612bf4c0ef8d Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Tue, 1 Sep 2026 07:11:28 +0000 Subject: [PATCH 1/3] feat(storage): implement 32 MiB replay buffer for appendable upload --- src/storage/src/storage/bidi_write.rs | 2 + .../src/storage/bidi_write/replay_buffer.rs | 300 ++++++++++++++++++ 2 files changed, 302 insertions(+) create mode 100644 src/storage/src/storage/bidi_write/replay_buffer.rs diff --git a/src/storage/src/storage/bidi_write.rs b/src/storage/src/storage/bidi_write.rs index 479a231a4c..d0156f30f5 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -19,6 +19,8 @@ pub(crate) mod connector; #[allow(dead_code)] mod redirect; #[allow(dead_code)] +pub(crate) mod replay_buffer; +#[allow(dead_code)] mod retry_redirect; #[allow(dead_code)] pub(crate) mod state; diff --git a/src/storage/src/storage/bidi_write/replay_buffer.rs b/src/storage/src/storage/bidi_write/replay_buffer.rs new file mode 100644 index 0000000000..df6441bceb --- /dev/null +++ b/src/storage/src/storage/bidi_write/replay_buffer.rs @@ -0,0 +1,300 @@ +// 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. + +//! Retains unacknowledged chunks, trims acknowledged data, and provides chunks +//! for resending upon stream reconnect. + +use crate::google::storage::v2::{ + BidiWriteObjectRequest, ChecksummedData, bidi_write_object_request::Data, +}; +use bytes::Bytes; +use std::collections::VecDeque; + +/// Defines the maximum capacity of the replay buffer in bytes (32 MiB). +pub const MAX_REPLAY_BUFFER_SIZE: usize = 32 * 1024 * 1024; + +/// Represents an unacknowledged data chunk retained in the [`ReplayBuffer`]. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ReplayChunk { + /// Holds the logical starting write offset of this chunk. + pub write_offset: i64, + /// Contains the raw payload bytes. + pub data: Bytes, + /// Stores the precomputed CRC32C checksum of [`data`][Self::data]. + pub crc32c: u32, +} + +impl ReplayChunk { + /// Creates a new replay chunk. + pub fn new(write_offset: i64, data: Bytes, crc32c: u32) -> Self { + Self { + write_offset, + data, + crc32c, + } + } + + /// Returns the ending byte offset (exclusive) of this chunk. + pub fn end_offset(&self) -> i64 { + self.write_offset + self.data.len() as i64 + } + + /// Converts this replay chunk into a [`BidiWriteObjectRequest`] for transmission. + pub fn to_request(&self) -> BidiWriteObjectRequest { + BidiWriteObjectRequest { + write_offset: self.write_offset, + data: Some(Data::ChecksummedData(ChecksummedData { + content: self.data.clone(), + crc32c: Some(self.crc32c), + })), + ..BidiWriteObjectRequest::default() + } + } +} + +/// Manages an in-memory FIFO queue of unacknowledged chunks up to [`MAX_REPLAY_BUFFER_SIZE`]. +#[derive(Debug, Default)] +pub struct ReplayBuffer { + queue: VecDeque, + current_size: usize, +} + +impl ReplayBuffer { + /// Creates a new, empty replay buffer. + pub fn new() -> Self { + Self { + queue: VecDeque::new(), + current_size: 0, + } + } + + /// Enqueues an unacknowledged [`ReplayChunk`] to the replay buffer. + pub fn push(&mut self, chunk: ReplayChunk) { + self.current_size += chunk.data.len(); + self.queue.push_back(chunk); + } + + /// Trims acknowledged chunks up to `persisted_size`. + /// + /// If `persisted_size` lands inside a chunk, that chunk is sliced in-place + /// and its CRC32C is recomputed for the unpersisted sub-slice only. + pub fn acknowledge(&mut self, persisted_size: i64) { + while let Some(front) = self.queue.front() + && front.end_offset() <= persisted_size + { + if let Some(chunk) = self.queue.pop_front() { + self.current_size -= chunk.data.len(); + } + } + + if let Some(front) = self.queue.front_mut() + && front.write_offset < persisted_size + { + let trimmed_bytes = (persisted_size - front.write_offset) as usize; + // SAFETY: persisted_size is guaranteed to be within the bounds of the chunk because front.write_offset < persisted_size and front.end_offset() > persisted_size. + front.data = front.data.slice(trimmed_bytes..); + front.write_offset = persisted_size; + front.crc32c = crc32c::crc32c(&front.data); + self.current_size -= trimmed_bytes; + } + } + + /// Returns `true` if the buffered byte count has reached or exceeded + /// [`MAX_REPLAY_BUFFER_SIZE`]. + pub fn is_full(&self) -> bool { + self.current_size >= MAX_REPLAY_BUFFER_SIZE + } + + /// Returns `true` if the replay buffer contains no chunks. + pub fn is_empty(&self) -> bool { + self.queue.is_empty() + } + + /// Returns the number of chunks currently held in the buffer. + pub fn len(&self) -> usize { + self.queue.len() + } + + /// Returns the total unpersisted byte count currently retained in the buffer. + pub fn current_size(&self) -> usize { + self.current_size + } + + /// Returns an iterator over the unpersisted [`ReplayChunk`]s in FIFO order for replay. + pub fn chunks_to_replay(&self) -> impl Iterator { + self.queue.iter() + } + + /// Clears all chunks from the buffer and resets byte tracking. + pub fn clear(&mut self) { + self.queue.clear(); + self.current_size = 0; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn new_buffer_state() { + // Arrange. + let mut buf = ReplayBuffer::new(); + + // Assert. + assert!(buf.is_empty()); + assert_eq!(buf.len(), 0); + assert_eq!(buf.current_size(), 0); + assert!(!buf.is_full()); + + // Act. + buf.acknowledge(100); + + // Assert. + assert!(buf.is_empty()); + } + + #[test] + fn push_and_acknowledge_full_chunks() { + // Arrange. + let mut buf = ReplayBuffer::new(); + let chunk1 = Bytes::from_static(b"hello "); + let chunk2 = Bytes::from_static(b"world!"); + + // Act. + buf.push(ReplayChunk::new(0, chunk1.clone(), crc32c::crc32c(&chunk1))); + buf.push(ReplayChunk::new(6, chunk2.clone(), crc32c::crc32c(&chunk2))); + + // Assert. + assert_eq!(buf.len(), 2); + assert_eq!(buf.current_size(), 12); + + // Act. + // Acknowledge partially up to 4 bytes (within chunk1) + buf.acknowledge(4); + + // Assert. + assert_eq!(buf.len(), 2); + assert_eq!(buf.current_size(), 8); + + let chunks: Vec<_> = buf.chunks_to_replay().cloned().collect(); + assert_eq!(chunks[0].write_offset, 4); + assert_eq!(chunks[0].data.as_ref(), b"o "); + assert_eq!(chunks[0].crc32c, crc32c::crc32c(b"o ")); + assert_eq!(chunks[1].write_offset, 6); + assert_eq!(chunks[1].data.as_ref(), b"world!"); + assert_eq!(chunks[1].crc32c, crc32c::crc32c(b"world!")); + + // Act. + // Acknowledge fully past chunk1 up to 10 (within chunk2) + buf.acknowledge(10); + + // Assert. + assert_eq!(buf.len(), 1); + assert_eq!(buf.current_size(), 2); + + let chunks: Vec<_> = buf.chunks_to_replay().cloned().collect(); + assert_eq!(chunks[0].write_offset, 10); + assert_eq!(chunks[0].data.as_ref(), b"d!"); + assert_eq!(chunks[0].crc32c, crc32c::crc32c(b"d!")); + + // Act. + // Acknowledge all remaining bytes. + buf.acknowledge(12); + + // Assert. + assert!(buf.is_empty()); + assert_eq!(buf.current_size(), 0); + } + + #[test] + fn acknowledge_duplicate_or_earlier_offset() { + // Arrange. + let mut buf = ReplayBuffer::new(); + let chunk = Bytes::from_static(b"abcdef"); + buf.push(ReplayChunk::new(10, chunk.clone(), crc32c::crc32c(&chunk))); + + // Act. + // Acknowledge offset earlier than front write_offset. + buf.acknowledge(5); + + // Assert. + assert_eq!(buf.len(), 1); + assert_eq!(buf.current_size(), 6); + + // Act. + // Acknowledge exact write_offset of the front chunk. + buf.acknowledge(10); + + // Assert. + assert_eq!(buf.len(), 1); + assert_eq!(buf.current_size(), 6); + } + + #[test] + fn is_full_threshold() { + // Arrange. + let mut buf = ReplayBuffer::new(); + let huge_chunk = Bytes::from(vec![0u8; MAX_REPLAY_BUFFER_SIZE]); + + // Act. + buf.push(ReplayChunk::new(0, huge_chunk, 0)); + + // Assert. + assert!(buf.is_full()); + + // Act. + buf.acknowledge(1); + + // Assert. + assert!(!buf.is_full()); + } + + #[test] + fn chunk_to_request_conversion() { + // Arrange. + let data = Bytes::from_static(b"replay data"); + let crc = crc32c::crc32c(&data); + let chunk = ReplayChunk::new(42, data.clone(), crc); + + // Act. + let req = chunk.to_request(); + + // Assert. + assert_eq!(req.write_offset, 42); + if let Some(Data::ChecksummedData(cd)) = req.data { + assert_eq!(cd.content, data); + assert_eq!(cd.crc32c, Some(crc)); + } else { + panic!("expected ChecksummedData"); + } + } + + #[test] + fn clear_resets_size_and_queue() { + // Arrange. + let mut buf = ReplayBuffer::new(); + let chunk = Bytes::from_static(b"test data"); + buf.push(ReplayChunk::new(0, chunk, 0)); + assert!(!buf.is_empty()); + assert!(buf.current_size() > 0); + + // Act. + buf.clear(); + + // Assert. + assert!(buf.is_empty()); + assert_eq!(buf.current_size(), 0); + } +} From 1926ec38c71ee48a8a7309a2b1af804907f07a35 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Tue, 1 Sep 2026 07:11:29 +0000 Subject: [PATCH 2/3] feat(storage): add stream reconnect support for appendable upload --- .../src/storage/bidi_write/connector.rs | 131 ++++++++++++++++++ 1 file changed, 131 insertions(+) diff --git a/src/storage/src/storage/bidi_write/connector.rs b/src/storage/src/storage/bidi_write/connector.rs index 04b935e8fa..620d424c94 100644 --- a/src/storage/src/storage/bidi_write/connector.rs +++ b/src/storage/src/storage/bidi_write/connector.rs @@ -86,6 +86,11 @@ where } } + #[cfg(test)] + pub(crate) fn set_spec_state(&mut self, state: AppendObjectSpecState) { + *self.spec.lock().expect("never poisoned") = state; + } + pub async fn connect_open( &mut self, req: crate::model_ext::OpenAppendableObjectRequest, @@ -155,6 +160,23 @@ where self.connect_attempt_loop().await } + /// Reconnects a broken or redirected bidirectional streaming write session. + /// + /// If `last_error` is a redirect error, this updates the internal routing token + /// and object spec before attempting reconnection. Reconnection attempts use + /// exponential backoff and retry policies configured in [`RequestOptions`]. + pub async fn reconnect( + &mut self, + last_error: Error, + ) -> Result<(BidiWriteObjectResponse, Connection)> { + if let Some(status) = gaxi::as_inner::as_inner::(&last_error) + { + let mut guard = self.spec.lock().expect("never poisoned"); + guard.handle_redirect(status.clone()); + } + self.connect_attempt_loop().await + } + async fn connect_attempt_loop( &mut self, ) -> Result<(BidiWriteObjectResponse, Connection)> { @@ -1061,4 +1083,113 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn reconnect_with_redirect_status_updates_spec() -> Result<()> { + // Arrange. + let (tx, rx) = tokio::sync::mpsc::channel::>(5); + let stream = TonicResponse::from(rx); + + let receivers = Arc::new(Mutex::new(Vec::new())); + let save = receivers.clone(); + let mut mock = MockTestClient::new(); + mock.expect_start() + .times(1) + .return_once(move |_, _, rx, _, _, params| { + assert!(params.contains("routing_token=new-token")); + save.lock().expect("never poisoned").push(rx); + Ok(Ok(stream)) + }); + let client = SharedMockClient::new(mock); + let mut connector = Connector::new(test_options(), client); + + // Pre-configure connector in Append state + let initial_spec = crate::google::storage::v2::AppendObjectSpec { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + generation: 123456, + routing_token: Some("old-token".into()), + write_handle: None, + ..Default::default() + }; + *connector.spec.lock().expect("never poisoned") = AppendObjectSpecState::Append { + spec: initial_spec, + initial_chunk: None, + }; + + let initial_response = BidiWriteObjectResponse { + write_status: Some( + crate::google::storage::v2::bidi_write_object_response::WriteStatus::PersistedSize( + 50, + ), + ), + ..Default::default() + }; + tx.send(Ok(initial_response.clone())).await?; + + let redirect_err = super::super::tests::redirect_error("new-token"); + + // Act. + let (resp, _conn) = connector.reconnect(redirect_err).await?; + + // Assert. + assert_eq!(resp, initial_response); + + let guard = connector.spec.lock().expect("never poisoned"); + if let AppendObjectSpecState::Append { spec: s, .. } = &*guard { + assert_eq!(s.routing_token.as_deref(), Some("new-token")); + assert_eq!(s.generation, 42); // from test redirect_status + } else { + panic!("Expected AppendObjectSpecState::Append"); + } + + Ok(()) + } + + #[tokio::test] + async fn reconnect_with_transient_error() -> Result<()> { + // Arrange. + let (tx, rx) = tokio::sync::mpsc::channel::>(5); + let stream = TonicResponse::from(rx); + + let mut mock = MockTestClient::new(); + mock.expect_start() + .times(1) + .return_once(move |_, _, _, _, _, _| Ok(Ok(stream))); + let client = SharedMockClient::new(mock); + let mut connector = Connector::new(test_options(), client); + + let initial_spec = crate::google::storage::v2::AppendObjectSpec { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + generation: 123456, + routing_token: Some("stable-token".into()), + write_handle: None, + ..Default::default() + }; + *connector.spec.lock().expect("never poisoned") = AppendObjectSpecState::Append { + spec: initial_spec, + initial_chunk: None, + }; + + let initial_response = BidiWriteObjectResponse { + write_status: Some( + crate::google::storage::v2::bidi_write_object_response::WriteStatus::PersistedSize( + 100, + ), + ), + ..Default::default() + }; + tx.send(Ok(initial_response.clone())).await?; + + let transient_err = super::super::tests::transient_error(); + + // Act. + let (resp, _conn) = connector.reconnect(transient_err).await?; + + // Assert. + assert_eq!(resp, initial_response); + + Ok(()) + } } From a12c61df51f823551bcf55315020df1e92bf3e7a Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Tue, 1 Sep 2026 07:11:29 +0000 Subject: [PATCH 3/3] feat(storage): integrate replay buffer and stream auto-reconnect into bidi write worker --- src/storage/src/storage/bidi_write/worker.rs | 476 +++++++++++++++---- src/storage/src/storage/transport.rs | 26 +- 2 files changed, 401 insertions(+), 101 deletions(-) diff --git a/src/storage/src/storage/bidi_write/worker.rs b/src/storage/src/storage/bidi_write/worker.rs index 1371bd7ee0..f1b01d968c 100644 --- a/src/storage/src/storage/bidi_write/worker.rs +++ b/src/storage/src/storage/bidi_write/worker.rs @@ -14,14 +14,19 @@ // TODO(#5716): Lift to shared bidi module -use super::connector::Connection; +use super::connector::{Connection, Connector}; +use super::replay_buffer::{ReplayBuffer, ReplayChunk}; use super::{Client, TonicStreaming}; use crate::Error; use crate::error::WriteError; -use crate::google::storage::v2::{BidiWriteObjectRequest, BidiWriteObjectResponse}; -use gaxi::grpc::tonic::Result as TonicResult; +use crate::google::storage::v2::{ + BidiWriteObjectRequest, BidiWriteObjectResponse, bidi_write_object_request::Data, + bidi_write_object_response::WriteStatus, +}; +use std::collections::VecDeque; use std::sync::Arc; +use gaxi::grpc::tonic::Result as TonicResult; use tokio::sync::mpsc::Receiver; use tokio::sync::oneshot; @@ -41,20 +46,67 @@ pub enum UploadIntent { ), } -/// The background worker that manages the live gRPC stream. +/// Tracks an in-flight flush or finalize request awaiting server confirmation. +#[derive(Debug)] +enum PendingRequest { + Flush { + target_offset: i64, + request: BidiWriteObjectRequest, + sender: oneshot::Sender>, + }, + Finalize { + target_offset: i64, + request: BidiWriteObjectRequest, + sender: oneshot::Sender>, + }, +} + +impl PendingRequest { + fn request(&self) -> BidiWriteObjectRequest { + match self { + PendingRequest::Flush { request, .. } => request.clone(), + PendingRequest::Finalize { request, .. } => request.clone(), + } + } + + fn is_satisfied(&self, response: &BidiWriteObjectResponse, persisted_size: i64) -> bool { + match self { + PendingRequest::Flush { target_offset, .. } => persisted_size >= *target_offset, + PendingRequest::Finalize { target_offset, .. } => { + matches!(response.write_status, Some(WriteStatus::Resource(_))) + && persisted_size >= *target_offset + } + } + } + + fn complete(self, response: crate::Result) { + match self { + PendingRequest::Flush { sender, .. } => { + let _ = sender.send(response); + } + PendingRequest::Finalize { sender, .. } => { + let _ = sender.send(response); + } + } + } +} + +/// The background worker that manages the live gRPC stream, unacknowledged chunk replay, +/// and automatic reconnection. 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. + connector: Connector, + replay_buffer: ReplayBuffer, + pending_requests: VecDeque, + /// Tracks if the client intends to complete the upload by sending a Finalize intent. finalized: bool, } impl Worker { - pub fn new(connector: super::connector::Connector) -> Self { + pub fn new(connector: Connector) -> Self { Self { - _connector: connector, - pending_flushes: std::collections::VecDeque::new(), + connector, + replay_buffer: ReplayBuffer::new(), + pending_requests: VecDeque::new(), finalized: false, } } @@ -75,7 +127,7 @@ where let error = loop { tokio::select! { m = rx.next_message() => { - match self.handle_response(m) { + match self.handle_response(m).await { // Successful end of stream, return without error. None => break None, // An unrecoverable error in the stream or its data, return @@ -84,7 +136,6 @@ where // 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))) => { @@ -92,12 +143,18 @@ where } } }, - intent = requests.recv() => { + intent = requests.recv(), if !self.replay_buffer.is_full() => { match intent { Some(intent) => { let request = self.process_intent(intent); if let Err(e) = tx.send(request).await { - break Some(Error::io(e)); + match self.reconnect(Error::io(e)).await { + Some(Ok(Some(connection))) => { + (rx, tx) = (connection.rx, connection.tx); + } + Some(Err(e)) => break Some(e), + _ => {} + } } } None => { @@ -121,14 +178,28 @@ where fn process_intent(&mut self, intent: UploadIntent) -> BidiWriteObjectRequest { match intent { - UploadIntent::Append(req) => req, + UploadIntent::Append(req) => { + if let Some(Data::ChecksummedData(ref cd)) = req.data { + let crc32c = cd.crc32c.unwrap_or_else(|| crc32c::crc32c(&cd.content)); + self.replay_buffer.push(ReplayChunk::new( + req.write_offset, + cd.content.clone(), + crc32c, + )); + } + req + } UploadIntent::Flush(req, sender) => { assert!( req.state_lookup, "state_lookup must be true for Flush intents" ); assert!(req.flush, "flush must be true for Flush intents"); - self.pending_flushes.push_back(sender); + self.pending_requests.push_back(PendingRequest::Flush { + target_offset: req.write_offset, + request: req.clone(), + sender, + }); req } UploadIntent::Finalize(req, sender) => { @@ -137,13 +208,112 @@ where req.finish_write, "finish_write must be true for Finalize intents" ); - self.pending_flushes.push_back(sender); self.finalized = true; + self.pending_requests.push_back(PendingRequest::Finalize { + target_offset: req.write_offset, + request: req.clone(), + sender, + }); req } } } + /// Handles an incoming response message or stream completion from the server. + /// + /// Returns `None` when the stream has terminated cleanly, `Some(Err(e))` if an + /// unrecoverable error occurred or reconnection failed, `Some(Ok(None))` if the + /// response message was processed successfully on the existing connection, or + /// `Some(Ok(Some(connection)))` if the connection was reconnected and replayed. + pub async fn handle_response( + &mut self, + message: TonicResult>, + ) -> Option>>> { + let response = match message { + Ok(Some(msg)) => msg, + Ok(None) => { + // If the stream is unexpectedly closed by the server before the client + // intends to finalize the upload, treat it as an error to trigger reconnect + // or prevent silent failures on subsequent client writes. + if !self.pending_requests.is_empty() || !self.finalized { + return self + .reconnect(Error::io("stream closed unexpectedly")) + .await; + } + return None; + } + Err(e) => return self.reconnect(Error::io(e)).await, + }; + self.handle_response_success(response); + Some(Ok(None)) + } + + /// Processes a successful [`BidiWriteObjectResponse`] from the server. + /// + /// Updates acknowledged offsets in the replay buffer and completes any matching + /// in-flight flush or finalize requests. + pub fn handle_response_success(&mut self, response: BidiWriteObjectResponse) { + let persisted_size = match response.write_status.as_ref() { + Some(WriteStatus::PersistedSize(s)) => *s, + Some(WriteStatus::Resource(r)) => r.size, + None => 0, + }; + + self.replay_buffer.acknowledge(persisted_size); + + let mut matched = false; + while let Some(front) = self.pending_requests.front() { + if front.is_satisfied(&response, persisted_size) { + let req = self.pending_requests.pop_front().unwrap(); + req.complete(Ok(response.clone())); + matched = true; + } else { + break; + } + } + + if !matched { + tracing::debug!( + "Received unprompted BidiWriteObjectResponse from server: {:?}", + response + ); + } + } + + async fn reconnect( + &mut self, + last_error: Error, + ) -> Option>>> { + let (initial_response, connection) = match self.connector.reconnect(last_error).await { + Ok(res) => res, + Err(e) => return Some(Err(e)), + }; + + // Process initial response from reconnected stream + let initial_persisted_size = match initial_response.write_status.as_ref() { + Some(WriteStatus::PersistedSize(s)) => *s, + Some(WriteStatus::Resource(r)) => r.size, + None => 0, + }; + self.replay_buffer.acknowledge(initial_persisted_size); + + // Replay all unpersisted chunks + for chunk in self.replay_buffer.chunks_to_replay() { + if let Err(e) = connection.tx.send(chunk.to_request()).await { + return Some(Err(Error::io(e.to_string()))); + } + } + + // Re-send pending flush / finalize requests + for pending in &self.pending_requests { + if let Err(e) = connection.tx.send(pending.request()).await { + return Some(Err(Error::io(e.to_string()))); + } + } + + Some(Ok(Some(connection))) + } + async fn wait_for_server_completion(&mut self, mut rx: C::Stream) -> Option { loop { match rx.next_message().await { @@ -161,8 +331,8 @@ where mut requests: Receiver, shared_error: Arc, ) { - for sender in self.pending_flushes.drain(..) { - let _ = sender.send(Err(Error::ser(Arc::clone(&shared_error)))); + for pending in self.pending_requests.drain(..) { + pending.complete(Err(Error::ser(Arc::clone(&shared_error)))); } // Drain remaining requests to notify pending flush/finalize intents if the stream failed. requests.close(); @@ -175,50 +345,18 @@ where } } } - - pub fn handle_response( - &mut self, - message: TonicResult>, - ) -> Option>>> { - let response = match message { - Ok(Some(msg)) => msg, - Ok(None) => { - // 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(Error::io("stream closed unexpectedly"))); - } - return None; - } - Err(e) => return Some(Err(Error::io(e))), - }; - self.handle_response_success(response); - - // TODO(#5716): Implement 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::super::tests::permanent_error; use super::*; use crate::google::storage::v2::{ BidiWriteObjectRequest, BidiWriteObjectResponse, bidi_write_object_response::WriteStatus, }; + use gaxi::grpc::tonic::Response as TonicResponse; + use gaxi::grpc::tonic::Result as TonicResult; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -230,9 +368,9 @@ mod tests { ); fn spawn_test_worker() -> TestWorkerContext { - let (request_tx, request_rx) = mpsc::channel(1); + let (request_tx, request_rx) = mpsc::channel(10); let (response_tx, response_rx) = mpsc::channel(10); - let (tx, rx) = mpsc::channel(1); + let (tx, rx) = mpsc::channel(10); let connection = Connection::new(request_tx, response_rx); let mut mock = MockTestClient::new(); @@ -273,6 +411,7 @@ mod tests { let flush_request = BidiWriteObjectRequest { flush: true, state_lookup: true, + write_offset: 100, ..Default::default() }; tx.send(UploadIntent::Flush(flush_request.clone(), flush_tx)) @@ -306,6 +445,7 @@ mod tests { let finalize_request = BidiWriteObjectRequest { flush: true, finish_write: true, + write_offset: 100, ..Default::default() }; tx.send(UploadIntent::Finalize( @@ -317,8 +457,13 @@ mod tests { let stream_req = request_rx.recv().await.unwrap(); assert!(stream_req.finish_write); + let object = crate::google::storage::v2::Object { + name: "test-obj".into(), + size: 100, + ..Default::default() + }; let server_resp = BidiWriteObjectResponse { - write_status: Some(WriteStatus::PersistedSize(100)), + write_status: Some(WriteStatus::Resource(object)), ..Default::default() }; response_tx.send(Ok(server_resp.clone())).await?; @@ -331,6 +476,115 @@ mod tests { Ok(()) } + #[tokio::test] + async fn run_reconnect_and_replay_unpersisted_chunks() -> anyhow::Result<()> { + // Arrange. + let (stream1_tx, mut stream1_rx) = mpsc::channel(10); + let (stream1_resp_tx, stream1_resp_rx) = mpsc::channel(10); + let conn1 = Connection::new(stream1_tx, stream1_resp_rx); + + let (captured_stream2_req_tx, mut captured_stream2_req_rx) = + mpsc::channel::>(1); + let (stream2_resp_tx, stream2_resp_rx) = mpsc::channel(10); + let stream2 = TonicResponse::from(stream2_resp_rx); + + let mut mock = MockTestClient::new(); + mock.expect_start() + .times(1) + .return_once(move |_, _, req_rx, _, _, _| { + let _ = captured_stream2_req_tx.try_send(req_rx); + Ok(Ok(stream2)) + }); + + let mut connector = mock_connector(mock); + let initial_spec = crate::google::storage::v2::AppendObjectSpec { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + generation: 0, + routing_token: None, + write_handle: None, + ..Default::default() + }; + connector.set_spec_state(super::super::state::AppendObjectSpecState::Append { + spec: initial_spec, + initial_chunk: None, + }); + + let worker = Worker::new(connector); + + let (intent_tx, intent_rx) = mpsc::channel(10); + let handle = tokio::spawn(worker.run(conn1, intent_rx)); + + // Append two 10-byte chunks + let chunk1 = bytes::Bytes::from_static(b"0123456789"); + let req1 = BidiWriteObjectRequest { + write_offset: 0, + data: Some(Data::ChecksummedData( + crate::google::storage::v2::ChecksummedData { + content: chunk1.clone(), + crc32c: Some(crc32c::crc32c(&chunk1)), + }, + )), + ..Default::default() + }; + let chunk2 = bytes::Bytes::from_static(b"abcdefghij"); + let req2 = BidiWriteObjectRequest { + write_offset: 10, + data: Some(Data::ChecksummedData( + crate::google::storage::v2::ChecksummedData { + content: chunk2.clone(), + crc32c: Some(crc32c::crc32c(&chunk2)), + }, + )), + ..Default::default() + }; + + // Act. + intent_tx.send(UploadIntent::Append(req1)).await?; + intent_tx.send(UploadIntent::Append(req2)).await?; + + // Assert. + // Ensure both chunks were dispatched on stream 1 and buffered for replay + let s1_req1 = stream1_rx.recv().await.unwrap(); + assert_eq!(s1_req1.write_offset, 0); + let s1_req2 = stream1_rx.recv().await.unwrap(); + assert_eq!(s1_req2.write_offset, 10); + + // Act. + // Simulate stream 1 failure by dropping response stream + drop(stream1_resp_tx); + + // Connector reconnects to stream 2; server initial message reports + // persisted_size = 10 (chunk 1 persisted) + let reconnect_initial = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(10)), + ..Default::default() + }; + stream2_resp_tx.send(Ok(reconnect_initial)).await?; + + // Assert. + // Verify that stream 2 received the reconnect opening handshake request + let mut stream2_req_rx = captured_stream2_req_rx.recv().await.unwrap(); + let initial_req = stream2_req_rx.recv().await.unwrap(); + assert!(initial_req.first_message.is_some()); + + // Verify that chunk 2 (unpersisted) is replayed over stream 2! + let replayed_req = stream2_req_rx.recv().await.unwrap(); + assert_eq!(replayed_req.write_offset, 10); + if let Some(Data::ChecksummedData(cd)) = replayed_req.data { + assert_eq!(cd.content, chunk2); + assert_eq!(cd.crc32c, Some(crc32c::crc32c(&chunk2))); + } else { + panic!("expected ChecksummedData"); + } + + drop(intent_tx); + tokio::task::yield_now().await; + drop(stream2_resp_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(); @@ -341,73 +595,108 @@ mod tests { Ok(()) } + fn setup_mock_worker_with_reconnect_error(err: Error) -> TestWorkerContext { + let (request_tx, request_rx) = mpsc::channel(10); + let (response_tx, response_rx) = mpsc::channel(10); + let (tx, rx) = mpsc::channel(10); + let connection = Connection::new(request_tx, response_rx); + + let mut mock = MockTestClient::new(); + mock.expect_start() + .return_once(move |_, _, _, _, _, _| Err(err)); + + let mut connector = mock_connector(mock); + let initial_spec = crate::google::storage::v2::AppendObjectSpec { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + generation: 0, + routing_token: None, + write_handle: None, + ..Default::default() + }; + connector.set_spec_state(super::super::state::AppendObjectSpecState::Append { + spec: initial_spec, + initial_chunk: None, + }); + + let worker = Worker::new(connector); + let handle = tokio::spawn(worker.run(connection, rx)); + + (handle, tx, request_rx, response_tx) + } + #[tokio::test] async fn run_server_closes_unexpectedly() -> anyhow::Result<()> { - let (handle, _tx, _request_rx, response_tx) = spawn_test_worker(); + // Arrange. + let (handle, tx, _request_rx, response_tx) = + setup_mock_worker_with_reconnect_error(permanent_error()); - // Close the stream from the server side unexpectedly. + // Act. + // Close the stream from the server side unexpectedly while upload is not finalized. drop(response_tx); + // Assert. let result = handle.await?; assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "cannot serialize the request the transport reports an error: stream closed unexpectedly" - ); + let err = result.unwrap_err().to_string(); + assert!(err.contains("cannot serialize the request")); + assert!(err.contains("PERMISSION_DENIED")); + drop(tx); Ok(()) } #[tokio::test] async fn run_stream_error_during_flush() -> anyhow::Result<()> { - let (handle, tx, mut request_rx, response_tx) = spawn_test_worker(); + // Arrange. + let (handle, tx, mut request_rx, response_tx) = + setup_mock_worker_with_reconnect_error(permanent_error()); let (flush_tx, flush_rx) = oneshot::channel(); let flush_request = BidiWriteObjectRequest { flush: true, state_lookup: true, + write_offset: 100, ..Default::default() }; + + // Act. 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 stream and simulate failed reconnect drop(response_tx); + // Assert. let received_resp = flush_rx.await?; assert!(received_resp.is_err()); - assert_eq!( - received_resp.unwrap_err().to_string(), - "cannot serialize the request the transport reports an error: stream closed unexpectedly" - ); + let err = received_resp.unwrap_err().to_string(); + assert!(err.contains("cannot serialize the request")); + assert!(err.contains("PERMISSION_DENIED")); let result = handle.await?; assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!(err.contains("cannot serialize the request")); + assert!(err.contains("PERMISSION_DENIED")); Ok(()) } #[tokio::test] async fn run_stream_error_then_queue_requests() -> anyhow::Result<()> { - let (request_tx, _request_rx) = mpsc::channel(10); - let (response_tx, response_rx) = mpsc::channel(10); - let (tx, rx) = mpsc::channel(10); - 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)); + // Arrange. + let (handle, tx, _request_rx, response_tx) = + setup_mock_worker_with_reconnect_error(permanent_error()); let (flush_tx1, flush_rx1) = oneshot::channel(); let (flush_tx2, flush_rx2) = oneshot::channel(); + // Act. // Drop the server response stream to simulate the remote network crash. - // The worker will wake up and eventually process this, triggering the drain. + // The worker will wake up and attempt reconnect, which fails and triggers draining. drop(response_tx); // Put requests into the channel immediately. Because it has capacity 10 @@ -415,6 +704,7 @@ mod tests { let valid_flush = || BidiWriteObjectRequest { flush: true, state_lookup: true, + write_offset: 100, ..Default::default() }; tx.send(UploadIntent::Flush(valid_flush(), flush_tx1)) @@ -422,26 +712,24 @@ mod tests { tx.send(UploadIntent::Flush(valid_flush(), flush_tx2)) .await?; + // Assert. let payload1 = flush_rx1.await.unwrap(); assert!(payload1.is_err()); - assert!( - payload1 - .unwrap_err() - .to_string() - .contains("stream closed unexpectedly") - ); + let err1 = payload1.unwrap_err().to_string(); + assert!(err1.contains("cannot serialize the request")); + assert!(err1.contains("PERMISSION_DENIED")); let payload2 = flush_rx2.await.unwrap(); assert!(payload2.is_err()); - assert!( - payload2 - .unwrap_err() - .to_string() - .contains("stream closed unexpectedly") - ); + let err2 = payload2.unwrap_err().to_string(); + assert!(err2.contains("cannot serialize the request")); + assert!(err2.contains("PERMISSION_DENIED")); let result = handle.await?; assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!(err.contains("cannot serialize the request")); + assert!(err.contains("PERMISSION_DENIED")); Ok(()) } diff --git a/src/storage/src/storage/transport.rs b/src/storage/src/storage/transport.rs index e0428d50c3..25dfaf20bd 100644 --- a/src/storage/src/storage/transport.rs +++ b/src/storage/src/storage/transport.rs @@ -1234,18 +1234,26 @@ mod tests { const OBJECT_NAME: &str = "test-object"; const BIND_ADDRESS: &str = "0.0.0.0:0"; + // Arrange. let guard = TestLayer::initialize(); let (tx, rx) = tokio::sync::mpsc::channel::>(10); + + // The first response is the initial handshake/metadata response expected by + // `connector.rs` immediately upon opening the stream, before any data is sent. 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?; + // The flush response must confirm persistence at or past the flush request's target + // offset (100 initial + 5 appended = 105) for the worker to satisfy the pending flush. + let flush_response = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(105)), + ..BidiWriteObjectResponse::default() + }; + let finalize_response = BidiWriteObjectResponse { write_status: Some(WriteStatus::Resource( storage_grpc_mock::google::storage::v2::Object { @@ -1261,12 +1269,12 @@ mod tests { let mut stream = req.into_inner(); tokio::spawn(async move { while let Some(Ok(msg)) = stream.recv().await { + // The second response is sent ONLY when the client explicitly requests a flush + // or finalize. `writer.append()` does not wait for a response. 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; + let _ = tx.send(Ok(flush_response.clone())).await; } } }); @@ -1280,6 +1288,8 @@ mod tests { .with_tracing() .build() .await?; + + // Act. let mut writer = client .reopen_appendable_object(BUCKET_NAME, OBJECT_NAME, 12345) .send() @@ -1288,6 +1298,8 @@ mod tests { writer.append(Bytes::from_static(b"hello")).await?; writer.flush().await?; let obj = writer.finalize().await?; + + // Assert. assert_eq!(obj.size, 105); let captured = TestLayer::capture(&guard); @@ -1297,7 +1309,7 @@ mod tests { .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, "flush", 12345, Some(105)); check_bidi_write_span_attributes(&captured, "finalize", 12345, Some(105)); Ok(())