feat(experimental): Add read resumption strategy (#1599)

* feat(experimental): Add read resumption strategy * add unit tests * minor fixes * resolving comments

Author: Pulkit0110

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