chore: import upstream snapshot with attribution

python/ray/serve/_private/replica_response_generator.py

@@ -0,0 +1,91 @@
+import asyncio
+import time
+from typing import Any, AsyncGenerator, Callable, Optional
+
+from ray.serve._private.proxy_response_generator import (
+    _ProxyResponseGeneratorBase,
+    swallow_cancelled,
+)
+from ray.serve._private.utils import calculate_remaining_timeout
+
+
+class ReplicaResponseGenerator(_ProxyResponseGeneratorBase):
+    """Generic wrapper that adds disconnect detection to any async generator.
+
+    This can be used to wrap any async generator and add timeout and disconnect
+    detection capabilities. When a disconnect is detected, the generator will
+    raise asyncio.CancelledError.
+    """
+
+    def __init__(
+        self,
+        async_generator: AsyncGenerator[Any, None],
+        *,
+        timeout_s: Optional[float] = None,
+        disconnected_task: Optional[asyncio.Task] = None,
+        result_callback: Optional[Callable[[Any], Any]] = None,
+    ):
+        super().__init__(
+            timeout_s=timeout_s,
+            disconnected_task=disconnected_task,
+            result_callback=result_callback,
+        )
+        self._async_generator = async_generator
+        self._done = False
+
+    async def __anext__(self):
+        if self._done:
+            raise StopAsyncIteration
+
+        try:
+            result = await self._get_next_result()
+            if self._result_callback is not None:
+                result = self._result_callback(result)
+            return result
+        except (StopAsyncIteration, asyncio.CancelledError) as e:
+            self._done = True
+            raise e from None
+        except Exception as e:
+            self._done = True
+            raise e from None
+
+    async def _await_response_anext(self) -> Any:
+        return await self._async_generator.__anext__()
+
+    async def _get_next_result(self) -> Any:
+        """Get the next result from the async generator with disconnect detection."""
+        # If there's no disconnect detection needed, use direct await to preserve
+        # cancellation propagation (important for gRPC cancellation)
+        remaining_timeout = calculate_remaining_timeout(
+            timeout_s=self._timeout_s,
+            start_time_s=self._start_time_s,
+            curr_time_s=time.time(),
+        )
+        if self._disconnected_task is None:
+            try:
+                return await asyncio.wait_for(
+                    self._await_response_anext(), timeout=remaining_timeout
+                )
+            except asyncio.TimeoutError:
+                raise TimeoutError()
+        # Otherwise use asyncio.wait for disconnect detection
+        next_result_task = asyncio.create_task(self._await_response_anext())
+        tasks = [next_result_task, self._disconnected_task]
+
+        done, _ = await asyncio.wait(
+            tasks,
+            return_when=asyncio.FIRST_COMPLETED,
+            timeout=remaining_timeout,
+        )
+
+        if next_result_task in done:
+            return next_result_task.result()
+        elif self._disconnected_task in done:
+            next_result_task.cancel()
+            next_result_task.add_done_callback(swallow_cancelled)
+            raise asyncio.CancelledError()
+        else:
+            # Timeout occurred
+            next_result_task.cancel()
+            next_result_task.add_done_callback(swallow_cancelled)
+            raise TimeoutError()