feat(experimental): Add read resumption strategy

google/cloud/storage/_experimental/asyncio/retry/base_strategy.py

@@ -31,13 +31,16 @@ def generate_requests(self, state: Any) -> Iterable[Any]:
         pass
 
     @abc.abstractmethod
-    def update_state_from_response(self, state: Any) -> None:
+    def update_state_from_response(self, response: Any, state: Any) -> None:
         """Updates the state based on a successful server response.
 
         This method is called for every message received from the server. It is
         responsible for processing the response and updating the shared state
         object.
 
+        :type response: Any
+        :param response: The response message received from the server.
+
         :type state: Any
         :param state: The shared state object for the operation, which will be
                       mutated by this method.

google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py

@@ -0,0 +1,71 @@
+from typing import Any, List, IO
+
+from google.cloud import _storage_v2 as storage_v2
+from google.cloud.storage.exceptions import DataCorruption
+from google.cloud.storage._experimental.asyncio.retry.base_strategy import (
+    _BaseResumptionStrategy,
+)
+
+class _DownloadState:
+    """A helper class to track the state of a single range download."""
+    def __init__(self, initial_offset: int, initial_length: int, user_buffer: IO[bytes]):
+        self.initial_offset = initial_offset
+        self.initial_length = initial_length
+        self.user_buffer = user_buffer
+        self.bytes_written = 0
+        self.next_expected_offset = initial_offset
+        self.is_complete = False
+
+
+class _ReadResumptionStrategy(_BaseResumptionStrategy):
+    """The concrete resumption strategy for bidi reads."""
+
+    def generate_requests(self, state: dict) -> List[storage_v2.ReadRange]:
+        """Generates new ReadRange requests for all incomplete downloads.
+
+        :type state: dict
+        :param state: A dictionary mapping a read_id to its corresponding
+                  _DownloadState object.
+        """
+        pending_requests = []
+        for read_id, read_state in state.items():
+            if not read_state.is_complete:
+                new_offset = read_state.initial_offset + read_state.bytes_written
+                new_length = read_state.initial_length - read_state.bytes_written
+
+                new_request = storage_v2.ReadRange(
+                    read_offset=new_offset,
+                    read_length=new_length,
+                    read_id=read_id,
+                )
+                pending_requests.append(new_request)
+        return pending_requests
+
+    def update_state_from_response(self, response: storage_v2.BidiReadObjectResponse, state: dict) -> None:
+        """Processes a server response, performs integrity checks, and updates state."""
+        for object_data_range in response.object_data_ranges:
+            read_id = object_data_range.read_range.read_id
+            read_state = state[read_id]
+
+            # Offset Verification
+            chunk_offset = object_data_range.read_range.read_offset
+            if chunk_offset != read_state.next_expected_offset:
+                raise DataCorruption(response, f"Offset mismatch for read_id {read_id}")
+
+            data = object_data_range.checksummed_data.content
+            chunk_size = len(data)
+            read_state.bytes_written += chunk_size
+            read_state.next_expected_offset += chunk_size
+            read_state.user_buffer.write(data)
+
+            # Final Byte Count Verification
+            if object_data_range.range_end:
+                read_state.is_complete = True
+                if read_state.initial_length != 0 and read_state.bytes_written != read_state.initial_length:
+                    raise DataCorruption(response, f"Byte count mismatch for read_id {read_id}")
+
+    async def recover_state_on_failure(self, error: Exception, state: Any) -> None:
+        """Handles BidiReadObjectRedirectError for reads."""
+        # This would parse the gRPC error details, extract the routing_token,
+        # and store it on the shared state object.
+        pass

tests/unit/asyncio/retry/test_reads_resumption_strategy.py

@@ -0,0 +1,206 @@
+# Copyright 2025 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
+#
+#     http://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.
+
+import io
+import unittest
+import pytest
+from google.cloud.storage.exceptions import DataCorruption
+
+from google.cloud import _storage_v2 as storage_v2
+from google.cloud.storage._experimental.asyncio.retry.reads_resumption_strategy import (
+    _DownloadState,
+    _ReadResumptionStrategy,
+)
+
+_READ_ID = 1
+
+
+class TestDownloadState(unittest.TestCase):
+    def test_initialization(self):
+        """Test that _DownloadState initializes correctly."""
+        initial_offset = 10
+        initial_length = 100
+        user_buffer = io.BytesIO()
+        state = _DownloadState(initial_offset, initial_length, user_buffer)
+
+        self.assertEqual(state.initial_offset, initial_offset)
+        self.assertEqual(state.initial_length, initial_length)
+        self.assertEqual(state.user_buffer, user_buffer)
+        self.assertEqual(state.bytes_written, 0)
+        self.assertEqual(state.next_expected_offset, initial_offset)
+        self.assertFalse(state.is_complete)
+
+
+class TestReadResumptionStrategy(unittest.TestCase):
+    def test_generate_requests_single_incomplete(self):
+        """Test generating a request for a single incomplete download."""
+        read_state = _DownloadState(0, 100, io.BytesIO())
+        read_state.bytes_written = 20
+        state = {_READ_ID: read_state}
+
+        read_strategy = _ReadResumptionStrategy()
+        requests = read_strategy.generate_requests(state)
+
+        self.assertEqual(len(requests), 1)
+        self.assertEqual(requests[0].read_offset, 20)
+        self.assertEqual(requests[0].read_length, 80)
+        self.assertEqual(requests[0].read_id, _READ_ID)
+
+    def test_generate_requests_multiple_incomplete(self):
+        """Test generating requests for multiple incomplete downloads."""
+        read_id2 = 2
+        read_state1 = _DownloadState(0, 100, io.BytesIO())
+        read_state1.bytes_written = 50
+        read_state2 = _DownloadState(200, 100, io.BytesIO())
+        state = {_READ_ID: read_state1, read_id2: read_state2}
+
+        read_strategy = _ReadResumptionStrategy()
+        requests = read_strategy.generate_requests(state)
+
+        self.assertEqual(len(requests), 2)
+        req1 = next(request for request in requests if request.read_id == _READ_ID)
+        req2 = next(request for request in requests if request.read_id == read_id2)
+
+        self.assertEqual(req1.read_offset, 50)
+        self.assertEqual(req1.read_length, 50)
+        self.assertEqual(req2.read_offset, 200)
+        self.assertEqual(req2.read_length, 100)
+
+    def test_generate_requests_with_complete(self):
+        """Test that no request is generated for a completed download."""
+        read_state = _DownloadState(0, 100, io.BytesIO())
+        read_state.is_complete = True
+        state = {_READ_ID: read_state}
+
+        read_strategy = _ReadResumptionStrategy()
+        requests = read_strategy.generate_requests(state)
+
+        self.assertEqual(len(requests), 0)
+
+    def test_generate_requests_empty_state(self):
+        """Test generating requests with an empty state."""
+        read_strategy = _ReadResumptionStrategy()
+        requests = read_strategy.generate_requests({})
+        self.assertEqual(len(requests), 0)
+
+    def test_update_state_processes_single_chunk_successfully(self):
+        """Test updating state from a successful response."""
+        buffer = io.BytesIO()
+        read_state = _DownloadState(0, 100, buffer)
+        state = {_READ_ID: read_state}
+        data = b"test_data"
+        read_strategy = _ReadResumptionStrategy()
+
+        response = storage_v2.BidiReadObjectResponse(
+            object_data_ranges=[
+                storage_v2.types.ObjectRangeData(
+                    read_range=storage_v2.ReadRange(read_id=_READ_ID, read_offset=0, read_length=len(data)),
+                    checksummed_data=storage_v2.ChecksummedData(content=data),
+                )
+            ]
+        )
+
+        read_strategy.update_state_from_response(response, state)
+
+        self.assertEqual(read_state.bytes_written, len(data))
+        self.assertEqual(read_state.next_expected_offset, len(data))
+        self.assertFalse(read_state.is_complete)
+        self.assertEqual(buffer.getvalue(), data)
+
+    def test_update_state_from_response_offset_mismatch(self):
+        """Test that an offset mismatch raises DataCorruption."""
+        read_state = _DownloadState(0, 100, io.BytesIO())
+        read_state.next_expected_offset = 10
+        state = {_READ_ID: read_state}
+        read_strategy = _ReadResumptionStrategy()
+
+        response = storage_v2.BidiReadObjectResponse(
+            object_data_ranges=[
+                storage_v2.types.ObjectRangeData(
+                    read_range=storage_v2.ReadRange(read_id=_READ_ID, read_offset=0, read_length=4),
+                    checksummed_data=storage_v2.ChecksummedData(content=b"data"),
+                )
+            ]
+        )
+
+        with pytest.raises(DataCorruption) as exc_info:
+            read_strategy.update_state_from_response(response, state)
+        assert "Offset mismatch" in str(exc_info.value)
+
+    def test_update_state_from_response_final_byte_count_mismatch(self):
+        """Test that a final byte count mismatch raises DataCorruption."""
+        read_state = _DownloadState(0, 100, io.BytesIO())
+        state = {_READ_ID: read_state}
+        read_strategy = _ReadResumptionStrategy()
+
+        response = storage_v2.BidiReadObjectResponse(
+            object_data_ranges=[
+                storage_v2.types.ObjectRangeData(
+                    read_range=storage_v2.ReadRange(read_id=_READ_ID, read_offset=0, read_length=4),
+                    checksummed_data=storage_v2.ChecksummedData(content=b"data"),
+                    range_end=True,
+                )
+            ]
+        )
+
+        with pytest.raises(DataCorruption) as exc_info:
+            read_strategy.update_state_from_response(response, state)
+        assert "Byte count mismatch" in str(exc_info.value)
+
+    def test_update_state_from_response_completes_download(self):
+        """Test that the download is marked complete on range_end."""
+        buffer = io.BytesIO()
+        data = b"test_data"
+        read_state = _DownloadState(0, len(data), buffer)
+        state = {_READ_ID: read_state}
+        read_strategy = _ReadResumptionStrategy()
+
+        response = storage_v2.BidiReadObjectResponse(
+            object_data_ranges=[
+                storage_v2.types.ObjectRangeData(
+                    read_range=storage_v2.ReadRange(read_id=_READ_ID, read_offset=0, read_length=len(data)),
+                    checksummed_data=storage_v2.ChecksummedData(content=data),
+                    range_end=True,
+                )
+            ]
+        )
+
+        read_strategy.update_state_from_response(response, state)
+
+        self.assertTrue(read_state.is_complete)
+        self.assertEqual(read_state.bytes_written, len(data))
+        self.assertEqual(buffer.getvalue(), data)
+
+    def test_update_state_from_response_completes_download_zero_length(self):
+        """Test completion for a download with initial_length of 0."""
+        buffer = io.BytesIO()
+        data = b"test_data"
+        read_state = _DownloadState(0, 0, buffer)
+        state = {_READ_ID: read_state}
+        read_strategy = _ReadResumptionStrategy()
+
+        response = storage_v2.BidiReadObjectResponse(
+            object_data_ranges=[
+                storage_v2.types.ObjectRangeData(
+                    read_range=storage_v2.ReadRange(read_id=_READ_ID, read_offset=0, read_length=len(data)),
+                    checksummed_data=storage_v2.ChecksummedData(content=data),
+                    range_end=True,
+                )
+            ]
+        )
+
+        read_strategy.update_state_from_response(response, state)
+
+        self.assertTrue(read_state.is_complete)
+        self.assertEqual(read_state.bytes_written, len(data))