From 6c47a4c167bcc2e30631009374103579e9bbb133 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 d61f0ae02d8c39ea8e60d661e364d78e3d1507dd 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 --- .../asyncio/retry/reads_strategy.py | 82 ------------------- 1 file changed, 82 deletions(-) delete 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 deleted file mode 100644 index 1ca8cd475..000000000 --- a/google/cloud/storage/_experimental/asyncio/retry/reads_strategy.py +++ /dev/null @@ -1,82 +0,0 @@ -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 3c74ca7d65bdb0312a020528b9a588ae0248e364 Mon Sep 17 00:00:00 2001 From: Pulkit Aggarwal Date: Mon, 3 Nov 2025 11:29:14 +0000 Subject: [PATCH 3/4] feat(experimental): Handle BidiReadObjectRedirectedError for bidi reads --- .../retry/reads_resumption_strategy.py | 6 ++++-- .../retry/test_reads_resumption_strategy.py | 20 +++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) 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 1ff038827..ab25a6170 100644 --- a/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py +++ b/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py @@ -5,6 +5,7 @@ from google.cloud.storage._experimental.asyncio.retry.base_strategy import ( _BaseResumptionStrategy, ) +from google.cloud._storage_v2.types.storage import BidiReadObjectRedirectedError class _DownloadState: """A helper class to track the state of a single range download.""" @@ -65,7 +66,8 @@ def update_state_from_response(self, response: storage_v2.BidiReadObjectResponse 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.""" + """Handles BidiReadObjectRedirectedError for reads.""" # This would parse the gRPC error details, extract the routing_token, # and store it on the shared state object. - pass + if isinstance(error, BidiReadObjectRedirectedError): + state['routing_token'] = error.routing_token diff --git a/tests/unit/asyncio/retry/test_reads_resumption_strategy.py b/tests/unit/asyncio/retry/test_reads_resumption_strategy.py index 7153844ad..f18329138 100644 --- a/tests/unit/asyncio/retry/test_reads_resumption_strategy.py +++ b/tests/unit/asyncio/retry/test_reads_resumption_strategy.py @@ -16,12 +16,14 @@ import unittest import pytest from google.cloud.storage.exceptions import DataCorruption +from google.api_core import exceptions from google.cloud import _storage_v2 as storage_v2 from google.cloud.storage._experimental.asyncio.retry.reads_resumption_strategy import ( _DownloadState, _ReadResumptionStrategy, ) +from google.cloud._storage_v2.types.storage import BidiReadObjectRedirectedError _READ_ID = 1 @@ -204,3 +206,21 @@ def test_update_state_from_response_completes_download_zero_length(self): self.assertTrue(read_state.is_complete) self.assertEqual(read_state.bytes_written, len(data)) + + async def test_recover_state_on_failure_handles_redirect(self): + """Verify recover_state_on_failure correctly extracts routing_token.""" + strategy = _ReadResumptionStrategy() + + state = {} + self.assertIsNone(state.get("routing_token")) + + dummy_token = "dummy-routing-token" + redirect_error = BidiReadObjectRedirectedError( + "Redirected by server", routing_token=dummy_token + ) + + final_error = exceptions.RetryError("Retry failed", cause=redirect_error) + + await strategy.recover_state_on_failure(final_error, state) + + self.assertEqual(state.get("routing_token"), dummy_token) From 667d7198f49b9f4e2b37d1187d2d7d180da70549 Mon Sep 17 00:00:00 2001 From: Pulkit Aggarwal Date: Tue, 4 Nov 2025 08:10:57 +0000 Subject: [PATCH 4/4] resolve comments --- .../_experimental/asyncio/retry/reads_resumption_strategy.py | 5 +++-- tests/unit/asyncio/retry/test_reads_resumption_strategy.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) 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 ab25a6170..9c82cf7f0 100644 --- a/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py +++ b/google/cloud/storage/_experimental/asyncio/retry/reads_resumption_strategy.py @@ -69,5 +69,6 @@ async def recover_state_on_failure(self, error: Exception, state: Any) -> None: """Handles BidiReadObjectRedirectedError for reads.""" # This would parse the gRPC error details, extract the routing_token, # and store it on the shared state object. - if isinstance(error, BidiReadObjectRedirectedError): - state['routing_token'] = error.routing_token + cause = getattr(error, "cause", error) + if isinstance(cause, BidiReadObjectRedirectedError): + state['routing_token'] = cause.routing_token diff --git a/tests/unit/asyncio/retry/test_reads_resumption_strategy.py b/tests/unit/asyncio/retry/test_reads_resumption_strategy.py index f18329138..fdde74f62 100644 --- a/tests/unit/asyncio/retry/test_reads_resumption_strategy.py +++ b/tests/unit/asyncio/retry/test_reads_resumption_strategy.py @@ -216,7 +216,7 @@ async def test_recover_state_on_failure_handles_redirect(self): dummy_token = "dummy-routing-token" redirect_error = BidiReadObjectRedirectedError( - "Redirected by server", routing_token=dummy_token + routing_token=dummy_token ) final_error = exceptions.RetryError("Retry failed", cause=redirect_error)