feat(experimental): Add bidi stream retry manager

Author: Pulkit0110

google/cloud/storage/_experimental/asyncio/retry/bidi_stream_retry_manager.py

@@ -0,0 +1,67 @@
+import asyncio
+from typing import Any, AsyncIterator, Callable
+
+from google.api_core import exceptions
+from google.cloud.storage._experimental.asyncio.retry.base_strategy import (
+    _BaseResumptionStrategy,
+)
+
+
+class _BidiStreamRetryManager:
+    """Manages the generic retry loop for a bidi streaming operation."""
+
+    def __init__(
+        self,
+        strategy: _BaseResumptionStrategy,
+        stream_opener: Callable[..., AsyncIterator[Any]],
+        retry_policy,
+    ):
+        """Initializes the retry manager."""
+        self._strategy = strategy
+        self._stream_opener = stream_opener
+        self._retry_policy = retry_policy
+
+    async def execute(self, initial_state: Any):
+        """
+        Executes the bidi operation with the configured retry policy.
+
+        This method implements a manual retry loop that provides the necessary
+        control points to manage state between attempts, which is not possible
+        with a simple retry decorator.
+        """
+        state = initial_state
+        retry_policy = self._retry_policy
+
+        while True:
+            try:
+                # 1. Generate requests based on the current state.
+                requests = self._strategy.generate_requests(state)
+
+                # 2. Open and consume the stream.
+                stream = self._stream_opener(requests, state)
+                async for response in stream:
+                    self._strategy.update_state_from_response(response, state)
+
+                # 3. If the stream completes without error, exit the loop.
+                return
+
+            except Exception as e:
+                # 4. If an error occurs, check if it's retriable.
+                if not retry_policy.predicate(e):
+                    # If not retriable, fail fast.
+                    raise
+
+                # 5. If retriable, allow the strategy to recover state.
+                #    This is where routing tokens are extracted or QueryWriteStatus is called.
+                try:
+                    await self._strategy.recover_state_on_failure(e, state)
+                except Exception as recovery_exc:
+                    # If state recovery itself fails, we must abort.
+                    raise exceptions.RetryError(
+                        "Failed to recover state after a transient error.",
+                        cause=recovery_exc,
+                    ) from recovery_exc
+
+                # 6. Use the policy to sleep and check for deadline expiration.
+                #    This will raise a RetryError if the deadline is exceeded.
+                await asyncio.sleep(await retry_policy.sleep(e))

tests/unit/asyncio/retry/test_bidi_stream_retry_manager.py

@@ -0,0 +1,111 @@
+import unittest
+from unittest import mock
+
+import pytest
+from google.api_core import exceptions
+from google.api_core.retry.retry_streaming_async import AsyncStreamingRetry
+
+from google.cloud.storage._experimental.asyncio.retry import manager
+from google.cloud.storage._experimental.asyncio.retry import strategy
+
+
+def _is_retriable(exc):
+    return isinstance(exc, exceptions.ServiceUnavailable)
+
+
+DEFAULT_TEST_RETRY = AsyncStreamingRetry(predicate=_is_retriable, deadline=1)
+
+
+class TestBidiStreamRetryManager(unittest.IsolatedAsyncioTestCase):
+    async def test_execute_success_on_first_try(self):
+        """Verify the manager correctly handles a stream that succeeds immediately."""
+        mock_strategy = mock.AsyncMock(spec=strategy._BaseResumptionStrategy)
+
+        async def mock_stream_opener(*args, **kwargs):
+            yield "response_1"
+
+        retry_manager = manager._BidiStreamRetryManager(
+            strategy=mock_strategy,
+            stream_opener=mock_stream_opener,
+            retry_policy=DEFAULT_TEST_RETRY,
+        )
+
+        await retry_manager.execute(initial_state={})
+
+        mock_strategy.generate_requests.assert_called_once()
+        mock_strategy.update_state_from_response.assert_called_once_with(
+            "response_1", {}
+        )
+        mock_strategy.recover_state_on_failure.assert_not_called()
+
+    async def test_execute_retries_and_succeeds(self):
+        """Verify the manager retries on a transient error and then succeeds."""
+        mock_strategy = mock.AsyncMock(spec=strategy._BaseResumptionStrategy)
+
+        attempt_count = 0
+
+        async def mock_stream_opener(*args, **kwargs):
+            nonlocal attempt_count
+            attempt_count += 1
+            if attempt_count == 1:
+                raise exceptions.ServiceUnavailable("Service is down")
+            else:
+                yield "response_2"
+
+        retry_manager = manager._BidiStreamRetryManager(
+            strategy=mock_strategy,
+            stream_opener=mock_stream_opener,
+            retry_policy=AsyncStreamingRetry(predicate=_is_retriable, initial=0.01),
+        )
+
+        await retry_manager.execute(initial_state={})
+
+        self.assertEqual(attempt_count, 2)
+        self.assertEqual(mock_strategy.generate_requests.call_count, 2)
+        mock_strategy.recover_state_on_failure.assert_called_once()
+        mock_strategy.update_state_from_response.assert_called_once_with(
+            "response_2", {}
+        )
+
+    async def test_execute_fails_after_deadline_exceeded(self):
+        """Verify the manager raises RetryError if the deadline is exceeded."""
+        mock_strategy = mock.AsyncMock(spec=strategy._BaseResumptionStrategy)
+
+        async def mock_stream_opener(*args, **kwargs):
+            raise exceptions.ServiceUnavailable("Service is always down")
+
+        # Use a very short deadline to make the test fast.
+        fast_retry = AsyncStreamingRetry(
+            predicate=_is_retriable, deadline=0.1, initial=0.05
+        )
+        retry_manager = manager._BidiStreamRetryManager(
+            strategy=mock_strategy,
+            stream_opener=mock_stream_opener,
+            retry_policy=fast_retry,
+        )
+
+        with pytest.raises(exceptions.RetryError, match="Deadline of 0.1s exceeded"):
+            await retry_manager.execute(initial_state={})
+
+        # Verify it attempted to recover state after each failure.
+        self.assertGreater(mock_strategy.recover_state_on_failure.call_count, 1)
+
+    async def test_execute_fails_immediately_on_non_retriable_error(self):
+        """Verify the manager aborts immediately on a non-retriable error."""
+        mock_strategy = mock.AsyncMock(spec=strategy._BaseResumptionStrategy)
+
+        async def mock_stream_opener(*args, **kwargs):
+            raise exceptions.PermissionDenied("Auth error")
+
+        retry_manager = manager._BidiStreamRetryManager(
+            strategy=mock_strategy,
+            stream_opener=mock_stream_opener,
+            retry_policy=DEFAULT_TEST_RETRY,
+        )
+
+        with pytest.raises(exceptions.PermissionDenied):
+            await retry_manager.execute(initial_state={})
+
+        # Verify that it did not try to recover or update state.
+        mock_strategy.recover_state_on_failure.assert_not_called()
+        mock_strategy.update_state_from_response.assert_not_called()