fix(storage): ensure bidi-gRPC stream is cleaned up on open failure and close (#18119)

- 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

Thank you for opening a Pull Request! Before submitting your PR, there
are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a
[bug/issue](https://github.com/googleapis/google-cloud-python/issues)
before writing your code! That way we can discuss the change, evaluate
designs, and agree on the general idea
- [x] Ensure the tests and linter pass
- [x] Code coverage does not decrease (if any source code was changed)
- [x] Appropriate docs were updated (if necessary)

Fixes b/547471318 🦕

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: zhixiangli

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,53 @@ 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, 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."""

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"