refactor(storage): clean up stream teardown, eliminate redundant state, and revert artificial diffs

Author: raj-prince

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

@@ -131,6 +131,19 @@ def record_completion(self, range_count: int, byte_count: int) -> None:
         self.pending_ranges = max(0, self.pending_ranges - range_count)
         self.pending_bytes = max(0, self.pending_bytes - byte_count)
 
+    async def close(self) -> None:
+        """Closes the underlying multiplexer and stream cleanly."""
+        try:
+            await self.multiplexer.close()
+        except Exception:
+            pass
+        try:
+            res = self.stream.close()
+            if inspect.isawaitable(res):
+                await res
+        except Exception:
+            pass
+
 
 class _StreamPool:
     """Manages a pool of _ManagedStream instances with dynamic scaling and least-loaded dispatch."""
@@ -196,16 +209,7 @@ async def _scale_up(self) -> None:
             new_worker = await self._stream_factory()
             async with self._lock:
                 if self._closed:
-                    try:
-                        await new_worker.multiplexer.close()
-                    except Exception:
-                        pass
-                    try:
-                        res = new_worker.stream.close()
-                        if inspect.isawaitable(res):
-                            await res
-                    except Exception:
-                        pass
+                    await new_worker.close()
                     return
                 self.workers.append(new_worker)
         except Exception as e:
@@ -228,16 +232,7 @@ async def close(self) -> None:
             for task in list(self._background_tasks):
                 task.cancel()
         for w in workers:
-            try:
-                await w.multiplexer.close()
-            except Exception:
-                pass
-            try:
-                res = w.stream.close()
-                if inspect.isawaitable(res):
-                    await res
-            except Exception:
-                pass
+            await w.close()
 
 
 class AsyncMultiRangeDownloader:
@@ -392,12 +387,6 @@ def __init__(
         self.full_obj_server_crc32c: Optional[int] = None
 
         self.stream_config = stream_config
-        self.min_connections = stream_config.min_connections if stream_config else 1
-        self.max_connections = stream_config.max_connections if stream_config else 1
-        self.target_io_depth = stream_config.target_io_depth if stream_config else 8
-        self.target_bytes = (
-            stream_config.target_bytes if stream_config else 8 * 1024 * 1024
-        )
         self._pool: Optional[_StreamPool] = None
         self._primary_worker: Optional[_ManagedStream] = None
         self._metadata: Optional[List[Tuple[str, str]]] = None
@@ -494,8 +483,6 @@ async def _do_open():
                 self.generation = self.read_obj_str.generation_number
             if self.read_obj_str.read_handle:
                 self.read_handle = self.read_obj_str.read_handle
-            if getattr(self.read_obj_str, "routing_token", None):
-                self._routing_token = self.read_obj_str.routing_token
             if self.read_obj_str.persisted_size is not None:
                 self.persisted_size = self.read_obj_str.persisted_size
             self.is_finalized = self.read_obj_str.is_finalized
@@ -511,14 +498,14 @@ async def _do_open():
         if self.stream_config is not None and self.stream_config.max_connections > 1:
             self._pool = _StreamPool(
                 stream_factory=self._create_new_stream_worker,
-                min_connections=self.min_connections,
-                max_connections=self.max_connections,
-                target_io_depth=self.target_io_depth,
-                target_bytes=self.target_bytes,
+                min_connections=self.stream_config.min_connections,
+                max_connections=self.stream_config.max_connections,
+                target_io_depth=self.stream_config.target_io_depth,
+                target_bytes=self.stream_config.target_bytes,
             )
             await self._pool.add_worker(self._primary_worker)
 
-            for _ in range(self.min_connections - 1):
+            for _ in range(self.stream_config.min_connections - 1):
                 worker = await self._create_new_stream_worker()
                 await self._pool.add_worker(worker)
         else:
@@ -545,8 +532,6 @@ async def _create_new_stream_worker(self) -> _ManagedStream:
             self.generation = stream.generation_number
         if stream.read_handle:
             self.read_handle = stream.read_handle
-        if getattr(stream, "routing_token", None):
-            self._routing_token = stream.routing_token
 
         mux = _StreamMultiplexer(stream)
         return _ManagedStream(stream, mux)
@@ -582,8 +567,6 @@ async def factory():
                 self.generation = stream.generation_number
             if stream.read_handle:
                 self.read_handle = stream.read_handle
-            if getattr(stream, "routing_token", None):
-                self._routing_token = stream.routing_token
             self.is_finalized = stream.is_finalized
             self.full_obj_server_crc32c = stream.full_obj_server_crc32c
 
@@ -665,13 +648,10 @@ async def _download_ranges_on_worker(
             # Heuristic to detect full object reads:
             # - Implicit full object read: start offset is 0 and length is 0 (read all).
             # - Explicit full object read: start offset is 0 and length matches the exact persisted size.
-            is_full_object_read = self.is_finalized and (
-                (offset == 0 and length == 0)
-                or (
-                    self.persisted_size is not None
-                    and offset == 0
-                    and length == self.persisted_size
-                )
+            is_full_object_read = (offset == 0 and length == 0) or (
+                self.persisted_size is not None
+                and offset == 0
+                and length == self.persisted_size
             )
             download_states[read_id] = _DownloadState(
                 initial_offset=offset,
@@ -685,9 +665,7 @@ async def _download_ranges_on_worker(
             "read_handle": self.read_handle,
             "routing_token": None,
             "enable_checksum": enable_checksum,
-            "full_obj_server_crc32c": self.full_obj_server_crc32c
-            if self.is_finalized
-            else None,
+            "full_obj_server_crc32c": self.full_obj_server_crc32c,
         }
 
         read_ids = set(download_states.keys())
@@ -787,8 +765,6 @@ async def generator():
             if initial_state.get("read_handle"):
                 self.read_handle = initial_state["read_handle"]
         finally:
-            if self._multiplexer is not None:
-                self._multiplexer.unregister(read_ids)
             if worker.multiplexer is not None:
                 worker.multiplexer.unregister(read_ids)
 

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

@@ -1029,10 +1029,10 @@ async def test_create_mrd_with_stream_config(self, mock_cls_stream):
             mock_client, "b", "o", stream_config=config
         )
 
-        assert mrd.min_connections == 2
-        assert mrd.max_connections == 4
-        assert mrd.target_io_depth == 8
-        assert mrd.target_bytes == 2 * 1024 * 1024
+        assert mrd.stream_config.min_connections == 2
+        assert mrd.stream_config.max_connections == 4
+        assert mrd.stream_config.target_io_depth == 8
+        assert mrd.stream_config.target_bytes == 2 * 1024 * 1024
         # Verified that 2 streams were opened initially for min_connections=2
         assert len(mrd._pool.workers) == 2
         await mrd.close()
@@ -1223,23 +1223,21 @@ async def test_routing_token_preservation_and_propagation(
         s1 = mock.MagicMock()
         s1.open = AsyncMock()
         s1.generation_number = 100
-        s1.routing_token = "token-abc"
         s1.read_handle = b"h1"
         s1.persisted_size = 1000
         s1.is_finalized = True
         s1.full_obj_server_crc32c = 12345
         mock_cls_async_read_object_stream.return_value = s1
 
         mrd = await AsyncMultiRangeDownloader.create_mrd(mock_client, "b", "o")
-        # Ensure routing token from initial stream was recorded
-        assert mrd._routing_token == "token-abc"
+        # Simulate a redirect having set _routing_token
+        mrd._routing_token = "token-abc"
 
         # Now when a new stream worker is opened via _create_new_stream_worker,
         # it should include routing_token in the metadata
         s2 = mock.MagicMock()
         s2.open = AsyncMock()
         s2.generation_number = 100
-        s2.routing_token = "token-xyz"
         s2.read_handle = b"h2"
         s2.persisted_size = 1000
         mock_cls_async_read_object_stream.return_value = s2
@@ -1251,7 +1249,5 @@ async def test_routing_token_preservation_and_propagation(
         assert ("x-goog-request-params", "routing_token=token-abc") in call_kwargs[
             "metadata"
         ]
-        # And after s2 opened, the token should update to s2's routing_token
-        assert mrd._routing_token == "token-xyz"
 
         await mrd.close()