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

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:
1. **Requests**: Concurrent tasks send range requests
(`BidiReadObjectRequest`) directly to the shared stream.
2. **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`.
3. **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.

Author: zhixiangli

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

@@ -0,0 +1,207 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ importannotations
+
+importasyncio
+importlogging
+fromtypingimportAwaitable, Callable, Dict, Optional, Set
+
+importgrpc
+
+fromgoogle.cloudimport_storage_v2
+fromgoogle.cloud.storage.asyncio.async_read_object_streamimport (
+_AsyncReadObjectStream,
+)
+
+logger=logging.getLogger(__name__)
+
+_DEFAULT_QUEUE_MAX_SIZE=100
+_DEFAULT_PUT_TIMEOUT_SECONDS=20.0
+
+
+class_StreamError:
+"""Wraps an error with the stream generation that produced it."""
+
+def__init__(self, exception: Exception, generation: int):
+self.exception=exception
+self.generation=generation
+
+
+class_StreamEnd:
+"""Signals the stream closed normally."""
+
+pass
+
+
+class_StreamMultiplexer:
+"""Multiplexes concurrent download tasks over a single bidi-gRPC stream.
+
+ Routes responses from a background recv loop to per-task asyncio.Queues
+ keyed by read_id. Coordinates stream reopening via generation-gated
+ locking.
+
+ A slow consumer on one task will slow down the entire shared connection
+ due to bounded queue backpressure propagating through gRPC flow control.
+ """
+
+def__init__(
+self,
+stream: _AsyncReadObjectStream,
+queue_max_size: int=_DEFAULT_QUEUE_MAX_SIZE,
+ ):
+self._stream=stream
+self._stream_generation: int=0
+self._queues: Dict[int, asyncio.Queue] = {}
+self._reopen_lock=asyncio.Lock()
+self._recv_task: Optional[asyncio.Task] =None
+self._queue_max_size=queue_max_size
+
+@property
+defstream_generation(self) ->int:
+returnself._stream_generation
+
+defregister(self, read_ids: Set[int]) ->asyncio.Queue:
+"""Register read_ids for a task and return its response queue."""
+queue=asyncio.Queue(maxsize=self._queue_max_size)
+forread_idinread_ids:
+self._queues[read_id] =queue
+returnqueue
+
+defunregister(self, read_ids: Set[int]) ->None:
+"""Remove read_ids from routing."""
+forread_idinread_ids:
+self._queues.pop(read_id, None)
+
+def_get_unique_queues(self) ->Set[asyncio.Queue]:
+returnset(self._queues.values())
+
+asyncdef_put_with_timeout(self, queue: asyncio.Queue, item) ->None:
+try:
+awaitasyncio.wait_for(
+queue.put(item), timeout=_DEFAULT_PUT_TIMEOUT_SECONDS
+ )
+exceptasyncio.TimeoutError:
+ifqueuenotinself._get_unique_queues():
+logger.debug("Dropped item for unregistered queue.")
+else:
+logger.warning(
+"Queue full for too long. Dropping item to prevent multiplexer hang."
+ )
+
+def_ensure_recv_loop(self) ->None:
+ifself._recv_taskisNoneorself._recv_task.done():
+self._recv_task=asyncio.create_task(self._recv_loop())
+
+def_stop_recv_loop(self) ->None:
+ifself._recv_taskandnotself._recv_task.done():
+self._recv_task.cancel()
+
+def_put_error_nowait(self, queue: asyncio.Queue, error: _StreamError) ->None:
+whileTrue:
+try:
+queue.put_nowait(error)
+break
+exceptasyncio.QueueFull:
+try:
+queue.get_nowait()
+exceptasyncio.QueueEmpty:
+pass
+
+asyncdef_recv_loop(self) ->None:
+try:
+whileTrue:
+response=awaitself._stream.recv()
+ifresponse==grpc.aio.EOF:
+sentinel=_StreamEnd()
+awaitasyncio.gather(
+*(
+self._put_with_timeout(queue, sentinel)
+forqueueinself._get_unique_queues()
+ )
+ )
+return
+
+ifresponse.object_data_ranges:
+queues_to_notify: Set[asyncio.Queue] =set()
+fordata_rangeinresponse.object_data_ranges:
+read_id=data_range.read_range.read_id
+queue=self._queues.get(read_id)
+ifqueue:
+queues_to_notify.add(queue)
+else:
+logger.warning(
+f"Received data for unregistered read_id: {read_id}"
+ )
+awaitasyncio.gather(
+*(
+self._put_with_timeout(queue, response)
+forqueueinqueues_to_notify
+ )
+ )
+else:
+awaitasyncio.gather(
+*(
+self._put_with_timeout(queue, response)
+forqueueinself._get_unique_queues()
+ )
+ )
+exceptasyncio.CancelledError:
+raise
+exceptExceptionase:
+logger.warning(f"Stream multiplexer recv loop failed: {e}", exc_info=True)
+error=_StreamError(e, self._stream_generation)
+forqueueinself._get_unique_queues():
+self._put_error_nowait(queue, error)
+
+asyncdefsend(self, request: _storage_v2.BidiReadObjectRequest) ->int:
+self._ensure_recv_loop()
+awaitself._stream.send(request)
+returnself._stream_generation
+
+asyncdefreopen_stream(
+self,
+broken_generation: int,
+stream_factory: Callable[[], Awaitable[_AsyncReadObjectStream]],
+ ) ->None:
+asyncwithself._reopen_lock:
+ifself._stream_generation!=broken_generation:
+return
+self._stop_recv_loop()
+ifself._recv_task:
+try:
+awaitself._recv_task
+except (asyncio.CancelledError, Exception):
+pass
+error=_StreamError(Exception("Stream reopening"), self._stream_generation)
+forqueueinself._get_unique_queues():
+self._put_error_nowait(queue, error)
+try:
+awaitself._stream.close()
+exceptException:
+pass
+self._stream=awaitstream_factory()
+self._stream_generation+=1
+self._ensure_recv_loop()
+
+asyncdefclose(self) ->None:
+self._stop_recv_loop()
+ifself._recv_task:
+try:
+awaitself._recv_task
+except (asyncio.CancelledError, Exception):
+pass
+error=_StreamError(Exception("Multiplexer closed"), self._stream_generation)
+forqueueinself._get_unique_queues():
+self._put_error_nowait(queue, error)