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

1import asyncio 

2import contextlib 

3import logging 

4from dataclasses import dataclass, field 

5 

6import grpc 

7from anyio import ( 

8 BrokenResourceError, 

9 EndOfStream, 

10) 

11from anyio.abc import ObjectStream 

12from jumpstarter_protocol import router_pb2 

13 

14logger = logging.getLogger(__name__) 

15 

16 

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) 

21 

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)}") 

30 

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 

36 

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 

42 

43 # Reference: https://grpc.github.io/grpc/python/grpc_asyncio.html#grpc.aio.StreamStreamCall.read 

44 if frame == grpc.aio.EOF: 

45 raise EndOfStream 

46 

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}") 

56 

57 return b"" 

58 

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() 

64 

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")