diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/_zero_copy_bidi_read.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/_zero_copy_bidi_read.py new file mode 100644 index 000000000000..fefa444a179e --- /dev/null +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/_zero_copy_bidi_read.py @@ -0,0 +1,354 @@ +# Copyright 2026 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. + +"""Zero-copy deserializer for ``BidiReadObjectResponse`` messages. + +The generated gRPC stub parses every response with proto-plus over upb, which +copies each data chunk twice (into the upb arena, then into a Python ``bytes`` +for ``checksummed_data.content``) on the event-loop thread. Data responses are +structurally tiny apart from the chunk itself, so this module walks the wire +format directly and exposes ``content`` as a ``memoryview`` of the raw message. +Responses carrying object metadata (the first message of a stream opened +without a read handle) or anything unexpected fall back to the generated +parser, so callers never see a difference other than the type of ``content``. + +:func:`wrapped_rpc` builds the matching ``BidiReadObject`` callable for a +transport; :class:`~google.cloud.storage.asyncio.async_read_object_stream._AsyncReadObjectStream` +prefers it over the generated one whenever :func:`is_supported` holds. +""" + +import os +import weakref + +import google_crc32c +from google.api_core import gapic_v1 + +from google.cloud import _storage_v2 +from google.cloud._storage_v2.services.storage.transports.base import ( + DEFAULT_CLIENT_INFO, +) +from google.cloud._storage_v2.services.storage.transports.grpc_asyncio import ( + StorageGrpcAsyncIOTransport, +) + +# Opt-in switch. Accepted spellings match ``_opentelemetry_tracing._parse_bool_env``. +_ENV_VAR = "GOOGLE_CLOUD_STORAGE_ZERO_COPY_BIDI_READ" +_ENABLED_VALUES = ("1", "true", "yes", "on") + +# Protobuf wire types (low 3 bits of a field tag): varint, 8-byte fixed, +# length-delimited, 4-byte fixed. Groups (3, 4) are left to the generated parser. +_WT_VARINT, _WT_I64, _WT_LEN, _WT_I32 = 0, 1, 2, 5 + + +class _ReadRange: + __slots__ = ("read_offset", "read_length", "read_id") + + def __init__(self): + self.read_offset = 0 + self.read_length = 0 + self.read_id = 0 + + +class _ChecksummedData: + __slots__ = ("content", "crc32c", "_has_crc32c") + + def __init__(self): + self.content = memoryview(b"") + self.crc32c = 0 + self._has_crc32c = False + + def HasField(self, name): + if name == "crc32c": + return self._has_crc32c + raise ValueError(f"Unknown field {name!r}") + + +class _ObjectRangeData: + __slots__ = ( + "checksummed_data", + "read_range", + "range_end", + "_has_checksummed_data", + "_has_read_range", + ) + + def __init__(self): + self.checksummed_data = _ChecksummedData() + self.read_range = _ReadRange() + self.range_end = False + self._has_checksummed_data = False + self._has_read_range = False + + def HasField(self, name): + if name == "checksummed_data": + return self._has_checksummed_data + if name == "read_range": + return self._has_read_range + raise ValueError(f"Unknown field {name!r}") + + +class _ZeroCopyBidiReadObjectResponse: + """Duck-typed stand-in for a ``BidiReadObjectResponse`` without metadata. + + Only the attributes the async read path uses are provided: + ``object_data_ranges``, ``read_handle``, ``metadata`` (always ``None``), + ``HasField`` and ``_pb`` (itself, so code that unwraps proto-plus messages + keeps working). + """ + + __slots__ = ("object_data_ranges", "read_handle", "metadata") + + def __init__(self, ranges, read_handle): + self.object_data_ranges = ranges + self.read_handle = read_handle + self.metadata = None + + @property + def _pb(self): + return self + + def HasField(self, name): + if name == "read_handle": + return self.read_handle is not None + if name == "metadata": + return False + raise ValueError(f"Unknown field {name!r}") + + +class _Malformed(ValueError): + """Raised when the wire bytes do not look like a well-formed message.""" + + +# A 10-byte varint carries 70 bits; like upb, keep only the low 64. +_UINT64_MASK = (1 << 64) - 1 +# A field tag (``field_number << 3 | wire_type``) must fit in 32 bits, as in upb. +_UINT32_MAX = (1 << 32) - 1 + + +def _varint(buf, i, end): + """Decode the varint at ``buf[i]``; it must end at or before ``end``. + + Like upb, accept up to ten bytes and keep only the low 64 bits. + """ + shift = 0 + result = 0 + while True: + b = buf[i] + i += 1 + result |= (b & 0x7F) << shift + if b < 0x80: + break + shift += 7 + if shift > 63: + raise _Malformed("varint longer than 10 bytes") + if i > end: + raise _Malformed("varint runs past its enclosing message") + return result & _UINT64_MASK, i + + +def _int64(value): + """Reinterpret the low 64 bits of a varint as a signed integer.""" + return value - (1 << 64) if value >> 63 else value + + +def _tag(buf, i, end): + """Decode a field tag at ``buf[i]`` into ``(field_number, wire_type, i)``. + + upb rejects field number 0, tags above 32 bits and tags encoded in more + than five bytes. + """ + start = i + tag, i = _varint(buf, i, end) + if tag < 8 or tag > _UINT32_MAX or i - start > 5: + raise _Malformed("invalid field tag") + return tag >> 3, tag & 7, i + + +def _length_prefixed(buf, i, end): + """Return ``(start, stop)`` of the length-delimited field body at ``i``.""" + n, i = _varint(buf, i, end) + stop = i + n + if stop > end: + raise _Malformed("field runs past its enclosing message") + return i, stop + + +def _skip(buf, i, wt, end): + if wt == _WT_VARINT: + return _varint(buf, i, end)[1] + if wt == _WT_LEN: + return _length_prefixed(buf, i, end)[1] + if wt == _WT_I64: + i += 8 + elif wt == _WT_I32: + i += 4 + else: + raise _Malformed(f"unsupported wire type {wt}") + if i > end: + raise _Malformed("field runs past its enclosing message") + return i + + +# The sub-message parsers fill in an existing object so that a singular +# embedded message occurring more than once is merged, as the generated +# parser does, instead of replacing what the earlier occurrence set. + + +def _parse_read_range(rr, buf, i, end): + while i < end: + f, wt, i = _tag(buf, i, end) + if wt == _WT_VARINT: + v, i = _varint(buf, i, end) + if f == 1: + rr.read_offset = _int64(v) + elif f == 2: + rr.read_length = _int64(v) + elif f == 3: + rr.read_id = _int64(v) + else: + i = _skip(buf, i, wt, end) + + +def _parse_checksummed(cd, buf, mv, i, end): + while i < end: + f, wt, i = _tag(buf, i, end) + if f == 1 and wt == _WT_LEN: + start, i = _length_prefixed(buf, i, end) + cd.content = mv[start:i] + elif f == 2 and wt == _WT_I32: + if i + 4 > end: + raise _Malformed("crc32c runs past its enclosing message") + cd.crc32c = int.from_bytes(buf[i : i + 4], "little") + cd._has_crc32c = True + i += 4 + else: + i = _skip(buf, i, wt, end) + + +def _parse_range_data(buf, mv, i, end): + rd = _ObjectRangeData() + while i < end: + f, wt, i = _tag(buf, i, end) + if wt == _WT_LEN: + start, i = _length_prefixed(buf, i, end) + if f == 1: + _parse_checksummed(rd.checksummed_data, buf, mv, start, i) + rd._has_checksummed_data = True + elif f == 2: + _parse_read_range(rd.read_range, buf, start, i) + rd._has_read_range = True + elif f == 3 and wt == _WT_VARINT: + v, i = _varint(buf, i, end) + rd.range_end = bool(v) + else: + i = _skip(buf, i, wt, end) + return rd + + +def _parse_handle(handle, buf, i, end): + while i < end: + f, wt, i = _tag(buf, i, end) + if f == 1 and wt == _WT_LEN: + start, i = _length_prefixed(buf, i, end) + handle = bytes(buf[start:i]) + else: + i = _skip(buf, i, wt, end) + return handle + + +_generated_deserialize = _storage_v2.BidiReadObjectResponse.deserialize + + +def deserialize(buf): + """Parse a serialized ``BidiReadObjectResponse``; chunk payloads alias ``buf``. + + Anything that is not a plain data response (object metadata present, or + bytes that do not parse cleanly) is handed to the generated parser, which + either produces the real message or raises its usual ``DecodeError``. + """ + try: + mv = memoryview(buf) + n = len(buf) + i = 0 + ranges = [] + handle = None + while i < n: + f, wt, i = _tag(buf, i, n) + if wt == _WT_LEN: + start, end = _length_prefixed(buf, i, n) + if f == 6: + ranges.append(_parse_range_data(buf, mv, start, end)) + elif f == 7: + handle = _parse_handle(handle or b"", buf, start, end) + elif f == 4: + return _generated_deserialize(buf) + i = end + else: + i = _skip(buf, i, wt, n) + if handle is not None: + # A real message: the handle is stored and passed back in the + # BidiReadObjectSpec of the next open, which only accepts the + # proto type. + handle = _storage_v2.BidiReadHandle(handle=handle) + return _ZeroCopyBidiReadObjectResponse(ranges, handle) + except Exception: + return _generated_deserialize(buf) + + +def is_supported(): + """Whether the zero-copy parser may be used in this process. + + The parser hands chunk payloads to the read path as ``memoryview`` objects + and checksum verification feeds them to ``google_crc32c``. Because + downstream sinks and releases of ``google-crc32c`` that only accept + ``bytes`` would fail on ``memoryview``, the zero-copy parser is enabled + only when ``GOOGLE_CLOUD_STORAGE_ZERO_COPY_BIDI_READ=1`` is set and + ``google_crc32c`` accepts ``memoryview`` inputs. + """ + if os.environ.get(_ENV_VAR, "").strip().lower() not in _ENABLED_VALUES: + return False + try: + google_crc32c.value(memoryview(b"\0")) + except Exception: + return False + return True + + +_rpcs = weakref.WeakKeyDictionary() + + +def wrapped_rpc(transport): + """Return the zero-copy ``BidiReadObject`` callable for ``transport``. + + The callable is the generated transport's stub with :func:`deserialize` + swapped in, wrapped exactly like ``transport._wrapped_methods`` entries + (error mapping, no default retry or timeout, the client-info header). + Returns ``None`` when the zero-copy path is unsupported or ``transport`` is + not the generated grpc_asyncio transport, in which case callers should use + the generated wrapped method. + """ + if not isinstance(transport, StorageGrpcAsyncIOTransport) or not is_supported(): + return None + rpc = _rpcs.get(transport) + if rpc is None: + stub = transport._logged_channel.stream_stream( + "/google.storage.v2.Storage/BidiReadObject", + request_serializer=_storage_v2.BidiReadObjectRequest.serialize, + response_deserializer=deserialize, + ) + rpc = gapic_v1.method_async.wrap_method( + stub, default_timeout=None, client_info=DEFAULT_CLIENT_INFO + ) + _rpcs[transport] = rpc + return rpc 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..a4860c17cbf2 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, ) @@ -447,6 +448,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 +524,42 @@ 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: + # Write the chunks received before the failure; a corrupt + # one is reported instead, as it would have been inline. + 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/async_read_object_stream.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py index d3c0780673c7..b13b221f45b2 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py @@ -19,6 +19,7 @@ from google.api_core.bidi_async import AsyncBidiRpc from google.cloud import _storage_v2 +from google.cloud.storage.asyncio import _zero_copy_bidi_read from google.cloud.storage.asyncio.async_abstract_object_stream import ( _AsyncAbstractObjectStream, ) @@ -76,9 +77,12 @@ def __init__( self._full_bucket_name = f"projects/_/buckets/{self.bucket_name}" - self.rpc = self.client._client._transport._wrapped_methods[ - self.client._client._transport.bidi_read_object - ] + transport = self.client._client._transport + # Zero-copy BidiReadObject responses when opted in and google-crc32c + # can verify memoryview payloads; otherwise the generated parser. + self.rpc = _zero_copy_bidi_read.wrapped_rpc(transport) + if self.rpc is None: + self.rpc = transport._wrapped_methods[transport.bidi_read_object] self.metadata = (("x-goog-request-params", f"bucket={self._full_bucket_name}"),) self.socket_like_rpc: Optional[AsyncBidiRpc] = None self._is_stream_open: bool = False 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..bdf64f9bbb09 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,69 @@ 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. 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( + "GOOGLE_CLOUD_STORAGE_CRC32C_OFFLOAD_MIN_BYTES", 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: Optional[concurrent.futures.Future] + + class _DownloadState: """A helper class to track the state of a single range download.""" @@ -47,6 +114,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 +206,158 @@ def update_state_from_response( # Checksum Verification checksummed_data = object_data_range.checksummed_data data = checksummed_data.content + server_checksum = None + checksum = 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 ( + 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 = _submit_crc32c(data) + if checksum 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, + ) + pending = state.get("pending_checksums") + if checksum 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. + """ + pending = state.get("pending_checksums") + while pending: + chunk = pending[0] + if chunk.checksum is not None: + if not chunk.checksum.done(): + if len(pending) <= max_pending: + return + await asyncio.wrap_future(chunk.checksum) + client_checksum = chunk.checksum.result() + 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 is not None: + chunk.checksum.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..23ca40120258 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,255 @@ 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) + + 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_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 +689,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 +793,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) diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py index 393fa7b2b709..608ff1178a88 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py @@ -113,6 +113,59 @@ def test_init_with_bucket_object_generation(mock_client, mock_async_bidi_rpc): assert read_obj_stream.rpc == rpc_sentinel +@mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") +@mock.patch( + "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" +) +@mock.patch( + "google.cloud.storage.asyncio.async_read_object_stream._zero_copy_bidi_read.wrapped_rpc" +) +def test_init_prefers_zero_copy_rpc(mock_wrapped_rpc, mock_client, mock_async_bidi_rpc): + # Arrange + mock_wrapped_rpc.return_value = mock.sentinel.zero_copy_rpc + mock_client._client._transport._wrapped_methods = {} + + # Act + read_obj_stream = _AsyncReadObjectStream( + client=mock_client, + bucket_name=_TEST_BUCKET_NAME, + object_name=_TEST_OBJECT_NAME, + ) + + # Assert + assert read_obj_stream.rpc is mock.sentinel.zero_copy_rpc + mock_wrapped_rpc.assert_called_once_with(mock_client._client._transport) + + +@mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") +@mock.patch( + "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" +) +@mock.patch( + "google.cloud.storage.asyncio.async_read_object_stream._zero_copy_bidi_read.wrapped_rpc" +) +def test_init_falls_back_to_generated_rpc( + mock_wrapped_rpc, mock_client, mock_async_bidi_rpc +): + # Arrange + mock_wrapped_rpc.return_value = None + rpc_sentinel = mock.sentinel.A + mock_client._client._transport.bidi_read_object = "bidi_read_object_rpc" + mock_client._client._transport._wrapped_methods = { + "bidi_read_object_rpc": rpc_sentinel, + } + + # Act + read_obj_stream = _AsyncReadObjectStream( + client=mock_client, + bucket_name=_TEST_BUCKET_NAME, + object_name=_TEST_OBJECT_NAME, + ) + + # Assert + assert read_obj_stream.rpc is rpc_sentinel + + @mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") @mock.patch( "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_zero_copy_bidi_read.py b/packages/google-cloud-storage/tests/unit/asyncio/test_zero_copy_bidi_read.py new file mode 100644 index 000000000000..8cbde30497e2 --- /dev/null +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_zero_copy_bidi_read.py @@ -0,0 +1,620 @@ +# Copyright 2026 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. + +"""Tests for the zero-copy ``BidiReadObjectResponse`` deserializer.""" + +import logging + +import google_crc32c +import grpc +import pytest +from google.auth.credentials import AnonymousCredentials +from google.protobuf.message import DecodeError + +from google.cloud import _storage_v2 +from google.cloud._storage_v2.services.storage.transports.grpc_asyncio import ( + StorageGrpcAsyncIOTransport, +) +from google.cloud.storage.asyncio import _zero_copy_bidi_read +from google.cloud.storage.asyncio.async_read_object_stream import ( + _AsyncReadObjectStream, +) + +PAYLOAD = bytes(range(256)) * 8 + +# Field numbers from google/storage/v2/storage.proto. +_RESPONSE_METADATA, _RESPONSE_RANGES, _RESPONSE_HANDLE = 4, 6, 7 +_RANGE_CHECKSUMMED, _RANGE_READ_RANGE, _RANGE_END = 1, 2, 3 +_CHECKSUMMED_CONTENT, _CHECKSUMMED_CRC32C = 1, 2 +_READ_RANGE_OFFSET, _READ_RANGE_LENGTH, _READ_RANGE_ID = 1, 2, 3 +_HANDLE_HANDLE = 1 +_UNKNOWN = 15 + +# --- minimal wire-format encoder ----------------------------------------------- + + +def _varint(value): + out = bytearray() + while True: + byte = value & 0x7F + value >>= 7 + if not value: + out.append(byte) + return bytes(out) + out.append(byte | 0x80) + + +def _tag(number, wire_type): + return _varint(number << 3 | wire_type) + + +def _varint_field(number, value): + return _tag(number, 0) + _varint(value) + + +def _i64_field(number, value): + return _tag(number, 1) + value.to_bytes(8, "little") + + +def _len_field(number, payload): + return _tag(number, 2) + _varint(len(payload)) + payload + + +def _i32_field(number, value): + return _tag(number, 5) + value.to_bytes(4, "little") + + +def _read_range(offset=4096, length=len(PAYLOAD), read_id=7, extra=b""): + return ( + _varint_field(_READ_RANGE_OFFSET, offset) + + _varint_field(_READ_RANGE_LENGTH, length) + + _varint_field(_READ_RANGE_ID, read_id) + + extra + ) + + +def _checksummed(content=PAYLOAD, crc32c=None, extra=b""): + body = _len_field(_CHECKSUMMED_CONTENT, content) + if crc32c is not None: + body += _i32_field(_CHECKSUMMED_CRC32C, crc32c) + return body + extra + + +def _range_data(checksummed=None, read_range=None, range_end=True, extra=b""): + body = b"" + if checksummed is not None: + body += _len_field(_RANGE_CHECKSUMMED, checksummed) + if read_range is not None: + body += _len_field(_RANGE_READ_RANGE, read_range) + if range_end: + body += _varint_field(_RANGE_END, 1) + return body + extra + + +def _response(*ranges, handle=None, extra=b""): + body = b"".join(_len_field(_RESPONSE_RANGES, r) for r in ranges) + if handle is not None: + body += _len_field(_RESPONSE_HANDLE, handle) + return body + extra + + +def _data_response(content=PAYLOAD, crc32c=None, handle=None): + return _response( + _range_data(_checksummed(content, crc32c), _read_range(length=len(content))), + handle=None if handle is None else _len_field(_HANDLE_HANDLE, handle), + ) + + +def _assert_same_as_generated(buf): + """Parse ``buf`` with both parsers and check the zero-copy one matches.""" + expected = _storage_v2.BidiReadObjectResponse.deserialize(buf) + actual = _zero_copy_bidi_read.deserialize(buf) + + assert isinstance(actual, _zero_copy_bidi_read._ZeroCopyBidiReadObjectResponse) + assert actual.metadata is None + assert not actual.HasField("metadata") + assert actual.HasField("read_handle") == expected._pb.HasField("read_handle") + if actual.HasField("read_handle"): + assert isinstance(actual.read_handle, _storage_v2.BidiReadHandle) + assert actual.read_handle.handle == expected.read_handle.handle + else: + assert actual.read_handle is None + + assert len(actual.object_data_ranges) == len(expected.object_data_ranges) + for got, want in zip(actual.object_data_ranges, expected.object_data_ranges): + assert got.HasField("read_range") == want._pb.HasField("read_range") + assert got.HasField("checksummed_data") == want._pb.HasField("checksummed_data") + assert got.read_range.read_offset == want.read_range.read_offset + assert got.read_range.read_length == want.read_range.read_length + assert got.read_range.read_id == want.read_range.read_id + assert got.range_end == want.range_end + assert isinstance(got.checksummed_data.content, memoryview) + assert bytes(got.checksummed_data.content) == want.checksummed_data.content + assert got.checksummed_data.HasField( + "crc32c" + ) == want.checksummed_data._pb.HasField("crc32c") + assert got.checksummed_data.crc32c == want.checksummed_data.crc32c + return actual + + +# --- deserialize --------------------------------------------------------------- + + +def test_deserialize_data_response_round_trip(): + crc = google_crc32c.value(PAYLOAD) + response = _assert_same_as_generated(_data_response(crc32c=crc)) + + [range_data] = response.object_data_ranges + assert range_data.read_range.read_offset == 4096 + assert range_data.read_range.read_length == len(PAYLOAD) + assert range_data.read_range.read_id == 7 + assert range_data.range_end is True + assert bytes(range_data.checksummed_data.content) == PAYLOAD + assert range_data.checksummed_data.crc32c == crc + + +def test_deserialize_matches_proto_plus_serialization(): + message = _storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + _storage_v2.ObjectRangeData( + checksummed_data=_storage_v2.ChecksummedData( + content=PAYLOAD, crc32c=google_crc32c.value(PAYLOAD) + ), + read_range=_storage_v2.ReadRange( + read_offset=4096, read_length=len(PAYLOAD), read_id=7 + ), + range_end=True, + ) + ], + read_handle=_storage_v2.BidiReadHandle(handle=b"opaque"), + ) + _assert_same_as_generated(_storage_v2.BidiReadObjectResponse.serialize(message)) + + +def test_deserialize_content_is_a_view_of_the_wire_buffer(): + buf = _data_response() + content = _zero_copy_bidi_read.deserialize(buf).object_data_ranges[0] + assert content.checksummed_data.content.obj is buf + + +def test_deserialize_without_crc32c_reports_no_field(): + response = _assert_same_as_generated(_data_response()) + assert response.object_data_ranges[0].checksummed_data.HasField("crc32c") is False + + +def test_deserialize_read_handle_is_a_bidi_read_handle(): + response = _assert_same_as_generated(_data_response(handle=b"opaque")) + + assert response.read_handle.handle == b"opaque" + # The SDK feeds a captured handle straight back into the next open. + spec = _storage_v2.BidiReadObjectSpec(read_handle=response.read_handle) + assert spec.read_handle.handle == b"opaque" + + +def test_deserialize_multiple_ranges(): + buf = _response( + _range_data(_checksummed(b"first"), _read_range(0, 5, 1), range_end=False), + _range_data(_checksummed(b"second"), _read_range(5, 6, 2)), + ) + response = _assert_same_as_generated(buf) + assert [bytes(r.checksummed_data.content) for r in response.object_data_ranges] == [ + b"first", + b"second", + ] + + +def test_deserialize_empty_message(): + response = _assert_same_as_generated(b"") + assert response.object_data_ranges == [] + assert response.read_handle is None + + +def test_deserialize_range_without_read_range_exposes_proto_defaults(): + # _StreamMultiplexer reads ``data_range.read_range.read_id`` without a + # HasField check, so the absent sub-message must look like an empty one. + response = _assert_same_as_generated(_response(_range_data(_checksummed()))) + [range_data] = response.object_data_ranges + assert range_data.HasField("read_range") is False + assert range_data.read_range.read_id == 0 + + +def test_deserialize_range_without_checksummed_data_exposes_proto_defaults(): + response = _assert_same_as_generated( + _response(_range_data(read_range=_read_range())) + ) + [range_data] = response.object_data_ranges + assert range_data.HasField("checksummed_data") is False + assert bytes(range_data.checksummed_data.content) == b"" + + +def test_deserialize_skips_unknown_top_level_fields(): + buf = _data_response() + ( + _varint_field(_UNKNOWN, 1) + + _i64_field(_UNKNOWN, 2) + + _i32_field(_UNKNOWN, 3) + + _len_field(_UNKNOWN, b"ignored") + ) + _assert_same_as_generated(buf) + + +def test_deserialize_skips_unknown_nested_fields(): + unknown = ( + _varint_field(_UNKNOWN, 1) + + _i64_field(_UNKNOWN, 2) + + _i32_field(_UNKNOWN, 3) + + _len_field(_UNKNOWN, b"ignored") + ) + buf = _response( + _range_data( + _checksummed(extra=unknown), + _read_range(extra=unknown), + extra=unknown, + ), + handle=_len_field(_HANDLE_HANDLE, b"opaque") + unknown, + ) + _assert_same_as_generated(buf) + + +@pytest.mark.parametrize( + "offset, length, read_id", + [ + pytest.param(-1, -1, -1, id="minus_one"), + pytest.param(-5, 0, -(1 << 55), id="negative"), + pytest.param(-(1 << 63), (1 << 63) - 1, 1, id="int64_extremes"), + ], +) +def test_deserialize_signed_int64_read_range(offset, length, read_id): + # read_offset, read_length and read_id are int64; negative values travel + # as ten-byte two's-complement varints. + message = _storage_v2.BidiReadObjectResponse( + object_data_ranges=[ + _storage_v2.ObjectRangeData( + read_range=_storage_v2.ReadRange( + read_offset=offset, read_length=length, read_id=read_id + ) + ) + ] + ) + response = _assert_same_as_generated( + _storage_v2.BidiReadObjectResponse.serialize(message) + ) + read_range = response.object_data_ranges[0].read_range + assert (read_range.read_offset, read_range.read_length, read_range.read_id) == ( + offset, + length, + read_id, + ) + + +def test_deserialize_over_long_varint_keeps_low_64_bits(): + # A ten-byte varint may carry up to 70 bits; the generated parser keeps + # the low 64 and interprets them as the field's type. + over_long = b"\xff" * 9 + b"\x7f" + read_range = _tag(_READ_RANGE_ID, 0) + over_long + response = _assert_same_as_generated(_response(_range_data(read_range=read_range))) + assert response.object_data_ranges[0].read_range.read_id == -1 + + +def test_deserialize_merges_split_sub_messages(): + # A singular embedded message that occurs twice is merged, so the second + # occurrence must not reset fields set by the first. + read_range = _len_field( + _RANGE_READ_RANGE, _varint_field(_READ_RANGE_OFFSET, 4096) + ) + _len_field(_RANGE_READ_RANGE, _varint_field(_READ_RANGE_ID, 7)) + checksummed = _len_field( + _RANGE_CHECKSUMMED, _len_field(_CHECKSUMMED_CONTENT, PAYLOAD) + ) + _len_field(_RANGE_CHECKSUMMED, _i32_field(_CHECKSUMMED_CRC32C, 1234)) + buf = _response( + _range_data(extra=read_range + checksummed), + handle=_len_field(_HANDLE_HANDLE, b"opaque"), + ) + _len_field(_RESPONSE_HANDLE, b"") + + response = _assert_same_as_generated(buf) + + [range_data] = response.object_data_ranges + assert range_data.read_range.read_offset == 4096 + assert range_data.read_range.read_id == 7 + assert bytes(range_data.checksummed_data.content) == PAYLOAD + assert range_data.checksummed_data.crc32c == 1234 + assert response.read_handle.handle == b"opaque" + + +def test_deserialize_metadata_response_uses_generated_parser(): + metadata = _storage_v2.Object(name="o", bucket="projects/_/buckets/b", size=10) + buf = _storage_v2.BidiReadObjectResponse.serialize( + _storage_v2.BidiReadObjectResponse(metadata=metadata) + ) + response = _zero_copy_bidi_read.deserialize(buf) + assert isinstance(response, _storage_v2.BidiReadObjectResponse) + assert response.metadata.size == 10 + + +def test_deserialize_unknown_group_uses_generated_parser(): + # Start-group/end-group wire types are valid protobuf that the zero-copy + # path does not understand; the generated parser keeps them as unknown fields. + buf = _data_response() + _tag(_UNKNOWN, 3) + _tag(_UNKNOWN, 4) + response = _zero_copy_bidi_read.deserialize(buf) + assert isinstance(response, _storage_v2.BidiReadObjectResponse) + assert response.object_data_ranges[0].checksummed_data.content == PAYLOAD + + +@pytest.mark.parametrize( + "buf", + [ + pytest.param( + # object_data_ranges (6 bytes) whose checksummed_data claims 8 bytes, + # followed by a well-formed read_handle. + bytes([0x32, 0x06, 0x0A, 0x08]) + b"abcd" + bytes([0x3A, 0x02, 0x0A, 0x00]), + id="inner_length_past_enclosing_message", + ), + pytest.param( + bytes([0x32]) + b"\x80" * 10 + b"\x00", id="varint_longer_than_10_bytes" + ), + pytest.param(bytes([0x32, 0x10, 0x0A, 0x00]), id="truncated_message"), + pytest.param( + _response(_range_data(_checksummed()))[:-3], id="truncated_tag_varint" + ), + pytest.param( + _response(_range_data(_checksummed(extra=_tag(_CHECKSUMMED_CRC32C, 5)))), + id="crc32c_past_enclosing_message", + ), + pytest.param( + _response(_range_data(read_range=_tag(_READ_RANGE_OFFSET, 0) + b"\x80")), + id="read_range_varint_runs_past_enclosing_message", + ), + pytest.param( + _response(_range_data(range_end=False, extra=_tag(_RANGE_END, 0) + b"\x80")) + + _varint_field(_UNKNOWN, 1), + id="range_end_varint_runs_past_enclosing_message", + ), + pytest.param( + _data_response() + _tag(_UNKNOWN, 1) + b"\x00" * 3, + id="truncated_fixed64", + ), + pytest.param( + _data_response() + _tag(_UNKNOWN, 5) + b"\x00", + id="truncated_fixed32", + ), + pytest.param(b"\x00\x00", id="field_number_zero"), + pytest.param( + _response(_range_data(extra=b"\x00\x00")), + id="field_number_zero_in_range_data", + ), + pytest.param( + _response(_range_data(_checksummed(extra=b"\x00\x00"))), + id="field_number_zero_in_checksummed_data", + ), + pytest.param( + _response(_range_data(read_range=_read_range(extra=b"\x00\x00"))), + id="field_number_zero_in_read_range", + ), + pytest.param( + _response(handle=b"\x00\x00"), + id="field_number_zero_in_read_handle", + ), + pytest.param( + # Wire type 6 ensures the pure-Python protobuf backend also raises + # DecodeError after _tag() rejects the >32-bit tag value. + _tag(1 << 29, 6) + b"\x01" + _data_response(), + id="field_number_too_large", + ), + pytest.param( + # Field 6 / wire type 6 padded to a six-byte varint: _tag() rejects + # i - start > 5 before inspecting the wire type, and both upb and + # the pure-Python backend raise DecodeError on the fallback call. + b"\xb6\x80\x80\x80\x80\x00" + _varint(2) + _varint_field(_RANGE_END, 1), + id="tag_longer_than_five_bytes", + ), + ], +) +def test_deserialize_malformed_input_raises_like_generated_parser(buf): + with pytest.raises(DecodeError): + _storage_v2.BidiReadObjectResponse.deserialize(buf) + with pytest.raises(DecodeError): + _zero_copy_bidi_read.deserialize(buf) + + +def test_has_field_rejects_unknown_names(): + response = _zero_copy_bidi_read.deserialize(_data_response(crc32c=1)) + [range_data] = response.object_data_ranges + for message in (response, range_data, range_data.checksummed_data): + with pytest.raises(ValueError): + message.HasField("no_such_field") + + +def test_zero_copy_response_pb_is_itself(): + # ReadsResumptionStrategy unwraps ``response._pb`` before calling HasField. + response = _zero_copy_bidi_read.deserialize(_data_response()) + assert response._pb is response + + +# --- is_supported -------------------------------------------------------------- + + +@pytest.mark.parametrize("value", [None, "", "0", "false", "No", " off "]) +def test_is_supported_disabled_unless_opted_in(monkeypatch, value): + assert _zero_copy_bidi_read._ENV_VAR == "GOOGLE_CLOUD_STORAGE_ZERO_COPY_BIDI_READ" + probed = [] + monkeypatch.setattr(google_crc32c, "value", lambda data: probed.append(data) or 0) + if value is None: + monkeypatch.delenv(_zero_copy_bidi_read._ENV_VAR, raising=False) + else: + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, value) + assert _zero_copy_bidi_read.is_supported() is False + assert probed == [] + + +def test_is_supported_requires_crc32c_to_accept_memoryview(monkeypatch): + probed_with = [] + + def bytes_only(data): + probed_with.append(type(data)) + raise TypeError("argument 1 must be read-only bytes-like object") + + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, "1") + monkeypatch.setattr(google_crc32c, "value", bytes_only) + assert _zero_copy_bidi_read.is_supported() is False + assert probed_with == [memoryview] + + +@pytest.mark.parametrize("value", ["1", "true", "Yes", " on "]) +def test_is_supported_when_opted_in_and_crc32c_accepts_memoryview(monkeypatch, value): + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, value) + monkeypatch.setattr(google_crc32c, "value", lambda data: len(data)) + assert _zero_copy_bidi_read.is_supported() is True + + +# --- wrapped_rpc --------------------------------------------------------------- + + +def _transport(target): + return StorageGrpcAsyncIOTransport( + channel=grpc.aio.insecure_channel(target), credentials=AnonymousCredentials() + ) + + +def test_wrapped_rpc_ignores_transports_it_does_not_know(monkeypatch): + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: True) + assert _zero_copy_bidi_read.wrapped_rpc(object()) is None + + +@pytest.mark.asyncio +async def test_wrapped_rpc_returns_none_when_unsupported(monkeypatch): + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: False) + transport = _transport("localhost:1") + try: + assert _zero_copy_bidi_read.wrapped_rpc(transport) is None + finally: + await transport.close() + + +@pytest.mark.asyncio +async def test_wrapped_rpc_is_built_once_per_transport(monkeypatch): + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: True) + transport = _transport("localhost:1") + other = _transport("localhost:1") + try: + rpc = _zero_copy_bidi_read.wrapped_rpc(transport) + assert rpc is not None + assert rpc is not transport._wrapped_methods[transport.bidi_read_object] + assert _zero_copy_bidi_read.wrapped_rpc(transport) is rpc + assert _zero_copy_bidi_read.wrapped_rpc(other) is not rpc + finally: + await transport.close() + await other.close() + + +# --- end to end through _AsyncReadObjectStream --------------------------------- + + +class _FakeStorageServer: + """In-process gRPC server that answers BidiReadObject with canned bytes.""" + + def __init__(self, responses): + self._responses = list(responses) + self.requests = [] + self.invocation_metadata = None + self._server = grpc.aio.server() + handler = grpc.method_handlers_generic_handler( + "google.storage.v2.Storage", + {"BidiReadObject": grpc.stream_stream_rpc_method_handler(self._handle)}, + ) + self._server.add_generic_rpc_handlers((handler,)) + self._port = self._server.add_insecure_port("localhost:0") + + @property + def target(self): + return f"localhost:{self._port}" + + async def _handle(self, request_iterator, context): + self.invocation_metadata = dict(context.invocation_metadata()) + async for request in request_iterator: + self.requests.append(_storage_v2.BidiReadObjectRequest.deserialize(request)) + yield self._responses.pop(0) + + async def __aenter__(self): + await self._server.start() + return self + + async def __aexit__(self, *exc_info): + await self._server.stop(None) + + +@pytest.fixture +def keep_google_logger_propagation(monkeypatch): + # Building a real GAPIC client runs api-core's client_logging + # initialization, which switches off propagation on the "google" logger + # and would hide log records from ``caplog`` in later test modules. + logger = logging.getLogger("google") + monkeypatch.setattr(logger, "propagate", logger.propagate) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("supported", [True, False]) +async def test_read_object_stream_end_to_end( + monkeypatch, keep_google_logger_propagation, supported +): + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: supported) + crc = google_crc32c.value(PAYLOAD) + first = _storage_v2.BidiReadObjectResponse.serialize( + _storage_v2.BidiReadObjectResponse( + metadata=_storage_v2.Object( + name="o", bucket="projects/_/buckets/b", generation=5, size=2048 + ), + read_handle=_storage_v2.BidiReadHandle(handle=b"opaque"), + ) + ) + second = _data_response(crc32c=crc) + + async with _FakeStorageServer([first, second]) as server: + transport = _transport(server.target) + client = _storage_v2.StorageAsyncClient(transport=transport) + stream = _AsyncReadObjectStream(client=client, bucket_name="b", object_name="o") + try: + await stream.open() + assert stream.generation_number == 5 + assert stream.persisted_size == 2048 + assert stream.read_handle.handle == b"opaque" + + await stream.send( + _storage_v2.BidiReadObjectRequest( + read_ranges=[ + _storage_v2.ReadRange( + read_offset=4096, read_length=len(PAYLOAD), read_id=7 + ) + ] + ) + ) + response = await stream.recv() + await stream.close() + finally: + await transport.close() + + [range_data] = response.object_data_ranges + content = range_data.checksummed_data.content + assert isinstance(content, memoryview if supported else bytes) + assert bytes(content) == PAYLOAD + assert range_data.checksummed_data.crc32c == crc + assert range_data.read_range.read_id == 7 + assert ( + isinstance(response, _zero_copy_bidi_read._ZeroCopyBidiReadObjectResponse) + is supported + ) + + assert server.requests[0].read_object_spec.object_ == "o" + assert server.requests[1].read_ranges[0].read_id == 7 + assert server.invocation_metadata["x-goog-request-params"] == ( + "bucket=projects/_/buckets/b" + ) + assert "gapic/" in server.invocation_metadata["x-goog-api-client"]