feat: implement AsyncMultiRangeDownloader with multiplexed bidi-gRPC stream support by zhixiangli · Pull Request #16528 · googleapis/google-cloud-python

zhixiangli · 2026-04-02T07:22:53Z

This PR implements AsyncMultiRangeDownloader with a new _StreamMultiplexer, enabling multiple concurrent range downloads to share a single bidirectional gRPC stream.

Before vs. After

Feature	Before	After (This PR)
Concurrency	Sequential or multiple connections	Concurrent over one connection
Overhead	High (multiple gRPC streams)	Low (multiplexed single stream)
Reliability	Per-stream retry logic	Unified generation-gated reopening

How it works

The system uses a background _StreamMultiplexer to manage the shared bidirectional stream:

Requests: Concurrent tasks send range requests (BidiReadObjectRequest) directly to the shared stream.
Multiplexing: A background Recv Loop listens for all responses. It uses the read_id in each response to route data to the correct task-specific asyncio.Queue.
Error Handling: If the stream breaks, a generation-gated lock ensures the stream is reopened only once. All active tasks receive a _StreamError and automatically retry using the new stream generation.

Key Changes:

_StreamMultiplexer: Background receiver loop for routing responses.
Generation-Gated Reopening: Coordinates stream recovery across concurrent tasks.
AsyncMultiRangeDownloader Integration: Full support for concurrent download_ranges calls.

gemini-code-assist

Code Review

This pull request introduces a _StreamMultiplexer to handle concurrent download tasks over a single bidirectional gRPC stream, replacing the previous locking mechanism in AsyncMultiRangeDownloader. The multiplexer routes responses to per-task asyncio queues based on read_id, allowing for better resource utilization. The lock parameter in download_ranges is now deprecated. Feedback focuses on improving the multiplexer's reliability and performance, specifically by using asyncio.gather to prevent head-of-line blocking during response broadcasting, ensuring the background receive loop terminates when no tasks are active, and adding error logging for observability.

packages/google-cloud-storage/google/cloud/storage/asyncio/_stream_multiplexer.py

chandra-siri · 2026-04-06T13:56:28Z

packages/google-cloud-storage/tests/system/test_zonal.py

+
+
+def test_mrd_concurrent_download(
+    storage_client, blobs_to_delete, event_loop, grpc_client


nit: can you use grpc_client_direct , although both works but we're generally interested in direct path

chandra-siri · 2026-04-06T14:39:22Z

packages/google-cloud-storage/google/cloud/storage/asyncio/_stream_multiplexer.py

+logger = logging.getLogger(__name__)
+
+_DEFAULT_QUEUE_MAX_SIZE = 100
+_DEFAULT_PUT_TIMEOUT = 20.0


nit: consider _DEFAULT_PUT_TIMEOUT_SECONDS

chandra-siri · 2026-04-06T14:41:06Z

packages/google-cloud-storage/google/cloud/storage/asyncio/_stream_multiplexer.py

+    keyed by read_id. Coordinates stream reopening via generation-gated
+    locking.
+
+    A slow consumer on one task will slow down the entire shared connection


does it makes sense @zhixiangli ?

chandra-siri · 2026-04-06T16:17:57Z

packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py

+                        if isinstance(item, _StreamEnd):
+                            if pending_read_ids:
+                                last_broken_generation = my_generation
+                                raise exceptions.ServiceUnavailable(


This may not be required. See my comment on _stream_multiplexer.py

chandra-siri · 2026-04-06T16:46:54Z

packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py


-                    pending_read_ids = {r.read_id for r in requests}
+                    my_generation = self._multiplexer.stream_generation


nit: how about renaming it to stream_generation ?

chandra-siri · 2026-04-07T11:56:31Z

packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py

        self, mock_cls_async_read_object_stream, mock_random_int
    ):
-        # Arrange


don't remove this. Please verify the same for all other tests.

chandra-siri · 2026-04-07T12:06:54Z

packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py

@@ -419,6 +422,8 @@ async def test_create_mrd_with_generation_number(
        mock_stream.generation_number = _TEST_GENERATION_NUMBER
        mock_stream.persisted_size = _TEST_OBJECT_SIZE
        mock_stream.read_handle = _TEST_READ_HANDLE
+        mock_stream.is_stream_open = True
+        mock_stream.recv = AsyncMock(side_effect=asyncio.Event().wait)


adding these lines have NO effect.

chandra-siri · 2026-04-07T12:08:27Z

packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py

Although you haven't changed code related to this line, I would appreciate if you can add assert on read_handle and object_size

chandra-siri · 2026-04-07T12:11:47Z

packages/google-cloud-storage/tests/unit/asyncio/test_stream_multiplexer.py

@@ -0,0 +1,503 @@
+# Copyright 2025 Google LLC


chandra-siri · 2026-04-07T12:14:43Z

packages/google-cloud-storage/tests/unit/asyncio/test_stream_multiplexer.py

+        assert isinstance(sentinel, _StreamEnd)
+
+
+class TestStreamMultiplexerInit:


please follow - go/unit-testing-practices?polyglot=python#structure for all tests in this file.

zhixiangli requested review from a team as code owners April 2, 2026 07:22

gemini-code-assist bot reviewed Apr 2, 2026

View reviewed changes

zhixiangli changed the title ~~feat: implement AsyncMultiRangeDownloader and integrate _StreamMultiplexer~~ feat: implement AsyncMultiRangeDownloader with multiplexed bidi-gRPC stream support Apr 2, 2026

zhixiangli force-pushed the zhixiangli/multiplexing-downloader branch from 13d5d08 to 774d691 Compare April 2, 2026 08:22

zhixiangli added 3 commits April 2, 2026 09:45

feat: add _StreamMultiplexer for asyncio bidi-gRPC streams

cf7c949

feat: integrate _StreamMultiplexer into AsyncMultiRangeDownloader

cb4e640

test: add system tests for concurrent AsyncMultiRangeDownloader

aed8682

zhixiangli force-pushed the zhixiangli/multiplexing-downloader branch from acfab40 to aed8682 Compare April 2, 2026 09:47

chandra-siri requested changes Apr 7, 2026

View reviewed changes

chandra-siri self-assigned this Apr 7, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: implement AsyncMultiRangeDownloader with multiplexed bidi-gRPC stream support#16528

feat: implement AsyncMultiRangeDownloader with multiplexed bidi-gRPC stream support#16528
zhixiangli wants to merge 3 commits intogoogleapis:mainfrom
zhixiangli:zhixiangli/multiplexing-downloader

zhixiangli commented Apr 2, 2026 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chandra-siri Apr 6, 2026

Uh oh!

chandra-siri Apr 6, 2026

Uh oh!

chandra-siri Apr 6, 2026

Uh oh!

chandra-siri Apr 6, 2026

Uh oh!

chandra-siri Apr 6, 2026

Uh oh!

chandra-siri Apr 7, 2026

Uh oh!

chandra-siri Apr 7, 2026

Uh oh!

chandra-siri Apr 7, 2026

Uh oh!

chandra-siri Apr 7, 2026

Uh oh!

chandra-siri Apr 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		def test_mrd_concurrent_download(
		storage_client, blobs_to_delete, event_loop, grpc_client


		pending_read_ids = {r.read_id for r in requests}
		my_generation = self._multiplexer.stream_generation

		assert isinstance(sentinel, _StreamEnd)


		class TestStreamMultiplexerInit:

Conversation

zhixiangli commented Apr 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Before vs. After

How it works

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

zhixiangli commented Apr 2, 2026 •

edited

Loading