From d366668b14f63ee56197035ccb02091fad90bf2b Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Tue, 1 Sep 2026 07:11:28 +0000 Subject: [PATCH 1/7] 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 168ea26dfd..46f40c2bf6 100644 --- a/src/storage/src/storage/bidi_write.rs +++ b/src/storage/src/storage/bidi_write.rs @@ -21,6 +21,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 63d99baac4ac80fce1b81963a9526b708d21bf33 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Tue, 8 Sep 2026 02:23:58 +0000 Subject: [PATCH 2/7] Incorporate reviewer feedback on variable and method naming for clarity --- .../src/storage/bidi_write/replay_buffer.rs | 82 ++++++++++--------- 1 file changed, 43 insertions(+), 39 deletions(-) diff --git a/src/storage/src/storage/bidi_write/replay_buffer.rs b/src/storage/src/storage/bidi_write/replay_buffer.rs index df6441bceb..b59237fed5 100644 --- a/src/storage/src/storage/bidi_write/replay_buffer.rs +++ b/src/storage/src/storage/bidi_write/replay_buffer.rs @@ -67,7 +67,7 @@ impl ReplayChunk { #[derive(Debug, Default)] pub struct ReplayBuffer { queue: VecDeque, - current_size: usize, + unpersisted_bytes: usize, } impl ReplayBuffer { @@ -75,13 +75,13 @@ impl ReplayBuffer { pub fn new() -> Self { Self { queue: VecDeque::new(), - current_size: 0, + unpersisted_bytes: 0, } } /// Enqueues an unacknowledged [`ReplayChunk`] to the replay buffer. pub fn push(&mut self, chunk: ReplayChunk) { - self.current_size += chunk.data.len(); + self.unpersisted_bytes += chunk.data.len(); self.queue.push_back(chunk); } @@ -89,12 +89,13 @@ impl ReplayBuffer { /// /// 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(); + pub fn ack(&mut self, persisted_size: i64) { + while let Some(front) = self.queue.front() { + if front.end_offset() <= persisted_size { + let chunk = self.queue.pop_front().expect("front chunk must exist"); + self.unpersisted_bytes -= chunk.data.len(); + } else { + break; } } @@ -106,14 +107,14 @@ impl ReplayBuffer { front.data = front.data.slice(trimmed_bytes..); front.write_offset = persisted_size; front.crc32c = crc32c::crc32c(&front.data); - self.current_size -= trimmed_bytes; + self.unpersisted_bytes -= 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 + self.unpersisted_bytes >= MAX_REPLAY_BUFFER_SIZE } /// Returns `true` if the replay buffer contains no chunks. @@ -122,13 +123,13 @@ impl ReplayBuffer { } /// Returns the number of chunks currently held in the buffer. - pub fn len(&self) -> usize { + pub fn num_chunks(&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 + pub fn unpersisted_bytes(&self) -> usize { + self.unpersisted_bytes } /// Returns an iterator over the unpersisted [`ReplayChunk`]s in FIFO order for replay. @@ -139,7 +140,7 @@ impl ReplayBuffer { /// Clears all chunks from the buffer and resets byte tracking. pub fn clear(&mut self) { self.queue.clear(); - self.current_size = 0; + self.unpersisted_bytes = 0; } } @@ -154,19 +155,19 @@ mod tests { // Assert. assert!(buf.is_empty()); - assert_eq!(buf.len(), 0); - assert_eq!(buf.current_size(), 0); + assert_eq!(buf.num_chunks(), 0); + assert_eq!(buf.unpersisted_bytes(), 0); assert!(!buf.is_full()); // Act. - buf.acknowledge(100); + buf.ack(100); // Assert. assert!(buf.is_empty()); } #[test] - fn push_and_acknowledge_full_chunks() { + fn push_and_ack_full_chunks() { // Arrange. let mut buf = ReplayBuffer::new(); let chunk1 = Bytes::from_static(b"hello "); @@ -177,16 +178,16 @@ mod tests { buf.push(ReplayChunk::new(6, chunk2.clone(), crc32c::crc32c(&chunk2))); // Assert. - assert_eq!(buf.len(), 2); - assert_eq!(buf.current_size(), 12); + assert_eq!(buf.num_chunks(), 2); + assert_eq!(buf.unpersisted_bytes(), 12); // Act. // Acknowledge partially up to 4 bytes (within chunk1) - buf.acknowledge(4); + buf.ack(4); // Assert. - assert_eq!(buf.len(), 2); - assert_eq!(buf.current_size(), 8); + assert_eq!(buf.num_chunks(), 2); + assert_eq!(buf.unpersisted_bytes(), 8); let chunks: Vec<_> = buf.chunks_to_replay().cloned().collect(); assert_eq!(chunks[0].write_offset, 4); @@ -198,11 +199,11 @@ mod tests { // Act. // Acknowledge fully past chunk1 up to 10 (within chunk2) - buf.acknowledge(10); + buf.ack(10); // Assert. - assert_eq!(buf.len(), 1); - assert_eq!(buf.current_size(), 2); + assert_eq!(buf.num_chunks(), 1); + assert_eq!(buf.unpersisted_bytes(), 2); let chunks: Vec<_> = buf.chunks_to_replay().cloned().collect(); assert_eq!(chunks[0].write_offset, 10); @@ -211,15 +212,15 @@ mod tests { // Act. // Acknowledge all remaining bytes. - buf.acknowledge(12); + buf.ack(12); // Assert. assert!(buf.is_empty()); - assert_eq!(buf.current_size(), 0); + assert_eq!(buf.unpersisted_bytes(), 0); } #[test] - fn acknowledge_duplicate_or_earlier_offset() { + fn ack_duplicate_or_earlier_offset() { // Arrange. let mut buf = ReplayBuffer::new(); let chunk = Bytes::from_static(b"abcdef"); @@ -227,19 +228,19 @@ mod tests { // Act. // Acknowledge offset earlier than front write_offset. - buf.acknowledge(5); + buf.ack(5); // Assert. - assert_eq!(buf.len(), 1); - assert_eq!(buf.current_size(), 6); + assert_eq!(buf.num_chunks(), 1); + assert_eq!(buf.unpersisted_bytes(), 6); // Act. // Acknowledge exact write_offset of the front chunk. - buf.acknowledge(10); + buf.ack(10); // Assert. - assert_eq!(buf.len(), 1); - assert_eq!(buf.current_size(), 6); + assert_eq!(buf.num_chunks(), 1); + assert_eq!(buf.unpersisted_bytes(), 6); } #[test] @@ -255,7 +256,7 @@ mod tests { assert!(buf.is_full()); // Act. - buf.acknowledge(1); + buf.ack(1); // Assert. assert!(!buf.is_full()); @@ -268,6 +269,9 @@ mod tests { let crc = crc32c::crc32c(&data); let chunk = ReplayChunk::new(42, data.clone(), crc); + // Assert end_offset calculation. + assert_eq!(chunk.end_offset(), 42 + data.len() as i64); + // Act. let req = chunk.to_request(); @@ -288,13 +292,13 @@ mod tests { let chunk = Bytes::from_static(b"test data"); buf.push(ReplayChunk::new(0, chunk, 0)); assert!(!buf.is_empty()); - assert!(buf.current_size() > 0); + assert!(buf.unpersisted_bytes() > 0); // Act. buf.clear(); // Assert. assert!(buf.is_empty()); - assert_eq!(buf.current_size(), 0); + assert_eq!(buf.unpersisted_bytes(), 0); } } From 7a720e5df703d41455bfa99f52ecd6a3209a17a7 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Tue, 8 Sep 2026 05:35:35 +0000 Subject: [PATCH 3/7] feat(storage): make ReplayBuffer capacity configurable --- .../src/storage/bidi_write/replay_buffer.rs | 91 +++++++++++++++---- 1 file changed, 75 insertions(+), 16 deletions(-) diff --git a/src/storage/src/storage/bidi_write/replay_buffer.rs b/src/storage/src/storage/bidi_write/replay_buffer.rs index b59237fed5..ef53aff719 100644 --- a/src/storage/src/storage/bidi_write/replay_buffer.rs +++ b/src/storage/src/storage/bidi_write/replay_buffer.rs @@ -21,8 +21,9 @@ use crate::google::storage::v2::{ 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; +/// Defines the default capacity of the [`ReplayBuffer`] in bytes (32 MiB). +// TODO(#5716): Remove once ReplayBuffer capacity is configured via CommonOptions. +pub const DEFAULT_REPLAY_BUFFER_SIZE: usize = 32 * 1024 * 1024; /// Represents an unacknowledged data chunk retained in the [`ReplayBuffer`]. #[derive(Clone, Debug, PartialEq, Eq)] @@ -36,7 +37,7 @@ pub struct ReplayChunk { } impl ReplayChunk { - /// Creates a new replay chunk. + /// Creates a new [`ReplayChunk`]. pub fn new(write_offset: i64, data: Bytes, crc32c: u32) -> Self { Self { write_offset, @@ -50,7 +51,7 @@ impl ReplayChunk { self.write_offset + self.data.len() as i64 } - /// Converts this replay chunk into a [`BidiWriteObjectRequest`] for transmission. + /// Converts this [`ReplayChunk`] into a [`BidiWriteObjectRequest`] for transmission. pub fn to_request(&self) -> BidiWriteObjectRequest { BidiWriteObjectRequest { write_offset: self.write_offset, @@ -63,23 +64,41 @@ impl ReplayChunk { } } -/// Manages an in-memory FIFO queue of unacknowledged chunks up to [`MAX_REPLAY_BUFFER_SIZE`]. -#[derive(Debug, Default)] +/// Manages an in-memory FIFO queue of unacknowledged chunks up to a configurable capacity. +#[derive(Debug)] pub struct ReplayBuffer { queue: VecDeque, unpersisted_bytes: usize, + capacity: usize, +} + +impl Default for ReplayBuffer { + fn default() -> Self { + Self::new() + } } impl ReplayBuffer { - /// Creates a new, empty replay buffer. + /// Creates a new, empty [`ReplayBuffer`] with the default capacity ([`DEFAULT_REPLAY_BUFFER_SIZE`]). pub fn new() -> Self { + Self::with_capacity(DEFAULT_REPLAY_BUFFER_SIZE) + } + + /// Creates a new, empty [`ReplayBuffer`] with a specified capacity in bytes. + pub fn with_capacity(capacity: usize) -> Self { Self { queue: VecDeque::new(), unpersisted_bytes: 0, + capacity, } } - /// Enqueues an unacknowledged [`ReplayChunk`] to the replay buffer. + /// Returns the configured capacity of the [`ReplayBuffer`] in bytes. + pub fn capacity(&self) -> usize { + self.capacity + } + + /// Enqueues an unacknowledged [`ReplayChunk`] to the [`ReplayBuffer`]. pub fn push(&mut self, chunk: ReplayChunk) { self.unpersisted_bytes += chunk.data.len(); self.queue.push_back(chunk); @@ -111,23 +130,22 @@ impl ReplayBuffer { } } - /// Returns `true` if the buffered byte count has reached or exceeded - /// [`MAX_REPLAY_BUFFER_SIZE`]. + /// Returns `true` if the buffered byte count has reached or exceeded [`capacity`][Self::capacity]. pub fn is_full(&self) -> bool { - self.unpersisted_bytes >= MAX_REPLAY_BUFFER_SIZE + self.unpersisted_bytes >= self.capacity } - /// Returns `true` if the replay buffer contains no chunks. + /// Returns `true` if the [`ReplayBuffer`] contains no chunks. pub fn is_empty(&self) -> bool { self.queue.is_empty() } - /// Returns the number of chunks currently held in the buffer. + /// Returns the number of chunks currently held in the [`ReplayBuffer`]. pub fn num_chunks(&self) -> usize { self.queue.len() } - /// Returns the total unpersisted byte count currently retained in the buffer. + /// Returns the total unpersisted byte count currently retained in the [`ReplayBuffer`]. pub fn unpersisted_bytes(&self) -> usize { self.unpersisted_bytes } @@ -137,7 +155,7 @@ impl ReplayBuffer { self.queue.iter() } - /// Clears all chunks from the buffer and resets byte tracking. + /// Clears all chunks from the [`ReplayBuffer`] and resets byte tracking. pub fn clear(&mut self) { self.queue.clear(); self.unpersisted_bytes = 0; @@ -247,7 +265,7 @@ mod tests { fn is_full_threshold() { // Arrange. let mut buf = ReplayBuffer::new(); - let huge_chunk = Bytes::from(vec![0u8; MAX_REPLAY_BUFFER_SIZE]); + let huge_chunk = Bytes::from(vec![0u8; DEFAULT_REPLAY_BUFFER_SIZE]); // Act. buf.push(ReplayChunk::new(0, huge_chunk, 0)); @@ -301,4 +319,45 @@ mod tests { assert!(buf.is_empty()); assert_eq!(buf.unpersisted_bytes(), 0); } + + #[test] + fn with_capacity_default() { + // Arrange & Act. + let buf = ReplayBuffer::new(); + + // Assert. + assert_eq!(buf.capacity(), DEFAULT_REPLAY_BUFFER_SIZE); + } + + #[test] + fn with_capacity_custom() { + // Arrange. + const CUSTOM_CAPACITY: usize = 64 * 1024 * 1024; // 64 MiB + + // Act. + let buf = ReplayBuffer::with_capacity(CUSTOM_CAPACITY); + + // Assert. + assert_eq!(buf.capacity(), CUSTOM_CAPACITY); + } + + #[test] + fn is_full_with_custom_capacity() { + // Arrange. + // Use a micro-capacity of 100 bytes for deterministic testing. + let mut buf = ReplayBuffer::with_capacity(100); + let chunk = Bytes::from(vec![0u8; 100]); + + // Act. + buf.push(ReplayChunk::new(0, chunk, 0)); + + // Assert. + assert!(buf.is_full()); + + // Act. + buf.ack(1); + + // Assert. + assert!(!buf.is_full()); + } } From 2814ea225cc9b6e1a8ed7cdf0a13d83e5d126de9 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Thu, 17 Sep 2026 12:57:17 +0000 Subject: [PATCH 4/7] Address reviewer nits in ReplayBuffer::ack - Drop the unreachable expect() by reading the front chunk length before popping it, removing the redundant panic branch. - Relabel the // SAFETY: comment as // Invariant:, since Bytes::slice is safe Rust and the comment documents an algorithmic invariant, not an unsafe contract. --- src/storage/src/storage/bidi_write/replay_buffer.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/storage/src/storage/bidi_write/replay_buffer.rs b/src/storage/src/storage/bidi_write/replay_buffer.rs index ef53aff719..258ef014cc 100644 --- a/src/storage/src/storage/bidi_write/replay_buffer.rs +++ b/src/storage/src/storage/bidi_write/replay_buffer.rs @@ -111,8 +111,8 @@ impl ReplayBuffer { pub fn ack(&mut self, persisted_size: i64) { while let Some(front) = self.queue.front() { if front.end_offset() <= persisted_size { - let chunk = self.queue.pop_front().expect("front chunk must exist"); - self.unpersisted_bytes -= chunk.data.len(); + self.unpersisted_bytes -= front.data.len(); + self.queue.pop_front(); } else { break; } @@ -122,7 +122,7 @@ impl ReplayBuffer { && 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. + // Invariant: 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); From c86091215947f77ae5867ecc4d225e3ce28249c7 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Tue, 1 Sep 2026 07:11:29 +0000 Subject: [PATCH 5/7] 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 a09503284e..25be6c19f8 100644 --- a/src/storage/src/storage/bidi_write/connector.rs +++ b/src/storage/src/storage/bidi_write/connector.rs @@ -84,6 +84,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, @@ -153,6 +158,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)> { @@ -1059,4 +1081,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 bcccd3b60b2be80654bcbb591005fb4b17f87d87 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Tue, 1 Sep 2026 07:11:29 +0000 Subject: [PATCH 6/7] 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 513c1d211c..4d43edb830 100644 --- a/src/storage/src/storage/bidi_write/worker.rs +++ b/src/storage/src/storage/bidi_write/worker.rs @@ -12,14 +12,19 @@ // See the License for the specific language governing permissions and // limitations under the License. -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; @@ -39,20 +44,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, } } @@ -73,7 +125,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 @@ -82,7 +134,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))) => { @@ -90,12 +141,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 => { @@ -119,14 +176,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) => { @@ -135,13 +206,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.ack(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.ack(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 { @@ -159,8 +329,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(); @@ -173,50 +343,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; @@ -228,9 +366,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(); @@ -271,6 +409,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)) @@ -304,6 +443,7 @@ mod tests { let finalize_request = BidiWriteObjectRequest { flush: true, finish_write: true, + write_offset: 100, ..Default::default() }; tx.send(UploadIntent::Finalize( @@ -315,8 +455,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?; @@ -329,6 +474,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(); @@ -339,73 +593,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 @@ -413,6 +702,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)) @@ -420,26 +710,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 4d7fb37b02..d46945cbc1 100644 --- a/src/storage/src/storage/transport.rs +++ b/src/storage/src/storage/transport.rs @@ -1232,18 +1232,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 { @@ -1259,12 +1267,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; } } }); @@ -1278,6 +1286,8 @@ mod tests { .with_tracing() .build() .await?; + + // Act. let mut writer = client .reopen_appendable_object(BUCKET_NAME, OBJECT_NAME, 12345) .send() @@ -1286,6 +1296,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); @@ -1295,7 +1307,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(()) From a3615748caeb87c55bf7fbb31aa4db63e0ee6931 Mon Sep 17 00:00:00 2001 From: Sharon Lynn Date: Thu, 17 Sep 2026 11:38:57 +0000 Subject: [PATCH 7/7] fix(storage): avoid replay buffer deadlock by injecting watermark flush The replay buffer gates the worker's intent branch on `is_full()`, but the buffer is only drained by `ack()`, which requires a server response. Nothing in the append path requested one, so ~32 MiB of appends without an explicit flush stalled the worker permanently and blocked the caller inside `append()`. Inject `flush` + `state_lookup` on the append that crosses a high watermark, and re-establish the invariant after reconnect. --- src/storage/src/storage/bidi_write/worker.rs | 301 ++++++++++++++++++- 1 file changed, 296 insertions(+), 5 deletions(-) diff --git a/src/storage/src/storage/bidi_write/worker.rs b/src/storage/src/storage/bidi_write/worker.rs index 4d43edb830..c36d52998f 100644 --- a/src/storage/src/storage/bidi_write/worker.rs +++ b/src/storage/src/storage/bidi_write/worker.rs @@ -14,7 +14,7 @@ use super::connector::{Connection, Connector}; use super::replay_buffer::{ReplayBuffer, ReplayChunk}; -use super::{Client, TonicStreaming}; +use super::{Client, MAX_WRITE_CHUNK_SIZE, TonicStreaming}; use crate::Error; use crate::error::WriteError; use crate::google::storage::v2::{ @@ -97,17 +97,50 @@ pub struct Worker { pending_requests: VecDeque, /// Tracks if the client intends to complete the upload by sending a Finalize intent. finalized: bool, + /// Tracks if a worker-initiated `state_lookup` request is awaiting a server response. + /// + /// The worker injects such a request when the replay buffer crosses its high + /// watermark, so that an `ack()` is guaranteed to arrive before the buffer fills. + self_flush_outstanding: bool, } impl Worker { pub fn new(connector: Connector) -> Self { + Self::with_replay_buffer(connector, ReplayBuffer::new()) + } + + /// Creates a [`Worker`] with a caller-provided [`ReplayBuffer`]. + /// + /// Tests use this to exercise the buffer capacity limits without buffering + /// [`DEFAULT_REPLAY_BUFFER_SIZE`][super::replay_buffer::DEFAULT_REPLAY_BUFFER_SIZE] bytes. + pub fn with_replay_buffer(connector: Connector, replay_buffer: ReplayBuffer) -> Self { Self { connector, - replay_buffer: ReplayBuffer::new(), + replay_buffer, pending_requests: VecDeque::new(), finalized: false, + self_flush_outstanding: false, } } + + /// Returns `true` if the replay buffer has crossed its high watermark and no + /// worker-initiated `state_lookup` is outstanding. + /// + /// The replay buffer only shrinks when [`ReplayBuffer::ack`] is called, which + /// requires a server response, and the server only responds when a request sets + /// `state_lookup`. Without injecting one, a caller that appends + /// [`ReplayBuffer::capacity`] bytes without an explicit flush would fill the buffer, + /// disable the intent branch of the worker loop, and stall forever. The watermark + /// leaves two chunks worth of headroom so the injected request is dispatched well + /// before [`ReplayBuffer::is_full`] trips. + fn needs_watermark_flush(&self) -> bool { + !self.self_flush_outstanding + && self.replay_buffer.unpersisted_bytes() + >= self + .replay_buffer + .capacity() + .saturating_sub(2 * MAX_WRITE_CHUNK_SIZE) + } } impl Worker @@ -176,7 +209,7 @@ where fn process_intent(&mut self, intent: UploadIntent) -> BidiWriteObjectRequest { match intent { - UploadIntent::Append(req) => { + UploadIntent::Append(mut 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( @@ -185,6 +218,14 @@ where crc32c, )); } + // Piggyback a `state_lookup` on this append once the replay buffer + // crosses its high watermark. Only `state_lookup` elicits a response, + // and only a response drives `ReplayBuffer::ack()`. + if self.needs_watermark_flush() { + req.flush = true; + req.state_lookup = true; + self.self_flush_outstanding = true; + } req } UploadIntent::Flush(req, sender) => { @@ -270,18 +311,25 @@ where } } - if !matched { + // A worker-initiated `state_lookup` has no entry in `pending_requests`, so its + // response legitimately matches nothing. Do not report it as unprompted. + if !matched && !self.self_flush_outstanding { tracing::debug!( "Received unprompted BidiWriteObjectResponse from server: {:?}", response ); } + self.self_flush_outstanding = false; } async fn reconnect( &mut self, last_error: Error, ) -> Option>>> { + // Any response the worker-initiated `state_lookup` was waiting on will never + // arrive on the dead stream. + self.self_flush_outstanding = false; + let (initial_response, connection) = match self.connector.reconnect(last_error).await { Ok(res) => res, Err(e) => return Some(Err(e)), @@ -309,6 +357,28 @@ where } } + // Replayed chunks set neither `flush` nor `state_lookup`. If the buffer is + // still above the watermark and no user request is pending, nothing would + // elicit a response on the new stream, so restore the invariant explicitly. + if self.pending_requests.is_empty() && self.needs_watermark_flush() { + let write_offset = self + .replay_buffer + .chunks_to_replay() + .last() + .map(|chunk| chunk.end_offset()) + .unwrap_or(initial_persisted_size); + let request = BidiWriteObjectRequest { + write_offset, + flush: true, + state_lookup: true, + ..BidiWriteObjectRequest::default() + }; + if let Err(e) = connection.tx.send(request).await { + return Some(Err(Error::io(e.to_string()))); + } + self.self_flush_outstanding = true; + } + Some(Ok(Some(connection))) } @@ -365,7 +435,20 @@ mod tests { mpsc::Sender>, ); + /// Defines the payload size of the synthetic appends used in the watermark tests. + const TEST_CHUNK_SIZE: usize = 512; + + /// Places the watermark (`capacity - 2 * MAX_WRITE_CHUNK_SIZE`) at exactly two test + /// chunks, so the second append is the one that crosses it. + const TEST_CAPACITY: usize = 2 * MAX_WRITE_CHUNK_SIZE + 2 * TEST_CHUNK_SIZE; + fn spawn_test_worker() -> TestWorkerContext { + spawn_test_worker_with_replay_capacity( + super::super::replay_buffer::DEFAULT_REPLAY_BUFFER_SIZE, + ) + } + + fn spawn_test_worker_with_replay_capacity(capacity: usize) -> TestWorkerContext { let (request_tx, request_rx) = mpsc::channel(10); let (response_tx, response_rx) = mpsc::channel(10); let (tx, rx) = mpsc::channel(10); @@ -375,12 +458,27 @@ mod tests { mock.expect_start().never(); let connector = mock_connector(mock); - let worker = Worker::new(connector); + let worker = Worker::with_replay_buffer(connector, ReplayBuffer::with_capacity(capacity)); let handle = tokio::spawn(worker.run(connection, rx)); (handle, tx, request_rx, response_tx) } + fn append_intent(write_offset: i64, len: usize) -> UploadIntent { + let content = bytes::Bytes::from(vec![b'x'; len]); + let crc32c = crc32c::crc32c(&content); + UploadIntent::Append(BidiWriteObjectRequest { + write_offset, + data: Some(Data::ChecksummedData( + crate::google::storage::v2::ChecksummedData { + content, + crc32c: Some(crc32c), + }, + )), + ..Default::default() + }) + } + #[tokio::test] async fn run_append() -> anyhow::Result<()> { let (handle, tx, mut request_rx, _response_tx) = spawn_test_worker(); @@ -787,4 +885,197 @@ mod tests { .await; assert!(handle.await.unwrap_err().is_panic()); } + + #[tokio::test] + async fn run_append_injects_watermark_flush() -> anyhow::Result<()> { + // Arrange. + let (handle, tx, mut request_rx, response_tx) = + spawn_test_worker_with_replay_capacity(TEST_CAPACITY); + + // Act. + tx.send(append_intent(0, TEST_CHUNK_SIZE)).await?; + tx.send(append_intent(TEST_CHUNK_SIZE as i64, TEST_CHUNK_SIZE)) + .await?; + + // Assert. + // The first append stays below the watermark and is dispatched untouched. + let first = request_rx.recv().await.unwrap(); + assert!(!first.flush, "{first:?}"); + assert!(!first.state_lookup, "{first:?}"); + + // The second append crosses the watermark, so the worker piggybacks a flush + // and a state_lookup on it. Only state_lookup elicits a server response. + let second = request_rx.recv().await.unwrap(); + assert!(second.flush, "{second:?}"); + assert!(second.state_lookup, "{second:?}"); + + drop(tx); + tokio::task::yield_now().await; + drop(response_tx); + handle.await??; + Ok(()) + } + + #[tokio::test] + async fn run_append_does_not_repeat_watermark_flush() -> anyhow::Result<()> { + // Arrange. + let (handle, tx, mut request_rx, response_tx) = + spawn_test_worker_with_replay_capacity(TEST_CAPACITY); + + // Act. + // Three appends, all above the watermark from the second one onwards, with no + // server response in between. + for i in 0..3 { + tx.send(append_intent((i * TEST_CHUNK_SIZE) as i64, TEST_CHUNK_SIZE)) + .await?; + } + + // Assert. + let _first = request_rx.recv().await.unwrap(); + let second = request_rx.recv().await.unwrap(); + assert!(second.state_lookup, "{second:?}"); + + // A state_lookup is already outstanding, so the third append is not flagged. + let third = request_rx.recv().await.unwrap(); + assert!(!third.flush, "{third:?}"); + assert!(!third.state_lookup, "{third:?}"); + + drop(tx); + tokio::task::yield_now().await; + drop(response_tx); + handle.await??; + Ok(()) + } + + #[tokio::test] + async fn run_append_past_capacity_without_explicit_flush() -> anyhow::Result<()> { + // Arrange. + // A capacity of four chunks, so the appends below overrun it several times over. + const APPEND_COUNT: usize = 16; + let (handle, tx, mut request_rx, response_tx) = + spawn_test_worker_with_replay_capacity(4 * TEST_CHUNK_SIZE); + + // A server that replies only when the client asks for it via state_lookup. + let server = tokio::spawn(async move { + let mut data_requests = 0_usize; + let mut persisted_size = 0_i64; + while let Some(request) = request_rx.recv().await { + if let Some(Data::ChecksummedData(cd)) = request.data.as_ref() { + data_requests += 1; + persisted_size = request.write_offset + cd.content.len() as i64; + } + if request.state_lookup { + let response = BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(persisted_size)), + ..Default::default() + }; + if response_tx.send(Ok(response)).await.is_err() { + break; + } + } + } + data_requests + }); + + // Act. + // Without the injected watermark flush the replay buffer fills, the intent + // branch of the worker loop is disabled forever and this block never returns. + let result = tokio::time::timeout(std::time::Duration::from_secs(10), async move { + for i in 0..APPEND_COUNT { + tx.send(append_intent((i * TEST_CHUNK_SIZE) as i64, TEST_CHUNK_SIZE)) + .await + .expect("the worker must keep accepting appends"); + } + drop(tx); + handle.await.expect("the worker task must not panic") + }) + .await + .expect("appends without an explicit flush must not deadlock"); + + // Assert. + result?; + assert_eq!(server.await?, APPEND_COUNT); + Ok(()) + } + + #[tokio::test] + async fn run_reconnect_restores_watermark_state_lookup() -> 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); + connector.set_spec_state(super::super::state::AppendObjectSpecState::Append { + spec: crate::google::storage::v2::AppendObjectSpec { + bucket: "projects/_/buckets/test-bucket".into(), + object: "test-object".into(), + ..Default::default() + }, + initial_chunk: None, + }); + + let worker = + Worker::with_replay_buffer(connector, ReplayBuffer::with_capacity(TEST_CAPACITY)); + let (intent_tx, intent_rx) = mpsc::channel(10); + let handle = tokio::spawn(worker.run(conn1, intent_rx)); + + // Fill the replay buffer up to the watermark on the first stream. + intent_tx.send(append_intent(0, TEST_CHUNK_SIZE)).await?; + intent_tx + .send(append_intent(TEST_CHUNK_SIZE as i64, TEST_CHUNK_SIZE)) + .await?; + let _ = stream1_rx.recv().await.unwrap(); + let _ = stream1_rx.recv().await.unwrap(); + + // Act. + // Break the first stream. The watermark state_lookup sent on it is lost. + drop(stream1_resp_tx); + + // The reconnected stream reports nothing persisted, so both chunks are replayed. + stream2_resp_tx + .send(Ok(BidiWriteObjectResponse { + write_status: Some(WriteStatus::PersistedSize(0)), + ..Default::default() + })) + .await?; + + // Assert. + 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()); + + // Replayed chunks carry data but neither flush nor state_lookup. + for i in 0..2 { + let replayed = stream2_req_rx.recv().await.unwrap(); + assert_eq!(replayed.write_offset, (i * TEST_CHUNK_SIZE) as i64); + assert!(!replayed.state_lookup, "{replayed:?}"); + } + + // The buffer is still at the watermark and no user request is pending, so the + // worker re-establishes the invariant with a standalone state_lookup. + let watermark_req = stream2_req_rx.recv().await.unwrap(); + assert!(watermark_req.flush, "{watermark_req:?}"); + assert!(watermark_req.state_lookup, "{watermark_req:?}"); + assert_eq!(watermark_req.write_offset, (2 * TEST_CHUNK_SIZE) as i64); + + drop(intent_tx); + tokio::task::yield_now().await; + drop(stream2_resp_tx); + handle.await??; + Ok(()) + } }