diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py index 32757bf754b7..cb1f79cfab10 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py @@ -40,6 +40,7 @@ _BidiStreamRetryManager, ) from google.cloud.storage.asyncio.retry.reads_resumption_strategy import ( + _CRC32C_MAX_PENDING, _DownloadState, _ReadResumptionStrategy, ) @@ -90,7 +91,7 @@ class AsyncMultiRangeDownloader: mrd = await AsyncMultiRangeDownloader.create_mrd( client, bucket_name="chandrasiri-rs", object_name="test_open9" ) - my_buff1 = open('my_fav_file.txt', 'wb') + my_buff1 = open("my_fav_file.txt", "wb") my_buff2 = BytesIO() my_buff3 = BytesIO() my_buff4 = any_object_which_provides_BytesIO_like_interface() @@ -394,6 +395,13 @@ async def download_ranges( :type enable_checksum: bool :param enable_checksum: (Optional) If True, checksums are verified for downloaded data. Defaults to True. + Chunks of ``GCS_PYTHON_CLIENT_CRC32C_OFFLOAD_MIN_BYTES`` bytes or + more (default 1 MiB) are verified on a worker thread and written to + their buffer once the checksum matches; chunks still waiting are + discarded if the call is cancelled or a mismatch is found, and a + retry requests them again. Setting + ``GCS_PYTHON_CLIENT_CRC32C_OFFLOAD_MIN_BYTES=0`` verifies every + chunk inline. :raises ValueError: if the underlying bidi-GRPC stream is not open. :raises ValueError: if the length of read_ranges is more than 1000. @@ -447,6 +455,7 @@ async def download_ranges( read_ids = set(download_states.keys()) queue = self._multiplexer.register(read_ids) + strategy = _ReadResumptionStrategy() try: attempt_count = 0 @@ -522,20 +531,45 @@ async def generator(): data_range.read_range.read_id ) yield item + # The strategy has consumed `item` by now: write the + # chunks whose checksums are done, and keep the + # backlog bounded. + await strategy.verify_pending_checksums( + state, _CRC32C_MAX_PENDING + ) + + # Write (or reject) every held chunk before the attempt + # ends, so a mismatch is raised inside the retry loop, + # as an inline one is. + await strategy.verify_pending_checksums(state) return generator() - strategy = _ReadResumptionStrategy() retry_manager = _BidiStreamRetryManager( strategy, send_and_recv_via_multiplexer ) try: - await retry_manager.execute(initial_state, retry_policy) + try: + await retry_manager.execute(initial_state, retry_policy) + except DataCorruption: + raise + except Exception: + # Chunks received before the stream error still have + # in-flight checksum futures in initial_state. Verify them + # before raising the stream error so an earlier corrupt + # chunk raises DataCorruption (matching inline verification) + # and valid chunks already received are written. + await strategy.verify_pending_checksums(initial_state) + raise + await strategy.verify_pending_checksums(initial_state) except DataCorruption: if self.is_stream_open: await self.close() raise + finally: + # Drop chunks still held, e.g. after cancellation. + strategy.discard_pending_checksums(initial_state) if initial_state.get("read_handle"): self.read_handle = initial_state["read_handle"] diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/reads_resumption_strategy.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/reads_resumption_strategy.py index 6cf17af19089..e4f89070a881 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/reads_resumption_strategy.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/reads_resumption_strategy.py @@ -12,8 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio +import collections +import concurrent.futures import logging -from typing import IO, Any, Dict, List +import os +from typing import IO, Any, Dict, List, NamedTuple, Optional import google_crc32c @@ -32,6 +36,71 @@ logger = logging.getLogger(__name__) +def _int_from_env(name: str, default: int) -> int: + """Reads an integer tuning knob from the environment; bad values are ignored.""" + try: + return int(os.environ.get(name, default)) + except ValueError: + return default + + +# Chunks at least this large are checksummed on a worker thread instead of the +# event-loop thread; smaller chunks are hashed inline, and 0 (or any value +# below 1) hashes every chunk inline. The 1 MiB default matches google_crc32c's +# GIL-release threshold, so offloaded hashes over ``bytes``-backed chunks +# (received chunks are) always run with the GIL released. +_CRC32C_OFFLOAD_MIN_BYTES = _int_from_env( + "GCS_PYTHON_CLIENT_CRC32C_OFFLOAD_MIN_BYTES", + getattr(google_crc32c, "GIL_RELEASE_THRESHOLD", 1024 * 1024), +) +# Upper bound on chunks held for verification per download_ranges() call. +# Chunks are written only once their checksum matches, so this keeps a slow +# worker from pinning an unbounded number of received messages. +_CRC32C_MAX_PENDING = 32 + + +def _new_crc32c_executor() -> concurrent.futures.ThreadPoolExecutor: + return concurrent.futures.ThreadPoolExecutor( + max_workers=2, thread_name_prefix="gcs-crc32c" + ) + + +# Process-wide pool shared by all clients; its threads start on the first submit(). +_crc32c_executor = _new_crc32c_executor() + + +def _reset_crc32c_executor() -> None: + """Replaces the parent's pool in a forked child, where it has no threads.""" + global _crc32c_executor + _crc32c_executor = _new_crc32c_executor() + + +if hasattr(os, "register_at_fork"): # pragma: no branch - not on Windows + os.register_at_fork(after_in_child=_reset_crc32c_executor) + + +def _submit_crc32c(data: Any) -> Optional[concurrent.futures.Future]: + """Starts hashing ``data`` on the pool; None if the pool refuses work.""" + try: + return _crc32c_executor.submit(google_crc32c.value, data) + except RuntimeError: + # Raised once interpreter shutdown has begun; hash inline instead. + return None + + +class _PendingChunk(NamedTuple): + """A received chunk that is not written to its buffer yet.""" + + read_id: int + data: Any + range_end: bool + response: Any + server_checksum: Optional[int] + # Worker-side CRC32C of ``data``; None if it was verified (or needs no + # verification) on receipt and only waits for the chunks ahead of it. + checksum_future: Optional[concurrent.futures.Future] + + class _DownloadState: """A helper class to track the state of a single range download.""" @@ -47,6 +116,9 @@ def __init__( self.initial_length = initial_length self.user_buffer = user_buffer self.bytes_written = 0 + # Offset of the next chunk to accept. Runs ahead of + # initial_offset + bytes_written while received chunks are held for + # checksum verification. self.next_expected_offset = initial_offset self.is_complete = False self.is_full_object_read = is_full_object_read @@ -136,64 +208,164 @@ def update_state_from_response( # Checksum Verification checksummed_data = object_data_range.checksummed_data data = checksummed_data.content + server_checksum = None + checksum_future = None if checksum_enabled and checksummed_data.HasField("crc32c"): server_checksum = checksummed_data.crc32c - client_checksum = google_crc32c.value(data) - 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}.", - ) + if ( + _CRC32C_OFFLOAD_MIN_BYTES > 0 + and len(data) >= _CRC32C_OFFLOAD_MIN_BYTES + and read_state.rolling_checksum is None + ): + # Hash off the loop thread; the chunk is written by + # verify_pending_checksums() once the checksum matches. + # Full-object reads are excluded: they already fold every + # chunk into rolling_checksum on this thread, so a second + # hash on a worker would only add CPU work. + checksum_future = _submit_crc32c(data) + if checksum_future is None: + client_checksum = google_crc32c.value(data) + 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}.", + ) + + read_state.next_expected_offset += len(data) + chunk = _PendingChunk( + read_id, + data, + object_data_range.range_end, + response, + server_checksum, + checksum_future, + ) + pending = state.get("pending_checksums") + if checksum_future is None and not pending: + self._write_chunk(state, chunk) + else: + # Hold the chunk; chunks are written in arrival order. + if pending is None: + pending = state["pending_checksums"] = collections.deque() + pending.append(chunk) + def _write_chunk(self, state: Dict[str, Any], chunk: _PendingChunk) -> None: + """Writes a verified chunk to its buffer and runs the end-of-range checks.""" + read_state = state["download_states"][chunk.read_id] + data = chunk.data + try: # Update State & Write Data - chunk_size = len(data) read_state.user_buffer.write(data) # Commit updates only after the write succeeds - if checksum_enabled and read_state.rolling_checksum is not None: + if ( + state.get("enable_checksum", True) + and read_state.rolling_checksum is not None + ): read_state.rolling_checksum.update(data) - read_state.bytes_written += chunk_size - read_state.next_expected_offset += chunk_size + read_state.bytes_written += len(data) + except BaseException: + # Keep every read resumable right after its written bytes. + self.discard_pending_checksums(state) + raise - # Final Byte Count & Full Object Checksum 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 - ): + # Final Byte Count & Full Object Checksum Verification + if chunk.range_end: + read_state.is_complete = True + if ( + read_state.initial_length != 0 + and read_state.bytes_written > read_state.initial_length + ): + raise DataCorruption( + chunk.response, + f"Byte count mismatch for read_id {chunk.read_id}. " + f"Expected {read_state.initial_length}, got {read_state.bytes_written}", + ) + + # Perform full-object checksum verification once the stream finishes. + if ( + read_state.is_full_object_read + and state.get("enable_checksum", True) + and read_state.rolling_checksum is not None + ): + full_obj_server_crc32c = state.get("full_obj_server_crc32c") + if full_obj_server_crc32c is not None: + # Use standard big-endian byte conversion to retrieve the rolling checksum value. + client_checksum = int.from_bytes( + read_state.rolling_checksum.digest(), + byteorder="big", + ) + if client_checksum != full_obj_server_crc32c: + raise DataCorruption( + chunk.response, + f"Full object checksum mismatch for read_id {chunk.read_id}. " + f"Server authoritative crc32c: {full_obj_server_crc32c}, client calculated rolling: {client_checksum}.", + ) + + async def verify_pending_checksums( + self, state: Dict[str, Any], max_pending: int = 0 + ) -> None: + """Writes held chunks, in arrival order, once their checksums match. + + Returns when at most ``max_pending`` chunks are still held; only an + unfinished checksum at the head of the queue is waited for. On a + mismatch, raises DataCorruption after discarding every held chunk + (see discard_pending_checksums()), so nothing from the corrupt chunk + on is written and a retry requests it again; a checksum that failed + on the worker is re-raised the same way. + """ + pending = state.get("pending_checksums") + while pending: + chunk = pending[0] + if chunk.checksum_future is not None: + try: + if not chunk.checksum_future.done(): + if len(pending) <= max_pending: + return + await asyncio.wrap_future(chunk.checksum_future) + client_checksum = chunk.checksum_future.result() + except BaseException: + self.discard_pending_checksums(state) + raise + if client_checksum != chunk.server_checksum: + self.discard_pending_checksums(state) raise DataCorruption( - response, - f"Byte count mismatch for read_id {read_id}. " - f"Expected {read_state.initial_length}, got {read_state.bytes_written}", + chunk.response, + f"Checksum mismatch for read_id {chunk.read_id}. " + f"Server sent {chunk.server_checksum}, client calculated {client_checksum}.", ) + pending.popleft() + self._write_chunk(state, chunk) - # Perform full-object checksum verification once the stream finishes. - if ( - read_state.is_full_object_read - and checksum_enabled - and read_state.rolling_checksum is not None - ): - full_obj_server_crc32c = state.get("full_obj_server_crc32c") - if full_obj_server_crc32c is not None: - # Use standard big-endian byte conversion to retrieve the rolling checksum value. - client_checksum = int.from_bytes( - read_state.rolling_checksum.digest(), - byteorder="big", - ) - if client_checksum != full_obj_server_crc32c: - raise DataCorruption( - response, - f"Full object checksum mismatch for read_id {read_id}. " - f"Server authoritative crc32c: {full_obj_server_crc32c}, client calculated rolling: {client_checksum}.", - ) + def discard_pending_checksums(self, state: Dict[str, Any]) -> None: + """Drops held chunks unwritten and rewinds each read to its written end. + + Checksums not yet started on the pool are cancelled. A retry requests + the dropped bytes again. + """ + pending = state.get("pending_checksums") + if pending: + for chunk in pending: + if chunk.checksum_future is not None: + chunk.checksum_future.cancel() + pending.clear() + for read_state in state["download_states"].values(): + read_state.next_expected_offset = ( + read_state.initial_offset + read_state.bytes_written + ) async def recover_state_on_failure(self, error: Exception, state: Any) -> None: - """Handles BidiReadObjectRedirectedError for reads.""" + """Handles BidiReadObjectRedirectedError, then writes held chunks. + + A retry resumes each read right after its written bytes, so held + chunks are verified and written first; a corrupt one raises + DataCorruption instead. + """ routing_token, read_handle = _handle_redirect(error) if routing_token: state["routing_token"] = routing_token if read_handle: state["read_handle"] = read_handle + await self.verify_pending_checksums(state) diff --git a/packages/google-cloud-storage/tests/unit/asyncio/retry/test_reads_resumption_strategy.py b/packages/google-cloud-storage/tests/unit/asyncio/retry/test_reads_resumption_strategy.py index 841ea655626e..dc7247b8ab8f 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/retry/test_reads_resumption_strategy.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/retry/test_reads_resumption_strategy.py @@ -13,14 +13,20 @@ # limitations under the License. import asyncio +import concurrent.futures import io +import os +import signal import unittest +import warnings +from unittest import mock import google_crc32c from google.api_core import exceptions from google.cloud import _storage_v2 as storage_v2 from google.cloud._storage_v2.types.storage import BidiReadObjectRedirectedError +from google.cloud.storage.asyncio.retry import reads_resumption_strategy from google.cloud.storage.asyncio.retry.reads_resumption_strategy import ( _DownloadState, _ReadResumptionStrategy, @@ -31,6 +37,27 @@ LOGGER_NAME = "google.cloud.storage.asyncio.retry.reads_resumption_strategy" +class TestIntFromEnv(unittest.TestCase): + def test_reads_integer(self): + with mock.patch.dict(os.environ, {"GCS_TEST_KNOB": "123"}): + self.assertEqual( + reads_resumption_strategy._int_from_env("GCS_TEST_KNOB", 7), 123 + ) + + def test_default_when_unset(self): + with mock.patch.dict(os.environ, {}, clear=True): + self.assertEqual( + reads_resumption_strategy._int_from_env("GCS_TEST_KNOB", 7), 7 + ) + + def test_default_for_invalid_values(self): + for value in ("", "abc", "1.5"): + with mock.patch.dict(os.environ, {"GCS_TEST_KNOB": value}): + self.assertEqual( + reads_resumption_strategy._int_from_env("GCS_TEST_KNOB", 7), 7 + ) + + class TestDownloadState(unittest.TestCase): def test_initialization(self): """Test that _DownloadState initializes correctly.""" @@ -309,6 +336,287 @@ def test_update_state_checksum_mismatch(self): with self.assertRaisesRegex(DataCorruption, "Checksum mismatch"): self.strategy.update_state_from_response(response, self.state) + def test_update_state_failed_write_leaves_read_resumable(self): + buffer = mock.Mock(spec=io.BytesIO) + buffer.write.side_effect = [OSError("disk full"), None] + read_state = self._add_download(_READ_ID, length=8, buffer=buffer) + response = self._create_response(b"data", _READ_ID, offset=0) + + with self.assertRaises(OSError): + self.strategy.update_state_from_response(response, self.state) + + self.assertEqual(read_state.bytes_written, 0) + self.assertEqual(read_state.next_expected_offset, 0) + self.strategy.update_state_from_response(response, self.state) + self.assertEqual(read_state.bytes_written, 4) + + # --- Offloaded Checksum Tests --- + + def _large_chunk(self): + return b"x" * reads_resumption_strategy._CRC32C_OFFLOAD_MIN_BYTES + + def _offload_to(self, futures): + """Make the strategy hand large chunks to ``futures`` instead of a pool.""" + executor = mock.Mock() + executor.submit.side_effect = list(futures) + patcher = mock.patch.object( + reads_resumption_strategy, "_crc32c_executor", executor + ) + patcher.start() + self.addCleanup(patcher.stop) + + def test_update_state_holds_large_chunk_until_verified(self): + """Large chunks are written only once their offloaded checksum matches.""" + content = self._large_chunk() + future = concurrent.futures.Future() + self._offload_to([future]) + read_state = self._add_download(_READ_ID, length=len(content)) + response = self._create_response(content, _READ_ID, offset=0, range_end=True) + + self.strategy.update_state_from_response(response, self.state) + + self.assertEqual(read_state.user_buffer.getvalue(), b"") + self.assertEqual(read_state.bytes_written, 0) + self.assertFalse(read_state.is_complete) + self.assertEqual(read_state.next_expected_offset, len(content)) + self.assertEqual(len(self.state["pending_checksums"]), 1) + self.assertIs(self.state["pending_checksums"][0].checksum_future, future) + + future.set_result(google_crc32c.value(content)) + asyncio.run(self.strategy.verify_pending_checksums(self.state)) + + self.assertEqual(read_state.user_buffer.getvalue(), content) + self.assertEqual(read_state.bytes_written, len(content)) + self.assertTrue(read_state.is_complete) + self.assertEqual(len(self.state["pending_checksums"]), 0) + + def test_verify_pending_checksums_mismatch_rewinds_to_corrupt_chunk(self): + """Nothing from the corrupt chunk on is written; a retry resumes there.""" + content = self._large_chunk() + size = len(content) + futures = [concurrent.futures.Future() for _ in range(3)] + for future, crc in zip(futures, (google_crc32c.value(content), 999999, 0)): + future.set_result(crc) + self._offload_to(futures) + read_state = self._add_download(_READ_ID, offset=10, length=3 * size) + for i in range(3): + response = self._create_response(content, _READ_ID, offset=10 + i * size) + self.strategy.update_state_from_response(response, self.state) + + with self.assertRaisesRegex(DataCorruption, "Checksum mismatch"): + asyncio.run(self.strategy.verify_pending_checksums(self.state)) + + self.assertEqual(read_state.user_buffer.getvalue(), content) + self.assertEqual(read_state.bytes_written, size) + self.assertEqual(read_state.next_expected_offset, 10 + size) + self.assertEqual(len(self.state["pending_checksums"]), 0) + (request,) = self.strategy.generate_requests(self.state) + self.assertEqual(request.read_offset, 10 + size) + self.assertEqual(request.read_length, 2 * size) + + def test_update_state_holds_small_chunk_behind_pending_one(self): + """Chunks are written in arrival order, even if verified out of order.""" + large, small = self._large_chunk(), b"tail" + future = concurrent.futures.Future() + self._offload_to([future]) + read_state = self._add_download(_READ_ID, length=len(large) + len(small)) + responses = [ + self._create_response(large, _READ_ID, offset=0), + self._create_response(small, _READ_ID, offset=len(large), range_end=True), + ] + for response in responses: + self.strategy.update_state_from_response(response, self.state) + + self.assertEqual(read_state.user_buffer.getvalue(), b"") + self.assertEqual(len(self.state["pending_checksums"]), 2) + + future.set_result(google_crc32c.value(large)) + asyncio.run(self.strategy.verify_pending_checksums(self.state)) + + self.assertEqual(read_state.user_buffer.getvalue(), large + small) + self.assertTrue(read_state.is_complete) + + def test_verify_pending_checksums_honors_max_pending(self): + """Finished chunks are written at once; unfinished ones only block + while more than max_pending are held.""" + content = self._large_chunk() + futures = [concurrent.futures.Future() for _ in range(3)] + futures[0].set_result(google_crc32c.value(content)) + self._offload_to(futures) + read_state = self._add_download(_READ_ID, length=3 * len(content)) + for i in range(3): + response = self._create_response(content, _READ_ID, offset=i * len(content)) + self.strategy.update_state_from_response(response, self.state) + self.assertEqual(len(self.state["pending_checksums"]), 3) + + asyncio.run(self.strategy.verify_pending_checksums(self.state, max_pending=2)) + self.assertEqual(len(self.state["pending_checksums"]), 2) + self.assertEqual(read_state.bytes_written, len(content)) + + for future in futures[1:]: + future.set_result(google_crc32c.value(content)) + asyncio.run(self.strategy.verify_pending_checksums(self.state)) + self.assertEqual(len(self.state["pending_checksums"]), 0) + self.assertEqual(read_state.user_buffer.getvalue(), content * 3) + + def test_verify_pending_checksums_without_pending_is_noop(self): + asyncio.run(self.strategy.verify_pending_checksums(self.state)) + self.assertNotIn("pending_checksums", self.state) + + def test_verify_pending_checksums_awaits_unfinished_future(self): + content = self._large_chunk() + future = concurrent.futures.Future() + self._offload_to([future]) + self._add_download(_READ_ID, length=len(content)) + response = self._create_response(content, _READ_ID, offset=0) + self.strategy.update_state_from_response(response, self.state) + + async def run(): + loop = asyncio.get_running_loop() + loop.call_later(0.01, future.set_result, google_crc32c.value(content)) + await self.strategy.verify_pending_checksums(self.state) + + asyncio.run(run()) + self.assertEqual(len(self.state["pending_checksums"]), 0) + + def test_verify_pending_checksums_cancels_queued_futures_on_mismatch(self): + """Work queued behind a corrupt chunk is dropped from the shared pool.""" + content = self._large_chunk() + corrupt, queued = concurrent.futures.Future(), concurrent.futures.Future() + corrupt.set_result(999999) + self._offload_to([corrupt, queued]) + self._add_download(_READ_ID, length=2 * len(content)) + for i in range(2): + response = self._create_response(content, _READ_ID, offset=i * len(content)) + self.strategy.update_state_from_response(response, self.state) + + with self.assertRaisesRegex(DataCorruption, "Checksum mismatch"): + asyncio.run(self.strategy.verify_pending_checksums(self.state)) + + self.assertTrue(queued.cancelled()) + self.assertEqual(len(self.state["pending_checksums"]), 0) + + def test_verify_pending_checksums_discards_queue_when_worker_raises(self): + content = self._large_chunk() + failed, queued = concurrent.futures.Future(), concurrent.futures.Future() + failed.set_exception(TypeError("worker failed")) + self._offload_to([failed, queued]) + read_state = self._add_download(_READ_ID, length=2 * len(content)) + for i in range(2): + response = self._create_response(content, _READ_ID, offset=i * len(content)) + self.strategy.update_state_from_response(response, self.state) + + with self.assertRaisesRegex(TypeError, "worker failed"): + asyncio.run(self.strategy.verify_pending_checksums(self.state)) + + self.assertTrue(queued.cancelled()) + self.assertEqual(len(self.state["pending_checksums"]), 0) + self.assertEqual(read_state.next_expected_offset, 0) + + def test_update_state_offload_disabled_by_zero_threshold(self): + content = self._large_chunk() + self._add_download(_READ_ID, length=len(content)) + with ( + mock.patch.object( + reads_resumption_strategy, "_CRC32C_OFFLOAD_MIN_BYTES", 0 + ), + mock.patch.object(reads_resumption_strategy, "_submit_crc32c") as submit, + ): + response = self._create_response(content, _READ_ID, offset=0) + self.strategy.update_state_from_response(response, self.state) + submit.assert_not_called() + self.assertFalse(self.state.get("pending_checksums")) + + def test_discard_pending_checksums_cancels_and_rewinds(self): + content = self._large_chunk() + future = concurrent.futures.Future() + self._offload_to([future]) + read_state = self._add_download(_READ_ID, offset=5, length=len(content) + 4) + responses = [ + self._create_response(content, _READ_ID, offset=5), + self._create_response(b"tail", _READ_ID, offset=5 + len(content)), + ] + for response in responses: + self.strategy.update_state_from_response(response, self.state) + + self.strategy.discard_pending_checksums(self.state) + + self.assertTrue(future.cancelled()) + self.assertEqual(len(self.state["pending_checksums"]), 0) + self.assertEqual(read_state.next_expected_offset, 5) + self.assertEqual(read_state.user_buffer.getvalue(), b"") + + def test_discard_pending_checksums_without_pending_is_noop(self): + self.strategy.discard_pending_checksums(self.state) + self.assertNotIn("pending_checksums", self.state) + + def test_verify_pending_checksums_checks_byte_count_of_held_chunk(self): + content = self._large_chunk() + self._add_download(_READ_ID, length=len(content) - 1) + response = self._create_response(content, _READ_ID, offset=0, range_end=True) + self.strategy.update_state_from_response(response, self.state) + + with self.assertRaisesRegex(DataCorruption, "Byte count mismatch"): + asyncio.run(self.strategy.verify_pending_checksums(self.state)) + + def test_update_state_hashes_inline_when_pool_rejects_work(self): + """After interpreter shutdown begins, submit() raises RuntimeError.""" + content = self._large_chunk() + executor = mock.Mock() + executor.submit.side_effect = RuntimeError( + "cannot schedule new futures after interpreter shutdown" + ) + patcher = mock.patch.object( + reads_resumption_strategy, "_crc32c_executor", executor + ) + patcher.start() + self.addCleanup(patcher.stop) + read_state = self._add_download(_READ_ID, length=2 * len(content)) + + response = self._create_response(content, _READ_ID, offset=0) + self.strategy.update_state_from_response(response, self.state) + self.assertEqual(read_state.user_buffer.getvalue(), content) + self.assertNotIn("pending_checksums", self.state) + + response = self._create_response( + content, _READ_ID, offset=len(content), crc=999999 + ) + with self.assertRaisesRegex(DataCorruption, "Checksum mismatch"): + self.strategy.update_state_from_response(response, self.state) + self.assertEqual(read_state.bytes_written, len(content)) + + def test_update_state_large_chunk_not_deferred_when_checksum_disabled(self): + content = self._large_chunk() + self.state["enable_checksum"] = False + self._add_download(_READ_ID, length=len(content)) + response = self._create_response(content, _READ_ID, offset=0, crc=999999) + + self.strategy.update_state_from_response(response, self.state) + + self.assertNotIn("pending_checksums", self.state) + + def test_update_state_full_object_read_hashes_large_chunk_inline(self): + """Full-object reads already hash every chunk into the rolling checksum + on the calling thread, so a worker-side hash would only be overhead.""" + content = self._large_chunk() + read_state = self._add_download( + _READ_ID, length=len(content), is_full_object_read=True + ) + response = self._create_response(content, _READ_ID, offset=0) + + self.strategy.update_state_from_response(response, self.state) + + self.assertNotIn("pending_checksums", self.state) + self.assertEqual(read_state.rolling_checksum._crc, google_crc32c.value(content)) + + def test_update_state_full_object_read_large_chunk_mismatch_raises_inline(self): + content = self._large_chunk() + self._add_download(_READ_ID, length=len(content), is_full_object_read=True) + response = self._create_response(content, _READ_ID, offset=0, crc=999999) + + with self.assertRaisesRegex(DataCorruption, "Checksum mismatch"): + self.strategy.update_state_from_response(response, self.state) + def test_update_state_final_byte_count_mismatch(self): """Test mismatch between expected length and actual bytes written on completion.""" self._add_download(_READ_ID, length=100) @@ -413,6 +721,53 @@ async def run(): # Token should remain unchanged self.assertEqual(self.state["routing_token"], "existing-token") + def test_recover_state_on_failure_raises_for_corrupt_pending_chunk(self): + """A corrupt chunk received before a stream error aborts instead of retrying, + without writing it; the redirect is still recorded.""" + content = self._large_chunk() + read_state = self._add_download(_READ_ID, length=len(content)) + response = self._create_response(content, _READ_ID, offset=0, crc=999999) + self.strategy.update_state_from_response(response, self.state) + token = "dummy-routing-token" + final_error = exceptions.Aborted( + "Retry failed", errors=[BidiReadObjectRedirectedError(routing_token=token)] + ) + + with self.assertRaisesRegex(DataCorruption, "Checksum mismatch"): + asyncio.run(self.strategy.recover_state_on_failure(final_error, self.state)) + + self.assertEqual(self.state["routing_token"], token) + self.assertEqual(read_state.user_buffer.getvalue(), b"") + (request,) = self.strategy.generate_requests(self.state) + self.assertEqual(request.read_offset, 0) + + def test_recover_state_on_failure_drains_pending_checksums(self): + """Verified chunks are written before the retry computes its offsets.""" + content = self._large_chunk() + future = concurrent.futures.Future() + self._offload_to([future]) + read_state = self._add_download(_READ_ID, length=2 * len(content)) + response = self._create_response(content, _READ_ID, offset=0) + self.strategy.update_state_from_response(response, self.state) + token = "dummy-routing-token" + final_error = exceptions.Aborted( + "Retry failed", errors=[BidiReadObjectRedirectedError(routing_token=token)] + ) + + async def run(): + asyncio.get_running_loop().call_later( + 0.01, future.set_result, google_crc32c.value(content) + ) + await self.strategy.recover_state_on_failure(final_error, self.state) + + asyncio.run(run()) + + self.assertEqual(len(self.state["pending_checksums"]), 0) + self.assertEqual(self.state["routing_token"], token) + self.assertEqual(read_state.user_buffer.getvalue(), content) + (request,) = self.strategy.generate_requests(self.state) + self.assertEqual(request.read_offset, len(content)) + def test_update_state_full_object_checksum_success(self): """Test that full object checksum verification succeeds on range_end.""" read_state = self._add_download( @@ -470,3 +825,39 @@ def test_update_state_full_object_checksum_mismatch_ignored_when_disabled(self): resp2 = self._create_response(b"data1", _READ_ID, offset=4, range_end=True) # Should NOT raise DataCorruption! self.strategy.update_state_from_response(resp2, self.state) + + +class TestCrc32cExecutor(unittest.TestCase): + def test_reset_replaces_the_pool(self): + pool = reads_resumption_strategy._crc32c_executor + self.addCleanup(setattr, reads_resumption_strategy, "_crc32c_executor", pool) + + reads_resumption_strategy._reset_crc32c_executor() + + new_pool = reads_resumption_strategy._crc32c_executor + self.addCleanup(new_pool.shutdown) + self.assertIsNot(new_pool, pool) + self.assertEqual(new_pool.submit(int, "7").result(), 7) + + @unittest.skipUnless(hasattr(os, "fork"), "requires os.fork") + def test_forked_child_gets_a_working_pool(self): + """The parent's pool has no threads in a forked child; work sent to it + would never run.""" + pool = reads_resumption_strategy._crc32c_executor + self.assertEqual(pool.submit(int, "7").result(), 7) + with warnings.catch_warnings(): + # Python 3.12+ warns about forking a process that has threads. + warnings.simplefilter("ignore", DeprecationWarning) + pid = os.fork() + if pid == 0: # pragma: NO COVER - the child exits without saving coverage + ok = False + try: + signal.alarm(30) # never leave the parent waiting on a hung child + child_pool = reads_resumption_strategy._crc32c_executor + ok = child_pool is not pool and ( + child_pool.submit(int, "7").result(timeout=10) == 7 + ) + finally: + os._exit(0 if ok else 1) + _, status = os.waitpid(pid, 0) + self.assertEqual(os.waitstatus_to_exitcode(status), 0) diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py index afe1dfd221d9..9fea72f414e9 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py @@ -13,6 +13,7 @@ # limitations under the License. import asyncio +import concurrent.futures from io import BytesIO from unittest import mock from unittest.mock import AsyncMock @@ -27,6 +28,7 @@ from google.cloud.storage.asyncio.async_multi_range_downloader import ( AsyncMultiRangeDownloader, ) +from google.cloud.storage.asyncio.retry import reads_resumption_strategy from google.cloud.storage.exceptions import DataCorruption _TEST_BUCKET_NAME = "test-bucket" @@ -328,10 +330,22 @@ async def test_download_ranges_raises_if_crc32c_c_extension_is_missing( ) @pytest.mark.asyncio + @pytest.mark.parametrize( + "test_data", + [ + pytest.param(b"some-data", id="inline"), + pytest.param( + b"x" * reads_resumption_strategy._CRC32C_OFFLOAD_MIN_BYTES, + id="offloaded", + ), + ], + ) @mock.patch( "google.cloud.storage.asyncio.retry.reads_resumption_strategy.google_crc32c.value" ) - async def test_download_ranges_raises_on_checksum_mismatch(self, mock_crc32c_value): + async def test_download_ranges_raises_on_checksum_mismatch( + self, mock_crc32c_value, test_data + ): from google.cloud.storage.asyncio._stream_multiplexer import _StreamMultiplexer from google.cloud.storage.asyncio.async_multi_range_downloader import ( AsyncMultiRangeDownloader, @@ -342,7 +356,6 @@ async def test_download_ranges_raises_on_checksum_mismatch(self, mock_crc32c_val spec=async_read_object_stream._AsyncReadObjectStream ) - test_data = b"some-data" server_checksum = 12345 mock_crc32c_value.return_value = 54321 @@ -367,15 +380,140 @@ async def test_download_ranges_raises_on_checksum_mismatch(self, mock_crc32c_val mrd._is_stream_open = True mrd._multiplexer = _StreamMultiplexer(mock_stream) + buffer = BytesIO() with pytest.raises(DataCorruption) as exc_info: with mock.patch( "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer", return_value=0, ): - await mrd.download_ranges([(0, len(test_data), BytesIO())]) + await mrd.download_ranges([(0, len(test_data), buffer)]) assert "Checksum mismatch" in str(exc_info.value) mock_crc32c_value.assert_called_once_with(test_data) + assert buffer.getvalue() == b"" + + def _mrd_with_held_chunk(self, content, recv_after_chunk): + """An MRD whose stream sends one offloaded chunk of ``content`` (whose + checksum future is returned) and then awaits ``recv_after_chunk()``.""" + from google.cloud.storage.asyncio._stream_multiplexer import _StreamMultiplexer + + future = concurrent.futures.Future() + submitted = asyncio.Event() + + def submit(fn, data): + submitted.set() + return future + + executor = mock.Mock() + executor.submit.side_effect = submit + patcher = mock.patch.object( + reads_resumption_strategy, "_crc32c_executor", executor + ) + response = _storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + _storage_v2.ObjectRangeData( + checksummed_data=_storage_v2.ChecksummedData( + content=content, crc32c=google_crc32c.value(content) + ), + read_range=_storage_v2.ReadRange( + read_id=0, read_offset=0, read_length=len(content) + ), + ) + ] + ) + responses = iter([response]) + + async def recv(): + item = next(responses, None) + if item is not None: + return item + return await recv_after_chunk(submitted, future) + + mock_stream = mock.AsyncMock( + spec=async_read_object_stream._AsyncReadObjectStream + ) + mock_stream.recv.side_effect = recv + mrd = AsyncMultiRangeDownloader(mock.MagicMock(), "bucket", "object") + mrd.read_obj_str = mock_stream + mrd._is_stream_open = True + mrd._multiplexer = _StreamMultiplexer(mock_stream) + return mrd, future, patcher + + @pytest.mark.asyncio + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer", + return_value=0, + ) + async def test_download_ranges_reports_held_corrupt_chunk_before_other_errors( + self, _ + ): + """A corrupt offloaded chunk is reported, and not written, even when a + non-retryable error arrives before its checksum is known.""" + content = b"x" * reads_resumption_strategy._CRC32C_OFFLOAD_MIN_BYTES + + async def fail_after_chunk(submitted, future): + await submitted.wait() # the chunk is held, its check still running + future.set_result(0) # ...and the worker finds a mismatch + raise exceptions.NotFound("object deleted") + + mrd, _, patcher = self._mrd_with_held_chunk(content, fail_after_chunk) + buffer = BytesIO() + with patcher, pytest.raises(DataCorruption) as exc_info: + await mrd.download_ranges([(0, 2 * len(content), buffer)]) + + assert "Checksum mismatch" in str(exc_info.value) + assert isinstance(exc_info.value.__context__, exceptions.NotFound) + assert buffer.getvalue() == b"" + + @pytest.mark.asyncio + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer", + return_value=0, + ) + async def test_download_ranges_writes_held_chunks_before_other_errors(self, _): + """Chunks received before a non-retryable error are written, as they + are when checked inline, and the error is raised.""" + content = b"x" * reads_resumption_strategy._CRC32C_OFFLOAD_MIN_BYTES + + async def fail_after_chunk(submitted, future): + await submitted.wait() + future.set_result(google_crc32c.value(content)) + raise exceptions.NotFound("object deleted") + + mrd, _, patcher = self._mrd_with_held_chunk(content, fail_after_chunk) + buffer = BytesIO() + with patcher, pytest.raises(exceptions.NotFound): + await mrd.download_ranges([(0, 2 * len(content), buffer)]) + + assert buffer.getvalue() == content + + @pytest.mark.asyncio + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader.generate_random_56_bit_integer", + return_value=0, + ) + async def test_download_ranges_cancellation_discards_held_chunks(self, _): + content = b"x" * reads_resumption_strategy._CRC32C_OFFLOAD_MIN_BYTES + + async def hang(submitted, future): + await asyncio.Event().wait() + + mrd, future, patcher = self._mrd_with_held_chunk(content, hang) + buffer = BytesIO() + with patcher: + task = asyncio.create_task( + mrd.download_ranges([(0, 2 * len(content), buffer)]) + ) + while not reads_resumption_strategy._crc32c_executor.submit.called: + await asyncio.sleep(0) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert future.cancelled() + assert buffer.getvalue() == b"" + await mrd.close() @mock.patch( "google.cloud.storage.asyncio.async_multi_range_downloader.AsyncMultiRangeDownloader.open", @@ -688,6 +826,7 @@ async def test_download_ranges_configures_full_object_read_state( mock_retry_manager = mock_retry_manager_cls.return_value mock_retry_manager.execute = AsyncMock() + mock_strategy_cls.return_value.verify_pending_checksums = AsyncMock() # Act # Implicit full read (0, 0) and explicit full read (0, persisted_size=100)