feat: implement _StreamMultiplexer register, recv loop, send, reopen, close

Author: zhixiangli

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

@@ -4,7 +4,7 @@
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#     https://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
@@ -46,8 +46,8 @@ class _StreamMultiplexer:
     """Multiplexes concurrent download tasks over a single bidi-gRPC stream.
 
     Routes responses from a background recv loop to per-task asyncio.Queues
-    keyed by read_id. Serializes sends via lock. Coordinates stream reopening
-    via generation-gated locking.
+    keyed by read_id. Coordinates stream reopening via generation-gated
+    locking.
 
     A slow consumer on one task will slow down the entire shared connection
     due to bounded queue backpressure propagating through gRPC flow control.
@@ -61,11 +61,118 @@ def __init__(
         self._stream = stream
         self._stream_generation: int = 0
         self._queues: Dict[int, asyncio.Queue] = {}
-        self._send_lock = asyncio.Lock()
         self._reopen_lock = asyncio.Lock()
         self._recv_task: Optional[asyncio.Task] = None
         self._queue_max_size = queue_max_size
 
     @property
     def stream_generation(self) -> int:
         return self._stream_generation
+
+    def register(self, read_ids: Set[int]) -> asyncio.Queue:
+        """Register read_ids for a task and return its response queue."""
+        queue = asyncio.Queue(maxsize=self._queue_max_size)
+        for read_id in read_ids:
+            self._queues[read_id] = queue
+        return queue
+
+    def unregister(self, read_ids: Set[int]) -> None:
+        """Remove read_ids from routing. Stops recv loop if no tasks remain."""
+        for read_id in read_ids:
+            self._queues.pop(read_id, None)
+
+    def _get_unique_queues(self) -> Set[asyncio.Queue]:
+        return set(self._queues.values())
+
+    def _ensure_recv_loop(self) -> None:
+        if self._recv_task is None or self._recv_task.done():
+            self._recv_task = asyncio.create_task(self._recv_loop())
+
+    def _stop_recv_loop(self) -> None:
+        if self._recv_task and not self._recv_task.done():
+            self._recv_task.cancel()
+
+    def _put_error_nowait(self, queue: asyncio.Queue, error: _StreamError) -> None:
+        while True:
+            try:
+                queue.put_nowait(error)
+                break
+            except asyncio.QueueFull:
+                try:
+                    queue.get_nowait()
+                except asyncio.QueueEmpty:
+                    pass
+
+    async def _recv_loop(self) -> None:
+        try:
+            while True:
+                response = await self._stream.recv()
+                if response is None:
+                    sentinel = _StreamEnd()
+                    for queue in self._get_unique_queues():
+                        await queue.put(sentinel)
+                    return
+
+                if response.object_data_ranges:
+                    queues_to_notify: Set[asyncio.Queue] = set()
+                    for data_range in response.object_data_ranges:
+                        read_id = data_range.read_range.read_id
+                        queue = self._queues.get(read_id)
+                        if queue:
+                            queues_to_notify.add(queue)
+                    for queue in queues_to_notify:
+                        await queue.put(response)
+                else:
+                    for queue in self._get_unique_queues():
+                        await queue.put(response)
+        except asyncio.CancelledError:
+            raise
+        except Exception as e:
+            error = _StreamError(e, self._stream_generation)
+            for queue in self._get_unique_queues():
+                self._put_error_nowait(queue, error)
+
+    async def send(self, request: _storage_v2.BidiReadObjectRequest) -> int:
+        self._ensure_recv_loop()
+        await self._stream.send(request)
+        return self._stream_generation
+
+    async def reopen_stream(
+        self,
+        broken_generation: int,
+        stream_factory: Callable[[], Awaitable[_AsyncReadObjectStream]],
+    ) -> None:
+        async with self._reopen_lock:
+            if self._stream_generation != broken_generation:
+                return
+            self._stop_recv_loop()
+            if self._recv_task:
+                try:
+                    await self._recv_task
+                except (asyncio.CancelledError, Exception):
+                    pass
+            error = _StreamError(
+                Exception("Stream reopening"), self._stream_generation
+            )
+            for queue in self._get_unique_queues():
+                self._put_error_nowait(queue, error)
+            try:
+                await self._stream.close()
+            except Exception:
+                pass
+            self._stream = await stream_factory()
+            self._stream_generation += 1
+            self._ensure_recv_loop()
+
+    async def close(self) -> None:
+        self._stop_recv_loop()
+        if self._recv_task:
+            try:
+                await self._recv_task
+            except (asyncio.CancelledError, Exception):
+                pass
+        error = _StreamError(
+            Exception("Multiplexer closed"), self._stream_generation
+        )
+        for queue in self._get_unique_queues():
+            self._put_error_nowait(queue, error)