fix(storage): preserve appended data when retrying a stream

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(

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

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(

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"])