feat: add timeout to queue.put in _StreamMultiplexer to prevent hang

Author: zhixiangli

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

@@ -26,6 +26,7 @@
 logger = logging.getLogger(__name__)
 
 _DEFAULT_QUEUE_MAX_SIZE = 100
+_DEFAULT_PUT_TIMEOUT = 20.0
 
 
 class _StreamError:
@@ -84,6 +85,17 @@ def unregister(self, read_ids: Set[int]) -> None:
     def _get_unique_queues(self) -> Set[asyncio.Queue]:
         return set(self._queues.values())
 
+    async def _put_with_timeout(self, queue: asyncio.Queue, item) -> None:
+        try:
+            await asyncio.wait_for(queue.put(item), timeout=_DEFAULT_PUT_TIMEOUT)
+        except asyncio.TimeoutError:
+            if queue not in self._get_unique_queues():
+                logger.debug("Dropped item for unregistered queue.")
+            else:
+                logger.warning(
+                    "Queue full for too long. Dropping item to prevent multiplexer hang."
+                )
+
     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())
@@ -109,8 +121,12 @@ async def _recv_loop(self) -> None:
                 response = await self._stream.recv()
                 if response is None:
                     sentinel = _StreamEnd()
-                    for queue in self._get_unique_queues():
-                        await queue.put(sentinel)
+                    await asyncio.gather(
+                        *(
+                            self._put_with_timeout(queue, sentinel)
+                            for queue in self._get_unique_queues()
+                        )
+                    )
                     return
 
                 if response.object_data_ranges:
@@ -120,11 +136,19 @@ async def _recv_loop(self) -> None:
                         queue = self._queues.get(read_id)
                         if queue:
                             queues_to_notify.add(queue)
-                    for queue in queues_to_notify:
-                        await queue.put(response)
+                    await asyncio.gather(
+                        *(
+                            self._put_with_timeout(queue, response)
+                            for queue in queues_to_notify
+                        )
+                    )
                 else:
-                    for queue in self._get_unique_queues():
-                        await queue.put(response)
+                    await asyncio.gather(
+                        *(
+                            self._put_with_timeout(queue, response)
+                            for queue in self._get_unique_queues()
+                        )
+                    )
         except asyncio.CancelledError:
             raise
         except Exception as e: