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

@@ -560,6 +560,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 +582,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 +608,58 @@ 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()
+        if retry_policy is None:
+            retry_policy = AsyncRetry(
+                predicate=_is_write_retryable, on_error=self._on_open_error
+            )
+        else:
+            original_on_error = retry_policy._on_error
+
+            def combined_on_error(exc):
+                self._on_open_error(exc)
+                if original_on_error:
+                    original_on_error(exc)
+
+            retry_policy = AsyncRetry(
+                predicate=_is_write_retryable,
+                initial=retry_policy._initial,
+                maximum=retry_policy._maximum,
+                multiplier=retry_policy._multiplier,
+                deadline=retry_policy._deadline,
+                on_error=combined_on_error,
+            )
+
+        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}"
+                )
+                self._is_stream_open = False
+                await self.open()
+
+            await self.write_obj_stream.close()
+            return self.persisted_size
+
+        try:
+            return await retry_policy(_do_close)()
+        finally:
+            self._is_stream_open = False
 
-        self._is_stream_open = False
-        return self.persisted_size
 
     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 +685,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 +716,62 @@ async def finalize(
                 ),
             )
 
-        try:
+        if retry_policy is None:
+            retry_policy = AsyncRetry(
+                predicate=_is_write_retryable, on_error=self._on_open_error
+            )
+        else:
+            original_on_error = retry_policy._on_error
+
+            def combined_on_error(exc):
+                self._on_open_error(exc)
+                if original_on_error:
+                    original_on_error(exc)
+
+            retry_policy = AsyncRetry(
+                predicate=_is_write_retryable,
+                initial=retry_policy._initial,
+                maximum=retry_policy._maximum,
+                multiplier=retry_policy._multiplier,
+                deadline=retry_policy._deadline,
+                on_error=combined_on_error,
+            )
+
+        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}"
+                )
+                self._is_stream_open = False
+                await self.open()
+
             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:
+                if self.write_obj_stream.is_stream_open:
+                    try:
+                        await self.write_obj_stream.close()
+                    except Exception as e:
+                        logger.warning(
+                            f"Error closing write stream during finalize cleanup. Got exception: {e}"
+                        )
+                else:
+                    try:
+                        await self.write_obj_stream.close()
+                    except Exception:
+                        pass
             self._is_stream_open = False
             self.offset = None
 

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

@@ -136,7 +136,7 @@ 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)
+            await writer.close(finalize_on_close=True, retry_policy=policy_to_pass)
 
             # If an exception was expected, this line should not be reached.
             if scenario["expected_error"] is not None:
@@ -235,6 +235,12 @@ async def test_bidi_writes(testbench):
             "instruction": "redirect-send-handle-and-token-tokenval",
             "expected_error": None,
         },
+        {
+            "name": "Retry exactly on finalize/close (503)",
+            "method": "storage.objects.insert",
+            "instruction": "return-503-after-10240K",
+            "expected_error": None,
+        },
     ]
 
     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()