From d8ae8bf913d94bb1507550c4f40b5344e188dc3b Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Fri, 2 Oct 2026 07:00:24 +0000 Subject: [PATCH 01/11] feat(storage): zero-copy BidiReadObjectResponse parsing for async reads The generated BidiReadObject stub parses every response with proto-plus over upb, which copies each data chunk twice on the event-loop thread before it reaches the caller's buffer. Add a wire-format parser for plain data responses that exposes `checksummed_data.content` as a memoryview of the received message, and let `_AsyncReadObjectStream` use a stub built with it (wrapped like the generated method: error mapping, no default retry or timeout, same client-info header). Responses carrying object metadata, or anything that does not parse cleanly, are handed to the generated parser, so callers keep seeing the real message or its usual DecodeError. CRC32C verification in the read path is unchanged: the memoryview flows through google_crc32c.value(), the user buffer write and the rolling checksum. Because google-crc32c releases up to 1.9.0 only accept bytes there, the fast path is used only when the installed google-crc32c can checksum a memoryview; GOOGLE_CLOUD_STORAGE_FAST_BIDI_READ=0 forces the generated parser as well. --- .../cloud/storage/asyncio/_fast_bidi_read.py | 325 +++++++++++ .../asyncio/async_read_object_stream.py | 10 +- .../asyncio/test_async_read_object_stream.py | 53 ++ .../tests/unit/asyncio/test_fast_bidi_read.py | 516 ++++++++++++++++++ 4 files changed, 901 insertions(+), 3 deletions(-) create mode 100644 packages/google-cloud-storage/google/cloud/storage/asyncio/_fast_bidi_read.py create mode 100644 packages/google-cloud-storage/tests/unit/asyncio/test_fast_bidi_read.py diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/_fast_bidi_read.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/_fast_bidi_read.py new file mode 100644 index 000000000000..9be3f0a61796 --- /dev/null +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/_fast_bidi_read.py @@ -0,0 +1,325 @@ +# 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, +) + +_ENV_VAR = "GOOGLE_CLOUD_STORAGE_FAST_BIDI_READ" +_DISABLED_VALUES = ("0", "false", "no", "off") + +_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 FastBidiReadObjectResponse: + """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.""" + + +def _varint(buf, i, end): + """Decode the varint at ``buf[i]``; it must end at or before ``end``.""" + 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, 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 + + +def _parse_read_range(buf, i, end): + rr = _ReadRange() + while i < end: + tag, i = _varint(buf, i, end) + f, wt = tag >> 3, tag & 7 + if wt == _WT_VARINT: + v, i = _varint(buf, i, end) + if f == 1: + rr.read_offset = v + elif f == 2: + rr.read_length = v + elif f == 3: + rr.read_id = v + else: + i = _skip(buf, i, wt, end) + return rr + + +def _parse_checksummed(buf, mv, i, end): + cd = _ChecksummedData() + while i < end: + tag, i = _varint(buf, i, end) + f, wt = tag >> 3, tag & 7 + 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) + return cd + + +def _parse_range_data(buf, mv, i, end): + rd = _ObjectRangeData() + while i < end: + tag, i = _varint(buf, i, end) + f, wt = tag >> 3, tag & 7 + if wt == _WT_LEN: + start, i = _length_prefixed(buf, i, end) + if f == 1: + rd.checksummed_data = _parse_checksummed(buf, mv, start, i) + rd._has_checksummed_data = True + elif f == 2: + rd.read_range = _parse_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(buf, i, end): + handle = b"" + while i < end: + tag, i = _varint(buf, i, end) + f, wt = tag >> 3, tag & 7 + 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) + # A real message: the handle is stored and passed back in the + # BidiReadObjectSpec of the next open, which only accepts the proto type. + return _storage_v2.BidiReadHandle(handle=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: + tag, i = _varint(buf, i, n) + f, wt = tag >> 3, tag & 7 + 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(buf, start, end) + elif f == 4: + return _generated_deserialize(buf) + i = end + else: + i = _skip(buf, i, wt, n) + return FastBidiReadObjectResponse(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``; releases of + google-crc32c that only accept ``bytes`` there would make every verified + read fail, so the generated parser is kept in that case. Setting + ``GOOGLE_CLOUD_STORAGE_FAST_BIDI_READ=0`` forces the generated parser too. + """ + if os.environ.get(_ENV_VAR, "").strip().lower() in _DISABLED_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 fast 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_read_object_stream.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py index d3c0780673c7..b536d3b3cf17 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 _fast_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 google-crc32c can verify + # memoryview payloads; otherwise the generated (proto-plus) parser. + self.rpc = _fast_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/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..88d9a0974125 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._fast_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._fast_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_fast_bidi_read.py b/packages/google-cloud-storage/tests/unit/asyncio/test_fast_bidi_read.py new file mode 100644 index 000000000000..677572cc5690 --- /dev/null +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_fast_bidi_read.py @@ -0,0 +1,516 @@ +# 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 _fast_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 fast one matches.""" + expected = _storage_v2.BidiReadObjectResponse.deserialize(buf) + actual = _fast_bidi_read.deserialize(buf) + + assert isinstance(actual, _fast_bidi_read.FastBidiReadObjectResponse) + 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 = _fast_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) + + +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 = _fast_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 fast path + # does not understand; the generated parser keeps them as unknown fields. + buf = _data_response() + _tag(_UNKNOWN, 3) + _tag(_UNKNOWN, 4) + response = _fast_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", + ), + ], +) +def test_deserialize_malformed_input_raises_like_generated_parser(buf): + with pytest.raises(DecodeError): + _storage_v2.BidiReadObjectResponse.deserialize(buf) + with pytest.raises(DecodeError): + _fast_bidi_read.deserialize(buf) + + +def test_has_field_rejects_unknown_names(): + response = _fast_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_fast_response_pb_is_itself(): + # ReadsResumptionStrategy unwraps ``response._pb`` before calling HasField. + response = _fast_bidi_read.deserialize(_data_response()) + assert response._pb is response + + +# --- is_supported -------------------------------------------------------------- + + +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.delenv(_fast_bidi_read._ENV_VAR, raising=False) + monkeypatch.setattr(google_crc32c, "value", bytes_only) + assert _fast_bidi_read.is_supported() is False + assert probed_with == [memoryview] + + +def test_is_supported_when_crc32c_accepts_memoryview(monkeypatch): + monkeypatch.delenv(_fast_bidi_read._ENV_VAR, raising=False) + monkeypatch.setattr(google_crc32c, "value", lambda data: len(data)) + assert _fast_bidi_read.is_supported() is True + + +@pytest.mark.parametrize("value", ["0", "false", "No", " off "]) +def test_is_supported_env_kill_switch(monkeypatch, value): + monkeypatch.setattr(google_crc32c, "value", lambda data: len(data)) + monkeypatch.setenv(_fast_bidi_read._ENV_VAR, value) + assert _fast_bidi_read.is_supported() is False + + +# --- 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(_fast_bidi_read, "is_supported", lambda: True) + assert _fast_bidi_read.wrapped_rpc(object()) is None + + +@pytest.mark.asyncio +async def test_wrapped_rpc_returns_none_when_unsupported(monkeypatch): + monkeypatch.setattr(_fast_bidi_read, "is_supported", lambda: False) + transport = _transport("localhost:1") + try: + assert _fast_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(_fast_bidi_read, "is_supported", lambda: True) + transport = _transport("localhost:1") + other = _transport("localhost:1") + try: + rpc = _fast_bidi_read.wrapped_rpc(transport) + assert rpc is not None + assert rpc is not transport._wrapped_methods[transport.bidi_read_object] + assert _fast_bidi_read.wrapped_rpc(transport) is rpc + assert _fast_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(_fast_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, _fast_bidi_read.FastBidiReadObjectResponse) 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"] From ff6c2f66bdd059d5ebe86f431da516822f033433 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Fri, 2 Oct 2026 10:24:51 +0000 Subject: [PATCH 02/11] fix(storage): match the generated parser on int64 sign, split sub-messages and invalid tags The zero-copy BidiReadObjectResponse parser is meant to be indistinguishable from the generated one apart from the type of ``content``; a comparison against upb found a few places where it was not: - ReadRange.read_offset/read_length/read_id are int64, so a ten-byte two's-complement varint must come back negative, and bits above 64 in an over-long varint are dropped rather than kept. - A singular embedded message that occurs twice in its parent is merged; the sub-message parsers now fill in the existing object, which also removes two throw-away allocations per ObjectRangeData. - Field number 0, tags above 32 bits and tags encoded in more than five bytes are malformed and now reach the generated parser's DecodeError. Tests cover each case against the generated parser. --- .../cloud/storage/asyncio/_fast_bidi_read.py | 81 ++++++++++------ .../tests/unit/asyncio/test_fast_bidi_read.py | 94 +++++++++++++++++++ 2 files changed, 146 insertions(+), 29 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/_fast_bidi_read.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/_fast_bidi_read.py index 9be3f0a61796..8c7102ae3c4c 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/_fast_bidi_read.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/_fast_bidi_read.py @@ -127,8 +127,15 @@ class _Malformed(ValueError): """Raised when the wire bytes do not look like a well-formed message.""" +_UINT64_MASK = (1 << 64) - 1 +_UINT32_MAX = (1 << 32) - 1 + + def _varint(buf, i, end): - """Decode the varint at ``buf[i]``; it must end at or before ``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: @@ -142,7 +149,25 @@ def _varint(buf, i, end): raise _Malformed("varint longer than 10 bytes") if i > end: raise _Malformed("varint runs past its enclosing message") - return result, i + 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): @@ -170,29 +195,29 @@ def _skip(buf, i, wt, end): return i -def _parse_read_range(buf, i, end): - rr = _ReadRange() +# 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: - tag, i = _varint(buf, i, end) - f, wt = tag >> 3, tag & 7 + f, wt, i = _tag(buf, i, end) if wt == _WT_VARINT: v, i = _varint(buf, i, end) if f == 1: - rr.read_offset = v + rr.read_offset = _int64(v) elif f == 2: - rr.read_length = v + rr.read_length = _int64(v) elif f == 3: - rr.read_id = v + rr.read_id = _int64(v) else: i = _skip(buf, i, wt, end) - return rr -def _parse_checksummed(buf, mv, i, end): - cd = _ChecksummedData() +def _parse_checksummed(cd, buf, mv, i, end): while i < end: - tag, i = _varint(buf, i, end) - f, wt = tag >> 3, tag & 7 + 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] @@ -204,21 +229,19 @@ def _parse_checksummed(buf, mv, i, end): i += 4 else: i = _skip(buf, i, wt, end) - return cd def _parse_range_data(buf, mv, i, end): rd = _ObjectRangeData() while i < end: - tag, i = _varint(buf, i, end) - f, wt = tag >> 3, tag & 7 + f, wt, i = _tag(buf, i, end) if wt == _WT_LEN: start, i = _length_prefixed(buf, i, end) if f == 1: - rd.checksummed_data = _parse_checksummed(buf, mv, start, i) + _parse_checksummed(rd.checksummed_data, buf, mv, start, i) rd._has_checksummed_data = True elif f == 2: - rd.read_range = _parse_read_range(buf, start, i) + _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) @@ -228,19 +251,15 @@ def _parse_range_data(buf, mv, i, end): return rd -def _parse_handle(buf, i, end): - handle = b"" +def _parse_handle(handle, buf, i, end): while i < end: - tag, i = _varint(buf, i, end) - f, wt = tag >> 3, tag & 7 + 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) - # A real message: the handle is stored and passed back in the - # BidiReadObjectSpec of the next open, which only accepts the proto type. - return _storage_v2.BidiReadHandle(handle=handle) + return handle _generated_deserialize = _storage_v2.BidiReadObjectResponse.deserialize @@ -260,19 +279,23 @@ def deserialize(buf): ranges = [] handle = None while i < n: - tag, i = _varint(buf, i, n) - f, wt = tag >> 3, tag & 7 + 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(buf, start, end) + 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 FastBidiReadObjectResponse(ranges, handle) except Exception: return _generated_deserialize(buf) diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_fast_bidi_read.py b/packages/google-cloud-storage/tests/unit/asyncio/test_fast_bidi_read.py index 677572cc5690..0f8ceb1dbe8b 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_fast_bidi_read.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_fast_bidi_read.py @@ -266,6 +266,70 @@ def test_deserialize_skips_unknown_nested_fields(): _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( @@ -322,6 +386,36 @@ def test_deserialize_unknown_group_uses_generated_parser(): _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): From 2b869dc410af1438adfae459b9dae3c4fa8c6995 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Fri, 2 Oct 2026 07:24:37 +0000 Subject: [PATCH 03/11] perf(storage): verify CRC32C of large read chunks off the event-loop thread update_state_from_response hashed every chunk inline on the asyncio loop thread, serialized with receive and parse. Chunks of at least 512 KiB are now submitted to a small shared thread pool (google_crc32c releases the GIL for read-only buffers of 1 MiB or more) and written immediately; download_ranges() awaits the pending results after each response (bounded to 32 outstanding) and fully before returning, raising DataCorruption on the first mismatch. Small chunks are still verified inline. On the 8-rank DCP load benchmark (test-c4-highmem, zonal bucket) read_data throughput per rank went from 1.08/1.10 GB/s to 1.39/1.46 GB/s with verification on; 1.57 GB/s is the same stack with verification off. --- .../asyncio/async_multi_range_downloader.py | 9 ++- .../retry/reads_resumption_strategy.py | 75 +++++++++++++++++-- .../retry/test_reads_resumption_strategy.py | 57 ++++++++++++++ .../test_async_multi_range_downloader.py | 1 + 4 files changed, 134 insertions(+), 8 deletions(-) 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..30afbce46d62 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,16 +524,21 @@ async def generator(): data_range.read_range.read_id ) yield item + # The strategy has consumed `item` by now; keep the + # offloaded checksum backlog bounded. + await strategy.verify_pending_checksums( + state, _CRC32C_MAX_PENDING + ) return generator() - strategy = _ReadResumptionStrategy() retry_manager = _BidiStreamRetryManager( strategy, send_and_recv_via_multiplexer ) try: await retry_manager.execute(initial_state, retry_policy) + await strategy.verify_pending_checksums(initial_state) except DataCorruption: if self.is_stream_open: await self.close() 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..78241bf4edea 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 concurrent.futures import logging -from typing import IO, Any, Dict, List +import os +import threading +from typing import IO, Any, Dict, List, Optional import google_crc32c @@ -31,6 +35,30 @@ ) logger = logging.getLogger(__name__) +# Chunks at least this large are checksummed on a worker thread instead of the +# event-loop thread (google_crc32c releases the GIL for read-only buffers of +# 1 MiB or more). Smaller chunks are hashed inline; the executor hand-off would +# cost about as much as the hash itself. +_CRC32C_OFFLOAD_MIN_BYTES = int( + os.environ.get("GOOGLE_CLOUD_STORAGE_CRC32C_OFFLOAD_MIN_BYTES", 512 * 1024) +) +# Upper bound on chunks awaiting verification per download_ranges() call, so +# a slow worker cannot pin an unbounded number of received messages. +_CRC32C_MAX_PENDING = 32 +_crc32c_executor: Optional[concurrent.futures.ThreadPoolExecutor] = None +_crc32c_executor_lock = threading.Lock() + + +def _get_crc32c_executor() -> concurrent.futures.ThreadPoolExecutor: + global _crc32c_executor + if _crc32c_executor is None: + with _crc32c_executor_lock: + if _crc32c_executor is None: + _crc32c_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=2, thread_name_prefix="gcs-crc32c" + ) + return _crc32c_executor + class _DownloadState: """A helper class to track the state of a single range download.""" @@ -139,13 +167,21 @@ def update_state_from_response( 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: + # Hash off the loop thread; verify_pending_checksums() + # raises DataCorruption before download_ranges() returns. + future = _get_crc32c_executor().submit(google_crc32c.value, data) + state.setdefault("pending_checksums", []).append( + (future, server_checksum, read_id, response) ) + else: + 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}.", + ) # Update State & Write Data chunk_size = len(data) @@ -190,6 +226,31 @@ def update_state_from_response( 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: + """Awaits offloaded chunk checksums until at most max_pending remain. + + Raises DataCorruption on the first mismatch. Chunks are verified in + arrival order, so the backlog is bounded without stalling the stream. + """ + pending = state.get("pending_checksums") + if not pending: + return + while len(pending) > max_pending: + future, server_checksum, read_id, response = pending[0] + if future.done(): + client_checksum = future.result() + else: + client_checksum = await asyncio.wrap_future(future) + del pending[0] + 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}.", + ) + async def recover_state_on_failure(self, error: Exception, state: Any) -> None: """Handles BidiReadObjectRedirectedError for reads.""" routing_token, read_handle = _handle_redirect(error) 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..6b58664a6421 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 @@ -21,6 +21,7 @@ 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, @@ -309,6 +310,62 @@ def test_update_state_checksum_mismatch(self): with self.assertRaisesRegex(DataCorruption, "Checksum mismatch"): self.strategy.update_state_from_response(response, self.state) + # --- Offloaded Checksum Tests --- + + def _large_chunk(self): + return b"x" * reads_resumption_strategy._CRC32C_OFFLOAD_MIN_BYTES + + def test_update_state_large_chunk_defers_checksum(self): + """Large chunks are written at once and verified by verify_pending_checksums.""" + content = self._large_chunk() + buffer = io.BytesIO() + self._add_download(_READ_ID, length=len(content), buffer=buffer) + response = self._create_response(content, _READ_ID, offset=0, range_end=True) + + self.strategy.update_state_from_response(response, self.state) + + self.assertEqual(buffer.getvalue(), content) + self.assertEqual(len(self.state["pending_checksums"]), 1) + asyncio.run(self.strategy.verify_pending_checksums(self.state)) + self.assertEqual(self.state["pending_checksums"], []) + + def test_verify_pending_checksums_raises_on_mismatch(self): + content = self._large_chunk() + 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) + + with self.assertRaisesRegex(DataCorruption, "Checksum mismatch"): + asyncio.run(self.strategy.verify_pending_checksums(self.state)) + + def test_verify_pending_checksums_honors_max_pending(self): + content = self._large_chunk() + 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) + asyncio.run(self.strategy.verify_pending_checksums(self.state)) + self.assertEqual(len(self.state["pending_checksums"]), 0) + + 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_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_final_byte_count_mismatch(self): """Test mismatch between expected length and actual bytes written on completion.""" self._add_download(_READ_ID, length=100) 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..baca33bd2b0f 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 @@ -688,6 +688,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) From f4d73cb6ffcb70aae9f950f901400755d4641a23 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Fri, 2 Oct 2026 11:04:07 +0000 Subject: [PATCH 04/11] fix(storage): verify offloaded checksums before retrying and keep full-object reads inline - Drain pending offloaded chunk checksums at the start of `recover_state_on_failure` so a corrupt chunk already written to the user buffer surfaces as `DataCorruption` before reopening the stream from the advanced offset. - Skip the worker-thread CRC32C offload when `rolling_checksum` is present (`is_full_object_read=True`), since full-object reads already fold every chunk into `rolling_checksum` on the event-loop thread. - Store `pending_checksums` in a `collections.deque` for O(1) `popleft` and cancel remaining queued futures on the first checksum mismatch. - Parse `GOOGLE_CLOUD_STORAGE_CRC32C_OFFLOAD_MIN_BYTES` via `_int_from_env` so malformed values fall back to the 512 KiB default and document the 1 MiB GIL-release threshold in `google_crc32c`. - Always acquire `_crc32c_executor_lock` in `_get_crc32c_executor`. --- .../retry/reads_resumption_strategy.py | 67 +++++++--- .../retry/test_reads_resumption_strategy.py | 121 +++++++++++++++++- .../test_async_multi_range_downloader.py | 16 ++- 3 files changed, 181 insertions(+), 23 deletions(-) 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 78241bf4edea..a899f1ae7f9a 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 @@ -13,6 +13,7 @@ # limitations under the License. import asyncio +import collections import concurrent.futures import logging import os @@ -35,12 +36,22 @@ ) 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 (google_crc32c releases the GIL for read-only buffers of -# 1 MiB or more). Smaller chunks are hashed inline; the executor hand-off would -# cost about as much as the hash itself. -_CRC32C_OFFLOAD_MIN_BYTES = int( - os.environ.get("GOOGLE_CLOUD_STORAGE_CRC32C_OFFLOAD_MIN_BYTES", 512 * 1024) +# event-loop thread; smaller chunks are hashed inline, since the executor +# hand-off would cost about as much as the hash itself. google_crc32c releases +# the GIL only while hashing read-only buffers of 1 MiB or more, so the workers +# run fully in parallel with the loop from that size up. +_CRC32C_OFFLOAD_MIN_BYTES = _int_from_env( + "GOOGLE_CLOUD_STORAGE_CRC32C_OFFLOAD_MIN_BYTES", 512 * 1024 ) # Upper bound on chunks awaiting verification per download_ranges() call, so # a slow worker cannot pin an unbounded number of received messages. @@ -51,13 +62,12 @@ def _get_crc32c_executor() -> concurrent.futures.ThreadPoolExecutor: global _crc32c_executor - if _crc32c_executor is None: - with _crc32c_executor_lock: - if _crc32c_executor is None: - _crc32c_executor = concurrent.futures.ThreadPoolExecutor( - max_workers=2, thread_name_prefix="gcs-crc32c" - ) - return _crc32c_executor + with _crc32c_executor_lock: + if _crc32c_executor is None: + _crc32c_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=2, thread_name_prefix="gcs-crc32c" + ) + return _crc32c_executor class _DownloadState: @@ -167,13 +177,20 @@ def update_state_from_response( if checksum_enabled and checksummed_data.HasField("crc32c"): server_checksum = checksummed_data.crc32c - if len(data) >= _CRC32C_OFFLOAD_MIN_BYTES: + if ( + len(data) >= _CRC32C_OFFLOAD_MIN_BYTES + and read_state.rolling_checksum is None + ): # Hash off the loop thread; verify_pending_checksums() # raises DataCorruption before download_ranges() returns. + # Full-object reads are excluded: they already fold every + # chunk into rolling_checksum on this thread below, so a + # second hash on a worker would only add CPU work. future = _get_crc32c_executor().submit(google_crc32c.value, data) - state.setdefault("pending_checksums", []).append( - (future, server_checksum, read_id, response) - ) + pending = state.get("pending_checksums") + if pending is None: + pending = state["pending_checksums"] = collections.deque() + pending.append((future, server_checksum, read_id, response)) else: client_checksum = google_crc32c.value(data) if server_checksum != client_checksum: @@ -231,8 +248,9 @@ async def verify_pending_checksums( ) -> None: """Awaits offloaded chunk checksums until at most max_pending remain. - Raises DataCorruption on the first mismatch. Chunks are verified in - arrival order, so the backlog is bounded without stalling the stream. + Raises DataCorruption on the first mismatch, cancelling any checksums + still queued behind it. Chunks are verified in arrival order, so the + backlog is bounded without stalling the stream. """ pending = state.get("pending_checksums") if not pending: @@ -243,8 +261,11 @@ async def verify_pending_checksums( client_checksum = future.result() else: client_checksum = await asyncio.wrap_future(future) - del pending[0] + pending.popleft() if server_checksum != client_checksum: + for entry in pending: + entry[0].cancel() + pending.clear() raise DataCorruption( response, f"Checksum mismatch for read_id {read_id}. " @@ -252,7 +273,13 @@ async def verify_pending_checksums( ) async def recover_state_on_failure(self, error: Exception, state: Any) -> None: - """Handles BidiReadObjectRedirectedError for reads.""" + """Verifies offloaded checksums, then handles BidiReadObjectRedirectedError. + + Chunks whose checksum is still pending have already been written to the + user buffer and advanced the resume offset, so a corrupt one must be + reported before the stream is reopened from that offset. + """ + await self.verify_pending_checksums(state) routing_token, read_handle = _handle_redirect(error) if routing_token: state["routing_token"] = routing_token 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 6b58664a6421..dd7a78ab092d 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,8 +13,11 @@ # limitations under the License. import asyncio +import concurrent.futures import io +import os import unittest +from unittest import mock import google_crc32c from google.api_core import exceptions @@ -32,6 +35,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.""" @@ -315,6 +339,16 @@ def test_update_state_checksum_mismatch(self): 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, "_get_crc32c_executor", return_value=executor + ) + patcher.start() + self.addCleanup(patcher.stop) + def test_update_state_large_chunk_defers_checksum(self): """Large chunks are written at once and verified by verify_pending_checksums.""" content = self._large_chunk() @@ -327,7 +361,7 @@ def test_update_state_large_chunk_defers_checksum(self): self.assertEqual(buffer.getvalue(), content) self.assertEqual(len(self.state["pending_checksums"]), 1) asyncio.run(self.strategy.verify_pending_checksums(self.state)) - self.assertEqual(self.state["pending_checksums"], []) + self.assertEqual(len(self.state["pending_checksums"]), 0) def test_verify_pending_checksums_raises_on_mismatch(self): content = self._large_chunk() @@ -356,6 +390,39 @@ 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_update_state_large_chunk_not_deferred_when_checksum_disabled(self): content = self._large_chunk() self.state["enable_checksum"] = False @@ -366,6 +433,28 @@ def test_update_state_large_chunk_not_deferred_when_checksum_disabled(self): 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) @@ -470,6 +559,36 @@ 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.""" + content = self._large_chunk() + 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) + + with self.assertRaisesRegex(DataCorruption, "Checksum mismatch"): + asyncio.run( + self.strategy.recover_state_on_failure( + exceptions.ServiceUnavailable("stream reset"), self.state + ) + ) + + def test_recover_state_on_failure_drains_pending_checksums(self): + """Verified chunks are drained before the redirect is applied.""" + content = self._large_chunk() + 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) + token = "dummy-routing-token" + final_error = exceptions.Aborted( + "Retry failed", errors=[BidiReadObjectRedirectedError(routing_token=token)] + ) + + asyncio.run(self.strategy.recover_state_on_failure(final_error, self.state)) + + self.assertEqual(len(self.state["pending_checksums"]), 0) + self.assertEqual(self.state["routing_token"], token) + def test_update_state_full_object_checksum_success(self): """Test that full object checksum verification succeeds on range_end.""" read_state = self._add_download( 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 baca33bd2b0f..fcd7ff68f9e0 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 @@ -27,6 +27,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 +329,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 +355,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 From 2cdc82f2f2a0ff88bcf0999ed9b78a34bf138e26 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Sat, 3 Oct 2026 14:36:22 +0000 Subject: [PATCH 05/11] fix(storage): write offloaded chunks only after their checksum verifies Chunks checksummed on the worker pool were written to the caller's buffer and counted in bytes_written before their CRC32C was known. A mismatch left the corrupt bytes in the buffer, a retry resumed after them, and an error or cancellation that arrived first hid the mismatch. Hold received chunks in arrival order and write each one only once its checksum matches. next_expected_offset still advances on receipt, so offsets are validated as before. On a mismatch, drop every held chunk and rewind each read to the end of its written bytes, so a retry asks for the dropped bytes again. download_ranges() writes or rejects the held chunks before an attempt ends and before another error propagates, and drops them on cancellation. --- .../asyncio/async_multi_range_downloader.py | 23 +- .../retry/reads_resumption_strategy.py | 207 ++++++++++++------ .../retry/test_reads_resumption_strategy.py | 161 ++++++++++++-- .../test_async_multi_range_downloader.py | 128 ++++++++++- 4 files changed, 421 insertions(+), 98 deletions(-) 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 30afbce46d62..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 @@ -524,12 +524,18 @@ async def generator(): data_range.read_range.read_id ) yield item - # The strategy has consumed `item` by now; keep the - # offloaded checksum backlog bounded. + # 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() retry_manager = _BidiStreamRetryManager( @@ -537,12 +543,23 @@ async def generator(): ) 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/retry/reads_resumption_strategy.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/retry/reads_resumption_strategy.py index a899f1ae7f9a..39d4adc27d65 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 @@ -18,7 +18,7 @@ import logging import os import threading -from typing import IO, Any, Dict, List, Optional +from typing import IO, Any, Dict, List, NamedTuple, Optional import google_crc32c @@ -53,8 +53,9 @@ def _int_from_env(name: str, default: int) -> int: _CRC32C_OFFLOAD_MIN_BYTES = _int_from_env( "GOOGLE_CLOUD_STORAGE_CRC32C_OFFLOAD_MIN_BYTES", 512 * 1024 ) -# Upper bound on chunks awaiting verification per download_ranges() call, so -# a slow worker cannot pin an unbounded number of received messages. +# 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 _crc32c_executor: Optional[concurrent.futures.ThreadPoolExecutor] = None _crc32c_executor_lock = threading.Lock() @@ -70,6 +71,19 @@ def _get_crc32c_executor() -> concurrent.futures.ThreadPoolExecutor: return _crc32c_executor +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.""" @@ -85,6 +99,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 @@ -174,6 +191,8 @@ 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 @@ -181,17 +200,13 @@ def update_state_from_response( len(data) >= _CRC32C_OFFLOAD_MIN_BYTES and read_state.rolling_checksum is None ): - # Hash off the loop thread; verify_pending_checksums() - # raises DataCorruption before download_ranges() returns. + # 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 below, so a - # second hash on a worker would only add CPU work. - future = _get_crc32c_executor().submit(google_crc32c.value, data) - pending = state.get("pending_checksums") - if pending is None: - pending = state["pending_checksums"] = collections.deque() - pending.append((future, server_checksum, read_id, response)) - else: + # chunk into rolling_checksum on this thread, so a second + # hash on a worker would only add CPU work. + checksum = _get_crc32c_executor().submit(google_crc32c.value, data) + if checksum is None: client_checksum = google_crc32c.value(data) if server_checksum != client_checksum: raise DataCorruption( @@ -200,88 +215,134 @@ def update_state_from_response( 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 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}", + ) - # 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 - ): - raise DataCorruption( - response, - f"Byte count mismatch for read_id {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", ) - - # 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( + 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}.", ) - 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}.", - ) async def verify_pending_checksums( self, state: Dict[str, Any], max_pending: int = 0 ) -> None: - """Awaits offloaded chunk checksums until at most max_pending remain. + """Writes held chunks, in arrival order, once their checksums match. - Raises DataCorruption on the first mismatch, cancelling any checksums - still queued behind it. Chunks are verified in arrival order, so the - backlog is bounded without stalling the stream. + 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") - if not pending: - return - while len(pending) > max_pending: - future, server_checksum, read_id, response = pending[0] - if future.done(): - client_checksum = future.result() - else: - client_checksum = await asyncio.wrap_future(future) + 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( + chunk.response, + f"Checksum mismatch for read_id {chunk.read_id}. " + f"Server sent {chunk.server_checksum}, client calculated {client_checksum}.", + ) pending.popleft() - if server_checksum != client_checksum: - for entry in pending: - entry[0].cancel() - pending.clear() - raise DataCorruption( - response, - f"Checksum mismatch for read_id {read_id}. " - f"Server sent {server_checksum}, client calculated {client_checksum}.", - ) + self._write_chunk(state, chunk) + + 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: - """Verifies offloaded checksums, then handles BidiReadObjectRedirectedError. + """Handles BidiReadObjectRedirectedError, then writes held chunks. - Chunks whose checksum is still pending have already been written to the - user buffer and advanced the resume offset, so a corrupt one must be - reported before the stream is reopened from that offset. + A retry resumes each read right after its written bytes, so held + chunks are verified and written first; a corrupt one raises + DataCorruption instead. """ - await self.verify_pending_checksums(state) 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 dd7a78ab092d..412e90826cac 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 @@ -334,6 +334,20 @@ 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): @@ -349,33 +363,84 @@ def _offload_to(self, futures): patcher.start() self.addCleanup(patcher.stop) - def test_update_state_large_chunk_defers_checksum(self): - """Large chunks are written at once and verified by verify_pending_checksums.""" + def test_update_state_holds_large_chunk_until_verified(self): + """Large chunks are written only once their offloaded checksum matches.""" content = self._large_chunk() - buffer = io.BytesIO() - self._add_download(_READ_ID, length=len(content), buffer=buffer) + 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(buffer.getvalue(), content) + 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_raises_on_mismatch(self): + 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() - 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) + 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() - self._add_download(_READ_ID, length=3 * len(content)) + 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) @@ -383,8 +448,13 @@ def test_verify_pending_checksums_honors_max_pending(self): 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)) @@ -423,6 +493,38 @@ def test_verify_pending_checksums_cancels_queued_futures_on_mismatch(self): 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_large_chunk_not_deferred_when_checksum_disabled(self): content = self._large_chunk() self.state["enable_checksum"] = False @@ -560,23 +662,31 @@ async def run(): 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.""" + """A corrupt chunk received before a stream error aborts instead of retrying, + without writing it; the redirect is still recorded.""" content = self._large_chunk() - self._add_download(_READ_ID, length=len(content)) + 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( - exceptions.ServiceUnavailable("stream reset"), self.state - ) - ) + 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 drained before the redirect is applied.""" + """Verified chunks are written before the retry computes its offsets.""" content = self._large_chunk() - self._add_download(_READ_ID, length=len(content)) + 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" @@ -584,10 +694,19 @@ def test_recover_state_on_failure_drains_pending_checksums(self): "Retry failed", errors=[BidiReadObjectRedirectedError(routing_token=token)] ) - asyncio.run(self.strategy.recover_state_on_failure(final_error, self.state)) + 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.""" 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 fcd7ff68f9e0..97b04f85418a 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 @@ -379,15 +380,140 @@ async def test_download_ranges_raises_on_checksum_mismatch( 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, "_get_crc32c_executor", return_value=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._get_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", From 1f2ef27718ffbfa034ee60837b7dc1c6dac2b643 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Sat, 3 Oct 2026 14:36:38 +0000 Subject: [PATCH 06/11] fix(storage): make the CRC32C executor fork- and shutdown-safe A forked child inherits the module-level pool but none of its worker threads, so offloaded checksums never finished there and the download hung. Reset the pool and its lock in the child after fork(). Once interpreter shutdown starts, ThreadPoolExecutor.submit() raises RuntimeError, which failed downloads still running in non-daemon threads. Hash those chunks inline instead. --- .../retry/reads_resumption_strategy.py | 25 ++++++- .../retry/test_reads_resumption_strategy.py | 73 +++++++++++++++++++ 2 files changed, 97 insertions(+), 1 deletion(-) 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 39d4adc27d65..d2e16156047b 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 @@ -71,6 +71,29 @@ def _get_crc32c_executor() -> concurrent.futures.ThreadPoolExecutor: return _crc32c_executor +def _reset_crc32c_executor() -> None: + """Forgets the parent's pool in a forked child, where it has no threads. + + The lock is replaced too, in case another thread held it during fork(). + """ + global _crc32c_executor, _crc32c_executor_lock + _crc32c_executor = None + _crc32c_executor_lock = threading.Lock() + + +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 _get_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.""" @@ -205,7 +228,7 @@ def update_state_from_response( # 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 = _get_crc32c_executor().submit(google_crc32c.value, data) + checksum = _submit_crc32c(data) if checksum is None: client_checksum = google_crc32c.value(data) if server_checksum != client_checksum: 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 412e90826cac..6ae1cc53bbe8 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 @@ -16,7 +16,10 @@ import concurrent.futures import io import os +import signal +import threading import unittest +import warnings from unittest import mock import google_crc32c @@ -525,6 +528,32 @@ def test_verify_pending_checksums_checks_byte_count_of_held_chunk(self): 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, "_get_crc32c_executor", return_value=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 @@ -765,3 +794,47 @@ 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_drops_the_pool_and_lock(self): + patcher = mock.patch.multiple( + reads_resumption_strategy, + _crc32c_executor=None, + _crc32c_executor_lock=threading.Lock(), + ) + patcher.start() + self.addCleanup(patcher.stop) + pool = reads_resumption_strategy._get_crc32c_executor() + self.addCleanup(pool.shutdown) + lock = reads_resumption_strategy._crc32c_executor_lock + + reads_resumption_strategy._reset_crc32c_executor() + + self.assertIsNot(reads_resumption_strategy._crc32c_executor_lock, lock) + new_pool = reads_resumption_strategy._get_crc32c_executor() + self.addCleanup(new_pool.shutdown) + self.assertIsNot(new_pool, pool) + + @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._get_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._get_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) From fb050b017b5c19494b2281c86fa420c31346d24d Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Sat, 3 Oct 2026 15:08:37 +0000 Subject: [PATCH 07/11] docs(storage): describe the google-crc32c GIL rule accurately google-crc32c releases the GIL only for chunks of 1 MiB or more that are backed by bytes, not for any read-only buffer. Received chunks are bytes or memoryview slices of bytes, so offloaded hashes still run in parallel. --- .../cloud/storage/asyncio/retry/reads_resumption_strategy.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 d2e16156047b..0a4e462413db 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 @@ -48,8 +48,9 @@ def _int_from_env(name: str, default: int) -> int: # Chunks at least this large are checksummed on a worker thread instead of the # event-loop thread; smaller chunks are hashed inline, since the executor # hand-off would cost about as much as the hash itself. google_crc32c releases -# the GIL only while hashing read-only buffers of 1 MiB or more, so the workers -# run fully in parallel with the loop from that size up. +# the GIL only while hashing chunks of 1 MiB or more backed by ``bytes`` +# (received chunks are), so the workers run fully in parallel with the loop +# from that size up. _CRC32C_OFFLOAD_MIN_BYTES = _int_from_env( "GOOGLE_CLOUD_STORAGE_CRC32C_OFFLOAD_MIN_BYTES", 512 * 1024 ) From c0d9d1ac29c85241c418d83b9cacf57859c48623 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Mon, 5 Oct 2026 14:01:11 +0000 Subject: [PATCH 08/11] perf(storage): raise the default CRC32C offload threshold to 1 MiB --- .../storage/asyncio/retry/reads_resumption_strategy.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) 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 0a4e462413db..cc46d1cef95a 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 @@ -46,13 +46,11 @@ def _int_from_env(name: str, default: int) -> int: # Chunks at least this large are checksummed on a worker thread instead of the -# event-loop thread; smaller chunks are hashed inline, since the executor -# hand-off would cost about as much as the hash itself. google_crc32c releases -# the GIL only while hashing chunks of 1 MiB or more backed by ``bytes`` -# (received chunks are), so the workers run fully in parallel with the loop -# from that size up. +# 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", 512 * 1024 + "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 From b203695c8c4a3104c7bca8fbf5ca41520be1fc60 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Tue, 6 Oct 2026 13:53:27 +0000 Subject: [PATCH 09/11] refactor(storage): rename _zero_copy_bidi_read and make it opt-in via GOOGLE_CLOUD_STORAGE_ZERO_COPY_BIDI_READ=1 --- ...t_bidi_read.py => _zero_copy_bidi_read.py} | 25 +++--- .../asyncio/async_read_object_stream.py | 8 +- .../asyncio/test_async_read_object_stream.py | 4 +- ...di_read.py => test_zero_copy_bidi_read.py} | 80 +++++++++++-------- 4 files changed, 64 insertions(+), 53 deletions(-) rename packages/google-cloud-storage/google/cloud/storage/asyncio/{_fast_bidi_read.py => _zero_copy_bidi_read.py} (92%) rename packages/google-cloud-storage/tests/unit/asyncio/{test_fast_bidi_read.py => test_zero_copy_bidi_read.py} (89%) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/_fast_bidi_read.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/_zero_copy_bidi_read.py similarity index 92% rename from packages/google-cloud-storage/google/cloud/storage/asyncio/_fast_bidi_read.py rename to packages/google-cloud-storage/google/cloud/storage/asyncio/_zero_copy_bidi_read.py index 8c7102ae3c4c..ee3cb1fd1dff 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/_fast_bidi_read.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/_zero_copy_bidi_read.py @@ -42,8 +42,8 @@ StorageGrpcAsyncIOTransport, ) -_ENV_VAR = "GOOGLE_CLOUD_STORAGE_FAST_BIDI_READ" -_DISABLED_VALUES = ("0", "false", "no", "off") +_ENV_VAR = "GOOGLE_CLOUD_STORAGE_ZERO_COPY_BIDI_READ" +_ENABLED_VALUES = ("1", "true", "yes", "on") _WT_VARINT, _WT_I64, _WT_LEN, _WT_I32 = 0, 1, 2, 5 @@ -95,7 +95,7 @@ def HasField(self, name): raise ValueError(f"Unknown field {name!r}") -class FastBidiReadObjectResponse: +class _ZeroCopyBidiReadObjectResponse: """Duck-typed stand-in for a ``BidiReadObjectResponse`` without metadata. Only the attributes the async read path uses are provided: @@ -296,7 +296,7 @@ def deserialize(buf): # BidiReadObjectSpec of the next open, which only accepts the # proto type. handle = _storage_v2.BidiReadHandle(handle=handle) - return FastBidiReadObjectResponse(ranges, handle) + return _ZeroCopyBidiReadObjectResponse(ranges, handle) except Exception: return _generated_deserialize(buf) @@ -305,12 +305,13 @@ 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``; releases of - google-crc32c that only accept ``bytes`` there would make every verified - read fail, so the generated parser is kept in that case. Setting - ``GOOGLE_CLOUD_STORAGE_FAST_BIDI_READ=0`` forces the generated parser too. + 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() in _DISABLED_VALUES: + if os.environ.get(_ENV_VAR, "").strip().lower() not in _ENABLED_VALUES: return False try: google_crc32c.value(memoryview(b"\0")) @@ -328,9 +329,9 @@ def wrapped_rpc(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 fast path is unsupported or ``transport`` is not - the generated grpc_asyncio transport, in which case callers should use the - generated wrapped method. + 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 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 b536d3b3cf17..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,7 +19,7 @@ from google.api_core.bidi_async import AsyncBidiRpc from google.cloud import _storage_v2 -from google.cloud.storage.asyncio import _fast_bidi_read +from google.cloud.storage.asyncio import _zero_copy_bidi_read from google.cloud.storage.asyncio.async_abstract_object_stream import ( _AsyncAbstractObjectStream, ) @@ -78,9 +78,9 @@ def __init__( self._full_bucket_name = f"projects/_/buckets/{self.bucket_name}" transport = self.client._client._transport - # Zero-copy BidiReadObject responses when google-crc32c can verify - # memoryview payloads; otherwise the generated (proto-plus) parser. - self.rpc = _fast_bidi_read.wrapped_rpc(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}"),) 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 88d9a0974125..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 @@ -118,7 +118,7 @@ def test_init_with_bucket_object_generation(mock_client, mock_async_bidi_rpc): "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" ) @mock.patch( - "google.cloud.storage.asyncio.async_read_object_stream._fast_bidi_read.wrapped_rpc" + "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 @@ -142,7 +142,7 @@ def test_init_prefers_zero_copy_rpc(mock_wrapped_rpc, mock_client, mock_async_bi "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" ) @mock.patch( - "google.cloud.storage.asyncio.async_read_object_stream._fast_bidi_read.wrapped_rpc" + "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 diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_fast_bidi_read.py b/packages/google-cloud-storage/tests/unit/asyncio/test_zero_copy_bidi_read.py similarity index 89% rename from packages/google-cloud-storage/tests/unit/asyncio/test_fast_bidi_read.py rename to packages/google-cloud-storage/tests/unit/asyncio/test_zero_copy_bidi_read.py index 0f8ceb1dbe8b..8cbde30497e2 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_fast_bidi_read.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_zero_copy_bidi_read.py @@ -26,7 +26,7 @@ from google.cloud._storage_v2.services.storage.transports.grpc_asyncio import ( StorageGrpcAsyncIOTransport, ) -from google.cloud.storage.asyncio import _fast_bidi_read +from google.cloud.storage.asyncio import _zero_copy_bidi_read from google.cloud.storage.asyncio.async_read_object_stream import ( _AsyncReadObjectStream, ) @@ -117,11 +117,11 @@ def _data_response(content=PAYLOAD, crc32c=None, handle=None): def _assert_same_as_generated(buf): - """Parse ``buf`` with both parsers and check the fast one matches.""" + """Parse ``buf`` with both parsers and check the zero-copy one matches.""" expected = _storage_v2.BidiReadObjectResponse.deserialize(buf) - actual = _fast_bidi_read.deserialize(buf) + actual = _zero_copy_bidi_read.deserialize(buf) - assert isinstance(actual, _fast_bidi_read.FastBidiReadObjectResponse) + 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") @@ -184,7 +184,7 @@ def test_deserialize_matches_proto_plus_serialization(): def test_deserialize_content_is_a_view_of_the_wire_buffer(): buf = _data_response() - content = _fast_bidi_read.deserialize(buf).object_data_ranges[0] + content = _zero_copy_bidi_read.deserialize(buf).object_data_ranges[0] assert content.checksummed_data.content.obj is buf @@ -335,16 +335,16 @@ def test_deserialize_metadata_response_uses_generated_parser(): buf = _storage_v2.BidiReadObjectResponse.serialize( _storage_v2.BidiReadObjectResponse(metadata=metadata) ) - response = _fast_bidi_read.deserialize(buf) + 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 fast path - # does not understand; the generated parser keeps them as unknown fields. + # 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 = _fast_bidi_read.deserialize(buf) + response = _zero_copy_bidi_read.deserialize(buf) assert isinstance(response, _storage_v2.BidiReadObjectResponse) assert response.object_data_ranges[0].checksummed_data.content == PAYLOAD @@ -422,26 +422,39 @@ def test_deserialize_malformed_input_raises_like_generated_parser(buf): with pytest.raises(DecodeError): _storage_v2.BidiReadObjectResponse.deserialize(buf) with pytest.raises(DecodeError): - _fast_bidi_read.deserialize(buf) + _zero_copy_bidi_read.deserialize(buf) def test_has_field_rejects_unknown_names(): - response = _fast_bidi_read.deserialize(_data_response(crc32c=1)) + 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_fast_response_pb_is_itself(): +def test_zero_copy_response_pb_is_itself(): # ReadsResumptionStrategy unwraps ``response._pb`` before calling HasField. - response = _fast_bidi_read.deserialize(_data_response()) + 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 = [] @@ -449,23 +462,17 @@ def bytes_only(data): probed_with.append(type(data)) raise TypeError("argument 1 must be read-only bytes-like object") - monkeypatch.delenv(_fast_bidi_read._ENV_VAR, raising=False) + monkeypatch.setenv(_zero_copy_bidi_read._ENV_VAR, "1") monkeypatch.setattr(google_crc32c, "value", bytes_only) - assert _fast_bidi_read.is_supported() is False + assert _zero_copy_bidi_read.is_supported() is False assert probed_with == [memoryview] -def test_is_supported_when_crc32c_accepts_memoryview(monkeypatch): - monkeypatch.delenv(_fast_bidi_read._ENV_VAR, raising=False) +@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 _fast_bidi_read.is_supported() is True - - -@pytest.mark.parametrize("value", ["0", "false", "No", " off "]) -def test_is_supported_env_kill_switch(monkeypatch, value): - monkeypatch.setattr(google_crc32c, "value", lambda data: len(data)) - monkeypatch.setenv(_fast_bidi_read._ENV_VAR, value) - assert _fast_bidi_read.is_supported() is False + assert _zero_copy_bidi_read.is_supported() is True # --- wrapped_rpc --------------------------------------------------------------- @@ -478,31 +485,31 @@ def _transport(target): def test_wrapped_rpc_ignores_transports_it_does_not_know(monkeypatch): - monkeypatch.setattr(_fast_bidi_read, "is_supported", lambda: True) - assert _fast_bidi_read.wrapped_rpc(object()) is None + 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(_fast_bidi_read, "is_supported", lambda: False) + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: False) transport = _transport("localhost:1") try: - assert _fast_bidi_read.wrapped_rpc(transport) is None + 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(_fast_bidi_read, "is_supported", lambda: True) + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: True) transport = _transport("localhost:1") other = _transport("localhost:1") try: - rpc = _fast_bidi_read.wrapped_rpc(transport) + 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 _fast_bidi_read.wrapped_rpc(transport) is rpc - assert _fast_bidi_read.wrapped_rpc(other) is not rpc + 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() @@ -558,7 +565,7 @@ def keep_google_logger_propagation(monkeypatch): async def test_read_object_stream_end_to_end( monkeypatch, keep_google_logger_propagation, supported ): - monkeypatch.setattr(_fast_bidi_read, "is_supported", lambda: supported) + monkeypatch.setattr(_zero_copy_bidi_read, "is_supported", lambda: supported) crc = google_crc32c.value(PAYLOAD) first = _storage_v2.BidiReadObjectResponse.serialize( _storage_v2.BidiReadObjectResponse( @@ -600,7 +607,10 @@ async def test_read_object_stream_end_to_end( assert bytes(content) == PAYLOAD assert range_data.checksummed_data.crc32c == crc assert range_data.read_range.read_id == 7 - assert isinstance(response, _fast_bidi_read.FastBidiReadObjectResponse) is supported + 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 From f8cbf4a64bb8d1ea524a88492af40e7d79926ff1 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Thu, 8 Oct 2026 14:36:40 +0000 Subject: [PATCH 10/11] docs(storage): comment the zero-copy parser's wire-type and bound constants Explain the opt-in switch's accepted spellings, the four protobuf wire types, the 64-bit varint mask and the 32-bit tag limit. --- .../google/cloud/storage/asyncio/_zero_copy_bidi_read.py | 5 +++++ 1 file changed, 5 insertions(+) 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 index ee3cb1fd1dff..fefa444a179e 100644 --- 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 @@ -42,9 +42,12 @@ 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 @@ -127,7 +130,9 @@ 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 From 49d5e29eb9a3a631a0b523a8c383d7c49c995fe9 Mon Sep 17 00:00:00 2001 From: yuxin00j Date: Fri, 9 Oct 2026 04:59:27 +0000 Subject: [PATCH 11/11] refactor(storage): create the CRC32C executor eagerly A ThreadPoolExecutor starts no threads until its first submit(), so creating it at import costs nothing and removes the lazy getter and its lock. The at-fork hook now just installs a fresh pool in the child, and the shutdown fallback in _submit_crc32c is unchanged. --- .../retry/reads_resumption_strategy.py | 31 +++++++------------ .../retry/test_reads_resumption_strategy.py | 27 ++++++---------- .../test_async_multi_range_downloader.py | 4 +-- 3 files changed, 23 insertions(+), 39 deletions(-) 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 cc46d1cef95a..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 @@ -17,7 +17,6 @@ import concurrent.futures import logging import os -import threading from typing import IO, Any, Dict, List, NamedTuple, Optional import google_crc32c @@ -56,28 +55,22 @@ def _int_from_env(name: str, default: int) -> int: # 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 -_crc32c_executor: Optional[concurrent.futures.ThreadPoolExecutor] = None -_crc32c_executor_lock = threading.Lock() -def _get_crc32c_executor() -> concurrent.futures.ThreadPoolExecutor: - global _crc32c_executor - with _crc32c_executor_lock: - if _crc32c_executor is None: - _crc32c_executor = concurrent.futures.ThreadPoolExecutor( - max_workers=2, thread_name_prefix="gcs-crc32c" - ) - return _crc32c_executor +def _new_crc32c_executor() -> concurrent.futures.ThreadPoolExecutor: + return concurrent.futures.ThreadPoolExecutor( + max_workers=2, thread_name_prefix="gcs-crc32c" + ) -def _reset_crc32c_executor() -> None: - """Forgets the parent's pool in a forked child, where it has no threads. +# Process-wide pool shared by all clients; its threads start on the first submit(). +_crc32c_executor = _new_crc32c_executor() - The lock is replaced too, in case another thread held it during fork(). - """ - global _crc32c_executor, _crc32c_executor_lock - _crc32c_executor = None - _crc32c_executor_lock = threading.Lock() + +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 @@ -87,7 +80,7 @@ def _reset_crc32c_executor() -> None: def _submit_crc32c(data: Any) -> Optional[concurrent.futures.Future]: """Starts hashing ``data`` on the pool; None if the pool refuses work.""" try: - return _get_crc32c_executor().submit(google_crc32c.value, data) + return _crc32c_executor.submit(google_crc32c.value, data) except RuntimeError: # Raised once interpreter shutdown has begun; hash inline instead. return None 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 6ae1cc53bbe8..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 @@ -17,7 +17,6 @@ import io import os import signal -import threading import unittest import warnings from unittest import mock @@ -361,7 +360,7 @@ def _offload_to(self, futures): executor = mock.Mock() executor.submit.side_effect = list(futures) patcher = mock.patch.object( - reads_resumption_strategy, "_get_crc32c_executor", return_value=executor + reads_resumption_strategy, "_crc32c_executor", executor ) patcher.start() self.addCleanup(patcher.stop) @@ -536,7 +535,7 @@ def test_update_state_hashes_inline_when_pool_rejects_work(self): "cannot schedule new futures after interpreter shutdown" ) patcher = mock.patch.object( - reads_resumption_strategy, "_get_crc32c_executor", return_value=executor + reads_resumption_strategy, "_crc32c_executor", executor ) patcher.start() self.addCleanup(patcher.stop) @@ -797,30 +796,22 @@ def test_update_state_full_object_checksum_mismatch_ignored_when_disabled(self): class TestCrc32cExecutor(unittest.TestCase): - def test_reset_drops_the_pool_and_lock(self): - patcher = mock.patch.multiple( - reads_resumption_strategy, - _crc32c_executor=None, - _crc32c_executor_lock=threading.Lock(), - ) - patcher.start() - self.addCleanup(patcher.stop) - pool = reads_resumption_strategy._get_crc32c_executor() - self.addCleanup(pool.shutdown) - lock = reads_resumption_strategy._crc32c_executor_lock + 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() - self.assertIsNot(reads_resumption_strategy._crc32c_executor_lock, lock) - new_pool = reads_resumption_strategy._get_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._get_crc32c_executor() + 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. @@ -830,7 +821,7 @@ def test_forked_child_gets_a_working_pool(self): ok = False try: signal.alarm(30) # never leave the parent waiting on a hung child - child_pool = reads_resumption_strategy._get_crc32c_executor() + child_pool = reads_resumption_strategy._crc32c_executor ok = child_pool is not pool and ( child_pool.submit(int, "7").result(timeout=10) == 7 ) 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 97b04f85418a..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 @@ -407,7 +407,7 @@ def submit(fn, data): executor = mock.Mock() executor.submit.side_effect = submit patcher = mock.patch.object( - reads_resumption_strategy, "_get_crc32c_executor", return_value=executor + reads_resumption_strategy, "_crc32c_executor", executor ) response = _storage_v2.BidiReadObjectResponse( object_data_ranges=[ @@ -504,7 +504,7 @@ async def hang(submitted, future): task = asyncio.create_task( mrd.download_ranges([(0, 2 * len(content), buffer)]) ) - while not reads_resumption_strategy._get_crc32c_executor().submit.called: + while not reads_resumption_strategy._crc32c_executor.submit.called: await asyncio.sleep(0) await asyncio.sleep(0) task.cancel()