Skip to content

feat(storage): multi-stream support in AsyncMultiRangeDownloader - #18537

Merged
raj-prince merged 5 commits into
googleapis:mainfrom
raj-prince:multi-stream-mrd
Nov 4, 2026
Merged

raj-prince merged 5 commits into
googleapis:mainfrom
raj-prince:multi-stream-mrd

Conversation

@raj-prince

@raj-prince raj-prince commented Oct 1, 2026 •

Copy link
Copy Markdown
Contributor

Description

This PR adds optional multi-stream support to AsyncMultiRangeDownloader (MRD), enabling downstream consumers like gcsfs to scale throughput beyond single-stream limits on high-bandwidth storage classes like RAPID.

By default, existing single-stream behavior is fully preserved with zero pooling or load-tracking overhead (no lock acquisition or pending range/byte accounting). Multi-stream is strictly opt-in via the stream_config parameter.

Summary of changes:

  • MRDStreamConfig: Added a configuration dataclass to bundle multi-stream settings (min_streams, max_streams, target_io_depth, target_bytes).
  • Strictly opt-in (stream_config=None by default): When stream_config is omitted, the pool is completely bypassed (_pool = None). Requests execute directly on the primary stream with no byte/range counting, locking, or load calculation.
  • Dynamic stream pool (_StreamPool): When enabled, range requests are distributed across streams using a hybrid selection strategy (least-loaded across base min_streams, first-fit across dynamically scaled streams). If load exceeds target capacity, additional streams are opened in the background in proportion to overall load without blocking active reads on already-open streams.
  • Routing state reuse: Newly opened streams reuse the object's existing read_handle and routing_token (with redirect-aware retry handling) to connect directly to the same backend without repeating lookup overhead.
  • Clean lifecycle management: Background scale-up tasks are tracked and cancelled when the pool or MRD is closed.

Benchmark Results (20 GiB download on RAPID)

Streams (min_streams = max_streams) I/O Size I/O Depth Elapsed Time Throughput (MiB/s) Throughput (GB/s)
1 8 MB 8 12.93 s 1,584.47 MiB/s 1.66 GB/s
2 8 MB 8 8.73 s 2,345.34 MiB/s 2.46 GB/s
4 8 MB 8 8.38 s 2,444.65 MiB/s 2.56 GB/s
8 8 MB 8 9.16 s 2,236.63 MiB/s 2.35 GB/s

test_reads.py execution

All 6 benchmark cases in tests/perf/microbenchmarks/time_based/reads/test_reads.py executed and PASSED (total elapsed time: 3m 41.69s):

Benchmark Case Bucket Type Pattern Processes Coroutines File Size Chunk Size Ranges Avg Throughput (MiB/s) Status
seq_regional_96p_1c_10240MiB_64KiB_1ranges Regional (JSON) Sequential 96 1 10 GiB 64 KiB 1 38.91 MiB/s PASSED
seq_rapid_96p_1c_10240MiB_64KiB_1ranges RAPID (MRD) Sequential 96 1 10 GiB 64 KiB 1 3,953.63 MiB/s PASSED
rand_regional_1p_1c_10240MiB_64KiB_1ranges Regional (JSON) Random 1 1 10 GiB 64 KiB 1 0.27 MiB/s PASSED
rand_regional_1p_16c_10240MiB_64KiB_1ranges Regional (JSON) Random 1 16 10 GiB 64 KiB 1 0.26 MiB/s PASSED
rand_rapid_1p_1c_10240MiB_64KiB_1ranges RAPID (MRD) Random 1 1 10 GiB 64 KiB 1 67.68 MiB/s PASSED
rand_rapid_1p_16c_10240MiB_64KiB_1ranges RAPID (MRD) Random 1 16 10 GiB 64 KiB 1 366.71 MiB/s PASSED

@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 multi-stream scaling and unfinalized object support in AsyncMultiRangeDownloader (MRD) by implementing an internal stream pool (_StreamPool) and stream manager (_ManagedStream). Feedback on these changes focuses on optimizing concurrency and resource management. Key recommendations include releasing the lock before awaiting the stream factory in acquire_stream to prevent blocking concurrent tasks, tracking background tasks to avoid garbage collection, initializing and updating self._routing_token to prevent AttributeError and preserve routing state, and properly handling pool closure to avoid resource leaks.

@raj-prince raj-prince changed the title feat(mrd): multi-stream support in MRD feat(storage): multi-stream support in MRD Oct 1, 2026
@raj-prince raj-prince changed the title feat(storage): multi-stream support in MRD feat(storage): multi-stream support in AsyncMultiRangeDownloader Oct 1, 2026
@raj-prince
raj-prince marked this pull request as ready for review October 1, 2026 19:09
@raj-prince
raj-prince requested a review from a team as a code owner October 1, 2026 19:09
@chandra-siri

Copy link
Copy Markdown
Contributor

/gcbrun(b183f7f)

@chandra-siri chandra-siri self-assigned this Oct 12, 2026
@chandra-siri

Copy link
Copy Markdown
Contributor

@raj-prince

some comments on descriptions

  • scale throughput beyond single-connection TCP/HTTP2 limits

Just mention single-stream (because each stream maps to 1 grpc channel , the underlying TCP connections per channel depends on the gRPC's xDS for direct path)

  • existing single-connection behavior

change this to single-stream

  • what do you mean by "tracking overhead" ?

  • "dispatched immediately to the least-loaded active stream."

Aren't we going with the hybrid approach ?

@raj-prince

Copy link
Copy Markdown
Contributor Author

/gcbrun

@chandra-siri

Copy link
Copy Markdown
Contributor

@raj-prince Please address my comment - b/575739236#comment2

@raj-prince

raj-prince commented Oct 29, 2026 •

Copy link
Copy Markdown
Contributor Author

@raj-prince Please address my comment - b/575739236#comment2

Great catch, I was not aware - opening stream with read-handle comes with default metadata. Fixed the override scenario.

Enabled all existing MRD systems test in test_rapid.py with multi-stream-configs.

@raj-prince

Copy link
Copy Markdown
Contributor Author

/gcbrun

1 similar comment
@raj-prince

Copy link
Copy Markdown
Contributor Author

/gcbrun

@raj-prince

Copy link
Copy Markdown
Contributor Author

/gcbrun

@raj-prince

Copy link
Copy Markdown
Contributor Author

/gcbrun

@raj-prince
raj-prince merged commit deb7f1e into googleapis:main Nov 4, 2026
53 of 54 checks passed
chandra-siri added a commit to chandra-siri/google-cloud-python that referenced this pull request Nov 4, 2026
parthea added a commit that referenced this pull request Nov 4, 2026
…em tests (#18980)

## Summary

Adds Rapid system tests for the opt-in zero-copy
`BidiReadObjectResponse` parser
(`GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ=1`) from #18627, as requested in
its review:

- `test_mrd_zero_copy_bidi_read`, a new test for behavior specific to
the parser: chunk types, buffer destinations and short reads.
- Every other test in `tests/system/test_rapid.py` that calls
`download_ranges` now runs with the parser off and on, and with
`accepts_memoryview` set to `False` and `True`.

They run with the existing Rapid system tests (and in manual RCU runs);
no CI changes.

## `test_mrd_zero_copy_bidi_read`

The test writes one 5 MiB appendable object, finalizes it, and reads it
back over a single `AsyncMultiRangeDownloader` stream with the parser
enabled:

- The stream uses the zero-copy `BidiReadObject` RPC, and the downloader
still gets the object size from the open response.
- A full-object read into a default `write()` destination receives only
`bytes` chunks.
- With `download_ranges(..., accepts_memoryview=True)`, chunks are
`memoryview`s when the installed `google-crc32c` accepts `memoryview`
and `bytes` otherwise (for example `google-crc32c` 1.9.0). The chunks
kept after the call still match the object data.
- One `download_ranges` call fills a `bytearray` and a writable
`memoryview` slice of a larger buffer, 2.5 MiB each. Guard bytes around
the slice catch out-of-range writes, and both buffers are cleared
afterwards, which fails if the SDK kept a view of either one.
- A read clamped at EOF into a buffer raises `BufferError` and leaves
the stream open. The buffer is cleared while the exception is still
held, which checks that the SDK released its view itself, and the next
read on the same downloader succeeds.

## Existing read tests

Two new fixtures: `zero_copy_bidi_read` sets
`GCS_PYTHON_CLIENT_ZERO_COPY_BIDI_READ` to `0` or `1` through
`monkeypatch`, so nothing leaks into later tests on the shared
session-scoped gRPC client, and `accepts_memoryview` is passed to every
`download_ranges` call. The 14 tests that call `download_ranges` take
both, on top of their `multi_stream_config` parameter from #18537, so
each runs in eight combinations: those tests go from 62 to 248 cases and
the file from 73 to 260. `test_mrd_zero_copy_bidi_read` keeps its own
settings. In the `python-sdk-zb-system-test` check, pytest takes 156 s
instead of 47 s on `main` (252 passed; the 8 skipped cases are the
read-handle test, which was already skipped).

## Testing

```
cd packages/google-cloud-storage
ruff check --select I tests/system/test_rapid.py
ruff format --check --target-version=py310 tests/system/test_rapid.py
flake8 tests/system/test_rapid.py
```

I ran the new test locally against an in-process `grpc.aio` server
implementing `BidiWriteObject` and `BidiReadObject` (responses of at
most 2 MiB, with and without an empty final response on a 2 MiB
boundary), once with `google-crc32c` 1.9.0 and once with a
`google-crc32c` build from `main` that accepts `memoryview`. Against the
same server I ran the 10 parametrized tests it can serve in all their
combinations, including both `multi_stream_config` values, checking that
each stream used the expected parser and each call passed the expected
`accepts_memoryview`; the cross-region, write-handle and
append-to-existing-object tests need a real bucket, and the read-handle
test is skipped. The `python-sdk-zb-system-test` check runs everything
against a real Rapid bucket.

---------

Co-authored-by: Anthonios Partheniou <partheniou@google.com>
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.

2 participants