perf(storage): verify CRC32C of large read chunks off the event-loop thread

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"]

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)