diff --git a/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py b/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py index e4e443815..8a06608b2 100644 --- a/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py +++ b/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. diff --git a/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py b/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py new file mode 100644 index 000000000..1ff038827 --- /dev/null +++ b/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 diff --git a/tests/unit/asyncio/retry/test_reads_resumption_strategy.py b/tests/unit/asyncio/retry/test_reads_resumption_strategy.py new file mode 100644 index 000000000..7153844ad --- /dev/null +++ b/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))