Repository navigation
Conversation
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.
… GOOGLE_CLOUD_STORAGE_ZERO_COPY_BIDI_READ=1
…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.
There was a problem hiding this comment.
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.
|
/gcbrun |
…e buffer destinations
… 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
left a comment
There was a problem hiding this comment.
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(1enables it). The code readsGCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ(0/false/no/offdisables it). zero_copy=was removed but is still accepted:AsyncMultiRangeDownloader.__init__swallows**kwargs, soAsyncMultiRangeDownloader(client, "b", "o", zero_copy=False)runs without error and has no effect. Consider rejecting unknown kwargs, or at least warning onzero_copy.- Old
wrapped_rpcsignature: the description showswrapped_rpc(transport, zero_copy, client_info); the code takes(transport, client_info). - Missing from the description:
accepts_memoryviewand the writable buffer destinations.write()now receives amemoryviewonly withaccepts_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.
|
/gcbrun |
Updated the description. This feature is opt-in. Previous commit renamed zero_copy to accept_memoryview, so keep **kwargs unchanged. |
Summary
An opt-in parser for
BidiReadObjectResponsedata messages that reads each message in place over thebytesobject gRPC hands to the deserializer, instead of copying every chunk twice on the event-loop thread (upbarena decode, then.contentinto a newbytes; 257 us per 2 MiB message vs. 5 us with this parser). Off by default; setGCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=1(ortrue/yes/on) to enable it..write()destination (io.BytesIO, binary files, custom streams) keep receivingbytes. A destination receives the chunk as amemoryviewof the received message only when both the parser is enabled and the caller passesdownload_ranges(..., accepts_memoryview=True); either switch alone leaves itbytes.bytearrayor a writablememoryviewslice) as the destination get each chunk copied straight into it, without an intermediatebytes.google-crc32c: when the installed release acceptsmemoryview(fix(crc32c): accept any contiguous bytes-like object in value() and extend() #18626), chunks are views of the wire buffer; withgoogle-crc32c <= 1.9.0they arebytesslices, which still saves theupbarena copy.GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=1,_AsyncReadObjectStream.recv()andDataCorruption.responsecarry the duck-typed_ZeroCopyBidiReadObjectResponseon data messages instead of the generatedproto-pluswrapper.Changes
asyncio/_zero_copy_bidi_read.py(new):deserialize(buf)walks the top-level wire tags of data responses (object_data_ranges,read_handle), sliceschecksummed_data.contentasmemoryview(buf)[start:i](orbuf[start:i]whengoogle-crc32conly takesbytes), and parsesReadRangeandBidiReadHandlesub-messages with the generated parser. Varint, int64, tag, and length-prefix handling matchupb.metadata, input that raises(_Malformed, IndexError, DecodeError), or a buffer that is not an immutablebytesobject is handed toBidiReadObjectResponse.deserialize, which returns the real message or raises its usualDecodeError.is_enabled()readsGCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ;is_supported()probes bothgoogle_crc32c.value(memoryview(...))andgoogle_crc32c.Checksum().update(memoryview(...))once per process.wrapped_rpc(transport)builds and caches the transport'sBidiReadObjectcallable withdeserializeswapped in and the same error mapping andx-goog-api-clientheader as_wrapped_methods.asyncio/async_read_object_stream.py:_AsyncReadObjectStreamuseswrapped_rpcwhen enabled, the generated method otherwise.asyncio/async_multi_range_downloader.py:download_ranges(..., accepts_memoryview=False):.write()destinations are wrapped in_BytesWriter, which convertsmemoryviewchunks tobytes, unless the caller opts in.bytes_to_read > 0) are wrapped in_BufferWriter, which copies each chunk into the firstbytes_to_readbytes of the buffer and releases its view whendownload_rangesexits.asyncio/retry/reads_resumption_strategy.py:update_state_from_responsechecksbytes_written + len(data) > initial_lengthbefore writing touser_buffer, so server over-delivery raisesDataCorruptionand 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 cachedvalue/Checksum.updateprobe,_BytesWriterand_BufferWritervalidation / retry / checksum-mismatch / full-object / view-release paths, an unpatchedgoogle-crc32ccheck, and an in-processgrpc.aioserver header parity test. 100% line and branch coverage of_zero_copy_bidi_read.pyunder bothupbandPROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python.Benchmark
GCE
c4-standard-192, Rapid bucket in the same zone, 8 processes x 16 concurrentAsyncMultiRangeDownloaderstreams,4 x 16 MiBranges perdownload_ranges()call into preallocated buffers (accepts_memoryview=True),google-crc32cwith #18626; 8 interleaved rounds per arm, aggregate GB/s:main