async_multi_range_downloader.py

# Copyright 2025 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 logging
from google.api_core import exceptions
from google.api_core.retry_async import AsyncRetry
from google.cloud.storage.asyncio.retry._helpers import _handle_redirect
from google.rpc import status_pb2

from typing import List, Optional, Tuple, Any, Dict

from ._utils import raise_if_no_fast_crc32c
from google.cloud.storage.asyncio.async_read_object_stream import (
    _AsyncReadObjectStream,
)
from google.cloud.storage.asyncio.async_grpc_client import (
    AsyncGrpcClient,
)
from google.cloud.storage.asyncio.retry.bidi_stream_retry_manager import (
    _BidiStreamRetryManager,
)
from google.cloud.storage.asyncio.retry.reads_resumption_strategy import (
    _ReadResumptionStrategy,
    _DownloadState,
)

from io import BytesIO
from google.cloud import _storage_v2
from google.cloud.storage._helpers import generate_random_56_bit_integer

_MAX_READ_RANGES_PER_BIDI_READ_REQUEST = 100
_BIDI_READ_REDIRECTED_TYPE_URL = (
    "type.googleapis.com/google.storage.v2.BidiReadObjectRedirectedError"
)

logger = logging.getLogger(__name__)


def _is_read_retryable(exc):
    """Predicate to determine if a read operation should be retried."""
    if isinstance(
        exc,
        (
            exceptions.InternalServerError,
            exceptions.ServiceUnavailable,
            exceptions.DeadlineExceeded,
            exceptions.TooManyRequests,
        ),
    ):
        return True

    if not isinstance(exc, exceptions.Aborted) or not exc.errors:
        return False

    try:
        grpc_error = exc.errors[0]
        trailers = grpc_error.trailing_metadata()
        if not trailers:
            return False

        status_details_bin = next(
            (v for k, v in trailers if k == "grpc-status-details-bin"), None
        )

        if not status_details_bin:
            return False

        status_proto = status_pb2.Status()
        status_proto.ParseFromString(status_details_bin)
        return any(
            detail.type_url == _BIDI_READ_REDIRECTED_TYPE_URL
            for detail in status_proto.details
        )
    except Exception as e:
        logger.error(f"Error parsing status_details_bin: {e}")
        return False


class AsyncMultiRangeDownloader:
    """Provides an interface for downloading multiple ranges of a GCS ``Object``
    concurrently.

    Example usage:

    .. code-block:: python

        client = AsyncGrpcClient()
        mrd = await AsyncMultiRangeDownloader.create_mrd(
            client, bucket_name="chandrasiri-rs", object_name="test_open9"
        )
        my_buff1 = open('my_fav_file.txt', 'wb')
        my_buff2 = BytesIO()
        my_buff3 = BytesIO()
        my_buff4 = any_object_which_provides_BytesIO_like_interface()
        await mrd.download_ranges(
            [
                # (start_byte, bytes_to_read, writeable_buffer)
                (0, 100, my_buff1),
                (100, 20, my_buff2),
                (200, 123, my_buff3),
                (300, 789, my_buff4),
            ]
        )

        # verify data in buffers...
        assert my_buff2.getbuffer().nbytes == 20


    """

    @classmethod
    async def create_mrd(
        cls,
        client: AsyncGrpcClient,
        bucket_name: str,
        object_name: str,
        generation: Optional[int] = None,
        read_handle: Optional[_storage_v2.BidiReadHandle] = None,
        retry_policy: Optional[AsyncRetry] = None,
        metadata: Optional[List[Tuple[str, str]]] = None,
        **kwargs,
    ) -> AsyncMultiRangeDownloader:
        """Initializes a MultiRangeDownloader and opens the underlying bidi-gRPC
        object for reading.

        :type client: :class:`~google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient`
        :param client: The asynchronous client to use for making API requests.

        :type bucket_name: str
        :param bucket_name: The name of the bucket containing the object.

        :type object_name: str
        :param object_name: The name of the object to be read.

        :type generation: int
        :param generation: (Optional) If present, selects a specific
                                  revision of this object.

        :type read_handle: _storage_v2.BidiReadHandle
        :param read_handle: (Optional) An existing handle for reading the object.
                            If provided, opening the bidi-gRPC connection will be faster.

        :type retry_policy: :class:`~google.api_core.retry_async.AsyncRetry`
        :param retry_policy: (Optional) The retry policy to use for the ``open`` operation.

        :type metadata: List[Tuple[str, str]]
        :param metadata: (Optional) The metadata to be sent with the ``open`` request.

        :rtype: :class:`~google.cloud.storage.asyncio.async_multi_range_downloader.AsyncMultiRangeDownloader`
        :returns: An initialized AsyncMultiRangeDownloader instance for reading.
        """
        mrd = cls(
            client,
            bucket_name,
            object_name,
            generation=generation,
            read_handle=read_handle,
            **kwargs,
        )
        await mrd.open(retry_policy=retry_policy, metadata=metadata)
        return mrd

    def __init__(
        self,
        client: AsyncGrpcClient,
        bucket_name: str,
        object_name: str,
        generation: Optional[int] = None,
        read_handle: Optional[_storage_v2.BidiReadHandle] = None,
        **kwargs,
    ) -> None:
        """Constructor for AsyncMultiRangeDownloader, clients are not adviced to
         use it directly. Instead it's adviced 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.

        :type bucket_name: str
        :param bucket_name: The name of the bucket containing the object.

        :type object_name: str
        :param object_name: The name of the object to be read.

        :type generation: int
        :param generation: (Optional) If present, selects a specific revision of
                                  this object.

        :type read_handle: _storage_v2.BidiReadHandle
        :param read_handle: (Optional) An existing read handle.
        """
        if "generation_number" in kwargs:
            if generation is not None:
                raise TypeError(
                    "Cannot set both 'generation' and 'generation_number'. "
                    "Use 'generation' for new code."
                )
            logger.warning(
                "'generation_number' is deprecated and will be removed in a future "
                "major release. Please use 'generation' instead."
            )
            generation = kwargs.pop("generation_number")

        raise_if_no_fast_crc32c()

        self.client = client
        self.bucket_name = bucket_name
        self.object_name = object_name
        self.generation = generation
        self.read_handle: Optional[_storage_v2.BidiReadHandle] = read_handle
        self.read_obj_str: Optional[_AsyncReadObjectStream] = None
        self._is_stream_open: bool = False
        self._routing_token: Optional[str] = None
        self._read_id_to_writable_buffer_dict = {}
        self._read_id_to_download_ranges_id = {}
        self._download_ranges_id_to_pending_read_ids = {}
        self.persisted_size: Optional[int] = None  # updated after opening the stream
        self._open_retries: int = 0

    async def __aenter__(self):
        """Opens the underlying bidi-gRPC connection to read from the object."""
        await self.open()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Closes the underlying bidi-gRPC connection."""
        if self.is_stream_open:
            await self.close()

    def _on_open_error(self, exc):
        """Extracts routing token and read handle on redirect error during open."""
        logger.warning(f"Error occurred while opening MRD: {exc}")
        routing_token, read_handle = _handle_redirect(exc)
        if routing_token:
            self._routing_token = routing_token
        if read_handle:
            self.read_handle = read_handle

    async def open(
        self,
        retry_policy: Optional[AsyncRetry] = None,
        metadata: Optional[List[Tuple[str, str]]] = None,
    ) -> None:
        """Opens the bidi-gRPC connection to read from the object."""
        if self._is_stream_open:
            raise ValueError("Underlying bidi-gRPC stream is already open")

        if retry_policy is None:
            def on_error_wrapper(exc):
                self._open_retries += 1
                self._on_open_error(exc)

            retry_policy = AsyncRetry(
                predicate=_is_read_retryable, on_error=on_error_wrapper
            )
        else:
            original_on_error = retry_policy._on_error

            def combined_on_error(exc):
                self._open_retries += 1
                self._on_open_error(exc)
                if original_on_error:
                    original_on_error(exc)

            retry_policy = AsyncRetry(
                predicate=_is_read_retryable,
                initial=retry_policy._initial,
                maximum=retry_policy._maximum,
                multiplier=retry_policy._multiplier,
                deadline=retry_policy._deadline,
                on_error=combined_on_error,
            )

        async def _do_open():
            current_metadata = list(metadata) if metadata else []

            # Cleanup stream from previous failed attempt, if any.
            if self.read_obj_str:
                if self.read_obj_str.is_stream_open:
                    try:
                        await self.read_obj_str.close()
                    except exceptions.GoogleAPICallError as e:
                        logger.warning(
                            f"Failed to close existing stream during resumption: {e}"
                        )
                self.read_obj_str = None
                self._is_stream_open = False

            self.read_obj_str = _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,
            )

            if self._routing_token:
                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
            )

            if self.read_obj_str.generation_number:
                self.generation = self.read_obj_str.generation_number
            if self.read_obj_str.read_handle:
                self.read_handle = self.read_obj_str.read_handle
            if self.read_obj_str.persisted_size is not None:
                self.persisted_size = self.read_obj_str.persisted_size

            self._is_stream_open = True

        await retry_policy(_do_open)()

    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,
    ) -> None:
        """Downloads multiple byte ranges from the object into the buffers
        provided by user with automatic retries.

        :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: (Optional) An asyncio lock to synchronize sends and recvs
            on the underlying bidi-GRPC stream. This is required when multiple
            coroutines are calling this method concurrently.

            i.e. Example usage with multiple coroutines:

            ```
            lock = asyncio.Lock()
            task1 = asyncio.create_task(mrd.download_ranges(ranges1, lock))
            task2 = asyncio.create_task(mrd.download_ranges(ranges2, lock))
            await asyncio.gather(task1, task2)

            ```

            If user want to call this method serially from multiple coroutines,
            then providing a lock is not necessary.

            ```
            await mrd.download_ranges(ranges1)
            await mrd.download_ranges(ranges2)

            # ... some other code code...

            ```

        :type retry_policy: :class:`~google.api_core.retry_async.AsyncRetry`
        :param retry_policy: (Optional) The retry policy to use for the operation.

        :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 not self._is_stream_open:
            raise ValueError("Underlying bidi-gRPC stream is not open")

        if lock is None:
            lock = asyncio.Lock()

        if retry_policy is None:
            retry_policy = AsyncRetry(predicate=_is_read_retryable)

        # Initialize Global State for Retry Strategy
        download_states = {}
        for read_range in read_ranges:
            read_id = generate_random_56_bit_integer()
            download_states[read_id] = _DownloadState(
                initial_offset=read_range[0],
                initial_length=read_range[1],
                user_buffer=read_range[2],
            )

        initial_state = {
            "download_states": download_states,
            "read_handle": self.read_handle,
            "routing_token": None,
        }

        # Track attempts to manage stream reuse
        attempt_count = 0

        def send_ranges_and_get_bytes(
            requests: List[_storage_v2.ReadRange],
            state: Dict[str, Any],
            metadata: Optional[List[Tuple[str, str]]] = None,
        ):
            async def generator():
                nonlocal attempt_count
                attempt_count += 1

                if attempt_count > 1:
                    logger.info(
                        f"Resuming download (attempt {attempt_count}) for {len(requests)} ranges."
                    )

                async with lock:
                    current_handle = state.get("read_handle")
                    current_token = state.get("routing_token")

                    # We reopen if it's a redirect (token exists) OR if this is a retry
                    # (not first attempt). This prevents trying to send data on a dead
                    # stream from a previous failed attempt.
                    should_reopen = (
                        (attempt_count > 1)
                        or (current_token is not None)
                        or (metadata is not None)
                    )

                    if should_reopen:
                        if current_token:
                            logger.info(
                                f"Re-opening stream with routing token: {current_token}"
                            )

                        self.read_obj_str = _AsyncReadObjectStream(
                            client=self.client.grpc_client,
                            bucket_name=self.bucket_name,
                            object_name=self.object_name,
                            generation_number=self.generation,
                            read_handle=current_handle,
                        )

                        # Inject routing_token into metadata if present
                        current_metadata = list(metadata) if metadata else []
                        if current_token:
                            current_metadata.append(
                                (
                                    "x-goog-request-params",
                                    f"routing_token={current_token}",
                                )
                            )

                        await self.read_obj_str.open(
                            metadata=current_metadata if current_metadata else None
                        )
                        self._is_stream_open = True

                    pending_read_ids = {r.read_id for r in requests}

                    # Send 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]
                        await self.read_obj_str.send(
                            _storage_v2.BidiReadObjectRequest(read_ranges=batch)
                        )

                    while pending_read_ids:
                        response = await self.read_obj_str.recv()
                        if response is None:
                            break
                        if response.object_data_ranges:
                            for data_range in response.object_data_ranges:
                                if data_range.range_end:
                                    pending_read_ids.discard(
                                        data_range.read_range.read_id
                                    )
                        yield response

            return generator()

        strategy = _ReadResumptionStrategy()
        retry_manager = _BidiStreamRetryManager(
            strategy, lambda r, s: send_ranges_and_get_bytes(r, s, metadata=metadata)
        )

        await retry_manager.execute(initial_state, retry_policy)

        if initial_state.get("read_handle"):
            self.read_handle = initial_state["read_handle"]

    async def close(self):
        """
        Closes the underlying bidi-gRPC connection.
        """
        if not self._is_stream_open:
            raise ValueError("Underlying bidi-gRPC stream is not open")

        if self.read_obj_str:
            await self.read_obj_str.close()
        self.read_obj_str = None
        self._is_stream_open = False

    @property
    def is_stream_open(self) -> bool:
        return self._is_stream_open