minor change

Author: Pulkit0110

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

@@ -63,7 +63,7 @@ def update_state_from_response(self, response: Any, state: Any) -> None:
         pass
 
     @abc.abstractmethod
-    async def recover_state_on_failure(self, error: Exception, state: Any) -> None:
+    def recover_state_on_failure(self, error: Exception, state: Any) -> None:
         """Prepares the state for the next retry attempt after a failure.
 
         This method is called when a retriable gRPC error occurs. It is

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

@@ -33,7 +33,6 @@ def __init__(
         stream_opener: Callable[[Iterable[Any], Any], AsyncIterator[Any]],
     ):
         """Initializes the retry manager.
-
         Args:
             strategy: The strategy for managing the state of a specific
                 bidi operation (e.g., reads or writes).
@@ -45,28 +44,35 @@ def __init__(
     async def execute(self, initial_state: Any, retry_policy: "AsyncRetry"):
         """
         Executes the bidi operation with the configured retry policy.
-
         Args:
             initial_state: An object containing all state for the operation.
             retry_policy: The `google.api_core.retry_async.AsyncRetry` object to
                 govern the retry behavior for this specific operation.
         """
         state = initial_state
 
+        def on_error(e: Exception):
+            """The single point of recovery logic."""
+            self._strategy.recover_state_on_failure(e, state)
+
         async def attempt():
+            """The core operation to be retried."""
             requests = self._strategy.generate_requests(state)
             stream = self._stream_opener(requests, state)
-            try:
-                async for response in stream:
-                    self._strategy.update_state_from_response(response, state)
-                return  # Successful completion of the stream.
-            except Exception as e:
-                if retry_policy._predicate(e):
-                    await self._strategy.recover_state_on_failure(e, state)
-                raise e
+            async for response in stream:
+                self._strategy.update_state_from_response(response, state)
+
+        # Correctly create a new retry instance with the on_error handler.
+        retry_with_error_handler = type(retry_policy)(
+            predicate=retry_policy._predicate,
+            initial=retry_policy._initial,
+            maximum=retry_policy._maximum,
+            multiplier=retry_policy._multiplier,
+            deadline=retry_policy._deadline,
+            on_error=on_error,
+        )
 
-        # Wrap the attempt function with the retry policy.
-        wrapped_attempt = retry_policy(attempt)
+        wrapped_attempt = retry_with_error_handler(attempt)
 
         # Execute the operation with retry.
         await wrapped_attempt()

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

@@ -90,7 +90,7 @@ def update_state_from_response(
                         response, f"Byte count mismatch for read_id {read_id}"
                     )
 
-    async def recover_state_on_failure(self, error: Exception, state: Any) -> None:
+    def recover_state_on_failure(self, error: Exception, state: Any) -> None:
         """Handles BidiReadObjectRedirectedError for reads."""
         # This would parse the gRPC error details, extract the routing_token,
         # and store it on the shared state object.

poc_bidi_retry_final.py

@@ -0,0 +1,120 @@
+# poc_bidi_retry_final.py
+
+import asyncio
+from unittest import mock
+from google.api_core import exceptions
+from google.api_core.retry_async import AsyncRetry
+
+# Assuming the retry components are in these locations
+# In a real scenario, these would be imported from the library
+from google.cloud.storage._experimental.asyncio.retry.bidi_stream_retry_manager import (
+    _BidiStreamRetryManager,
+)
+from google.cloud.storage._experimental.asyncio.retry.base_strategy import (
+    _BaseResumptionStrategy,
+)
+
+
+class ReadResumptionStrategy(_BaseResumptionStrategy):
+    """
+    A concrete implementation of the strategy for bidi reads.
+    This is a simplified version for the POC.
+    """
+
+    def __init__(self):
+        self.state = {"offset": 0, "remaining_bytes": float("inf")}
+
+    def generate_requests(self, state):
+        print(f"[Strategy] Generating request with state: {state}")
+        # In a real scenario, this yields ReadObjectRequest protos
+        yield {"read_offset": state["offset"]}
+
+    def handle_response(self, response):
+        # In a real scenario, this is a ReadObjectResponse proto
+        chunk = response.get("chunk", b"")
+        self.state["offset"] += len(chunk)
+        print(f"[Strategy] Handled response, new state: {self.state}")
+        return response
+
+    async def recover_state_on_failure(self, error, state):
+        print(f"[Strategy] Recovering state from error: {error}. Current state: {state}")
+        # For reads, the offset is already updated, so we just return the current state
+        return self.state
+
+
+# --- Simulation Setup ---
+
+# A mock stream that fails once mid-stream
+ATTEMPT_COUNT = 0
+STREAM_CONTENT = [
+    [{"chunk": b"part_one"}, {"chunk": b"part_two"}, exceptions.ServiceUnavailable("Network error")],
+    [{"chunk": b"part_three"}, {"chunk": b"part_four"}],
+]
+
+
+async def mock_stream_opener(requests, state):
+    """
+    A mock stream opener that simulates a failing and then succeeding stream.
+    """
+    global ATTEMPT_COUNT
+    print(f"\n--- Stream Attempt {ATTEMPT_COUNT + 1} ---")
+    # Consume the request iterator (in a real scenario, this sends requests to gRPC)
+    _ = [req for req in requests]
+    print(f"Mock stream opened with state: {state}")
+
+    content_for_this_attempt = STREAM_CONTENT[ATTEMPT_COUNT]
+    ATTEMPT_COUNT += 1
+
+    for item in content_for_this_attempt:
+        await asyncio.sleep(0.01)  # Simulate network latency
+        if isinstance(item, Exception):
+            print(f"!!! Stream yielding an error: {item} !!!")
+            raise item
+        else:
+            print(f"Stream yielding chunk of size: {len(item.get('chunk', b''))}")
+            yield item
+
+
+async def main():
+    """
+    Main function to run the POC.
+    """
+    print("--- Starting Bidi Read Retry POC ---")
+
+    # 1. Define a retry policy
+    retry_policy = AsyncRetry(
+        predicate=lambda e: isinstance(e, exceptions.ServiceUnavailable),
+        deadline=30.0,
+        initial=0.1,  # Start with a short wait
+    )
+
+    # 2. Instantiate the strategy and retry manager
+    strategy = ReadResumptionStrategy()
+    retry_manager = _BidiStreamRetryManager(
+        strategy=strategy, stream_opener=mock_stream_opener
+    )
+
+    # 3. Execute the operation
+    print("\nExecuting the retry manager...")
+    final_stream_iterator = await retry_manager.execute(
+        initial_state={"offset": 0}, retry_policy=retry_policy
+    )
+
+    # 4. Consume the final, successful stream
+    all_content = b""
+    print("\n--- Consuming Final Stream ---")
+    async for response in final_stream_iterator:
+        chunk = response.get("chunk", b"")
+        all_content += chunk
+        print(f"Received chunk: {chunk.decode()}. Total size: {len(all_content)}")
+
+    print("\n--- POC Finished ---")
+    print(f"Final downloaded content: {all_content.decode()}")
+    print(f"Total attempts made: {ATTEMPT_COUNT}")
+    assert all_content == b"part_onepart_twopart_threepart_four"
+    assert ATTEMPT_COUNT == 2
+    print("\nAssertion passed: Content correctly assembled across retries.")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())