fix(storage): add retry support for finalize and close in AsyncAppendableObjectWriter

Fixes: b/532527637

Author: sahusneha2004

packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py

@@ -299,21 +299,11 @@ def _on_open_error(self, exc):
             if redirect_proto.generation:
                 self.generation = redirect_proto.generation
 
-    async def open(
-        self,
-        retry_policy: Optional[AsyncRetry] = None,
-        metadata: Optional[List[Tuple[str, str]]] = None,
-    ) -> None:
-        """Opens the underlying bidi-gRPC stream.
-
-        :raises ValueError: If the stream is already open.
-
-        """
-        if self._is_stream_open:
-            raise ValueError("Underlying bidi-gRPC stream is already open")
-
+    def _merge_retry_policy(
+        self, retry_policy: Optional[AsyncRetry] = None
+    ) -> AsyncRetry:
         if retry_policy is None:
-            retry_policy = AsyncRetry(
+            return AsyncRetry(
                 predicate=_is_write_retryable, on_error=self._on_open_error
             )
         else:
@@ -324,7 +314,7 @@ def combined_on_error(exc):
                 if original_on_error:
                     original_on_error(exc)
 
-            retry_policy = AsyncRetry(
+            return AsyncRetry(
                 predicate=_is_write_retryable,
                 initial=retry_policy._initial,
                 maximum=retry_policy._maximum,
@@ -333,6 +323,21 @@ def combined_on_error(exc):
                 on_error=combined_on_error,
             )
 
+    async def open(
+        self,
+        retry_policy: Optional[AsyncRetry] = None,
+        metadata: Optional[List[Tuple[str, str]]] = None,
+    ) -> None:
+        """Opens the underlying bidi-gRPC stream.
+
+        :raises ValueError: If the stream is already open.
+
+        """
+        if self._is_stream_open:
+            raise ValueError("Underlying bidi-gRPC stream is already open")
+
+        retry_policy = self._merge_retry_policy(retry_policy)
+
         async def _do_open():
             current_metadata = list(metadata) if metadata else []
 
@@ -560,6 +565,7 @@ async def close(
         self,
         finalize_on_close=False,
         full_object_checksum: Optional[int] = None,
+        retry_policy: Optional[AsyncRetry] = None,
     ) -> Union[int, _storage_v2.Object]:
         """Closes the underlying bidi-gRPC stream.
 
@@ -581,6 +587,9 @@ async def close(
                 crc32c_int = google_crc32c.value(data)
                 print(crc32c_int)
 
+        :type retry_policy: :class:`~google.api_core.retry_async.AsyncRetry`
+        :param retry_policy: (Optional) The retry policy to use for the operation.
+
         rtype: Union[int, _storage_v2.Object]
         returns: Updated `self.persisted_size` by default after closing the
             bidi-gRPC stream. However, if `finalize_on_close=True` is passed,
@@ -604,15 +613,47 @@ async def close(
             )
 
         if finalize_on_close:
-            return await self.finalize(full_object_checksum=full_object_checksum)
+            return await self.finalize(
+                full_object_checksum=full_object_checksum,
+                retry_policy=retry_policy,
+            )
 
-        await self.write_obj_stream.close()
+        retry_policy = self._merge_retry_policy(retry_policy)
 
-        self._is_stream_open = False
-        return self.persisted_size
+        attempt_count = 0
+
+        async def _do_close():
+            nonlocal attempt_count
+            attempt_count += 1
+
+            if attempt_count > 1:
+                logger.info(
+                    f"Re-opening the stream for close retry attempt: {attempt_count}"
+                )
+                expected_offset = self.offset
+                self._is_stream_open = False
+                await self.open()
+                if (
+                    self.offset is not None
+                    and expected_offset is not None
+                    and self.offset < expected_offset
+                ):
+                    raise exceptions.InternalServerError(
+                        f"Unrecoverable data loss during reconnect. Expected offset {expected_offset}, got {self.offset}"
+                    )
+
+            await self.write_obj_stream.close()
+            return self.persisted_size
+
+        try:
+            return await retry_policy(_do_close)()
+        finally:
+            self._is_stream_open = False
 
     async def finalize(
-        self, full_object_checksum: Optional[int] = None
+        self,
+        full_object_checksum: Optional[int] = None,
+        retry_policy: Optional[AsyncRetry] = None,
     ) -> _storage_v2.Object:
         """Finalizes the Appendable Object.
 
@@ -638,6 +679,9 @@ async def finalize(
                 crc32c_int = google_crc32c.value(data)
                 print(crc32c_int)
 
+        :type retry_policy: :class:`~google.api_core.retry_async.AsyncRetry`
+        :param retry_policy: (Optional) The retry policy to use for the operation.
+
         rtype: google.cloud.storage_v2.types.Object
         returns: The finalized object resource.
 
@@ -666,14 +710,46 @@ async def finalize(
                 ),
             )
 
-        try:
+        retry_policy = self._merge_retry_policy(retry_policy)
+
+        attempt_count = 0
+
+        async def _do_finalize():
+            nonlocal attempt_count
+            attempt_count += 1
+
+            if attempt_count > 1:
+                logger.info(
+                    f"Re-opening the stream for finalize retry attempt: {attempt_count}"
+                )
+                expected_offset = self.offset
+                self._is_stream_open = False
+                await self.open()
+                if (
+                    self.offset is not None
+                    and expected_offset is not None
+                    and self.offset < expected_offset
+                ):
+                    raise exceptions.InternalServerError(
+                        f"Unrecoverable data loss during reconnect. Expected offset {expected_offset}, got {self.offset}"
+                    )
+
             await self.write_obj_stream.send(finalize_req)
             response = await self.write_obj_stream.recv()
             self.object_resource = response.resource
             self.persisted_size = self.object_resource.size
             return self.object_resource
+
+        try:
+            return await retry_policy(_do_finalize)()
         finally:
-            await self.write_obj_stream.close()
+            if self.write_obj_stream:
+                try:
+                    await self.write_obj_stream.close()
+                except Exception as e:
+                    logger.debug(
+                        f"Stream close during finalize cleanup resulted in: {e}"
+                    )
             self._is_stream_open = False
             self.offset = None
 

packages/google-cloud-storage/tests/conformance/test_bidi_writes.py

@@ -136,24 +136,27 @@ def on_retry_error(exc):
                 CONTENT, metadata=fault_injection_metadata, retry_policy=policy_to_pass
             )
             # await writer.finalize()
-            await writer.close(finalize_on_close=True)
+            f_o_c = scenario.get("finalize_on_close", True)
+            await writer.close(finalize_on_close=f_o_c, retry_policy=policy_to_pass)
 
             # If an exception was expected, this line should not be reached.
             if scenario["expected_error"] is not None:
                 raise AssertionError(
                     f"Expected exception {scenario['expected_error']} was not raised."
                 )
 
-            # 4. Verify the object content.
-            read_request = storage_v2.ReadObjectRequest(
-                bucket=f"projects/_/buckets/{bucket_name}",
-                object=object_name,
-            )
-            read_stream = await gapic_client.read_object(request=read_request)
-            data = b""
-            async for chunk in read_stream:
-                data += chunk.checksummed_data.content
-            assert data == CONTENT
+            # 4. Verify the object content if applicable.
+            if not scenario.get("skip_verification"):
+                read_request = storage_v2.ReadObjectRequest(
+                    bucket=f"projects/_/buckets/{bucket_name}",
+                    object=object_name,
+                )
+                read_stream = await gapic_client.read_object(request=read_request)
+                data = b""
+                async for chunk in read_stream:
+                    data += chunk.checksummed_data.content
+                assert data == CONTENT
+
             if scenario["expected_error"] is None:
                 # Scenarios like 503, 500, smarter resumption, and redirects
                 # SHOULD trigger at least one retry attempt.
@@ -235,6 +238,26 @@ async def test_bidi_writes(testbench):
             "instruction": "redirect-send-handle-and-token-tokenval",
             "expected_error": None,
         },
+        {
+            "name": "Retry exactly on finalize/close (Redirect Error)",
+            "method": "storage.objects.insert",
+            "instruction": "redirect-send-handle-and-token-mytoken-on-finish-write",
+            "expected_error": None,
+        },
+        {
+            "name": "Retry exactly on finalize/close (503)",
+            "method": "storage.objects.insert",
+            "instruction": "return-503-on-finish-write",
+            "expected_error": None,
+        },
+        {
+            "name": "Retry exactly on close (finalize_on_close=False) (503)",
+            "method": "storage.objects.insert",
+            "instruction": "return-503-on-half-close",
+            "expected_error": None,
+            "finalize_on_close": False,
+            "skip_verification": True,
+        },
     ]
 
     try:

packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py

@@ -560,7 +560,9 @@ async def test_close_with_checksum_and_finalize(self, mock_appendable_writer):
 
         checksum = 12345678
         await writer.close(finalize_on_close=True, full_object_checksum=checksum)
-        writer.finalize.assert_awaited_once_with(full_object_checksum=checksum)
+        writer.finalize.assert_awaited_once_with(
+            full_object_checksum=checksum, retry_policy=None
+        )
 
     @pytest.mark.asyncio
     async def test_close_with_checksum_without_finalize_raises(
@@ -625,3 +627,134 @@ async def test_finalize_mismatch_closes_stream(self, mock_appendable_writer):
         # Assert stream was closed and local state reset despite exception
         mock_appendable_writer["mock_stream"].close.assert_awaited()
         assert not writer._is_stream_open
+
+    @pytest.mark.asyncio
+    async def test_finalize_retry_on_transient_error(self, mock_appendable_writer):
+        writer = self._make_one(mock_appendable_writer["mock_client"])
+        writer._is_stream_open = True
+        writer.write_obj_stream = mock_appendable_writer["mock_stream"]
+
+        resource = storage_type.Object(size=999)
+        mock_appendable_writer["mock_stream"].recv.side_effect = [
+            exceptions.InternalServerError("500 Transient Error"),
+            storage_type.BidiWriteObjectResponse(resource=resource),
+        ]
+
+        res = await writer.finalize()
+
+        assert res == resource
+        assert writer.persisted_size == 999
+        assert mock_appendable_writer["mock_stream"].send.await_count == 2
+        assert not writer._is_stream_open
+
+    @pytest.mark.asyncio
+    async def test_finalize_custom_retry_policy(self, mock_appendable_writer):
+        from google.api_core.retry_async import AsyncRetry
+
+        writer = self._make_one(mock_appendable_writer["mock_client"])
+        writer._is_stream_open = True
+        writer.write_obj_stream = mock_appendable_writer["mock_stream"]
+
+        custom_policy = AsyncRetry(
+            predicate=lambda exc: isinstance(exc, exceptions.InternalServerError)
+        )
+        resource = storage_type.Object(size=999)
+        mock_appendable_writer[
+            "mock_stream"
+        ].recv.return_value = storage_type.BidiWriteObjectResponse(resource=resource)
+
+        res = await writer.finalize(retry_policy=custom_policy)
+        assert res == resource
+
+    @pytest.mark.asyncio
+    async def test_close_with_finalize_and_custom_retry_policy(
+        self, mock_appendable_writer
+    ):
+        from google.api_core.retry_async import AsyncRetry
+
+        writer = self._make_one(mock_appendable_writer["mock_client"])
+        writer._is_stream_open = True
+        writer.finalize = AsyncMock()
+
+        custom_policy = AsyncRetry(predicate=lambda exc: False)
+        await writer.close(finalize_on_close=True, retry_policy=custom_policy)
+        writer.finalize.assert_awaited_once_with(
+            full_object_checksum=None,
+            retry_policy=custom_policy,
+        )
+
+    @pytest.mark.asyncio
+    async def test_close_retry_on_transient_error(self, mock_appendable_writer):
+        writer = self._make_one(mock_appendable_writer["mock_client"])
+        writer._is_stream_open = True
+        writer.write_obj_stream = mock_appendable_writer["mock_stream"]
+
+        resource = storage_type.Object(size=999)
+        mock_appendable_writer["mock_stream"].recv.side_effect = [
+            exceptions.InternalServerError("500 Transient Error"),
+            storage_type.BidiWriteObjectResponse(resource=resource),
+        ]
+
+        res = await writer.close(finalize_on_close=True)
+
+        assert res == resource
+        assert writer.persisted_size == 999
+        assert mock_appendable_writer["mock_stream"].send.await_count == 2
+        assert not writer._is_stream_open
+
+    @pytest.mark.asyncio
+    async def test_finalize_retry_on_redirect_error(self, mock_appendable_writer):
+        writer = self._make_one(mock_appendable_writer["mock_client"])
+        writer._is_stream_open = True
+        writer.write_obj_stream = mock_appendable_writer["mock_stream"]
+
+        redirect = BidiWriteObjectRedirectedError(
+            routing_token="rt1",
+            write_handle=storage_type.BidiWriteHandle(handle=b"h1"),
+        )
+        exc = exceptions.Aborted("aborted", errors=[redirect])
+        
+        resource = storage_type.Object(size=999)
+        mock_appendable_writer["mock_stream"].recv.side_effect = [
+            exc,
+            storage_type.BidiWriteObjectResponse(resource=resource),
+        ]
+        
+        writer.open = mock.AsyncMock()
+
+        res = await writer.finalize()
+
+        assert res == resource
+        assert writer.persisted_size == 999
+        assert mock_appendable_writer["mock_stream"].send.await_count == 2
+        assert writer._routing_token == "rt1"
+        assert writer.write_handle.handle == b"h1"
+        writer.open.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    async def test_close_retry_on_redirect_error(self, mock_appendable_writer):
+        writer = self._make_one(mock_appendable_writer["mock_client"])
+        writer._is_stream_open = True
+        writer.write_obj_stream = mock_appendable_writer["mock_stream"]
+
+        redirect = BidiWriteObjectRedirectedError(
+            routing_token="rt2",
+            write_handle=storage_type.BidiWriteHandle(handle=b"h2"),
+        )
+        exc = exceptions.Aborted("aborted", errors=[redirect])
+
+        mock_appendable_writer["mock_stream"].close.side_effect = [
+            exc,
+            None,
+        ]
+        
+        writer.open = mock.AsyncMock()
+        writer.persisted_size = 999
+
+        res = await writer.close()
+
+        assert res == 999
+        assert mock_appendable_writer["mock_stream"].close.await_count == 2
+        assert writer._routing_token == "rt2"
+        assert writer.write_handle.handle == b"h2"
+        writer.open.assert_awaited_once()