Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/streams/router.py: 34%
53 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:20 +0200
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:20 +0200
1import asyncio
2import contextlib
3import logging
4from dataclasses import dataclass, field
6import grpc
7from anyio import (
8 BrokenResourceError,
9 EndOfStream,
10)
11from anyio.abc import ObjectStream
12from jumpstarter_protocol import router_pb2
14logger = logging.getLogger(__name__)
17@dataclass(kw_only=True, slots=True)
18class RouterStream(ObjectStream[bytes]):
19 context: grpc.aio.StreamStreamCall | grpc._cython.cygrpc._ServicerContext
20 cls: type = field(init=False)
22 def __post_init__(self):
23 match self.context:
24 case grpc.aio.StreamStreamCall():
25 self.cls = router_pb2.StreamRequest
26 case grpc._cython.cygrpc._ServicerContext():
27 self.cls = router_pb2.StreamResponse
28 case _:
29 raise ValueError(f"RouterStream: invalid context type: {type(self.context)}")
31 async def send(self, payload: bytes) -> None:
32 try:
33 await self.context.write(self.cls(payload=payload))
34 except grpc.aio.AioRpcError as e:
35 raise BrokenResourceError from e
37 async def receive(self) -> bytes:
38 try:
39 frame = await self.context.read()
40 except grpc.aio.AioRpcError as e:
41 raise BrokenResourceError from e
43 # Reference: https://grpc.github.io/grpc/python/grpc_asyncio.html#grpc.aio.StreamStreamCall.read
44 if frame == grpc.aio.EOF:
45 raise EndOfStream
47 match frame.frame_type:
48 case router_pb2.FRAME_TYPE_DATA:
49 return frame.payload
50 case router_pb2.FRAME_TYPE_GOAWAY:
51 raise EndOfStream
52 case router_pb2.FRAME_TYPE_PING:
53 pass
54 case _:
55 logger.debug(f"RouterStream: unrecognized frame ignored: {frame}")
57 return b""
59 async def send_eof(self):
60 with contextlib.suppress(grpc.aio.AioRpcError, asyncio.exceptions.InvalidStateError):
61 await self.context.write(self.cls(frame_type=router_pb2.FRAME_TYPE_GOAWAY))
62 if isinstance(self.context, grpc.aio.StreamStreamCall):
63 await self.context.done_writing()
65 async def aclose(self):
66 with contextlib.suppress(grpc.aio.AioRpcError, asyncio.exceptions.InvalidStateError):
67 await self.send_eof()
68 if isinstance(self.context, grpc._cython.cygrpc._ServicerContext):
69 await self.context.abort(grpc.StatusCode.ABORTED, "RouterStream: aclose")