From aae6babf0f9a06872994532ef5790ca55b8b42f4 Mon Sep 17 00:00:00 2001 From: zhixiangli Date: Fri, 14 Aug 2026 14:37:32 +0000 Subject: [PATCH 1/2] fix(storage): ensure bidi-gRPC stream is cleaned up on open failure and close - Ensure socket_like_rpc is closed if open() or the initial recv() fails or is cancelled in _AsyncReadObjectStream.open() - Add _close_socket_like_rpc helper to safely handle RPC closure - Only call requests_done() in close() if socket_like_rpc is active, and ensure stream state and RPC are closed even if requests_done() fails - Add unit tests for open failures/cancellations and close behavior on inactive or failing streams --- .../asyncio/async_read_object_stream.py | 74 +++++++---- .../asyncio/test_async_read_object_stream.py | 119 ++++++++++++++++++ 2 files changed, 166 insertions(+), 27 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py index 8fd98d623571..3f976afc1fcf 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio +import logging from typing import List, Optional, Tuple from google.api_core.bidi_async import AsyncBidiRpc @@ -22,6 +24,8 @@ ) from google.cloud.storage.asyncio.async_grpc_client import AsyncGrpcClient +logger = logging.getLogger(__name__) + class _AsyncReadObjectStream(_AsyncAbstractObjectStream): """Class representing a gRPC bidi-stream for reading data from a GCS ``Object``. @@ -126,40 +130,56 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None: initial_request=self.first_bidi_read_req, metadata=current_metadata, ) - await self.socket_like_rpc.open() # this is actually 1 send - response = await self.socket_like_rpc.recv() - # populated only in the first response of bidi-stream and when opened - # without using `read_handle` - if hasattr(response, "metadata") and response.metadata: - if self.generation_number is None: - self.generation_number = response.metadata.generation - # update persisted size - self.persisted_size = response.metadata.size - self.object_metadata = response.metadata - if ( - hasattr(response.metadata, "finalize_time") - and response.metadata.finalize_time - and response.metadata.finalize_time.second > 0 - ): - self.is_finalized = True + try: + await self.socket_like_rpc.open() # this is actually 1 send + response = await self.socket_like_rpc.recv() + # populated only in the first response of bidi-stream and when opened + # without using `read_handle` + if hasattr(response, "metadata") and response.metadata: + if self.generation_number is None: + self.generation_number = response.metadata.generation + # update persisted size + self.persisted_size = response.metadata.size + self.object_metadata = response.metadata if ( - hasattr(response.metadata, "checksums") - and response.metadata.checksums + hasattr(response.metadata, "finalize_time") + and response.metadata.finalize_time + and response.metadata.finalize_time.second > 0 ): - self.full_obj_server_crc32c = response.metadata.checksums.crc32c - - if response and response.read_handle: - self.read_handle = response.read_handle - - self._is_stream_open = True + self.is_finalized = True + if ( + hasattr(response.metadata, "checksums") + and response.metadata.checksums + ): + self.full_obj_server_crc32c = response.metadata.checksums.crc32c + + if response and response.read_handle: + self.read_handle = response.read_handle + + self._is_stream_open = True + except asyncio.CancelledError: + await self._close_socket_like_rpc() + raise + except Exception: + await self._close_socket_like_rpc() + raise + + async def _close_socket_like_rpc(self) -> None: + try: + await self.socket_like_rpc.close() + except Exception as exc: + logger.debug("Error while closing the read bidi-gRPC stream: %s", exc) async def close(self) -> None: """Closes the bidi-gRPC connection.""" if not self._is_stream_open: raise ValueError("Stream is not open") - await self.requests_done() - await self.socket_like_rpc.close() - self._is_stream_open = False + try: + if self.socket_like_rpc.is_active: + await self.requests_done() + finally: + self._is_stream_open = False + await self._close_socket_like_rpc() async def requests_done(self): """Signals that all requests have been sent.""" diff --git a/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py b/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py index a8f64422765e..393fa7b2b709 100644 --- a/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py +++ b/packages/google-cloud-storage/tests/unit/asyncio/test_async_read_object_stream.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio from unittest import mock from unittest.mock import AsyncMock import pytest +from google.api_core.exceptions import Aborted from google.cloud import _storage_v2 from google.cloud.storage.asyncio import async_read_object_stream @@ -203,6 +205,123 @@ async def test_close(mock_client, mock_cls_async_bidi_rpc): assert not read_obj_stream.is_stream_open +@mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") +@mock.patch( + "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" +) +@pytest.mark.asyncio +async def test_open_closes_rpc_when_first_recv_fails( + mock_client, mock_cls_async_bidi_rpc +): + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=False + ) + socket_like_rpc = mock_cls_async_bidi_rpc.return_value + socket_like_rpc.recv = AsyncMock( + side_effect=Aborted("Idle stream has been closed.") + ) + + with pytest.raises(Aborted): + await read_obj_stream.open() + + socket_like_rpc.open.assert_awaited_once() + socket_like_rpc.close.assert_awaited_once() + assert not read_obj_stream.is_stream_open + + +@mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") +@mock.patch( + "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" +) +@pytest.mark.parametrize("cancelled_await", ["open", "recv"]) +@pytest.mark.asyncio +async def test_open_closes_rpc_when_cancelled( + mock_client, mock_cls_async_bidi_rpc, cancelled_await +): + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=False + ) + socket_like_rpc = mock_cls_async_bidi_rpc.return_value + setattr( + socket_like_rpc, + cancelled_await, + AsyncMock(side_effect=asyncio.CancelledError), + ) + + with pytest.raises(asyncio.CancelledError): + await read_obj_stream.open() + + socket_like_rpc.close.assert_awaited_once() + assert not read_obj_stream.is_stream_open + + +@mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") +@mock.patch( + "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" +) +@pytest.mark.asyncio +async def test_open_propagates_close_failure_from_failed_open( + mock_client, mock_cls_async_bidi_rpc +): + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=False + ) + socket_like_rpc = mock_cls_async_bidi_rpc.return_value + socket_like_rpc.recv = AsyncMock( + side_effect=Aborted("Idle stream has been closed.") + ) + socket_like_rpc.close = AsyncMock(side_effect=RuntimeError("close blew up")) + + with pytest.raises(Aborted): + await read_obj_stream.open() + + assert not read_obj_stream.is_stream_open + + +@mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") +@mock.patch( + "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" +) +@pytest.mark.asyncio +async def test_close_closes_rpc_when_requests_done_fails( + mock_client, mock_cls_async_bidi_rpc +): + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=True + ) + socket_like_rpc = read_obj_stream.socket_like_rpc + read_obj_stream.requests_done = AsyncMock( + side_effect=Aborted("Idle stream has been closed.") + ) + + with pytest.raises(Aborted): + await read_obj_stream.close() + + socket_like_rpc.close.assert_awaited_once() + assert not read_obj_stream.is_stream_open + + +@mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") +@mock.patch( + "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" +) +@pytest.mark.asyncio +async def test_close_skips_requests_done_on_inactive_rpc( + mock_client, mock_cls_async_bidi_rpc +): + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=True + ) + read_obj_stream.socket_like_rpc.is_active = False + read_obj_stream.requests_done = AsyncMock() + + await read_obj_stream.close() + + read_obj_stream.requests_done.assert_not_called() + read_obj_stream.socket_like_rpc.close.assert_awaited_once() + assert not read_obj_stream.is_stream_open + + @mock.patch("google.cloud.storage.asyncio.async_read_object_stream.AsyncBidiRpc") @mock.patch( "google.cloud.storage.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" From dab4c3801d3e4657169cf2cd35001ca05754457e Mon Sep 17 00:00:00 2001 From: Zhixiang Li Date: Fri, 14 Aug 2026 22:43:18 +0800 Subject: [PATCH 2/2] Update packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../google/cloud/storage/asyncio/async_read_object_stream.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py index 3f976afc1fcf..d3c0780673c7 100644 --- a/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py +++ b/packages/google-cloud-storage/google/cloud/storage/asyncio/async_read_object_stream.py @@ -157,10 +157,7 @@ async def open(self, metadata: Optional[List[Tuple[str, str]]] = None) -> None: self.read_handle = response.read_handle self._is_stream_open = True - except asyncio.CancelledError: - await self._close_socket_like_rpc() - raise - except Exception: + except (asyncio.CancelledError, Exception): await self._close_socket_like_rpc() raise