diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/_stream_pool.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/_stream_pool.py new file mode 100644 index 000000000000..e3c2b283bfa0 --- /dev/null +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/_stream_pool.py @@ -0,0 +1,306 @@ +# 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__ import annotations + +import asyncio +import inspect +import logging +import math +from dataclasses import dataclass +from io import BytesIO +from typing import Awaitable, Callable, Dict, List, Optional, Tuple + +from google.cloud.storage.asyncio._stream_multiplexer import _StreamMultiplexer +from google.cloud.storage.asyncio.async_read_object_stream import ( + _AsyncReadObjectStream, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class MultiStreamConfig: + """Configuration parameters for multi-stream download scaling and load balancing. + + :type min_streams: int + :param min_streams: Minimum number of streams opened at initialization. Defaults to 1. + + :type max_streams: int + :param max_streams: Maximum concurrent streams allowed. Defaults to 8. + + :type target_io_depth: int + :param target_io_depth: Target outstanding requests per stream before scaling up. Defaults to 8. + + :type target_bytes: int + :param target_bytes: Target outstanding bytes per stream before scaling up. Defaults to 8 MB. + """ + + min_streams: int = 1 + max_streams: int = 8 + target_io_depth: int = 8 + target_bytes: int = 8 * 1024 * 1024 + + +class _PooledStream: + """Wraps an active stream and its multiplexer with local load counters.""" + + def __init__( + self, + stream: _AsyncReadObjectStream, + multiplexer: _StreamMultiplexer, + ): + self.stream = stream + self.multiplexer = multiplexer + self.pending_ranges: int = 0 + self.pending_bytes: int = 0 + + def calculate_load(self, target_io_depth: int, target_bytes: int) -> float: + """Returns normalized load metric combining request count and bytes.""" + u_req = self.pending_ranges / target_io_depth if target_io_depth > 0 else 0.0 + u_bytes = self.pending_bytes / target_bytes if target_bytes > 0 else 0.0 + return 0.5 * (u_req + u_bytes) + + def is_fully_loaded(self, target_io_depth: int, target_bytes: int) -> bool: + """Returns True if normalized load is at or above 1.0.""" + return self.calculate_load(target_io_depth, target_bytes) >= 1.0 + + def record_request(self, range_count: int, byte_count: int) -> None: + self.pending_ranges += range_count + self.pending_bytes += byte_count + + def record_completion(self, range_count: int, byte_count: int) -> None: + self.pending_ranges = max(0, self.pending_ranges - range_count) + self.pending_bytes = max(0, self.pending_bytes - byte_count) + + async def close(self) -> None: + """Closes the underlying multiplexer and stream cleanly.""" + try: + await self.multiplexer.close() + except Exception as e: + logger.warning(f"Failed to close stream multiplexer: {e}") + try: + res = self.stream.close() + if inspect.isawaitable(res): + await res + except Exception as e: + logger.warning(f"Failed to close underlying read stream: {e}") + + +class _StreamPool: + """Manages a pool of ``_PooledStream`` instances with dynamic scaling, + hybrid stream selection, and contiguous range distribution. + + Each stream's normalized load is computed as an equal-weighted average of + its request-count utilization and byte utilization relative to target + thresholds:: + + load = 0.5 * (pending_ranges / target_io_depth + pending_bytes / target_bytes) + + A stream is considered *fully loaded* when ``load >= 1.0``. + + **Hybrid Stream Selection Strategy:** + + 1. **Least-Loaded across Base Streams (``streams[:min_streams]``):** + Incoming ranges are first balanced across the initial ``min_streams`` + base streams by selecting the stream with the lowest normalized load, + as long as that stream is not yet fully loaded (``load < 1.0``). + 2. **First-Fit across Scaled Streams (``streams[min_streams:]``):** + Once all base streams are fully loaded (``load >= 1.0``), dynamically + scaled streams are checked in creation order and the first stream with + ``load < 1.0`` is selected. This packs work onto earlier scaled streams + rather than spreading small amounts of load across all scaled streams. + 3. **Least-Loaded Fallback:** + If every open stream in the pool is already at or above target capacity + (``load >= 1.0``), the stream with the lowest normalized load across the + entire pool is selected. + + **Contiguous Range Distribution:** + When distributing a batch of ranges via :meth:`distribute_ranges`, + consecutive ranges are assigned to the currently selected stream until its + normalized load reaches ``1.0``, before moving to the next stream (or + triggering background scale-up for remaining overflow ranges). + """ + + def __init__( + self, + stream_factory: Callable[[], Awaitable[_PooledStream]], + multi_stream_config: Optional[MultiStreamConfig] = None, + ): + config = multi_stream_config or MultiStreamConfig() + self._stream_factory = stream_factory + self.min_streams = max(1, config.min_streams) + self.max_streams = max(self.min_streams, config.max_streams) + self.target_io_depth = config.target_io_depth + self.target_bytes = config.target_bytes + + self.streams: List[_PooledStream] = [] + self._lock = asyncio.Lock() + self._pending_scale_ups: int = 0 + self._closed = False + self._background_tasks: set[asyncio.Task] = set() + + async def add_stream(self, pooled_stream: _PooledStream) -> None: + async with self._lock: + self.streams.append(pooled_stream) + + def _acquire_stream(self) -> _PooledStream: + base_streams = self.streams[: self.min_streams] + best_base = min( + base_streams, + key=lambda s: s.calculate_load(self.target_io_depth, self.target_bytes), + ) + if not best_base.is_fully_loaded(self.target_io_depth, self.target_bytes): + return best_base + + for stream in self.streams[self.min_streams :]: + if not stream.is_fully_loaded(self.target_io_depth, self.target_bytes): + return stream + + return min( + self.streams, + key=lambda s: s.calculate_load(self.target_io_depth, self.target_bytes), + ) + + def _maybe_scale_up(self, extra_load: float = 0.0) -> List[asyncio.Task]: + total_load = ( + sum( + s.calculate_load(self.target_io_depth, self.target_bytes) + for s in self.streams + ) + + extra_load + ) + desired_streams = math.ceil(total_load) + planned_streams = len(self.streams) + self._pending_scale_ups + needed_scale_ups = max( + 0, min(desired_streams, self.max_streams) - planned_streams + ) + + tasks = [] + for _ in range(needed_scale_ups): + self._pending_scale_ups += 1 + task = asyncio.create_task(self._scale_up()) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + tasks.append(task) + return tasks + + async def distribute_ranges( + self, + read_ranges: List[Tuple[int, int, BytesIO]], + allow_scale_up: bool = True, + ) -> Tuple[ + Dict[_PooledStream, List[Tuple[int, int, BytesIO]]], + List[Tuple[int, int, BytesIO]], + List[asyncio.Task], + ]: + """Distributes ``read_ranges`` across open streams and triggers background + scale-up when all currently open streams become fully loaded. + + :type read_ranges: List[Tuple[int, int, BytesIO]] + :param read_ranges: List of ``(offset, length, buffer)`` tuples to distribute. + + :type allow_scale_up: bool + :param allow_scale_up: If True, triggers background scale-up and defers + remaining overflow ranges when all open streams reach full capacity. + + :rtype: Tuple[Dict[_PooledStream, List[Tuple[int, int, BytesIO]]], List[Tuple[int, int, BytesIO]], List[asyncio.Task]] + :returns: A 3-tuple ``(assignments, overflow, scale_tasks)``: + - ``assignments``: Dictionary mapping each selected ``_PooledStream`` + to the contiguous list of ``read_ranges`` assigned to it for immediate + download. + - ``overflow``: List of remaining ``read_ranges`` that could not fit on + the currently open streams because all streams reached full load and + background scale-up tasks were started. Empty if all ranges were + assigned immediately. + - ``scale_tasks``: List of background ``asyncio.Task`` objects created + to open new streams for ``overflow``. Empty if no scale-up was + triggered for overflow. + """ + async with self._lock: + if self._closed: + raise ValueError("Pool is closed") + if not self.streams: + raise ValueError("No streams available in stream pool") + + assignments: Dict[_PooledStream, List[Tuple[int, int, BytesIO]]] = {} + current_stream: Optional[_PooledStream] = None + + for idx, read_range in enumerate(read_ranges): + if current_stream is None or current_stream.is_fully_loaded( + self.target_io_depth, self.target_bytes + ): + current_stream = self._acquire_stream() + if ( + allow_scale_up + and len(read_ranges) > 1 + and all( + s.is_fully_loaded(self.target_io_depth, self.target_bytes) + for s in self.streams + ) + ): + overflow = read_ranges[idx:] + overflow_bytes = sum(length for _, length, _ in overflow) + u_req = ( + len(overflow) / self.target_io_depth + if self.target_io_depth > 0 + else 0.0 + ) + u_bytes = ( + overflow_bytes / self.target_bytes + if self.target_bytes > 0 + else 0.0 + ) + scale_tasks = self._maybe_scale_up( + extra_load=0.5 * (u_req + u_bytes) + ) + if scale_tasks: + return assignments, overflow, scale_tasks + + current_stream.record_request(1, read_range[1]) + assignments.setdefault(current_stream, []).append(read_range) + + if allow_scale_up: + self._maybe_scale_up() + return assignments, [], [] + + async def _scale_up(self) -> None: + try: + new_stream = await self._stream_factory() + async with self._lock: + if self._closed: + await new_stream.close() + return + self.streams.append(new_stream) + except Exception as e: + logger.warning(f"Failed to scale up MRD stream: {e}") + finally: + async with self._lock: + self._pending_scale_ups = max(0, self._pending_scale_ups - 1) + + def release_stream( + self, pooled_stream: _PooledStream, range_count: int, req_bytes: int + ) -> None: + pooled_stream.record_completion(range_count, req_bytes) + + async def close(self) -> None: + async with self._lock: + self._closed = True + self._pending_scale_ups = 0 + streams = list(self.streams) + self.streams.clear() + for task in list(self._background_tasks): + task.cancel() + for s in streams: + await s.close() diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py index 32757bf754b7..79f483b12f15 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_multi_range_downloader.py @@ -15,6 +15,7 @@ from __future__ import annotations import asyncio +import inspect import logging from io import BytesIO from typing import Any, Dict, List, Optional, Tuple @@ -29,6 +30,11 @@ _StreamError, _StreamMultiplexer, ) +from google.cloud.storage.asyncio._stream_pool import ( + MultiStreamConfig, + _PooledStream, + _StreamPool, +) from google.cloud.storage.asyncio.async_grpc_client import ( AsyncGrpcClient, ) @@ -120,6 +126,7 @@ async def create_mrd( read_handle: Optional[_storage_v2.BidiReadHandle] = None, retry_policy: Optional[AsyncRetry] = None, metadata: Optional[List[Tuple[str, str]]] = None, + multi_stream_config: Optional[MultiStreamConfig] = None, **kwargs, ) -> AsyncMultiRangeDownloader: """Initializes a MultiRangeDownloader and opens the underlying bidi-gRPC @@ -148,6 +155,11 @@ async def create_mrd( :type metadata: List[Tuple[str, str]] :param metadata: (Optional) The metadata to be sent with the ``open`` request. + :type multi_stream_config: Optional[MultiStreamConfig] + :param multi_stream_config: (Optional) Configuration dataclass grouping all multi-stream + parameters. If None, multi-stream is disabled and a single + stream is used. + :rtype: :class:`~google.cloud.storage.asyncio.async_multi_range_downloader.AsyncMultiRangeDownloader` :returns: An initialized AsyncMultiRangeDownloader instance for reading. """ @@ -157,6 +169,7 @@ async def create_mrd( object_name, generation=generation, read_handle=read_handle, + multi_stream_config=multi_stream_config, **kwargs, ) await mrd.open(retry_policy=retry_policy, metadata=metadata) @@ -169,10 +182,11 @@ def __init__( object_name: str, generation: Optional[int] = None, read_handle: Optional[_storage_v2.BidiReadHandle] = None, + multi_stream_config: Optional[MultiStreamConfig] = None, **kwargs, ) -> None: - """Constructor for AsyncMultiRangeDownloader, clients are not adviced to - use it directly. Instead it's adviced to use the classmethod `create_mrd`. + """Constructor for AsyncMultiRangeDownloader, clients are not advised to + use it directly. Instead it's advised to use the classmethod `create_mrd`. :type client: :class:`~google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient` :param client: The asynchronous client to use for making API requests. @@ -189,6 +203,11 @@ def __init__( :type read_handle: _storage_v2.BidiReadHandle :param read_handle: (Optional) An existing read handle. + + :type multi_stream_config: Optional[MultiStreamConfig] + :param multi_stream_config: (Optional) Configuration dataclass grouping all multi-stream + parameters. If None, multi-stream is disabled and a single + stream is used. """ if "generation_number" in kwargs: if generation is not None: @@ -216,6 +235,12 @@ def __init__( self.is_finalized: bool = False self.full_obj_server_crc32c: Optional[int] = None + self.multi_stream_config = multi_stream_config + self._pool: Optional[_StreamPool] = None + self._primary_stream: Optional[_PooledStream] = None + self._metadata: Optional[List[Tuple[str, str]]] = None + self._open_retry_policy: Optional[AsyncRetry] = None + async def __aenter__(self): """Opens the underlying bidi-gRPC connection to read from the object.""" await self.open() @@ -298,7 +323,6 @@ async def _do_open(): current_metadata.append( ("x-goog-request-params", f"routing_token={self._routing_token}") ) - self._routing_token = None await self.read_obj_str.open( metadata=current_metadata if current_metadata else None @@ -315,15 +339,92 @@ async def _do_open(): self._is_stream_open = True + self._open_retry_policy = retry_policy + self._metadata = list(metadata) if metadata else [] await retry_policy(_do_open)() self._multiplexer = _StreamMultiplexer(self.read_obj_str) + self._primary_stream = _PooledStream(self.read_obj_str, self._multiplexer) + + if ( + self.multi_stream_config is not None + and self.multi_stream_config.max_streams > 1 + ): + self._pool = _StreamPool( + stream_factory=self._create_pooled_stream, + multi_stream_config=self.multi_stream_config, + ) + await self._pool.add_stream(self._primary_stream) + + if self.multi_stream_config.min_streams > 1: + extra_streams = await asyncio.gather( + *( + self._create_pooled_stream() + for _ in range(self.multi_stream_config.min_streams - 1) + ), + return_exceptions=True, + ) + first_exc = next( + (s for s in extra_streams if isinstance(s, Exception)), None + ) + if first_exc is not None: + for s in extra_streams: + if not isinstance(s, Exception): + await s.close() + await self._pool.close() + self._pool = None + raise first_exc + + for stream in extra_streams: + await self._pool.add_stream(stream) + else: + self._pool = None + + async def _create_pooled_stream(self) -> _PooledStream: + """Opens an additional pooled stream using current routing and read_handle.""" + stream: Optional[_AsyncReadObjectStream] = None + + async def _do_open_pooled(): + nonlocal stream + if stream is not None and getattr(stream, "is_stream_open", False): + try: + await stream.close() + except exceptions.GoogleAPICallError as e: + logger.warning( + f"Failed to close existing pooled stream during resumption: {e}" + ) + + current_metadata = list(self._metadata) if self._metadata else [] + if self._routing_token: + current_metadata.append( + ("x-goog-request-params", f"routing_token={self._routing_token}") + ) + + stream = _AsyncReadObjectStream( + client=self.client.grpc_client, + bucket_name=self.bucket_name, + object_name=self.object_name, + generation_number=self.generation, + read_handle=self.read_handle, + ) + await stream.open(metadata=current_metadata if current_metadata else None) + + if stream.generation_number: + self.generation = stream.generation_number + if stream.read_handle: + self.read_handle = stream.read_handle + + await self._open_retry_policy(_do_open_pooled)() - def _create_stream_factory(self, state, metadata): + mux = _StreamMultiplexer(stream) + return _PooledStream(stream, mux) + + def _create_stream_factory(self, state, metadata, pooled_stream=None): """Create a factory that opens a new stream with current routing state.""" + target_stream = pooled_stream or self._primary_stream async def factory(): - current_handle = state.get("read_handle") - current_token = state.get("routing_token") + current_handle = state.get("read_handle") or self.read_handle + current_token = state.get("routing_token") or self._routing_token stream = _AsyncReadObjectStream( client=self.client.grpc_client, @@ -348,23 +449,29 @@ async def factory(): self.generation = stream.generation_number if stream.read_handle: self.read_handle = stream.read_handle + if current_token: + self._routing_token = current_token self.is_finalized = stream.is_finalized self.full_obj_server_crc32c = stream.full_obj_server_crc32c self.read_obj_str = stream + if target_stream is not None: + target_stream.stream = stream + if target_stream is None or target_stream == self._primary_stream: + self.read_obj_str = stream self._is_stream_open = True return stream return factory - async def download_ranges( + async def _download_ranges_on_stream( self, + pooled_stream: _PooledStream, read_ranges: List[Tuple[int, int, BytesIO]], - lock: asyncio.Lock = None, - retry_policy: Optional[AsyncRetry] = None, - metadata: Optional[List[Tuple[str, str]]] = None, - enable_checksum: bool = True, + retry_policy: AsyncRetry, + metadata: Optional[List[Tuple[str, str]]], + enable_checksum: bool, ) -> None: """Downloads multiple byte ranges from the object into the buffers provided by user with automatic retries. @@ -440,13 +547,13 @@ async def download_ranges( initial_state = { "download_states": download_states, "read_handle": self.read_handle, - "routing_token": None, + "routing_token": self._routing_token, "enable_checksum": enable_checksum, "full_obj_server_crc32c": self.full_obj_server_crc32c, } read_ids = set(download_states.keys()) - queue = self._multiplexer.register(read_ids) + queue = pooled_stream.multiplexer.register(read_ids) try: attempt_count = 0 @@ -473,14 +580,16 @@ async def generator(): broken_gen = ( last_broken_generation if attempt_count > 1 - else self._multiplexer.stream_generation + else pooled_stream.multiplexer.stream_generation ) - stream_factory = self._create_stream_factory(state, metadata) - await self._multiplexer.reopen_stream( + stream_factory = self._create_stream_factory( + state, metadata, pooled_stream=pooled_stream + ) + await pooled_stream.multiplexer.reopen_stream( broken_gen, stream_factory ) - stream_generation = self._multiplexer.stream_generation + stream_generation = pooled_stream.multiplexer.stream_generation # Send Requests pending_read_ids = {r.read_id for r in requests} @@ -489,7 +598,7 @@ async def generator(): ): batch = requests[i : i + _MAX_READ_RANGES_PER_BIDI_READ_REQUEST] try: - await self._multiplexer.send( + await pooled_stream.multiplexer.send( _storage_v2.BidiReadObjectRequest(read_ranges=batch) ) except Exception: @@ -539,9 +648,122 @@ async def generator(): if initial_state.get("read_handle"): self.read_handle = initial_state["read_handle"] + if initial_state.get("routing_token"): + self._routing_token = initial_state["routing_token"] finally: - if self._multiplexer is not None: - self._multiplexer.unregister(read_ids) + if pooled_stream.multiplexer is not None: + pooled_stream.multiplexer.unregister(read_ids) + + async def download_ranges( + self, + read_ranges: List[Tuple[int, int, BytesIO]], + lock: asyncio.Lock = None, + retry_policy: Optional[AsyncRetry] = None, + metadata: Optional[List[Tuple[str, str]]] = None, + enable_checksum: bool = True, + ) -> None: + """Downloads multiple byte ranges from the object into the buffers + provided by user with automatic retries across a managed multi-stream pool. + + :type read_ranges: List[Tuple[int, int, "BytesIO"]] + :param read_ranges: A list of tuples, where each tuple represents a + combination of byte_range and writeable buffer in format - + (`start_byte`, `bytes_to_read`, `writeable_buffer`). Buffer has + to be provided by the user, and user has to make sure appropriate + memory is available in the application to avoid out-of-memory crash. + + Special cases: + if the value of `bytes_to_read` is 0, it'll be interpreted as + download all contents until the end of the file from `start_byte`. + Examples: + * (0, 0, buffer) : downloads 0 to end , i.e. entire object. + * (100, 0, buffer) : downloads from 100 to end. + + :type lock: asyncio.Lock + :param lock: (Deprecated) This parameter is deprecated and has no effect. + + :type retry_policy: :class:`~google.api_core.retry_async.AsyncRetry` + :param retry_policy: (Optional) The retry policy to use for the operation. + + :type metadata: List[Tuple[str, str]] + :param metadata: (Optional) The metadata to be sent with the request. + + :type enable_checksum: bool + :param enable_checksum: (Optional) If True, checksums are verified for downloaded data. Defaults to True. + + :raises ValueError: if the underlying bidi-GRPC stream is not open. + :raises ValueError: if the length of read_ranges is more than 1000. + :raises DataCorruption: if a checksum mismatch is detected while reading data. + + """ + + if len(read_ranges) > 1000: + raise ValueError( + "Invalid input - length of read_ranges cannot be more than 1000" + ) + + if enable_checksum: + raise_if_no_fast_crc32c() + + if not self._is_stream_open: + raise ValueError("Underlying bidi-gRPC stream is not open") + + if self.read_obj_str is None or self._multiplexer is None: + raise ValueError( + "Underlying bidi-gRPC stream or multiplexer is not initialized" + ) + + if retry_policy is None: + retry_policy = AsyncRetry(predicate=_is_read_retryable) + + if self._primary_stream is None: + self._primary_stream = _PooledStream(self.read_obj_str, self._multiplexer) + + if self._pool is not None: + + async def _download_batch( + stream: _PooledStream, + batch: List[Tuple[int, int, BytesIO]], + ) -> None: + batch_bytes = sum(length for _, length, _ in batch) + try: + await self._download_ranges_on_stream( + stream, batch, retry_policy, metadata, enable_checksum + ) + finally: + if self._pool is not None: + self._pool.release_stream(stream, len(batch), batch_bytes) + + async def _download_overflow() -> None: + await asyncio.gather(*scale_tasks) + extra_assignments, _, _ = await self._pool.distribute_ranges( + overflow, allow_scale_up=False + ) + await asyncio.gather( + *(_download_batch(s, b) for s, b in extra_assignments.items()) + ) + + assignments, overflow, scale_tasks = await self._pool.distribute_ranges( + read_ranges + ) + tasks = [_download_batch(s, b) for s, b in assignments.items()] + if overflow: + tasks.append(_download_overflow()) + + # Fast path: await directly to avoid asyncio.gather() Task wrapper overhead + # when all ranges fit onto a single stream without overflow. + if len(tasks) == 1: + await tasks[0] + elif tasks: + await asyncio.gather(*tasks) + else: + await self._download_ranges_on_stream( + self._primary_stream, + read_ranges, + retry_policy, + metadata, + enable_checksum, + ) async def close(self): """ @@ -550,16 +772,24 @@ async def close(self): if not self._is_stream_open: raise ValueError("Underlying bidi-gRPC stream is not open") + if self._pool: + await self._pool.close() + self._pool = None + if self._multiplexer: await self._multiplexer.close() self._multiplexer = None if self.read_obj_str: try: - await self.read_obj_str.close() - except (asyncio.CancelledError, exceptions.GoogleAPICallError): + if getattr(self.read_obj_str, "is_stream_open", True): + res = self.read_obj_str.close() + if inspect.isawaitable(res): + await res + except (ValueError, asyncio.CancelledError, exceptions.GoogleAPICallError): pass self.read_obj_str = None + self._primary_stream = None self._is_stream_open = False @property diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py index afe1dfd221d9..91cfca3e8dc5 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_multi_range_downloader.py @@ -37,6 +37,42 @@ class TestAsyncMultiRangeDownloader: + @staticmethod + def _make_mock_client(): + mock_client = mock.MagicMock() + mock_client.grpc_client = mock.AsyncMock() + return mock_client + + @staticmethod + def _make_mock_stream( + generation=_TEST_GENERATION_NUMBER, + persisted_size=_TEST_OBJECT_SIZE, + read_handle=_TEST_READ_HANDLE, + is_finalized=True, + open_side_effect=None, + ): + s = mock.MagicMock() + s.open = AsyncMock(side_effect=open_side_effect) + s.close = AsyncMock() + s.generation_number = generation + s.persisted_size = persisted_size + s.read_handle = read_handle + s.is_finalized = is_finalized + s.full_obj_server_crc32c = 12345 if is_finalized else None + s.object_metadata = mock.Mock() + return s + + @staticmethod + def _make_dummy_pooled_stream(): + from google.cloud.storage.asyncio._stream_pool import ( + _PooledStream, + ) + + return _PooledStream( + mock.MagicMock(close=AsyncMock()), + mock.MagicMock(close=AsyncMock()), + ) + def create_read_ranges(self, num_ranges): ranges = [] for i in range(num_ranges): @@ -865,3 +901,590 @@ def test_is_read_retryable_predicate(self): assert _is_read_retryable(exceptions.NotFound("not found")) is False assert _is_read_retryable(exceptions.PermissionDenied("denied")) is False assert _is_read_retryable(exceptions.InvalidArgument("invalid")) is False + + def test_pooled_stream_load(self): + stream = self._make_dummy_pooled_stream() + + # Initially zero load + assert stream.calculate_load(target_io_depth=10, target_bytes=1000) == 0.0 + + # Add 5 ranges and 500 bytes -> 5/10 = 0.5, 500/1000 = 0.5 -> load = 0.5*0.5 + 0.5*0.5 = 0.5 + stream.record_request(5, 500) + assert stream.pending_ranges == 5 + assert stream.pending_bytes == 500 + assert stream.calculate_load(target_io_depth=10, target_bytes=1000) == 0.5 + + # Release + stream.record_completion(3, 300) + assert stream.pending_ranges == 2 + assert stream.pending_bytes == 200 + assert stream.calculate_load(target_io_depth=10, target_bytes=1000) == 0.2 + + @mock.patch("google.cloud.storage.asyncio._stream_pool.logger") + @pytest.mark.asyncio + async def test_pooled_stream_close_logs_warnings_on_error(self, mock_logger): + from google.cloud.storage.asyncio._stream_pool import ( + _PooledStream, + ) + + mock_stream = mock.MagicMock( + close=AsyncMock(side_effect=RuntimeError("stream close error")) + ) + mock_mux = mock.MagicMock( + close=AsyncMock(side_effect=RuntimeError("mux close error")) + ) + pooled = _PooledStream(mock_stream, mock_mux) + + await pooled.close() + + assert mock_logger.warning.call_count == 2 + + @pytest.mark.asyncio + async def test_stream_pool_scale_up_and_dispatch(self): + from google.cloud.storage.asyncio._stream_pool import ( + MultiStreamConfig, + _StreamPool, + ) + + async def stream_factory(): + return self._make_dummy_pooled_stream() + + pool = _StreamPool( + stream_factory=stream_factory, + multi_stream_config=MultiStreamConfig( + min_streams=1, + max_streams=2, + target_io_depth=2, + target_bytes=100, + ), + ) + initial_stream = await stream_factory() + await pool.add_stream(initial_stream) + + # 1. Distribute with small load (1 range, 10 bytes) -> load < 1.0, stays on initial stream + assignments, overflow, scale_tasks = await pool.distribute_ranges( + [(0, 10, BytesIO())] + ) + assert list(assignments.keys()) == [initial_stream] + assert overflow == [] + assert scale_tasks == [] + assert len(pool.streams) == 1 + + # 2. Add enough load to exceed target load > 1.0 -> triggers background scale up + assignments2, _, _ = await pool.distribute_ranges([(10, 100, BytesIO())]) + assert list(assignments2.keys()) == [initial_stream] + + # Allow event loop to run background scale-up task + await asyncio.sleep(0.01) + assert len(pool.streams) == 2 + s2 = pool.streams[1] + assert s2 != initial_stream + + # 3. Next distribute selects idle stream s2 + assignments3, _, _ = await pool.distribute_ranges([(110, 10, BytesIO())]) + assert list(assignments3.keys()) == [s2] + + # 4. Release stream capacity and close + pool.release_stream(initial_stream, 2, 110) + pool.release_stream(s2, 1, 10) + await pool.close() + assert len(pool.streams) == 0 + + @pytest.mark.asyncio + async def test_stream_pool_proportional_scale_up(self): + from google.cloud.storage.asyncio._stream_pool import ( + MultiStreamConfig, + _StreamPool, + ) + + async def stream_factory(): + return self._make_dummy_pooled_stream() + + pool = _StreamPool( + stream_factory=stream_factory, + multi_stream_config=MultiStreamConfig( + min_streams=1, + max_streams=5, + target_io_depth=2, + target_bytes=100, + ), + ) + initial_stream = await stream_factory() + await pool.add_stream(initial_stream) + + # Huge burst: 6 ranges of 50 bytes each (total load = 3.0) -> triggers 2 concurrent background scale-ups + ranges = [(i * 50, 50, BytesIO()) for i in range(6)] + assignments, overflow, scale_tasks = await pool.distribute_ranges(ranges) + assert list(assignments.keys()) == [initial_stream] + assert len(assignments[initial_stream]) == 2 + assert len(overflow) == 4 + assert len(scale_tasks) == 2 + assert pool._pending_scale_ups == 2 + assert len(pool._background_tasks) == 2 + + # Allow background scale-up tasks to finish + await asyncio.gather(*scale_tasks) + assert len(pool.streams) == 3 + assert pool._pending_scale_ups == 0 + + # Burst exceeding max_streams -> capped by remaining headroom (5 - 3 = 2) + more_ranges = [(i * 50, 50, BytesIO()) for i in range(10)] + _, _, scale_tasks2 = await pool.distribute_ranges(more_ranges) + assert len(scale_tasks2) == 2 + assert pool._pending_scale_ups == 2 + + await asyncio.gather(*scale_tasks2) + assert len(pool.streams) == 5 + assert pool._pending_scale_ups == 0 + + await pool.close() + assert len(pool.streams) == 0 + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_create_mrd_with_multi_stream_config(self, mock_cls_stream): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + MultiStreamConfig, + ) + + mock_client = self._make_mock_client() + mock_cls_stream.side_effect = [ + self._make_mock_stream(read_handle=b"h1"), + self._make_mock_stream(read_handle=b"h2"), + ] + + config = MultiStreamConfig( + min_streams=2, + max_streams=4, + target_io_depth=8, + target_bytes=2 * 1024 * 1024, + ) + + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, "b", "o", multi_stream_config=config + ) + + assert mrd.multi_stream_config.min_streams == 2 + assert mrd.multi_stream_config.max_streams == 4 + assert mrd.multi_stream_config.target_io_depth == 8 + assert mrd.multi_stream_config.target_bytes == 2 * 1024 * 1024 + # Verified that 2 streams were opened concurrently for min_streams=2 + assert len(mrd._pool.streams) == 2 + await mrd.close() + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_mrd_open_min_streams_failure_cleanup(self, mock_cls_stream): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + MultiStreamConfig, + ) + + mock_client = self._make_mock_client() + mock_cls_stream.side_effect = [ + self._make_mock_stream(), + self._make_mock_stream( + open_side_effect=exceptions.PermissionDenied("Stream failed") + ), + ] + + config = MultiStreamConfig(min_streams=2, max_streams=4) + + with pytest.raises(exceptions.PermissionDenied, match="Stream failed"): + await AsyncMultiRangeDownloader.create_mrd( + mock_client, "b", "o", multi_stream_config=config + ) + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_mrd_download_ranges_triggers_pool_scaling(self, mock_cls_stream): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + MultiStreamConfig, + ) + + mock_client = self._make_mock_client() + mock_cls_stream.side_effect = [ + self._make_mock_stream(), + self._make_mock_stream(), + ] + + config = MultiStreamConfig( + min_streams=1, + max_streams=3, + target_io_depth=1, + target_bytes=50, + ) + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, "b", "o", multi_stream_config=config + ) + assert len(mrd._pool.streams) == 1 + + with mock.patch.object(mrd, "_download_ranges_on_stream", new=AsyncMock()): + # Download ranges with load > 1.0 (2 ranges, 100 bytes) + await mrd.download_ranges([(0, 50, BytesIO()), (50, 50, BytesIO())]) + await asyncio.sleep(0.01) + # Pool dynamically scaled up to 2 streams + assert len(mrd._pool.streams) == 2 + + await mrd.close() + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_create_mrd_single_stream_bypass(self, mock_cls_stream): + mock_client = self._make_mock_client() + mock_cls_stream.return_value = self._make_mock_stream() + + # Default create_mrd without multi_stream_config -> single-stream bypass + mrd = await AsyncMultiRangeDownloader.create_mrd(mock_client, "b", "o") + assert mrd.multi_stream_config is None + assert mrd._pool is None + + # download_ranges should use _primary_stream directly without pool + with mock.patch.object( + mrd, "_download_ranges_on_stream", new=AsyncMock() + ) as mock_dl: + await mrd.download_ranges([(0, 50, BytesIO())]) + assert mock_dl.call_count == 1 + assert mock_dl.call_args[0][0] == mrd._primary_stream + + await mrd.close() + + @pytest.mark.asyncio + async def test_stream_pool_closed_and_cancellation(self): + from google.cloud.storage.asyncio._stream_pool import ( + MultiStreamConfig, + _StreamPool, + ) + + scale_up_started = asyncio.Event() + scale_up_finish = asyncio.Event() + + async def slow_factory(): + scale_up_started.set() + await scale_up_finish.wait() + return self._make_dummy_pooled_stream() + + pool = _StreamPool( + stream_factory=slow_factory, + multi_stream_config=MultiStreamConfig( + min_streams=1, + max_streams=2, + target_io_depth=1, + target_bytes=10, + ), + ) + initial_s = self._make_dummy_pooled_stream() + await pool.add_stream(initial_s) + + # Trigger scale up + await pool.distribute_ranges([(0, 20, BytesIO())]) + await scale_up_started.wait() + assert len(pool._background_tasks) == 1 + + # Close pool while scale up task is pending + await pool.close() + assert pool._closed is True + assert len(pool.streams) == 0 + + # Background task should be cancelled + with pytest.raises(asyncio.CancelledError): + await next( + iter(pool._background_tasks) + ) if pool._background_tasks else asyncio.sleep(0) + + # distribute_ranges on closed pool must raise ValueError + with pytest.raises(ValueError, match="Pool is closed"): + await pool.distribute_ranges([(0, 10, BytesIO())]) + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_routing_token_preservation_and_propagation( + self, mock_cls_async_read_object_stream + ): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + MultiStreamConfig, + ) + + def make_redirect_error(token: str, handle: bytes): + redirect = _storage_v2.BidiReadObjectRedirectedError( + routing_token=token, + read_handle=_storage_v2.BidiReadHandle(handle=handle), + ) + status = status_pb2.Status() + detail = status.details.add() + detail.type_url = ( + "type.googleapis.com/google.storage.v2.BidiReadObjectRedirectedError" + ) + detail.value = _storage_v2.BidiReadObjectRedirectedError.serialize(redirect) + grpc_error = mock.MagicMock() + grpc_error.trailing_metadata.return_value = [ + ("grpc-status-details-bin", status.SerializeToString()) + ] + return exceptions.Aborted("redirect", errors=[grpc_error]) + + mock_client = self._make_mock_client() + + # 1. Primary stream fails on first open with redirect ("token-abc"), succeeds on retry + s1_failed = self._make_mock_stream( + open_side_effect=make_redirect_error("token-abc", b"h1") + ) + s1_ok = self._make_mock_stream(read_handle=b"h1") + # 2. Second stream opened via min_streams=2 receives "token-abc", fails once with redirect ("token-xyz"), succeeds on retry + s2_failed = self._make_mock_stream( + open_side_effect=make_redirect_error("token-xyz", b"h2") + ) + s2_ok = self._make_mock_stream(read_handle=b"h2") + + mock_cls_async_read_object_stream.side_effect = [ + s1_failed, + s1_ok, + s2_failed, + s2_ok, + ] + + config = MultiStreamConfig(min_streams=2, max_streams=4) + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, "b", "o", multi_stream_config=config + ) + + # Verify s1_ok received routing_token=token-abc on retry + assert ("x-goog-request-params", "routing_token=token-abc") in ( + s1_ok.open.call_args.kwargs["metadata"] + ) + # Verify s2_failed inherited routing_token=token-abc preserved from primary open + assert ("x-goog-request-params", "routing_token=token-abc") in ( + s2_failed.open.call_args.kwargs["metadata"] + ) + # Verify s2_ok retried with the updated routing_token=token-xyz + assert ("x-goog-request-params", "routing_token=token-xyz") in ( + s2_ok.open.call_args.kwargs["metadata"] + ) + assert mrd._routing_token == "token-xyz" + + await mrd.close() + + @pytest.mark.asyncio + async def test_stream_pool_selection_min_streams_and_further(self): + from google.cloud.storage.asyncio._stream_pool import ( + MultiStreamConfig, + _StreamPool, + ) + + pool = _StreamPool( + stream_factory=lambda: self._make_dummy_pooled_stream(), + multi_stream_config=MultiStreamConfig( + min_streams=2, + max_streams=4, + target_io_depth=2, + target_bytes=100, + ), + ) + + s1 = self._make_dummy_pooled_stream() + s2 = self._make_dummy_pooled_stream() + await pool.add_stream(s1) + await pool.add_stream(s2) + + # s1 has some load, s2 has 0 load -> selects least loaded among min_streams (s2) + s1.record_request(1, 10) + assert s1.calculate_load(2, 100) == 0.3 + assert s2.calculate_load(2, 100) == 0.0 + + assignments, _, _ = await pool.distribute_ranges( + [(0, 10, BytesIO())], allow_scale_up=False + ) + assert list(assignments.keys()) == [s2] + + # Fully load both s1 and s2 (load >= 1.0) + s1.record_request(2, 100) + s2.record_request(3, 100) + assert s1.is_fully_loaded(2, 100) + assert s2.is_fully_loaded(2, 100) + + # Add further streams s3 and s4 (beyond min_streams=2) + s3 = self._make_dummy_pooled_stream() + s4 = self._make_dummy_pooled_stream() + await pool.add_stream(s3) + await pool.add_stream(s4) + + # Give s3 some load (0.35), s4 has 0.0 load + s3.record_request(1, 20) + assert s3.calculate_load(2, 100) == 0.35 + assert s4.calculate_load(2, 100) == 0.0 + + # Since base streams are fully loaded, selects FIRST not fully loaded further stream (s3, not s4!) + assignments_further, _, _ = await pool.distribute_ranges( + [(0, 10, BytesIO())], allow_scale_up=False + ) + assert list(assignments_further.keys()) == [s3] + + # Now fully load s3 + s3.record_request(3, 100) + assert s3.is_fully_loaded(2, 100) + + # Next distribute chooses s4 (the first not fully loaded further stream now) + assignments_s4, _, _ = await pool.distribute_ranges( + [(0, 10, BytesIO())], allow_scale_up=False + ) + assert list(assignments_s4.keys()) == [s4] + + # Now fully load s4 + s4.record_request(3, 100) + assert s4.is_fully_loaded(2, 100) + + # When all streams are fully loaded, fallback selects the least loaded overall + # Let s3 have load 2.0 and s1, s2, s4 have load 3.0 + s3.pending_ranges = 4 + s3.pending_bytes = 200 + s1.pending_ranges = 6 + s1.pending_bytes = 300 + s2.pending_ranges = 6 + s2.pending_bytes = 300 + s4.pending_ranges = 6 + s4.pending_bytes = 300 + + fallback_assignments, _, _ = await pool.distribute_ranges( + [(0, 10, BytesIO())], allow_scale_up=False + ) + assert list(fallback_assignments.keys()) == [s3] + + await pool.close() + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_mrd_download_ranges_sharding_across_streams(self, mock_cls_stream): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + MultiStreamConfig, + ) + + mock_client = self._make_mock_client() + s1 = self._make_mock_stream(read_handle=b"h1") + s2 = self._make_mock_stream(read_handle=b"h2") + s3 = self._make_mock_stream(read_handle=b"h3") + + open_stream_started = asyncio.Event() + + async def s3_open_after_s1_starts(*args, **kwargs): + # Prove s1/s2 start downloading before s3 finishes opening + await asyncio.wait_for(open_stream_started.wait(), timeout=1.0) + + s3.open = AsyncMock(side_effect=s3_open_after_s1_starts) + mock_cls_stream.side_effect = [s1, s2, s3] + + # Each range of 50 bytes with target_io_depth=2, target_bytes=100 has load = 0.5. + # So 2 ranges (100 bytes) saturate a stream (load = 1.0). + config = MultiStreamConfig( + min_streams=2, + max_streams=3, + target_io_depth=2, + target_bytes=100, + ) + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, "b", "o", multi_stream_config=config + ) + assert len(mrd._pool.streams) == 2 + + dispatched_shards = [] + + async def record_shard(pooled_stream, sub_ranges, *args, **kwargs): + open_stream_started.set() + dispatched_shards.append((pooled_stream, list(sub_ranges))) + + with mock.patch.object( + mrd, "_download_ranges_on_stream", side_effect=record_shard + ): + # Pass 6 ranges of 50 bytes each (total load = 3.0 across max_streams=3). + # Open streams (s1, s2) start immediately while s3 opens in parallel for overflow. + ranges = [(i * 50, 50, BytesIO()) for i in range(6)] + await mrd.download_ranges(ranges) + + assert len(mrd._pool.streams) == 3 + assert len(dispatched_shards) == 3 + assert [len(sub_ranges) for _, sub_ranges in dispatched_shards] == [2, 2, 2] + + # Verify all counters were released back to 0 after completion + for ps in mrd._pool.streams: + assert ps.pending_ranges == 0 + assert ps.pending_bytes == 0 + + await mrd.close() + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_mrd_download_ranges_fills_current_stream_until_fully_loaded( + self, mock_cls_stream + ): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + MultiStreamConfig, + ) + + mock_client = self._make_mock_client() + s1 = self._make_mock_stream(read_handle=b"h1") + s2 = self._make_mock_stream(read_handle=b"h2") + mock_cls_stream.side_effect = [s1, s2] + + # Each 25-byte range has load = 0.5*(1/4) + 0.5*(25/100) = 0.25. + # So 4 ranges (100 bytes) saturate a stream (load = 1.0). + config = MultiStreamConfig( + min_streams=2, + max_streams=4, + target_io_depth=4, + target_bytes=100, + ) + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, "b", "o", multi_stream_config=config + ) + assert len(mrd._pool.streams) == 2 + + dispatched_shards = [] + + async def record_shard(pooled_stream, sub_ranges, *args, **kwargs): + dispatched_shards.append((pooled_stream, list(sub_ranges))) + + with mock.patch.object( + mrd, "_download_ranges_on_stream", side_effect=record_shard + ): + # 1. Small batch (2 ranges, total load = 0.5 < 1.0): stays together on stream 0 + r_small = [(0, 25, BytesIO()), (25, 25, BytesIO())] + await mrd.download_ranges(r_small) + assert len(dispatched_shards) == 1 + assert dispatched_shards[0] == (mrd._pool.streams[0], r_small) + + dispatched_shards.clear() + + # 2. Larger batch (6 ranges, total load = 1.5): + # First 4 ranges fill stream 0 to load 1.0; next 2 ranges go contiguously to stream 1 + r_large = [(i * 25, 25, BytesIO()) for i in range(6)] + await mrd.download_ranges(r_large) + assert len(dispatched_shards) == 2 + assert dispatched_shards[0] == (mrd._pool.streams[0], r_large[:4]) + assert dispatched_shards[1] == (mrd._pool.streams[1], r_large[4:]) + + await mrd.close() + + @pytest.mark.asyncio + async def test_download_ranges_uninitialized_stream_or_multiplexer_raises(self): + mock_client = self._make_mock_client() + mrd = AsyncMultiRangeDownloader(mock_client, "b", "o") + mrd._is_stream_open = True + mrd.read_obj_str = None + mrd._multiplexer = None + + with pytest.raises( + ValueError, + match="Underlying bidi-gRPC stream or multiplexer is not initialized", + ): + await mrd.download_ranges([(0, 10, BytesIO())])