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..314da3ba8a42 --- /dev/null +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/_stream_pool.py @@ -0,0 +1,377 @@ +# 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 + + def __post_init__(self) -> None: + if self.min_streams < 1: + raise ValueError("min_streams must be >= 1") + if self.max_streams < self.min_streams: + raise ValueError("max_streams must be >= min_streams") + if self.target_io_depth <= 0: + raise ValueError("target_io_depth must be > 0") + if self.target_bytes <= 0: + raise ValueError("target_bytes must be > 0") + + +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`` (or a water-filling target load cap when + all streams are saturated and no further scale-up is possible), 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 = config.min_streams + self.max_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() + + def effective_range_bytes( + self, + offset: int, + length: int, + persisted_size: Optional[int] = None, + ) -> int: + """Returns normalized byte length for load accounting, resolving + ``length == 0`` ("read to end of object") using ``persisted_size`` or + ``target_bytes`` as a fallback. + """ + if length > 0: + return length + if persisted_size is not None: + return max(0, persisted_size - offset) + return self.target_bytes + + def _batch_load( + self, + ranges: List[Tuple[int, int, BytesIO]], + persisted_size: Optional[int] = None, + ) -> float: + total_bytes = sum( + self.effective_range_bytes(offset, length, persisted_size) + for offset, length, _ in ranges + ) + return 0.5 * ( + len(ranges) / self.target_io_depth + total_bytes / self.target_bytes + ) + + 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, + persisted_size: Optional[int] = None, + ) -> 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. + + :type persisted_size: Optional[int] + :param persisted_size: Known persisted object size used to estimate byte + load when ``bytes_to_read == 0`` (read to end of object). + + :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 or are in flight. Empty if all + ranges were assigned immediately. + - ``scale_tasks``: List of background ``asyncio.Task`` objects currently + opening new streams for ``overflow``. Empty if no scale-up was + triggered or in flight 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") + + can_overflow = ( + allow_scale_up + and (len(read_ranges) > 1 or bool(self._background_tasks)) + and ( + len(self.streams) < self.max_streams or bool(self._background_tasks) + ) + ) + existing_load = sum( + s.calculate_load(self.target_io_depth, self.target_bytes) + for s in self.streams + ) + load_cap = ( + 1.0 + if can_overflow + else max( + 1.0, + (existing_load + self._batch_load(read_ranges, persisted_size)) + / len(self.streams), + ) + ) + + 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.calculate_load( + self.target_io_depth, self.target_bytes + ) + >= load_cap + and ( + can_overflow + or any( + s.calculate_load(self.target_io_depth, self.target_bytes) + < load_cap + for s in self.streams + ) + ) + ): + current_stream = self._acquire_stream() + if can_overflow and all( + s.is_fully_loaded(self.target_io_depth, self.target_bytes) + for s in self.streams + ): + overflow = read_ranges[idx:] + self._maybe_scale_up( + extra_load=self._batch_load(overflow, persisted_size) + ) + return assignments, overflow, list(self._background_tasks) + + req_bytes = self.effective_range_bytes( + read_range[0], read_range[1], persisted_size + ) + current_stream.record_request(1, req_bytes) + assignments.setdefault(current_stream, []).append(read_range) + + if allow_scale_up: + self._maybe_scale_up() + return assignments, [], [] + + async def _scale_up(self) -> None: + new_stream: Optional[_PooledStream] = None + added = False + try: + new_stream = await self._stream_factory() + async with self._lock: + if not self._closed: + self.streams.append(new_stream) + added = True + except Exception as e: + logger.warning(f"Failed to scale up MRD stream: {e}") + finally: + if new_stream is not None and not added: + await new_stream.close() + 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() + bg_tasks = list(self._background_tasks) + for task in bg_tasks: + task.cancel() + if bg_tasks: + await asyncio.gather(*bg_tasks, return_exceptions=True) + 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 74b9581333d1..654759f4aa57 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 typing import IO, Any, Dict, List, Optional, Tuple, Union @@ -28,6 +29,11 @@ _StreamError, _StreamMultiplexer, ) +from google.cloud.storage.asyncio._stream_pool import ( + MultiStreamConfig, + _PooledStream, + _StreamPool, +) from google.cloud.storage.asyncio.async_grpc_client import ( AsyncGrpcClient, ) @@ -175,6 +181,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 @@ -203,6 +210,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. """ @@ -212,6 +224,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) @@ -224,10 +237,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. @@ -244,6 +258,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: @@ -271,6 +290,13 @@ 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 + self._user_open_on_error: Optional[Any] = None + def _trace_span( self, name: str, @@ -314,6 +340,9 @@ async def open( raise ValueError("Underlying bidi-gRPC stream is already open") async with self._trace_span("Storage.AsyncMultiRangeDownloader.open"): + self._user_open_on_error = ( + retry_policy._on_error if retry_policy is not None else None + ) if retry_policy is None: def on_error_wrapper(exc): @@ -388,14 +417,121 @@ 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, BaseException)), + None, + ) + if first_exc is not None: + for s in extra_streams: + if not isinstance(s, BaseException): + await s.close() + await self._pool.close() + self._pool = None + self._multiplexer = None + self.read_obj_str = None + self._primary_stream = None + self._is_stream_open = False + 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 read_handle.""" + stream: Optional[_AsyncReadObjectStream] = None + pooled_routing_token: Optional[str] = None + user_on_error = self._user_open_on_error + + def _on_pooled_open_error(exc): + nonlocal pooled_routing_token + self._open_retries += 1 + logger.warning(f"Error occurred while opening MRD: {exc}") + routing_token, read_handle = _handle_redirect(exc) + if routing_token: + pooled_routing_token = routing_token + if read_handle: + self.read_handle = read_handle + if user_on_error: + user_on_error(exc) + + base_policy = self._open_retry_policy + pooled_retry_policy = AsyncRetry( + predicate=_is_open_retryable, + initial=base_policy._initial, + maximum=base_policy._maximum, + multiplier=base_policy._multiplier, + deadline=base_policy._deadline, + on_error=_on_pooled_open_error, + ) + + async def _do_open_pooled(): + nonlocal stream, pooled_routing_token + 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 pooled_routing_token: + current_metadata.append( + ("x-goog-request-params", f"routing_token={pooled_routing_token}") + ) + pooled_routing_token = None + + 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) - def _create_stream_factory(self, state, metadata): + if stream.generation_number: + self.generation = stream.generation_number + if stream.read_handle: + self.read_handle = stream.read_handle + + await pooled_retry_policy(_do_open_pooled)() + + 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_handle = state.get("read_handle") or self.read_handle current_token = state.get("routing_token") stream = _AsyncReadObjectStream( @@ -414,6 +550,7 @@ async def factory(): f"routing_token={current_token}", ) ) + state["routing_token"] = None await stream.open(metadata=current_metadata if current_metadata else None) @@ -421,16 +558,218 @@ async def factory(): self.generation = stream.generation_number if stream.read_handle: self.read_handle = stream.read_handle - self.is_finalized = stream.is_finalized - self.full_obj_server_crc32c = stream.full_obj_server_crc32c - - self.read_obj_str = stream + if stream.object_metadata is not None: + self.is_finalized = stream.is_finalized + self.full_obj_server_crc32c = stream.full_obj_server_crc32c + if target_stream is not None: + target_stream.stream = stream + if target_stream is None or target_stream is self._primary_stream: + self.read_obj_str = stream self._is_stream_open = True return stream return factory + async def _download_ranges_on_stream( + self, + pooled_stream: _PooledStream, + read_ranges: List[Tuple[int, int, Union[IO[bytes], bytearray, memoryview]]], + retry_policy: AsyncRetry, + metadata: Optional[List[Tuple[str, str]]], + enable_checksum: bool, + accepts_memoryview: bool = False, + ) -> None: + """Downloads multiple byte ranges from the object into the buffers + provided by user with automatic retries on a single pooled stream. + """ + 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 retry_policy is None: + retry_policy = AsyncRetry(predicate=_is_read_retryable) + + # Initialize Global State for Retry Strategy + download_states = {} + read_ids = set() + try: + for read_range in read_ranges: + read_id = generate_random_56_bit_integer() + # Unpack tuple into self-documenting variable names to improve readability. + offset, length, user_buffer = read_range + + # Heuristic to detect full object reads: + # - Implicit full object read: start offset is 0 and length is 0 (read all). + # - Explicit full object read: start offset is 0 and length matches the exact persisted size. + is_full_object_read = (offset == 0 and length == 0) or ( + self.persisted_size is not None + and offset == 0 + and length == self.persisted_size + ) + download_states[read_id] = _DownloadState( + initial_offset=offset, + initial_length=length, + user_buffer=_wrap_user_buffer( + user_buffer, length, accepts_memoryview + ), + is_full_object_read=is_full_object_read, + ) + + initial_state = { + "download_states": download_states, + "read_handle": self.read_handle, + "routing_token": None, + "enable_checksum": enable_checksum, + "full_obj_server_crc32c": self.full_obj_server_crc32c, + } + + read_ids = set(download_states.keys()) + queue = pooled_stream.multiplexer.register(read_ids) + strategy = _ReadResumptionStrategy() + + attempt_count = 0 + last_broken_generation = None + + def send_and_recv_via_multiplexer( + requests: List[_storage_v2.ReadRange], + state: Dict[str, Any], + ): + async def generator(): + nonlocal attempt_count, last_broken_generation + attempt_count += 1 + + if attempt_count > 1: + logger.info( + f"Resuming download (attempt {attempt_count}) for {len(requests)} ranges." + ) + + # Reopen stream if needed + should_reopen = ( + attempt_count > 1 and last_broken_generation is not None + ) or (attempt_count == 1 and metadata is not None) + if should_reopen: + broken_gen = ( + last_broken_generation + if attempt_count > 1 + else pooled_stream.multiplexer.stream_generation + ) + stream_factory = self._create_stream_factory( + state, metadata, pooled_stream=pooled_stream + ) + await pooled_stream.multiplexer.reopen_stream( + broken_gen, stream_factory + ) + + stream_generation = pooled_stream.multiplexer.stream_generation + + # Send Requests + pending_read_ids = {r.read_id for r in requests} + for i in range( + 0, len(requests), _MAX_READ_RANGES_PER_BIDI_READ_REQUEST + ): + batch = requests[i : i + _MAX_READ_RANGES_PER_BIDI_READ_REQUEST] + try: + await pooled_stream.multiplexer.send( + _storage_v2.BidiReadObjectRequest(read_ranges=batch) + ) + except Exception: + last_broken_generation = stream_generation + raise + + # Receive Responses + while pending_read_ids: + item = await queue.get() + + if isinstance(item, _StreamEnd): + if pending_read_ids: + last_broken_generation = stream_generation + raise exceptions.ServiceUnavailable( + "Stream ended with pending read_ids" + ) + break + + if isinstance(item, _StreamError): + if item.generation < stream_generation: + continue # stale error, skip + last_broken_generation = item.generation + raise item.exception + + # Track completion + if item.object_data_ranges: + for data_range in item.object_data_ranges: + if data_range.range_end: + pending_read_ids.discard( + data_range.read_range.read_id + ) + yield item + # The strategy has consumed `item` by now: write the + # chunks whose checksums are done, and keep the + # backlog bounded. + await strategy.verify_pending_checksums( + state, _CRC32C_MAX_PENDING + ) + + # Write (or reject) every held chunk before the attempt + # ends, so a mismatch is raised inside the retry loop, + # as an inline one is. + await strategy.verify_pending_checksums(state) + + return generator() + + retry_manager = _BidiStreamRetryManager( + strategy, send_and_recv_via_multiplexer + ) + + try: + try: + await retry_manager.execute(initial_state, retry_policy) + except DataCorruption: + raise + except Exception: + # Chunks received before the stream error still have + # in-flight checksum futures in initial_state. Verify them + # before raising the stream error so an earlier corrupt + # chunk raises DataCorruption (matching inline verification) + # and valid chunks already received are written. + await strategy.verify_pending_checksums(initial_state) + raise + await strategy.verify_pending_checksums(initial_state) + except DataCorruption: + if self.is_stream_open: + await self.close() + raise + finally: + # Drop chunks still held, e.g. after cancellation. + strategy.discard_pending_checksums(initial_state) + + if initial_state.get("read_handle"): + self.read_handle = initial_state["read_handle"] + + for download_state in download_states.values(): + if ( + isinstance(download_state.user_buffer, _BufferWriter) + and download_state.bytes_written != download_state.initial_length + ): + raise BufferError( + f"Buffer destination expected {download_state.initial_length} " + f"bytes at offset {download_state.initial_offset}, got " + f"{download_state.bytes_written}" + ) + finally: + if pooled_stream.multiplexer is not None: + pooled_stream.multiplexer.unregister(read_ids) + for download_state in download_states.values(): + if isinstance(download_state.user_buffer, _BufferWriter): + download_state.user_buffer.close() + async def download_ranges( self, read_ranges: List[Tuple[int, int, Union[IO[bytes], bytearray, memoryview]]], @@ -441,7 +780,7 @@ async def download_ranges( accepts_memoryview: bool = False, ) -> None: """Downloads multiple byte ranges from the object into the buffers - provided by user with automatic retries. + provided by user with automatic retries across a managed multi-stream pool. :type read_ranges: List[Tuple[int, int, Union[IO[bytes], bytearray, memoryview]]] :param read_ranges: A list of tuples, where each tuple represents a @@ -531,6 +870,11 @@ async def download_ranges( 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" + ) + async with self._trace_span( "Storage.AsyncMultiRangeDownloader.downloadRanges", attributes={"gcp.storage.range.count": len(read_ranges)}, @@ -538,181 +882,114 @@ async def download_ranges( if retry_policy is None: retry_policy = AsyncRetry(predicate=_is_read_retryable) - # Initialize Global State for Retry Strategy - download_states = {} - read_ids = set() - try: - for read_range in read_ranges: - read_id = generate_random_56_bit_integer() - # Unpack tuple into self-documenting variable names to improve readability. - offset, length, user_buffer = read_range - - # Heuristic to detect full object reads: - # - Implicit full object read: start offset is 0 and length is 0 (read all). - # - Explicit full object read: start offset is 0 and length matches the exact persisted size. - is_full_object_read = (offset == 0 and length == 0) or ( - self.persisted_size is not None - and offset == 0 - and length == self.persisted_size - ) - download_states[read_id] = _DownloadState( - initial_offset=offset, - initial_length=length, - user_buffer=_wrap_user_buffer( - user_buffer, length, accepts_memoryview - ), - is_full_object_read=is_full_object_read, - ) - - initial_state = { - "download_states": download_states, - "read_handle": self.read_handle, - "routing_token": None, - "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) - strategy = _ReadResumptionStrategy() + if self._primary_stream is None: + self._primary_stream = _PooledStream( + self.read_obj_str, self._multiplexer + ) - attempt_count = 0 - last_broken_generation = None + if self._pool is not None: + pool = self._pool + persisted_size = self.persisted_size - def send_and_recv_via_multiplexer( - requests: List[_storage_v2.ReadRange], - state: Dict[str, Any], - ): - async def generator(): - nonlocal attempt_count, last_broken_generation - attempt_count += 1 + async def _run_concurrently(awaitables) -> None: + running = [asyncio.ensure_future(t) for t in awaitables] + try: + await asyncio.gather(*running) + except BaseException: + for t in running: + t.cancel() + await asyncio.gather(*running, return_exceptions=True) + raise - if attempt_count > 1: - logger.info( - f"Resuming download (attempt {attempt_count}) for {len(requests)} ranges." - ) + async def _download_batch( + stream: _PooledStream, + batch: List[ + Tuple[int, int, Union[IO[bytes], bytearray, memoryview]] + ], + ) -> None: + batch_bytes = sum( + pool.effective_range_bytes(offset, length, persisted_size) + for offset, length, _ in batch + ) + try: + await self._download_ranges_on_stream( + stream, + batch, + retry_policy, + metadata, + enable_checksum, + accepts_memoryview=accepts_memoryview, + ) + finally: + pool.release_stream(stream, len(batch), batch_bytes) - # Reopen stream if needed - should_reopen = ( - attempt_count > 1 and last_broken_generation is not None - ) or (attempt_count == 1 and metadata is not None) - if should_reopen: - broken_gen = ( - last_broken_generation - if attempt_count > 1 - else self._multiplexer.stream_generation - ) - stream_factory = self._create_stream_factory( - state, metadata - ) - await self._multiplexer.reopen_stream( - broken_gen, stream_factory + async def _download_overflow() -> None: + remaining_overflow = list(overflow) + pending_scale_tasks = set(scale_tasks) + overflow_tasks: List[asyncio.Task] = [] + try: + while remaining_overflow and pending_scale_tasks: + _, pending_scale_tasks = await asyncio.wait( + pending_scale_tasks, + return_when=asyncio.FIRST_COMPLETED, ) - - stream_generation = self._multiplexer.stream_generation - - # Send Requests - pending_read_ids = {r.read_id for r in requests} - for i in range( - 0, len(requests), _MAX_READ_RANGES_PER_BIDI_READ_REQUEST - ): - batch = requests[ - i : i + _MAX_READ_RANGES_PER_BIDI_READ_REQUEST - ] - try: - await self._multiplexer.send( - _storage_v2.BidiReadObjectRequest(read_ranges=batch) + if pool._closed: + break + if pending_scale_tasks: + ( + extra_assignments, + remaining_overflow, + _, + ) = await pool.distribute_ranges( + remaining_overflow, + allow_scale_up=True, + persisted_size=persisted_size, ) - except Exception: - last_broken_generation = stream_generation - raise - - # Receive Responses - while pending_read_ids: - item = await queue.get() - - if isinstance(item, _StreamEnd): - if pending_read_ids: - last_broken_generation = stream_generation - raise exceptions.ServiceUnavailable( - "Stream ended with pending read_ids" + for s, b in extra_assignments.items(): + overflow_tasks.append( + asyncio.create_task(_download_batch(s, b)) ) - break - if isinstance(item, _StreamError): - if item.generation < stream_generation: - continue # stale error, skip - last_broken_generation = item.generation - raise item.exception - - # Track completion - if item.object_data_ranges: - for data_range in item.object_data_ranges: - if data_range.range_end: - pending_read_ids.discard( - data_range.read_range.read_id - ) - yield item - # The strategy has consumed `item` by now: write the - # chunks whose checksums are done, and keep the - # backlog bounded. - await strategy.verify_pending_checksums( - state, _CRC32C_MAX_PENDING + if remaining_overflow and not pool._closed: + extra_assignments, _, _ = await pool.distribute_ranges( + remaining_overflow, + allow_scale_up=False, + persisted_size=persisted_size, ) + for s, b in extra_assignments.items(): + overflow_tasks.append( + asyncio.create_task(_download_batch(s, b)) + ) - # Write (or reject) every held chunk before the attempt - # ends, so a mismatch is raised inside the retry loop, - # as an inline one is. - await strategy.verify_pending_checksums(state) - - return generator() + await asyncio.gather(*overflow_tasks) + except BaseException: + for t in overflow_tasks: + t.cancel() + await asyncio.gather(*overflow_tasks, return_exceptions=True) + raise - retry_manager = _BidiStreamRetryManager( - strategy, send_and_recv_via_multiplexer + assignments, overflow, scale_tasks = await pool.distribute_ranges( + read_ranges, persisted_size=persisted_size + ) + 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 _run_concurrently(tasks) + else: + await self._download_ranges_on_stream( + self._primary_stream, + read_ranges, + retry_policy, + metadata, + enable_checksum, + accepts_memoryview=accepts_memoryview, ) - - try: - try: - await retry_manager.execute(initial_state, retry_policy) - except DataCorruption: - raise - except Exception: - # Chunks received before the stream error still have - # in-flight checksum futures in initial_state. Verify them - # before raising the stream error so an earlier corrupt - # chunk raises DataCorruption (matching inline verification) - # and valid chunks already received are written. - await strategy.verify_pending_checksums(initial_state) - raise - await strategy.verify_pending_checksums(initial_state) - except DataCorruption: - if self.is_stream_open: - await self.close() - raise - finally: - # Drop chunks still held, e.g. after cancellation. - strategy.discard_pending_checksums(initial_state) - - if initial_state.get("read_handle"): - self.read_handle = initial_state["read_handle"] - - for download_state in download_states.values(): - if ( - isinstance(download_state.user_buffer, _BufferWriter) - and download_state.bytes_written - != download_state.initial_length - ): - raise BufferError( - f"Buffer destination expected {download_state.initial_length} " - f"bytes at offset {download_state.initial_offset}, got " - f"{download_state.bytes_written}" - ) - finally: - if self._multiplexer is not None: - self._multiplexer.unregister(read_ids) - for download_state in download_states.values(): - if isinstance(download_state.user_buffer, _BufferWriter): - download_state.user_buffer.close() async def close(self): """ @@ -722,16 +999,28 @@ async def close(self): raise ValueError("Underlying bidi-gRPC stream is not open") async with self._trace_span("Storage.AsyncMultiRangeDownloader.close"): + 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/system/test_rapid.py b/packages/google-cloud-storage/tests/system/test_rapid.py index 038203cb5103..8a92a1477e45 100644 --- a/packages/google-cloud-storage/tests/system/test_rapid.py +++ b/packages/google-cloud-storage/tests/system/test_rapid.py @@ -35,12 +35,26 @@ from google.cloud.storage.asyncio.async_grpc_client import AsyncGrpcClient from google.cloud.storage.asyncio.async_multi_range_downloader import ( AsyncMultiRangeDownloader, + MultiStreamConfig, ) from google.cloud.storage.blob import ( ObjectContexts, ObjectCustomContextPayload, ) +_MULTI_STREAM_CONFIGS = [ + pytest.param(None, id="single_stream"), + pytest.param( + MultiStreamConfig( + min_streams=2, + max_streams=4, + target_io_depth=2, + target_bytes=1024 * 1024, + ), + id="multi_stream", + ), +] + PREPROD_GRPC_ENDPOINT = "storage-preprod-test-grpc.googleusercontent.com:443" # Run system test for either Rapid or RCU. But not both => XOR pytestmark = pytest.mark.skipif( @@ -176,6 +190,7 @@ def _get_equal_dist(a: int, b: int) -> tuple[int, int]: @pytest.mark.skipif( RCU_SYSTEM_TESTS, reason="X regions reads/writes for RCU not supported in SDK yet" ) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) @pytest.mark.parametrize( "object_size", [ @@ -188,6 +203,7 @@ def test_basic_wrd_x_region( storage_client, blobs_to_delete, object_size, + multi_stream_config, event_loop, grpc_client, ): @@ -207,9 +223,15 @@ async def _run(): assert int(object_metadata.checksums.crc32c) == object_checksum buffer = BytesIO() - mrd = AsyncMultiRangeDownloader(grpc_client, _CROSS_REGION_BUCKET, object_name) + mrd = AsyncMultiRangeDownloader( + grpc_client, + _CROSS_REGION_BUCKET, + object_name, + multi_stream_config=multi_stream_config, + ) + expected_retries = multi_stream_config.min_streams if multi_stream_config else 1 async with mrd: - assert mrd._open_retries == 1 + assert mrd._open_retries == expected_retries # (0, 0) means read the whole object await mrd.download_ranges([(0, 0, buffer)]) assert mrd.persisted_size == object_size @@ -226,6 +248,7 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) @pytest.mark.parametrize( "object_size", [ @@ -243,6 +266,7 @@ def test_basic_wrd( blobs_to_delete, attempt_direct_path, object_size, + multi_stream_config, event_loop, grpc_clients, ): @@ -266,7 +290,10 @@ async def _run(): buffer = BytesIO() async with AsyncMultiRangeDownloader( - grpc_client, _BUCKET_UNDER_TEST, object_name + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, ) as mrd: # (0, 0) means read the whole object await mrd.download_ranges([(0, 0, buffer)]) @@ -284,6 +311,7 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) @pytest.mark.parametrize( "object_size", [ @@ -293,7 +321,12 @@ async def _run(): ], ) def test_basic_wrd_in_slices( - storage_client, blobs_to_delete, object_size, event_loop, grpc_client + storage_client, + blobs_to_delete, + object_size, + multi_stream_config, + event_loop, + grpc_client, ): object_name = f"test_basic_wrd-{str(uuid.uuid4())}" @@ -316,7 +349,12 @@ async def _run(): assert object_metadata.size == object_size assert int(object_metadata.checksums.crc32c) == object_checksum - mrd = AsyncMultiRangeDownloader(grpc_client, _BUCKET_UNDER_TEST, object_name) + mrd = AsyncMultiRangeDownloader( + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, + ) buffer = BytesIO() await mrd.open() # (0, 0) means read the whole object @@ -336,6 +374,7 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) @pytest.mark.parametrize( "flush_interval", [ @@ -349,6 +388,7 @@ def test_wrd_with_non_default_flush_interval( storage_client, blobs_to_delete, flush_interval, + multi_stream_config, event_loop, grpc_client, ): @@ -375,7 +415,12 @@ async def _run(): assert object_metadata.size == object_size assert int(object_metadata.checksums.crc32c) == object_checksum - mrd = AsyncMultiRangeDownloader(grpc_client, _BUCKET_UNDER_TEST, object_name) + mrd = AsyncMultiRangeDownloader( + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, + ) buffer = BytesIO() await mrd.open() # (0, 0) means read the whole object @@ -537,8 +582,13 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) def test_read_unfinalized_appendable_object( - storage_client, blobs_to_delete, event_loop, grpc_client_direct + storage_client, + blobs_to_delete, + multi_stream_config, + event_loop, + grpc_client_direct, ): object_name = f"read_unfinalized_appendable_object-{str(uuid.uuid4())[:4]}" @@ -554,7 +604,12 @@ async def _run(): await writer.append(_BYTES_TO_UPLOAD) await writer.flush() - mrd = AsyncMultiRangeDownloader(grpc_client, _BUCKET_UNDER_TEST, object_name) + mrd = AsyncMultiRangeDownloader( + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, + ) buffer = BytesIO() await mrd.open() assert mrd.persisted_size == len(_BYTES_TO_UPLOAD) @@ -575,7 +630,8 @@ async def _run(): @pytest.mark.skip(reason="Flaky test b/478129078") -def test_mrd_open_with_read_handle(event_loop, grpc_client_direct): +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) +def test_mrd_open_with_read_handle(event_loop, grpc_client_direct, multi_stream_config): object_name = f"test_read_handl-{str(uuid.uuid4())[:4]}" async def _run(): @@ -590,7 +646,10 @@ async def _run(): await writer.close() mrd = AsyncMultiRangeDownloader( - grpc_client_direct, _BUCKET_UNDER_TEST, object_name + grpc_client_direct, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, ) await mrd.open() read_handle = mrd.read_handle @@ -598,7 +657,11 @@ async def _run(): # Open a new MRD using the `read_handle` obtained above new_mrd = AsyncMultiRangeDownloader( - grpc_client_direct, _BUCKET_UNDER_TEST, object_name, read_handle=read_handle + grpc_client_direct, + _BUCKET_UNDER_TEST, + object_name, + read_handle=read_handle, + multi_stream_config=multi_stream_config, ) await new_mrd.open() # persisted_size not set when opened with read_handle @@ -614,7 +677,10 @@ async def _run(): event_loop.run_until_complete(_run()) -def test_mrd_open_with_read_handle_over_cloud_path(event_loop, grpc_client): +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) +def test_mrd_open_with_read_handle_over_cloud_path( + event_loop, grpc_client, multi_stream_config +): object_name = f"test_read_handl-{str(uuid.uuid4())[:4]}" async def _run(): @@ -628,14 +694,23 @@ async def _run(): await writer.append(_BYTES_TO_UPLOAD) await writer.close() - mrd = AsyncMultiRangeDownloader(grpc_client, _BUCKET_UNDER_TEST, object_name) + mrd = AsyncMultiRangeDownloader( + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, + ) await mrd.open() read_handle = mrd.read_handle await mrd.close() # Open a new MRD using the `read_handle` obtained above new_mrd = AsyncMultiRangeDownloader( - grpc_client, _BUCKET_UNDER_TEST, object_name, read_handle=read_handle + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + read_handle=read_handle, + multi_stream_config=multi_stream_config, ) await new_mrd.open() # persisted_size is set regardless of whether we use read_handle or not @@ -652,8 +727,13 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) def test_wrd_open_with_write_handle( - event_loop, grpc_client_direct, storage_client, blobs_to_delete + event_loop, + grpc_client_direct, + storage_client, + blobs_to_delete, + multi_stream_config, ): object_name = f"test_write_handl-{str(uuid.uuid4())[:4]}" @@ -690,7 +770,10 @@ async def _run(): # 4. Verify the data was written correctly by reading it back mrd = AsyncMultiRangeDownloader( - grpc_client_direct, _BUCKET_UNDER_TEST, object_name + grpc_client_direct, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, ) buffer = BytesIO() await mrd.open() @@ -710,8 +793,13 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) def test_read_unfinalized_appendable_object_with_generation( - storage_client, blobs_to_delete, event_loop, grpc_client_direct + storage_client, + blobs_to_delete, + multi_stream_config, + event_loop, + grpc_client_direct, ): object_name = f"read_unfinalized_appendable_object-{str(uuid.uuid4())[:4]}" grpc_client = grpc_client_direct @@ -720,7 +808,11 @@ async def _run(): async def _read_and_verify(expected_content, generation=None): """Helper to read object content and verify against expected.""" mrd = AsyncMultiRangeDownloader( - grpc_client, _BUCKET_UNDER_TEST, object_name, generation + grpc_client, + _BUCKET_UNDER_TEST, + object_name, + generation, + multi_stream_config=multi_stream_config, ) buffer = BytesIO() await mrd.open() @@ -1077,6 +1169,7 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) @pytest.mark.parametrize( "ranges_desc, chunk_ranges", [ @@ -1093,6 +1186,7 @@ def test_mrd_concurrent_download( grpc_client_direct, ranges_desc, chunk_ranges, + multi_stream_config, ): """ Test that mrd can handle concurrent `download_ranges` calls correctly. @@ -1116,7 +1210,10 @@ async def _run(): await writer.close(finalize_on_close=True) async with AsyncMultiRangeDownloader( - grpc_client_direct, _BUCKET_UNDER_TEST, object_name + grpc_client_direct, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, ) as mrd: tasks = [] ranges_to_fetch = [] @@ -1159,8 +1256,13 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) def test_mrd_concurrent_download_cancellation( - storage_client, blobs_to_delete, event_loop, grpc_client_direct + storage_client, + blobs_to_delete, + multi_stream_config, + event_loop, + grpc_client_direct, ): """ Test task cancellation / abort mid-stream. @@ -1184,7 +1286,10 @@ async def _run(): await writer.close(finalize_on_close=True) async with AsyncMultiRangeDownloader( - grpc_client_direct, _BUCKET_UNDER_TEST, object_name + grpc_client_direct, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, ) as mrd: tasks = [] num_chunks = 40 @@ -1225,8 +1330,13 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) def test_mrd_concurrent_download_out_of_bounds( - storage_client, blobs_to_delete, event_loop, grpc_client_direct + storage_client, + blobs_to_delete, + multi_stream_config, + event_loop, + grpc_client_direct, ): """ Test out-of-bounds & edge ranges concurrent with valid requests. @@ -1250,7 +1360,10 @@ async def _run(): await writer.close(finalize_on_close=True) async with AsyncMultiRangeDownloader( - grpc_client_direct, _BUCKET_UNDER_TEST, object_name + grpc_client_direct, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, ) as mrd: valid_buffer = BytesIO() valid_task = asyncio.create_task( @@ -1279,6 +1392,7 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) @pytest.mark.parametrize( "read_start, read_length, enable_checksum", [ @@ -1295,6 +1409,7 @@ def test_mrd_checksum_validation( read_start, read_length, enable_checksum, + multi_stream_config, ): """ Tests full downloads with specified offset, length, and enable_checksum toggle on finalized objects. @@ -1316,7 +1431,10 @@ async def _run(): await writer.close(finalize_on_close=True) async with AsyncMultiRangeDownloader( - grpc_client_direct, _BUCKET_UNDER_TEST, object_name + grpc_client_direct, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, ) as mrd: buffer = BytesIO() await mrd.download_ranges( @@ -1334,8 +1452,13 @@ async def _run(): event_loop.run_until_complete(_run()) +@pytest.mark.parametrize("multi_stream_config", _MULTI_STREAM_CONFIGS) def test_mrd_checksum_unfinalized_appendable_skipped( - storage_client, blobs_to_delete, event_loop, grpc_client_direct + storage_client, + blobs_to_delete, + multi_stream_config, + event_loop, + grpc_client_direct, ): """ Verifies that live, unfinalized appendable objects skip the full-object checksum check @@ -1356,7 +1479,10 @@ async def _run(): # Download the unfinalized appendable object with enable_checksum=True async with AsyncMultiRangeDownloader( - grpc_client_direct, _BUCKET_UNDER_TEST, object_name + grpc_client_direct, + _BUCKET_UNDER_TEST, + object_name, + multi_stream_config=multi_stream_config, ) as mrd: buffer = BytesIO() # Since it's unfinalized, it should skip the checksum check without raising 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 8bda5c6d8d8b..74befb8eda2e 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 @@ -81,6 +81,42 @@ def _reprobe_crc32c_memoryview_support(monkeypatch): 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): @@ -1525,3 +1561,1073 @@ 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, + ) + + # Arrange: primary and 2nd stream succeed, 3rd stream fails during open() + mock_client = self._make_mock_client() + primary_stream = self._make_mock_stream() + second_stream = self._make_mock_stream() + failing_stream = self._make_mock_stream( + open_side_effect=exceptions.PermissionDenied("Stream failed") + ) + mock_cls_stream.side_effect = [ + primary_stream, + second_stream, + failing_stream, + ] + + config = MultiStreamConfig(min_streams=3, max_streams=4) + mrd = AsyncMultiRangeDownloader( + mock_client, "b", "o", multi_stream_config=config + ) + + # Act & Assert + with pytest.raises(exceptions.PermissionDenied, match="Stream failed"): + await mrd.open() + + assert mrd._is_stream_open is False + assert mrd._pool is None + assert mrd.read_obj_str is None + assert mrd._multiplexer is None + assert mrd._primary_stream is None + primary_stream.close.assert_awaited_once() + second_stream.close.assert_awaited_once() + + @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() + + @pytest.mark.parametrize( + "config_kwargs", + [ + pytest.param(None, id="default-no-config"), + pytest.param( + {"min_streams": 1, "max_streams": 1}, id="explicit-max-streams-one" + ), + ], + ) + @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, config_kwargs + ): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + MultiStreamConfig, + ) + + # Arrange + mock_client = self._make_mock_client() + mock_cls_stream.return_value = self._make_mock_stream() + config = MultiStreamConfig(**config_kwargs) if config_kwargs else None + + # Act + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, "b", "o", multi_stream_config=config + ) + + # Assert + assert mrd._pool is None + 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 + task = next(iter(pool._background_tasks)) + + # 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 and awaited by close() + with pytest.raises(asyncio.CancelledError): + await task + assert task.cancelled() + + # 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, + ] + + from google.api_core.retry_async import AsyncRetry + + on_err_mock = mock.MagicMock() + custom_retry = AsyncRetry(initial=0.01, on_error=on_err_mock) + config = MultiStreamConfig(min_streams=2, max_streams=4) + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, + "b", + "o", + retry_policy=custom_retry, + multi_stream_config=config, + ) + + # Verify custom on_error was invoked for both primary and pooled stream redirects + assert on_err_mock.call_count == 2 + + # Verify s1_ok received routing_token=token-abc on retry and cleared mrd._routing_token + assert ("x-goog-request-params", "routing_token=token-abc") in ( + s1_ok.open.call_args.kwargs["metadata"] + ) + assert mrd._routing_token is None + + # Verify s2_failed started without a stale routing_token, then s2_ok retried with routing_token=token-xyz + assert s2_failed.open.call_args.kwargs["metadata"] is None + assert ("x-goog-request-params", "routing_token=token-xyz") in ( + s2_ok.open.call_args.kwargs["metadata"] + ) + assert mrd.read_handle == b"h2" + assert mrd._routing_token is None + + 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())]) + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_reopen_non_primary_stream_preserves_primary_stream_and_metadata( + self, mock_cls_stream + ): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + MultiStreamConfig, + ) + + mock_client = self._make_mock_client() + primary_stream = self._make_mock_stream(read_handle=b"h1", is_finalized=True) + non_primary_initial = self._make_mock_stream( + read_handle=b"h2", is_finalized=False + ) + non_primary_initial.object_metadata = None + non_primary_initial.full_obj_server_crc32c = None + + replacement_stream = self._make_mock_stream( + read_handle=b"h3", is_finalized=False + ) + # Reopening with a read_handle may not return response.metadata + replacement_stream.object_metadata = None + replacement_stream.full_obj_server_crc32c = None + + mock_cls_stream.side_effect = [ + primary_stream, + non_primary_initial, + replacement_stream, + ] + + config = MultiStreamConfig(min_streams=2, max_streams=4) + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, "b", "o", multi_stream_config=config + ) + assert mrd.read_obj_str is primary_stream + assert mrd._primary_stream.stream is primary_stream + assert mrd.is_finalized is True + assert mrd.full_obj_server_crc32c == 12345 + + non_primary_pooled = mrd._pool.streams[1] + assert non_primary_pooled is not mrd._primary_stream + assert non_primary_pooled.stream is non_primary_initial + + # Reopen the non-primary stream via _create_stream_factory + factory = mrd._create_stream_factory( + state={"read_handle": b"h2", "routing_token": "tok"}, + metadata=None, + pooled_stream=non_primary_pooled, + ) + reopened = await factory() + + assert reopened is replacement_stream + assert non_primary_pooled.stream is replacement_stream + assert mrd.read_obj_str is mrd._primary_stream.stream + assert mrd.read_obj_str is primary_stream + assert mrd.is_finalized is True + assert mrd.full_obj_server_crc32c == 12345 + + await mrd.close() + + @pytest.mark.parametrize( + ("kwargs", "expected_msg"), + [ + pytest.param( + {"min_streams": 0}, + "min_streams must be >= 1", + id="min-streams-below-one", + ), + pytest.param( + {"min_streams": 4, "max_streams": 2}, + "max_streams must be >= min_streams", + id="max-below-min-streams", + ), + pytest.param( + {"target_io_depth": 0}, + "target_io_depth must be > 0", + id="non-positive-target-io-depth", + ), + pytest.param( + {"target_bytes": -1}, + "target_bytes must be > 0", + id="non-positive-target-bytes", + ), + ], + ) + def test_multi_stream_config_validation(self, kwargs, expected_msg): + from google.cloud.storage.asyncio._stream_pool import MultiStreamConfig + + with pytest.raises(ValueError, match=expected_msg): + MultiStreamConfig(**kwargs) + + @pytest.mark.asyncio + async def test_stream_pool_empty_raises_and_scale_up_error_logged(self): + from google.cloud.storage.asyncio._stream_pool import ( + MultiStreamConfig, + _StreamPool, + ) + + async def failing_factory(): + raise RuntimeError("open failed") + + pool = _StreamPool( + stream_factory=failing_factory, + multi_stream_config=MultiStreamConfig( + min_streams=1, max_streams=2, target_io_depth=1, target_bytes=10 + ), + ) + + with pytest.raises(ValueError, match="No streams available in stream pool"): + await pool.distribute_ranges([(0, 10, BytesIO())]) + + await pool.add_stream(self._make_dummy_pooled_stream()) + await pool.distribute_ranges([(0, 20, BytesIO())]) + await asyncio.sleep(0.01) + + assert pool._pending_scale_ups == 0 + assert len(pool.streams) == 1 + await pool.close() + + @pytest.mark.asyncio + async def test_stream_pool_scale_up_cancellation_closes_new_stream(self): + from google.cloud.storage.asyncio._stream_pool import ( + MultiStreamConfig, + _StreamPool, + ) + + # Arrange + factory_returning = asyncio.Event() + new_pooled_stream = self._make_dummy_pooled_stream() + new_pooled_stream.close = AsyncMock() + + async def factory(): + factory_returning.set() + await asyncio.sleep(0) + return new_pooled_stream + + pool = _StreamPool( + stream_factory=factory, + multi_stream_config=MultiStreamConfig( + min_streams=1, + max_streams=2, + target_io_depth=1, + target_bytes=10, + ), + ) + await pool.add_stream(self._make_dummy_pooled_stream()) + + # Act: hold pool._lock while factory() finishes, then cancel the pending scale-up task + await pool.distribute_ranges([(0, 20, BytesIO())]) + await factory_returning.wait() + async with pool._lock: + await asyncio.sleep(0) + scale_task = next(iter(pool._background_tasks)) + scale_task.cancel() + + await pool.close() + + # Assert + new_pooled_stream.close.assert_awaited_once() + assert pool._pending_scale_ups == 0 + + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_download_ranges_cancels_siblings_on_failure_and_data_corruption( + self, mock_cls_stream + ): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + MultiStreamConfig, + ) + + # Arrange: s1 (base), s2 (fast scale-up), s3 (slow scale-up) + 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") + + s2_overflow_started = asyncio.Event() + scale_up_gate = asyncio.Event() + + async def slow_s3_open(*args, **kwargs): + await scale_up_gate.wait() + + s3.open = AsyncMock(side_effect=slow_s3_open) + mock_cls_stream.side_effect = [s1, s2, s3] + + 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 + ) + + overflow_batch_cancelled = asyncio.Event() + + async def side_effect_dl(pooled_stream, batch, *args, **kwargs): + if pooled_stream.stream is s1: + await s2_overflow_started.wait() + await mrd.close() + raise DataCorruption(None, "checksum mismatch") + else: + s2_overflow_started.set() + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + overflow_batch_cancelled.set() + raise + + # Act & Assert + with mock.patch.object( + mrd, "_download_ranges_on_stream", side_effect=side_effect_dl + ): + ranges = [(i * 50, 50, BytesIO()) for i in range(3)] + with pytest.raises(DataCorruption, match="checksum mismatch"): + await mrd.download_ranges(ranges) + + assert overflow_batch_cancelled.is_set() + assert mrd.is_stream_open is False + assert mrd._pool is None + + @pytest.mark.asyncio + async def test_stream_pool_contiguous_distribution_when_saturated(self): + from google.cloud.storage.asyncio._stream_pool import ( + MultiStreamConfig, + _StreamPool, + ) + + # Arrange + pool = _StreamPool( + stream_factory=lambda: self._make_dummy_pooled_stream(), + multi_stream_config=MultiStreamConfig( + min_streams=2, + max_streams=2, + 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) + ranges = [(i * 50, 50, BytesIO()) for i in range(8)] + + # Act + assignments, overflow, scale_tasks = await pool.distribute_ranges(ranges) + + # Assert: contiguous slices without range-by-range interleaving + assert overflow == [] + assert scale_tasks == [] + assert assignments[s1] == ranges[:4] + assert assignments[s2] == ranges[4:] + + await pool.close() + + @pytest.mark.parametrize( + ("persisted_size", "offset", "expected_effective_bytes"), + [ + pytest.param(1000, 200, 800, id="known-persisted-size"), + pytest.param(None, 0, 500, id="unknown-persisted-size-fallback"), + ], + ) + @mock.patch( + "google.cloud.storage.asyncio.async_multi_range_downloader._AsyncReadObjectStream" + ) + @pytest.mark.asyncio + async def test_zero_length_range_load_accounting_and_release( + self, mock_cls_stream, persisted_size, offset, expected_effective_bytes + ): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + MultiStreamConfig, + ) + + # Arrange + mock_client = self._make_mock_client() + s1 = self._make_mock_stream(read_handle=b"h1", persisted_size=persisted_size) + s2 = self._make_mock_stream(read_handle=b"h2", persisted_size=persisted_size) + mock_cls_stream.side_effect = [s1, s2] + + config = MultiStreamConfig( + min_streams=1, + max_streams=2, + target_io_depth=2, + target_bytes=500, + ) + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, "b", "o", multi_stream_config=config + ) + recorded_loads = [] + + async def inspect_load(pooled_stream, batch, *args, **kwargs): + recorded_loads.append( + (pooled_stream.pending_ranges, pooled_stream.pending_bytes) + ) + + # Act + with mock.patch.object( + mrd, "_download_ranges_on_stream", side_effect=inspect_load + ): + await mrd.download_ranges([(offset, 0, BytesIO())]) + + # Assert + assert recorded_loads == [(1, expected_effective_bytes)] + 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_download_overflow_incremental_dispatch_and_inflight_reuse( + 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") + + s2_dispatched = asyncio.Event() + + async def s3_waits_for_s2_dispatch(*args, **kwargs): + # Prove s2 receives and starts downloading overflow BEFORE s3 finishes opening! + await asyncio.wait_for(s2_dispatched.wait(), timeout=1.0) + + s3.open = AsyncMock(side_effect=s3_waits_for_s2_dispatch) + mock_cls_stream.side_effect = [s1, s2, s3] + + config = MultiStreamConfig( + min_streams=1, + max_streams=3, + target_io_depth=2, + target_bytes=100, + ) + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, "b", "o", multi_stream_config=config + ) + + dispatched = [] + + async def record_dl(pooled_stream, batch, *args, **kwargs): + dispatched.append((pooled_stream, list(batch))) + if pooled_stream.stream is s2: + s2_dispatched.set() + + with mock.patch.object( + mrd, "_download_ranges_on_stream", side_effect=record_dl + ): + # 6 ranges (2 per stream across s1, s2, s3). + # s2 opens immediately, so overflow ranges r2..r3 are dispatched to s2 + # without head-of-line blocking on s3.open! + ranges = [(i * 50, 50, BytesIO()) for i in range(6)] + await mrd.download_ranges(ranges) + + assert len(mrd._pool.streams) == 3 + assert len(dispatched) == 3 + 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_concurrent_multi_stream_stress_reliability(self, mock_cls_stream): + from google.cloud.storage.asyncio.async_multi_range_downloader import ( + MultiStreamConfig, + ) + + mock_client = self._make_mock_client() + + def make_stream(*args, **kwargs): + s = self._make_mock_stream(persisted_size=10_000) + + async def delayed_open(*a, **kw): + await asyncio.sleep(0.002) + + s.open = AsyncMock(side_effect=delayed_open) + return s + + mock_cls_stream.side_effect = make_stream + + config = MultiStreamConfig( + min_streams=2, + max_streams=6, + target_io_depth=4, + target_bytes=200, + ) + mrd = await AsyncMultiRangeDownloader.create_mrd( + mock_client, "b", "o", multi_stream_config=config + ) + + completed_ranges = 0 + + async def simulated_dl(pooled_stream, batch, *args, **kwargs): + nonlocal completed_ranges + await asyncio.sleep(0.001) + for offset, length, buf in batch: + size = length if length > 0 else (10_000 - offset) + buf.write(b"x" * min(size, 16)) + completed_ranges += len(batch) + + with mock.patch.object( + mrd, "_download_ranges_on_stream", side_effect=simulated_dl + ): + # Run 12 concurrent download_ranges calls with varying batch sizes and zero-length ranges + concurrent_calls = [] + expected_ranges = 0 + for call_idx in range(12): + batch = [ + (i * 50, 50 if (i + call_idx) % 7 != 0 else 0, BytesIO()) + for i in range(1, (call_idx % 8) + 2) + ] + expected_ranges += len(batch) + concurrent_calls.append(mrd.download_ranges(batch)) + + await asyncio.gather(*concurrent_calls) + + assert completed_ranges == expected_ranges + assert 2 < len(mrd._pool.streams) <= 6 + for ps in mrd._pool.streams: + assert ps.pending_ranges == 0 + assert ps.pending_bytes == 0 + + await mrd.close()