wehub-resource-sync
106 lines
4.1 KiB
Python
106 lines
4.1 KiB
Python
import asyncio
|
|
from typing import Any, AsyncIterator, Optional, TypeVar
|
|
|
|
from ray._common.utils import get_or_create_event_loop
|
|
|
|
T = TypeVar("T")
|
|
|
|
# Sentinel placed on the queue to signal that the request stream has ended.
|
|
_STREAM_END = object()
|
|
|
|
|
|
class gRPCDIReceiveStream(AsyncIterator[T]):
|
|
"""Exposes a native gRPC request iterator to user code for direct ingress.
|
|
|
|
The native iterator is bound to the replica's server event loop, but user code
|
|
may run on a separate event loop. When the loops differ, a fetch task drains the
|
|
iterator on the server loop and hands messages to user code across loops; when
|
|
they're the same, user code iterates the native iterator directly.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
request_iterator: AsyncIterator[T],
|
|
user_event_loop: asyncio.AbstractEventLoop,
|
|
*,
|
|
cancel_event: Optional[asyncio.Event] = None,
|
|
):
|
|
self._request_iterator = request_iterator
|
|
self._user_event_loop = user_event_loop
|
|
self._same_event_loop = user_event_loop is get_or_create_event_loop()
|
|
# Set when the client disconnects/errors mid-stream so the consumer's
|
|
# gRPCInputStream reports is_cancelled() and ends gracefully.
|
|
self._cancel_event = cancel_event
|
|
# Lazily created so it binds to the user-code event loop (where it is
|
|
# always accessed), never the server loop.
|
|
self._queue: Optional[asyncio.Queue] = None
|
|
self._fetch_task: Optional[asyncio.Task] = None
|
|
|
|
@property
|
|
def _message_queue(self) -> asyncio.Queue:
|
|
if self._queue is None:
|
|
self._queue = asyncio.Queue()
|
|
return self._queue
|
|
|
|
def _put(self, item: Any):
|
|
self._message_queue.put_nowait(item)
|
|
|
|
def start(self) -> Optional[asyncio.Task]:
|
|
"""Start draining the native iterator. Must be called on the server loop."""
|
|
# We don't create another task to consume the input stream if we run on the
|
|
# same event loop.
|
|
if self._same_event_loop:
|
|
return
|
|
|
|
self._fetch_task = asyncio.ensure_future(self._fetch_until_done())
|
|
return self._fetch_task
|
|
|
|
def cancel(self):
|
|
"""Stop draining the native iterator (e.g. if the consumer finished)."""
|
|
if self._fetch_task is not None:
|
|
self._fetch_task.cancel()
|
|
|
|
async def _fetch_until_done(self):
|
|
"""Drive the native iterator on the server loop, forwarding each message.
|
|
|
|
A stream error (e.g. client disconnect) is treated as cancellation: the
|
|
cancel event is set and the stream ends gracefully -- mirroring the proxy
|
|
path (`receive_grpc_messages`) -- rather than surfacing the raw gRPC error
|
|
to user code. A sentinel is always enqueued at the end so the consumer
|
|
terminates.
|
|
"""
|
|
try:
|
|
async for message in self._request_iterator:
|
|
# Stop draining if the consumer cancelled (e.g. user code called
|
|
# input_stream.cancel()).
|
|
if self._cancel_event is not None and self._cancel_event.is_set():
|
|
break
|
|
self._user_event_loop.call_soon_threadsafe(self._put, message)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
if self._cancel_event is not None:
|
|
self._user_event_loop.call_soon_threadsafe(self._cancel_event.set)
|
|
finally:
|
|
self._user_event_loop.call_soon_threadsafe(self._put, _STREAM_END)
|
|
|
|
def __aiter__(self) -> "gRPCDIReceiveStream":
|
|
return self
|
|
|
|
async def __anext__(self) -> T:
|
|
"""Return the next request message. Runs on the user-code loop."""
|
|
if self._same_event_loop:
|
|
try:
|
|
return await anext(self._request_iterator)
|
|
except (asyncio.CancelledError, StopAsyncIteration):
|
|
raise
|
|
except Exception:
|
|
if self._cancel_event is not None:
|
|
self._cancel_event.set()
|
|
raise StopAsyncIteration
|
|
|
|
item = await self._message_queue.get()
|
|
if item is _STREAM_END:
|
|
raise StopAsyncIteration
|
|
return item
|