[yuxin00j]

Summary

An opt-in parser for BidiReadObjectResponse that reads the message in place over the bytes object gRPC hands to the deserializer, so _AsyncReadObjectStream passes each chunk's checksummed_data.content to the read path as a memoryview of the received message instead of the two copies the generated parser makes on the event-loop thread (the upb parse, then .content into a new bytes; 257 us per 2 MiB message, 5 us with this parser). Off by default. Requires a google-crc32c release that accepts memoryview (#18626); with an older release the generated parser is used automatically.

Changes

• asyncio/_zero_copy_bidi_read.py (new): deserialize(buf) walks the wire format of data responses (object_data_ranges, read_handle) over memoryview(buf) and returns stand-ins with the attributes the read path uses. A response carrying metadata, bytes that do not parse cleanly, or a buffer that is not an immutable bytes object is handed to BidiReadObjectResponse.deserialize, which returns the real message or raises its usual DecodeError. Varint, int64 and tag handling match upb.
• Opt-in per downloader: AsyncMultiRangeDownloader.create_mrd(..., zero_copy=None). True asks for the parser, False keeps the generated one, None (default) defers to GOOGLE_CLOUD_STORAGE_ZERO_COPY_BIDI_READ (1/true/yes/on). Either way the parser is used only when the installed google-crc32c can hash a memoryview (is_supported() probe). The flag reaches every stream the downloader opens, including on resumption.
• wrapped_rpc(transport, zero_copy, client_info): the transport's BidiReadObject stub with the deserializer swapped in, wrapped like the generated _wrapped_methods entry (same error mapping and x-goog-api-client header, built from the client's ClientInfo), cached per transport.
• download_ranges contract: with zero_copy on, writeable_buffer.write() receives a read-only memoryview of the received message, which is immutable and stays alive as long as the view does; a buffer may keep the view and copy from it after write() returns, and a buffer that needs bytes must call bytes(data). io.BytesIO and binary files work unchanged.

Tests

tests/unit/asyncio/test_zero_copy_bidi_read.py: field-by-field comparison with the generated parser on encoder-built and proto-plus-serialized messages (including unknown fields of every wire type at each nesting level), 16 malformed inputs on which both parsers raise DecodeError, fallback cases (metadata, non-bytes buffers), is_enabled() precedence (argument over environment, both gated by the probe), wrapped_rpc caching, and an in-process grpc.aio server driving _AsyncReadObjectStream with the parser on and off and a custom ClientInfo (user agent and gapic version reach the server). test_async_multi_range_downloader.py / test_async_read_object_stream.py cover the argument plumbing. 100% line and branch coverage of the new module, on the upb and the pure-Python protobuf backends.

nox -s unit lint cover

Benchmark

c4-standard-192 (us-central1-c), Rapid bucket in the same zone, 16 concurrent streams per process, 4 x 16 MiB per download_ranges call into preallocated buffers, CRC32C verification on, google-crc32c from the companion PR; median of 3 interleaved runs, aggregate GB/s:

Processes	main	this PR
1	1.55	2.62
8	9.76	11.29

[gemini-code-assist]

Code Review

This pull request introduces a zero-copy deserializer for BidiReadObjectResponse messages in _zero_copy_bidi_read.py to optimize performance by avoiding extra data copies on the event-loop thread. It integrates this parser into the async read path (async_multi_range_downloader.py and async_read_object_stream.py) with a configurable zero_copy option, and adds comprehensive unit tests. Feedback on the changes highlights a potential IndexError in the _varint decoding loop when parsing truncated inputs, suggesting a robust bounds check at the start of each iteration.

[raj-prince]

/gcbrun

[zhixiangli]

Thanks for this. The throughput gains are compelling and the parity tests against the generated parser are thorough. Inline comments below; two PR-level points:

1. The title and description are out of date relative to the code (since d9e019e and a297106):

• Default on/off: the title says "opt-in" and the description says "Off by default", but is_enabled() turns the parser on by default.
• Env var: the description documents GOOGLE_CLOUD_STORAGE_ZERO_COPY_BIDI_READ (1 enables it). The code reads GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ (0/false/no/off disables it).
• zero_copy= was removed but is still accepted: AsyncMultiRangeDownloader.__init__ swallows **kwargs, so AsyncMultiRangeDownloader(client, "b", "o", zero_copy=False) runs without error and has no effect. Consider rejecting unknown kwargs, or at least warning on zero_copy.
• Old wrapped_rpc signature: the description shows wrapped_rpc(transport, zero_copy, client_info); the code takes (transport, client_info).
• Missing from the description: accepts_memoryview and the writable buffer destinations. write() now receives a memoryview only with accepts_memoryview=True.

2. On-by-default changes what callers get back. _AsyncReadObjectStream.recv() and DataCorruption.response now return _ZeroCopyBidiReadObjectResponse instead of the proto-plus message for every async reader. If that's intended, it would be good to call it out in the description and release notes, together with the env-var kill switch.

[zhixiangli]

High: server overflow into a buffer destination writes past bytes_to_read, or raises the wrong error.

_ReadResumptionStrategy.update_state_from_response calls user_buffer.write(data) before the byte-count check, and only checks on range_end. That was harmless for BytesIO, but _BufferWriter writes into the caller's memory, and self._view covers the whole buffer rather than [:length]. With length=10 and a 12-byte chunk:

• bytearray(10): ValueError: memoryview assignment: lvalue and rvalue have different structures instead of DataCorruption. download_ranges only closes the stream on DataCorruption, so it stays open.
• bytearray(b"B" * 20): becomes b"AAAAAAAAAAAABBBBBBBB" (bytes 10–11 overwritten). DataCorruption is raised only at range_end; a non-final chunk raises nothing.

Suggested fix: in reads_resumption_strategy.py, check before writing:

chunk_size = len(data)
if (
    read_state.initial_length != 0
    and read_state.bytes_written + chunk_size > 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 + chunk_size}",
    )
read_state.user_buffer.write(data)

Also limit the view to [:length] as a second guard (see the next comment).

Test: a 12-byte response for (0, 10, bytearray(10)) and for (0, 10, bytearray(b"B" * 20)) should raise DataCorruption, close the stream, and leave buf[10:] == b"B" * 10.

[zhixiangli]

The view on the caller's buffer is never released.

Neither view nor self._view is ever .release()d. When download_ranges raises, the traceback keeps its frame (download_states -> _BufferWriter -> _view) alive, so the caller can't resize their buffer while handling the error:

buf = bytearray(4)
try:
    await mrd.download_ranges([(0, 4, buf), (0, 0, bytearray(4))])  # 2nd range invalid
except ValueError:
    buf.clear()  # BufferError: Existing exports of data: object cannot be re-sized

The same thing happens inside except DataCorruption:.

Suggested change

	view = memoryview(buffer)
	if view.readonly or not view.c_contiguous:
	view.release()
	raise ValueError(
	"Buffer destination must be a writable, C-contiguous buffer"
	)
	nbytes = view.nbytes
	if nbytes < length:
	view.release()
	raise ValueError(
	f"Buffer destination ({nbytes} bytes) is smaller than "
	f"bytes_to_read ({length})"
	)
	self._view = view.cast("B")
	self._pos = 0
	view = memoryview(buffer)
	try:
	if view.readonly or not view.c_contiguous:
	raise ValueError(
	"Buffer destination must be a writable, C-contiguous buffer"
	)
	nbytes = view.nbytes
	if nbytes < length:
	raise ValueError(
	f"Buffer destination ({nbytes} bytes) is smaller than "
	f"bytes_to_read ({length})"
	)
	self._view = view.cast("B")[:length]
	finally:
	view.release()
	self._pos = 0

Then add def close(self): self._view.release(). In download_ranges, move the read_ranges loop inside the existing try: and close every _BufferWriter in its finally:.

Test: buf.clear() succeeds inside except ValueError: and inside except DataCorruption:.

[zhixiangli]

The x-goog-api-client header now depends on the zero-copy setting.

AsyncGrpcClient builds the transport as transport_cls(channel=channel) without client_info, so the generated _wrapped_methods entry (and every other RPC on the client) sends the default header, while this path sends the caller's header. Measured:

client_info passed to AsyncGrpcClient	zero-copy off	zero-copy on
api_core.client_info.ClientInfo("gcsfs/1")	gl-python/… gax/… gapic/3.16.0 pb/7.36.2	same
gapic_v1.client_info.ClientInfo("gcsfs/1")	gl-python/… gax/… gapic/3.16.0 pb/7.36.2	gcsfs/1 gcloud-python/3.16.0 … gccl/3.16.0

Two related points:

• Cache ignores client_info: _rpcs is keyed by transport only, so wrapped_rpc(t, None) followed by wrapped_rpc(t, other) returns the first callable and ignores other.
• No kind=: this calls gapic_v1.method_async.wrap_method directly, whereas the transport's _wrap_method also passes kind= when api-core supports it.

Suggestion: pass the client info to the transport in AsyncGrpcClient (converting an api_core ClientInfo to gapic_v1.client_info.ClientInfo), and build this callable with transport._wrap_method(...) using the transport's own client info. Both parsers then send the same metadata, and the cache key and the client_info= plumbing go away.

Test: an end-to-end test parametrized over zero-copy on/off that asserts the same x-goog-api-client value reaches the server.

[zhixiangli]

Length prefixes longer than 5 bytes are accepted, unlike the generated parser.

_tag rejects tags longer than 5 bytes, but _length_prefixed uses _varint's 10-byte, low-64-bit result unchecked. No exception is raised, so the fallback to the generated parser never runs:

input	generated (upb)	zero-copy
b"\x32\x80\x80\x80\x80\x80\x00" (6-byte length 0)	DecodeError	parses
b"\x32" + b"\x80"*9 + b"\x02" (length 2**64, masked to 0)	DecodeError	parses
b"\x32" + b"\x80"*9 + b"\x00" (10-byte length 0)	DecodeError	parses

Suggested change

	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 _length_prefixed(buf, i, end):
	"""Return ``(start, stop)`` of the length-delimited field body at ``i``.
	Like upb, the length must be encoded in at most five bytes.
	"""
	start = i
	n, i = _varint(buf, i, end)
	if i - start > 5:
	raise _Malformed("field length longer than 5 bytes")
	stop = i + n
	if stop > end:
	raise _Malformed("field runs past its enclosing message")
	return i, stop

Test: add the three inputs to test_deserialize_malformed_input_raises_like_generated_parser.

[zhixiangli]

Also probe Checksum.update, which whole-object reads use.

Whole-object reads ((0, 0, …) / (0, persisted_size, …)) also call rolling_checksum.update(data), but the probe only checks value. The companion crc32c change updates both entry points together, so real installs agree. The test fixtures here, however, patch only google_crc32c.value, so a whole-object read through this path fails with TypeError: argument 2 must be read-only bytes-like object, not memoryview. The new MRD tests don't catch this because _TEST_OBJECT_SIZE is 1 MiB and they read 18 bytes, so is_full_object_read is always False.

Suggested change

	def is_supported():
	"""Whether the installed ``google_crc32c`` can hash ``memoryview`` inputs.
	When true, :func:`deserialize` exposes chunk payloads as ``memoryview``
	slices of the gRPC wire buffer; when false, it slices ``bytes`` instead so
	checksum verification remains compatible with older ``google_crc32c``
	releases while still avoiding the intermediate upb arena copy.
	"""
	try:
	google_crc32c.value(memoryview(b"\0"))
	except Exception:
	return False
	return True
	def is_supported():
	"""Whether the installed ``google_crc32c`` can hash ``memoryview`` inputs.
	When true, :func:`deserialize` exposes chunk payloads as ``memoryview``
	slices of the gRPC wire buffer; when false, it slices ``bytes`` instead so
	checksum verification remains compatible with older ``google_crc32c``
	releases while still avoiding the intermediate upb arena copy. Both
	``value`` and ``Checksum.update`` (used for full-object reads) must
	accept a ``memoryview``.
	"""
	probe = memoryview(b"\0")
	try:
	google_crc32c.value(probe)
	google_crc32c.Checksum().update(probe)
	except Exception:
	return False
	return True

With this change, the tests that patch only google_crc32c.value (the autouse fixture in test_zero_copy_bidi_read.py and the two new MRD tests) need to patch Checksum.update too. With that, all tests in the three files pass. Please also add (0, 0, sink) and (0, persisted_size, bytearray(persisted_size)) cases with enable_checksum=True.

[zhixiangli]

Docstring nit: "The SDK drops its reference after write() returns" isn't quite accurate. The parsed response still holds checksummed_data.content, and _StreamMultiplexer._recv_loop / _BidiStreamRetryManager.execute keep the latest response alive while awaiting the next message. Memory use is the same as with the generated parser, so a wording fix is enough:

Suggested change

	:type accepts_memoryview: bool
	:param accepts_memoryview: (Optional) If True, ``write()`` destinations
	may receive each chunk as a read-only :class:`memoryview` of the
	received gRPC message instead of ``bytes``. The SDK drops its
	reference after ``write()`` returns; the message (the chunk plus
	a few bytes of framing) is freed when the destination drops the
	view. Destinations must still accept ``bytes``, which the generated
	parser produces. No effect on buffer destinations. Defaults to
	False.
	:type accepts_memoryview: bool
	:param accepts_memoryview: (Optional) If True, ``write()`` destinations
	may receive each chunk as a read-only :class:`memoryview` of the
	received gRPC message instead of ``bytes``. The view keeps that
	message (the chunk plus a few bytes of framing) alive while the
	destination or the SDK's most recently processed response
	references it; use ``bytes(data)`` to keep the data without the
	message. Destinations must still accept ``bytes``, which the
	generated parser produces. No effect on buffer destinations.
	Defaults to False.

[zhixiangli]

nit: the signature still says read_ranges: List[Tuple[int, int, BytesIO]] (and :type read_ranges: says "BytesIO"), so type checkers will reject the bytearray / memoryview destinations documented here. Suggest Union[IO[bytes], bytearray, memoryview] for the third element, and lock: Optional[asyncio.Lock] = None. _AsyncReadObjectStream.__init__ also gained client_info=None without a type hint or :param: entry; that may become moot with the client-info comment.

[zhixiangli]

nit: _ZeroCopyBidiReadObjectResponse has _pb -> self, but _ObjectRangeData and _ChecksummedData don't. With the generated parser, response.object_data_ranges[0]._pb works; here it raises AttributeError. Similarly, checksummed_data._pb.HasField("_crc32c") (the proto3 optional synthetic oneof) returns a bool on the real message but raises ValueError here. Nothing in this repo relies on either, but these objects are exposed via recv() and DataCorruption.response.

[zhixiangli]

Test gaps for buffer destinations. Resuming into a buffer after a retry works today (the retry asks for offset 10, length 8, and the buffer ends up correct), but nothing locks it in:

• Retry into a buffer: the first stream writes 10 of 18 bytes into bytearray(18) and then raises exceptions.Aborted; the second stream sends the rest at read_offset=10. Assert bytes(raw) == data and that the second _AsyncReadObjectStream(...) call (from _create_stream_factory) received client_info=mock_client._client_info. Today only the open() call is asserted.
• Checksum mismatch: a CRC mismatch on a zero-copy response into a bytearray raises DataCorruption, leaves the buffer untouched, and closes the stream.