feat: add _StreamMultiplexer skeleton with sentinel types

Author: zhixiangli

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

@@ -0,0 +1,71 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# 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
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Awaitable, Callable, Dict, Optional, Set
+
+from google.cloud import _storage_v2
+from google.cloud.storage.asyncio.async_read_object_stream import (
+    _AsyncReadObjectStream,
+)
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_QUEUE_MAX_SIZE = 100
+
+
+class _StreamError:
+    """Wraps a retryable error with the stream generation that produced it."""
+
+    def __init__(self, exception: Exception, generation: int):
+        self.exception = exception
+        self.generation = generation
+
+
+class _StreamEnd:
+    """Signals the stream closed normally."""
+
+    pass
+
+
+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.
+
+    A slow consumer on one task will slow down the entire shared connection
+    due to bounded queue backpressure propagating through gRPC flow control.
+    """
+
+    def __init__(
+        self,
+        stream: _AsyncReadObjectStream,
+        queue_max_size: int = _DEFAULT_QUEUE_MAX_SIZE,
+    ):
+        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

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

@@ -0,0 +1,52 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# 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
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import pytest
+from unittest.mock import AsyncMock, MagicMock
+
+from google.cloud.storage.asyncio._stream_multiplexer import (
+    _StreamMultiplexer,
+    _StreamError,
+    _StreamEnd,
+    _DEFAULT_QUEUE_MAX_SIZE,
+)
+
+
+class TestSentinelTypes:
+    def test_stream_error_stores_exception_and_generation(self):
+        exc = ValueError("test")
+        error = _StreamError(exc, generation=3)
+        assert error.exception is exc
+        assert error.generation == 3
+
+    def test_stream_end_is_instantiable(self):
+        sentinel = _StreamEnd()
+        assert isinstance(sentinel, _StreamEnd)
+
+
+class TestStreamMultiplexerInit:
+    def test_init_sets_stream_and_defaults(self):
+        mock_stream = AsyncMock()
+        mux = _StreamMultiplexer(mock_stream)
+        assert mux._stream is mock_stream
+        assert mux.stream_generation == 0
+        assert mux._queues == {}
+        assert mux._recv_task is None
+        assert mux._queue_max_size == _DEFAULT_QUEUE_MAX_SIZE
+
+    def test_init_custom_queue_size(self):
+        mock_stream = AsyncMock()
+        mux = _StreamMultiplexer(mock_stream, queue_max_size=50)
+        assert mux._queue_max_size == 50