From 9d70ea80c068d6638ff873bcb29434074913c62c Mon Sep 17 00:00:00 2001 From: Pulkit Aggarwal Date: Mon, 3 Nov 2025 07:26:07 +0000 Subject: [PATCH 1/4] feat(experimental): Add read resumption strategy --- .../asyncio/retry/reads_strategy.py | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) create mode 100644 google/cloud/storage/_experimental/asyncio/retry/reads_strategy.py diff --git a/google/cloud/storage/_experimental/asyncio/retry/reads_strategy.py b/google/cloud/storage/_experimental/asyncio/retry/reads_strategy.py new file mode 100644 index 000000000..1ca8cd475 --- /dev/null +++ b/google/cloud/storage/_experimental/asyncio/retry/reads_strategy.py @@ -0,0 +1,82 @@ +import abc +from typing import Any, List + +from google.api_core import exceptions +from google.cloud import _storage_v2 as storage_v2 +from google.cloud.storage.exceptions import DataCorruption +from google_crc32c import Checksum +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, initial_length, user_buffer): + 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.""" + 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 exceptions.DataCorruption(f"Offset mismatch for read_id {read_id}") + + # Checksum Verification + checksummed_data = object_data_range.checksummed_data + data = checksummed_data.content + server_checksum = checksummed_data.crc32c + + client_crc32c = Checksum(data).digest() + client_checksum = int.from_bytes(client_crc32c, "big") + if server_checksum != client_checksum: + raise DataCorruption( + response, + f"Checksum mismatch for read_id {read_id}. " + f"Server sent {server_checksum}, client calculated {client_checksum}.", + ) + + 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 exceptions.DataCorruption(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 From 961464d5e407040b5e35540274250b47720a29a3 Mon Sep 17 00:00:00 2001 From: Pulkit Aggarwal Date: Mon, 3 Nov 2025 09:28:26 +0000 Subject: [PATCH 2/4] add unit tests --- ...rategy.py => reads_resumption_strategy.py} | 20 +- .../retry/test_reads_resumption_strategy.py | 203 ++++++++++++++++++ 2 files changed, 206 insertions(+), 17 deletions(-) rename google/cloud/storage/_experimental/asyncio/retry/{reads_strategy.py => reads_resumption_strategy.py} (76%) create mode 100644 tests/unit/asyncio/retry/test_reads_resumption_strategy.py diff --git a/google/cloud/storage/_experimental/asyncio/retry/reads_strategy.py b/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py similarity index 76% rename from google/cloud/storage/_experimental/asyncio/retry/reads_strategy.py rename to google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py index 1ca8cd475..a469b8abc 100644 --- a/google/cloud/storage/_experimental/asyncio/retry/reads_strategy.py +++ b/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py @@ -1,7 +1,6 @@ import abc from typing import Any, List -from google.api_core import exceptions from google.cloud import _storage_v2 as storage_v2 from google.cloud.storage.exceptions import DataCorruption from google_crc32c import Checksum @@ -48,22 +47,9 @@ def update_state_from_response(self, response: storage_v2.BidiReadObjectResponse # Offset Verification chunk_offset = object_data_range.read_range.read_offset if chunk_offset != read_state.next_expected_offset: - raise exceptions.DataCorruption(f"Offset mismatch for read_id {read_id}") - - # Checksum Verification - checksummed_data = object_data_range.checksummed_data - data = checksummed_data.content - server_checksum = checksummed_data.crc32c - - client_crc32c = Checksum(data).digest() - client_checksum = int.from_bytes(client_crc32c, "big") - if server_checksum != client_checksum: - raise DataCorruption( - response, - f"Checksum mismatch for read_id {read_id}. " - f"Server sent {server_checksum}, client calculated {client_checksum}.", - ) + 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 @@ -73,7 +59,7 @@ def update_state_from_response(self, response: storage_v2.BidiReadObjectResponse 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 exceptions.DataCorruption(f"Byte count mismatch for read_id {read_id}") + 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.""" 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..d36e55f44 --- /dev/null +++ b/tests/unit/asyncio/retry/test_reads_resumption_strategy.py @@ -0,0 +1,203 @@ +# 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, +) + + +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 = {1: 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, 1) + + def test_generate_requests_multiple_incomplete(self): + """Test generating requests for multiple incomplete downloads.""" + read_state1 = _DownloadState(0, 100, io.BytesIO()) + read_state1.bytes_written = 50 + read_state2 = _DownloadState(200, 100, io.BytesIO()) + state = {1: read_state1, 2: 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 == 1) + req2 = next(request for request in requests if request.read_id == 2) + + 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 = {1: 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_from_response_success(self): + """Test updating state from a successful response.""" + buffer = io.BytesIO() + read_state = _DownloadState(0, 100, buffer) + state = {1: 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=1, 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 = {1: read_state} + read_strategy = _ReadResumptionStrategy() + + response = storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + storage_v2.types.ObjectRangeData( + read_range=storage_v2.ReadRange(read_id=1, 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 = {1: read_state} + read_strategy = _ReadResumptionStrategy() + + response = storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + storage_v2.types.ObjectRangeData( + read_range=storage_v2.ReadRange(read_id=1, 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 = {1: read_state} + read_strategy = _ReadResumptionStrategy() + + response = storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + storage_v2.types.ObjectRangeData( + read_range=storage_v2.ReadRange(read_id=1, 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 = {1: read_state} + read_strategy = _ReadResumptionStrategy() + + response = storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + storage_v2.types.ObjectRangeData( + read_range=storage_v2.ReadRange(read_id=1, 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)) From 538a1b9202d99bb09079276837e3da8fa946213a Mon Sep 17 00:00:00 2001 From: Pulkit Aggarwal Date: Mon, 3 Nov 2025 09:47:43 +0000 Subject: [PATCH 3/4] minor fixes --- .../storage/_experimental/asyncio/retry/base_strategy.py | 2 +- .../asyncio/retry/reads_resumption_strategy.py | 6 ++---- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py b/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py index e4e443815..2f1248f60 100644 --- a/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py +++ b/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py @@ -31,7 +31,7 @@ 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 diff --git a/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py b/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py index a469b8abc..805e99dc1 100644 --- a/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py +++ b/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py @@ -1,16 +1,14 @@ -import abc -from typing import Any, List +from typing import Any, List, IO from google.cloud import _storage_v2 as storage_v2 from google.cloud.storage.exceptions import DataCorruption -from google_crc32c import Checksum 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, initial_length, user_buffer): + 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 From a0727a41fb3071ab81c7b289c801cf0b6a7af704 Mon Sep 17 00:00:00 2001 From: Pulkit Aggarwal Date: Wed, 5 Nov 2025 06:14:51 +0000 Subject: [PATCH 4/4] resolving comments --- .../asyncio/retry/base_strategy.py | 3 ++ .../retry/reads_resumption_strategy.py | 7 +++- .../retry/test_reads_resumption_strategy.py | 37 ++++++++++--------- 3 files changed, 29 insertions(+), 18 deletions(-) diff --git a/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py b/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py index 2f1248f60..8a06608b2 100644 --- a/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py +++ b/google/cloud/storage/_experimental/asyncio/retry/base_strategy.py @@ -38,6 +38,9 @@ def update_state_from_response(self, response: Any, state: Any) -> None: 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 index 805e99dc1..1ff038827 100644 --- a/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py +++ b/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py @@ -21,7 +21,12 @@ 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.""" + """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: diff --git a/tests/unit/asyncio/retry/test_reads_resumption_strategy.py b/tests/unit/asyncio/retry/test_reads_resumption_strategy.py index d36e55f44..7153844ad 100644 --- a/tests/unit/asyncio/retry/test_reads_resumption_strategy.py +++ b/tests/unit/asyncio/retry/test_reads_resumption_strategy.py @@ -23,6 +23,8 @@ _ReadResumptionStrategy, ) +_READ_ID = 1 + class TestDownloadState(unittest.TestCase): def test_initialization(self): @@ -45,7 +47,7 @@ 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 = {1: read_state} + state = {_READ_ID: read_state} read_strategy = _ReadResumptionStrategy() requests = read_strategy.generate_requests(state) @@ -53,21 +55,22 @@ def test_generate_requests_single_incomplete(self): 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, 1) + 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 = {1: read_state1, 2: read_state2} + 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 == 1) - req2 = next(request for request in requests if request.read_id == 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) @@ -78,7 +81,7 @@ 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 = {1: read_state} + state = {_READ_ID: read_state} read_strategy = _ReadResumptionStrategy() requests = read_strategy.generate_requests(state) @@ -91,18 +94,18 @@ def test_generate_requests_empty_state(self): requests = read_strategy.generate_requests({}) self.assertEqual(len(requests), 0) - def test_update_state_from_response_success(self): + 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 = {1: read_state} + 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=1, read_offset=0, read_length=len(data)), + read_range=storage_v2.ReadRange(read_id=_READ_ID, read_offset=0, read_length=len(data)), checksummed_data=storage_v2.ChecksummedData(content=data), ) ] @@ -119,13 +122,13 @@ 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 = {1: read_state} + 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=1, read_offset=0, read_length=4), + read_range=storage_v2.ReadRange(read_id=_READ_ID, read_offset=0, read_length=4), checksummed_data=storage_v2.ChecksummedData(content=b"data"), ) ] @@ -138,13 +141,13 @@ def test_update_state_from_response_offset_mismatch(self): 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 = {1: read_state} + 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=1, read_offset=0, read_length=4), + 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, ) @@ -160,13 +163,13 @@ def test_update_state_from_response_completes_download(self): buffer = io.BytesIO() data = b"test_data" read_state = _DownloadState(0, len(data), buffer) - state = {1: read_state} + 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=1, read_offset=0, read_length=len(data)), + 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, ) @@ -184,13 +187,13 @@ def test_update_state_from_response_completes_download_zero_length(self): buffer = io.BytesIO() data = b"test_data" read_state = _DownloadState(0, 0, buffer) - state = {1: read_state} + 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=1, read_offset=0, read_length=len(data)), + 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, )