From 42ff57de72044e39d208f0e000d23a454ad2ea60 Mon Sep 17 00:00:00 2001 From: Shubham Padkonde Date: Sun, 4 Oct 2026 17:26:34 +0530 Subject: [PATCH] fix(storage): resume appends using buffer-relative offsets --- .../asyncio/async_appendable_object_writer.py | 13 +++- .../retry/writes_resumption_strategy.py | 12 ++-- .../retry/test_writes_resumption_strategy.py | 18 ++++++ .../test_async_appendable_object_writer.py | 61 +++++++++++++++++++ 4 files changed, 98 insertions(+), 6 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index 1af54d6ffe43..17c2927c8f6a 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py @@ -423,7 +423,8 @@ async def append( :type enable_checksum: bool :param enable_checksum: (Optional) If True, calculates and checks checksums for each chunk. Defaults to True. - :raises ValueError: If the stream is not open. + :raises ValueError: If the stream is not open, or recovery requires an + offset outside the current append buffer. """ if not self._is_stream_open: raise ValueError("Stream is not open. Call open() before append().") @@ -472,7 +473,14 @@ async def generator(): write_state.write_handle = self.write_handle write_state.routing_token = None - write_state.user_buffer.seek(write_state.persisted_size) + buffer_offset = ( + write_state.persisted_size - write_state.buffer_start_offset + ) + if not 0 <= buffer_offset <= len(data): + raise ValueError( + "Cannot resume upload: persisted offset is outside the current append buffer." + ) + write_state.user_buffer.seek(buffer_offset) write_state.bytes_sent = write_state.persisted_size write_state.bytes_since_last_flush = 0 self.bytes_appended_since_last_flush = 0 @@ -514,6 +522,7 @@ async def generator(): write_state.persisted_size = self.persisted_size # offset is set during `open()` call. write_state.bytes_sent = self.offset or 0 + write_state.buffer_start_offset = write_state.bytes_sent write_state.bytes_since_last_flush = self.bytes_appended_since_last_flush retry_manager = _BidiStreamRetryManager( diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/writes_resumption_strategy.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/writes_resumption_strategy.py index bec0b31c5407..5f55cea8b767 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/writes_resumption_strategy.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/writes_resumption_strategy.py @@ -48,6 +48,8 @@ def __init__( ): self.chunk_size = chunk_size self.user_buffer = user_buffer + # Absolute object offset corresponding to the beginning of user_buffer. + self.buffer_start_offset: int = 0 self.persisted_size: int = 0 # Bytes sent to the server (it may be unpersisted), # i.e. latest object size = persisted_size + some more bytes. @@ -76,7 +78,7 @@ def generate_requests( write_state: _WriteState = state["write_state"] requests = [] - # The buffer should already be seeked to the correct position (persisted_size) + # The buffer should already be seeked to the buffer-relative persisted position # by the `recover_state_on_failure` method before this is called. while not write_state.is_finalized: chunk = write_state.user_buffer.read(write_state.chunk_size) @@ -149,7 +151,9 @@ async def recover_state_on_failure( write_state.write_handle = redirect_proto.write_handle # We must assume any data sent beyond 'persisted_size' was lost. - # Reset the user buffer to the last known good byte confirmed by the server. - write_state.user_buffer.seek(write_state.persisted_size) - write_state.bytes_sent = write_state.persisted_size + # The last response may predate this append. Reopening the stream will + # obtain and validate the authoritative persisted size before sending. + resume_offset = max(write_state.persisted_size, write_state.buffer_start_offset) + write_state.user_buffer.seek(resume_offset - write_state.buffer_start_offset) + write_state.bytes_sent = resume_offset write_state.bytes_since_last_flush = 0 diff --git a/packages/google-cloud-storage/tests/unit/asyncio/retry/test_writes_resumption_strategy.py b/packages/google-cloud-storage/tests/unit/asyncio/retry/test_writes_resumption_strategy.py index 2af036f80cb2..fe3383d27b17 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/retry/test_writes_resumption_strategy.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/retry/test_writes_resumption_strategy.py @@ -245,6 +245,24 @@ async def test_generate_requests_after_failure_and_recovery(self, strategy): # Tests for update_state_from_response # ------------------------------------------------------------------------- + @pytest.mark.asyncio + async def test_recover_later_append_uses_buffer_relative_offset(self, strategy): + write_state = _WriteState(250, io.BytesIO(b"C" * 1000), 10**9) + write_state.buffer_start_offset = 2000 + write_state.bytes_sent = 2500 + write_state.persisted_size = 2250 + write_state.user_buffer.seek(500) + state = {"write_state": write_state} + + await strategy.recover_state_on_failure( + exceptions.ServiceUnavailable("disconnected"), state + ) + requests = strategy.generate_requests(state) + + assert [request.write_offset for request in requests] == [2250, 2500, 2750] + assert b"".join(r.checksummed_data.content for r in requests) == b"C" * 750 + assert write_state.bytes_sent == 3000 + def test_update_state_from_response_all_fields(self, strategy): """Verify all fields from a BidiWriteObjectResponse update the state.""" write_state = _WriteState( diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index e6fcd90f0b84..3f372aaaaccf 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py @@ -507,6 +507,67 @@ async def mock_execute(state, policy): mock_writer_open.assert_awaited() assert writer.persisted_size == 5 + @pytest.mark.asyncio + @pytest.mark.parametrize("known_size", [0, 2000, 2250]) + @pytest.mark.parametrize("reopened_size", [2000, 2250, 3000]) + async def test_later_append_retry_preserves_data( + self, mock_appendable_writer, known_size, reopened_size + ): + from google.api_core.retry_async import AsyncRetry + + writer = self._make_one(mock_appendable_writer["mock_client"]) + writer._is_stream_open = True + writer.offset = 2000 + writer.persisted_size = known_size + stream = mock_appendable_writer["mock_stream"] + writer.write_obj_stream = stream + stream.send.side_effect = [exceptions.ServiceUnavailable("disconnected"), None] + data = bytes(range(250)) * 4 + + async def reopen(metadata=None): + writer.persisted_size = reopened_size + writer.offset = reopened_size + writer._is_stream_open = True + + with mock.patch.object(writer, "open", side_effect=reopen) as open_mock: + await writer.append(data, retry_policy=AsyncRetry(initial=0, maximum=0)) + + open_mock.assert_awaited_once() + resent = [call.args[0] for call in stream.send.await_args_list[1:]] + assert ( + b"".join(r.checksummed_data.content for r in resent) + == data[reopened_size - 2000 :] + ) + if resent: + assert resent[0].write_offset == reopened_size + assert writer.offset == 3000 + + @pytest.mark.asyncio + @pytest.mark.parametrize("reopened_size", [1999, 3001]) + async def test_later_append_retry_rejects_unrecoverable_offset( + self, mock_appendable_writer, reopened_size + ): + from google.api_core.retry_async import AsyncRetry + + writer = self._make_one(mock_appendable_writer["mock_client"]) + writer._is_stream_open = True + writer.offset = 2000 + writer.persisted_size = 2000 + stream = mock_appendable_writer["mock_stream"] + writer.write_obj_stream = stream + stream.send.side_effect = [exceptions.ServiceUnavailable("disconnected"), None] + + async def reopen(metadata=None): + writer.persisted_size = reopened_size + writer._is_stream_open = True + + with mock.patch.object(writer, "open", side_effect=reopen): + with pytest.raises(ValueError, match="outside the current append buffer"): + await writer.append( + b"C" * 1000, retry_policy=AsyncRetry(initial=0, maximum=0) + ) + assert stream.send.await_count == 1 + @pytest.mark.asyncio async def test_append_unimplemented_string_raises(self, mock_appendable_writer): writer = self._make_one(mock_appendable_writer["mock_client"])