Skip to content

feat(storage): opt-in zero-copy BidiReadObjectResponse parsing for async reads - #18627

Open
yuxin00j wants to merge 11 commits into
googleapis:mainfrom
yuxin00j:storage-zero-copy-read
Open

yuxin00j wants to merge 11 commits into
googleapis:mainfrom
yuxin00j:storage-zero-copy-read

Conversation

@yuxin00j

@yuxin00j yuxin00j commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor

Summary

An opt-in parser for BidiReadObjectResponse data messages that reads each message in place over the bytes object gRPC hands to the deserializer, instead of copying every chunk twice on the event-loop thread (upb arena decode, then .content into a new bytes; 257 us per 2 MiB message vs. 5 us with this parser). Off by default; set GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=1 (or true/yes/on) to enable it.

  • Existing callers passing a .write() destination (io.BytesIO, binary files, custom streams) keep receiving bytes. A destination receives the chunk as a memoryview of the received message only when both the parser is enabled and the caller passes download_ranges(..., accepts_memoryview=True); either switch alone leaves it bytes.
  • Callers passing a writable buffer (bytearray or a writable memoryview slice) as the destination get each chunk copied straight into it, without an intermediate bytes.
  • Works with any google-crc32c: when the installed release accepts memoryview (fix(crc32c): accept any contiguous bytes-like object in value() and extend() #18626), chunks are views of the wire buffer; with google-crc32c <= 1.9.0 they are bytes slices, which still saves the upb arena copy.
  • When opted in via GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=1, _AsyncReadObjectStream.recv() and DataCorruption.response carry the duck-typed _ZeroCopyBidiReadObjectResponse on data messages instead of the generated proto-plus wrapper.

Changes

  • asyncio/_zero_copy_bidi_read.py (new):
    • deserialize(buf) walks the top-level wire tags of data responses (object_data_ranges, read_handle), slices checksummed_data.content as memoryview(buf)[start:i] (or buf[start:i] when google-crc32c only takes bytes), and parses ReadRange and BidiReadHandle sub-messages with the generated parser. Varint, int64, tag, and length-prefix handling match upb.
    • A response carrying metadata, input that raises (_Malformed, IndexError, DecodeError), 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.
    • is_enabled() reads GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ; is_supported() probes both google_crc32c.value(memoryview(...)) and google_crc32c.Checksum().update(memoryview(...)) once per process.
    • wrapped_rpc(transport) builds and caches the transport's BidiReadObject callable with deserialize swapped in and the same error mapping and x-goog-api-client header as _wrapped_methods.
  • asyncio/async_read_object_stream.py: _AsyncReadObjectStream uses wrapped_rpc when enabled, the generated method otherwise.
  • asyncio/async_multi_range_downloader.py:
    • download_ranges(..., accepts_memoryview=False): .write() destinations are wrapped in _BytesWriter, which converts memoryview chunks to bytes, unless the caller opts in.
    • Writable, C-contiguous buffer-protocol destinations (bytes_to_read > 0) are wrapped in _BufferWriter, which copies each chunk into the first bytes_to_read bytes of the buffer and releases its view when download_ranges exits.
  • asyncio/retry/reads_resumption_strategy.py: update_state_from_response checks bytes_written + len(data) > initial_length before writing to user_buffer, so server over-delivery raises DataCorruption and closes the stream before touching the caller's buffer.

Tests

tests/unit/asyncio/test_zero_copy_bidi_read.py, test_async_read_object_stream.py, test_async_multi_range_downloader.py: parametrized equivalence against the generated parser (_pb / HasField, unknown fields, split sub-messages, over-long varints and length prefixes), malformed and fallback inputs, the environment switch and cached value / Checksum.update probe, _BytesWriter and _BufferWriter validation / retry / checksum-mismatch / full-object / view-release paths, an unpatched google-crc32c check, and an in-process grpc.aio server header parity test. 100% line and branch coverage of _zero_copy_bidi_read.py under both upb and PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python.

nox -s unit lint cover

Benchmark

GCE c4-standard-192, Rapid bucket in the same zone, 8 processes x 16 concurrent AsyncMultiRangeDownloader streams, 4 x 16 MiB ranges per download_ranges() call into preallocated buffers (accepts_memoryview=True), google-crc32c with #18626; 8 interleaved rounds per arm, aggregate GB/s:

Arm Mean Median Min - Max
main 9.47 9.48 8.07 - 10.51
this PR, parser enabled 11.24 (+18.6%) 11.07 (+16.7%) 9.62 - 13.73

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.
…sages 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.
…stants

Explain the opt-in switch's accepted spellings, the four protobuf wire types, the 64-bit varint mask and the 32-bit tag limit.
Until now the only way to opt a downloader into the zero-copy
BidiReadObjectResponse parser was the GOOGLE_CLOUD_STORAGE_ZERO_COPY_BIDI_READ
environment variable, which _AsyncReadObjectStream re-reads on every open.
A library such as gcsfs that wants memoryview chunks therefore had to set a
process-wide variable, which also switches every other downloader in the
process (and in its children) to memoryview payloads on their next reopen.

create_mrd() and the constructor now take zero_copy: True asks for the parser,
False keeps the generated one, None (the default) defers to the environment
variable as before. The flag is passed to every stream the downloader opens,
including the ones created on resumption. _zero_copy_bidi_read.is_supported()
is now only the google-crc32c memoryview probe; the new is_enabled(zero_copy)
combines it with the argument or the environment, and wrapped_rpc() takes the
flag.

download_ranges() documents what a buffer receives: bytes, or with zero_copy
a read-only memoryview of the received message, which is immutable and stays
alive for as long as the view does, so a buffer may keep the view and copy
from it after write() returns.

Test plan: from packages/google-cloud-storage, with the upb and the pure-Python
protobuf backends:

    PYTHONPATH=$PWD python -m pytest tests/unit/asyncio/test_zero_copy_bidi_read.py \
        tests/unit/asyncio/test_async_read_object_stream.py \
        tests/unit/asyncio/test_async_multi_range_downloader.py
    PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python PYTHONPATH=$PWD python -m pytest \
        tests/unit/asyncio/test_zero_copy_bidi_read.py

102 and 55 tests pass; new tests cover is_enabled() precedence (argument over
environment, both gated by is_supported()), wrapped_rpc() with an explicit
flag, and the flag reaching _AsyncReadObjectStream from create_mrd().

Authored with an AI assistant.
…he zero-copy parser

wrapped_rpc built its BidiReadObject callable with DEFAULT_CLIENT_INFO, while
the generated method uses the ClientInfo the transport was created with. A
caller passing its own ClientInfo therefore sent a different x-goog-api-client
header on BidiReadObject only. AsyncGrpcClient now remembers the ClientInfo it
configures, AsyncMultiRangeDownloader hands it to each stream, and wrapped_rpc
uses it (the transport default when none is given).

deserialize() aliases chunk payloads to the buffer gRPC hands it, which is only
safe for an immutable bytes object. Anything else now goes to the generated
parser instead of being aliased.

Test plan: from packages/google-cloud-storage, upb and pure-Python protobuf:

    PYTHONPATH=$PWD python -m pytest tests/unit/asyncio/test_zero_copy_bidi_read.py \
        tests/unit/asyncio/test_async_read_object_stream.py \
        tests/unit/asyncio/test_async_multi_range_downloader.py \
        tests/unit/asyncio/test_async_grpc_client.py
    PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python PYTHONPATH=$PWD python -m pytest \
        tests/unit/asyncio/test_zero_copy_bidi_read.py

126 and 57 pass. The fake-server test now runs both parser paths with a custom
ClientInfo and asserts its user agent and gapic version reach the server; a
new test feeds bytearray and memoryview buffers to deserialize().

Authored with an AI assistant.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread packages/google-cloud-storage/tests/unit/asyncio/test_zero_copy_bidi_read.py Outdated
@raj-prince

Copy link
Copy Markdown
Contributor

/gcbrun

@yuxin00j
yuxin00j requested a review from a team as a code owner October 20, 2026 09:04
… write() destinations

download_ranges(accepts_memoryview=True) passes each chunk to a write()
destination as the read-only memoryview the zero-copy parser produced
instead of materializing bytes first. The default stays False, so existing
callers keep receiving bytes. Buffer destinations are unaffected.

Also accept a plain google.api_core.client_info.ClientInfo in wrapped_rpc:
AsyncGrpcClient takes one (gcsfs passes it), but gapic_v1's wrap_method
needs the gapic_v1 subclass for its x-goog-api-client header, so the
zero-copy stream failed to open with AttributeError. Such an object now
falls back to the transport default, which is the header main sends.
When a read ends on a chunk boundary the server's last message carries
range_end and a crc32c but no content field, so the zero-copy parser left
_ChecksummedData.content at its memoryview(b"") default even in the mode
that slices bytes for google-crc32c releases that cannot hash views. The
inline checksum then raised TypeError on stock google-crc32c 1.9.0, which
is what failed the zonal system tests (test_basic_wrd for 10 and 20 MiB
objects) since the parser became the default. Empty content is now b"" in
both modes, as the generated parser returns it.

@zhixiangli zhixiangli left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@yuxin00j

Copy link
Copy Markdown
Contributor Author

/gcbrun

@yuxin00j

Copy link
Copy Markdown
Contributor Author

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.

Updated the description. This feature is opt-in. Previous commit renamed zero_copy to accept_memoryview, so keep **kwargs unchanged.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants