From ce980a53c5b98c9817dedd3de410dce031834b4c Mon Sep 17 00:00:00 2001 From: Sneha Sahu Date: Thu, 16 Jul 2026 10:19:46 +0000 Subject: [PATCH] fix(storage): add retry support for finalize and close in AsyncAppendableObjectWriter Fixes: b/532527637 --- .../asyncio/async_appendable_object_writer.py | 112 ++++++++++++++- .../tests/conformance/test_bidi_writes.py | 45 ++++-- .../test_async_appendable_object_writer.py | 135 +++++++++++++++++- 3 files changed, 273 insertions(+), 19 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py index ce954ed3d7c2..b343294ecd6b 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_appendable_object_writer.py +++ b/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 diff --git a/packages/google-cloud-storage/tests/conformance/test_bidi_writes.py b/packages/google-cloud-storage/tests/conformance/test_bidi_writes.py index b98d8e221048..5490a419ec84 100644 --- a/packages/google-cloud-storage/tests/conformance/test_bidi_writes.py +++ b/packages/google-cloud-storage/tests/conformance/test_bidi_writes.py @@ -136,7 +136,8 @@ 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: @@ -144,16 +145,18 @@ def on_retry_error(exc): 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: diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py index a51b0961eb8d..c483b65d79a2 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/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()