Coverage for core / src / sensorkit / backend / event.py: 89%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5import collections 

6import contextlib 

7import itertools 

8from typing import TYPE_CHECKING, ClassVar, override 

9 

10import uuid_utils.compat as uuid 

11from loguru import logger 

12from pydantic import Field, ValidationError 

13 

14from sensorkit.backend.base import SpecialProperty 

15from sensorkit.common.model import ModelRegistry, RegistryBaseModel 

16 

17if TYPE_CHECKING: 

18 from sensorkit.backend.base import StreamContext 

19 

20 

21class Event(RegistryBaseModel): 

22 """Base SensorKit event model.""" 

23 

24 event_model: str | None = None 

25 event_id: uuid.UUID = Field(default_factory=uuid.uuid7) 

26 

27 model_config: ClassVar[dict] = {"frozen": True} 

28 registry: ClassVar[ModelRegistry[Event]] = ModelRegistry( 

29 discriminator="event_model", 

30 default_tag="UnknownEvent", 

31 ) 

32 

33 @classmethod 

34 def model_registry(cls): 

35 return cls.registry 

36 

37 def timestamp(self): 

38 """Return the millisecond UTC timestamp encoded in the UUIDv7 event_id.""" 

39 return self.event_id.int >> 80 

40 

41 

42class UnknownEvent(Event, extra="allow"): 

43 """Represents an event for which no corresponding model is known.""" 

44 

45 

46class EventMultiplexer: 

47 """Distributes parsed events to per-type and wildcard subscriber queues.""" 

48 

49 def __init__(self): 

50 self._all_queues: set[asyncio.Queue[Event]] = set() 

51 self._single_queues: dict[str, set[asyncio.Queue[Event]]] = collections.defaultdict(set) 

52 

53 async def wait_ready(self): 

54 """Wait until events are being received. A plain multiplexer has no source to wait for.""" 

55 

56 async def parse_event(self, json: bytes): 

57 """Deserialise a JSON event payload and enqueue it to all matching subscriber queues.""" 

58 # Parse the event and narrow its type. 

59 try: 

60 model = Event.model_validate_json(json) 

61 except Exception: 

62 logger.exception("Failed to parse event") 

63 raise 

64 

65 # Notify queues registered for this event. 

66 single = self._single_queues.get(model.event_model, ()) 

67 

68 for queue in itertools.chain(self._all_queues, single): 

69 queue.put_nowait(model) 

70 

71 @contextlib.contextmanager 

72 def event_queue(self, event: type[Event]): 

73 """Context manager yielding a Queue that receives events of the specified type.""" 

74 queue: asyncio.Queue[Event] = asyncio.Queue() 

75 name = event.model_tag() 

76 

77 try: 

78 self._single_queues[name].add(queue) 

79 yield queue 

80 finally: 

81 self._single_queues[name].discard(queue) 

82 

83 @contextlib.contextmanager 

84 def all_events(self): 

85 """Context manager yielding a Queue that receives all events regardless of type.""" 

86 queue: asyncio.Queue[Event] = asyncio.Queue() 

87 

88 try: 

89 self._all_queues.add(queue) 

90 yield queue 

91 finally: 

92 self._all_queues.discard(queue) 

93 

94 

95class EventStreamConsumer(EventMultiplexer): 

96 """An EventMultiplexer that consumes events from a given stream.""" 

97 

98 def __init__(self, context: StreamContext): 

99 super().__init__() 

100 self._context = context 

101 self._startup = asyncio.get_running_loop().create_future() 

102 

103 async def start(self): 

104 """Start the background event consumer task and wait until the stream is ready.""" 

105 self._task = asyncio.create_task(self._event_consumer_task()) 

106 await self._startup 

107 

108 @override 

109 async def wait_ready(self): 

110 """Wait until the stream subscription is live, raising if it failed to start. 

111 

112 Blocks indefinitely until `start` is called, so callers must hold a multiplexer that is 

113 already starting. 

114 """ 

115 await self._startup 

116 

117 async def _event_consumer_task(self): 

118 try: 

119 stream = await self._context.consume(SpecialProperty.EVENTS) 

120 except BaseException as e: 

121 self._startup.set_exception(e) 

122 raise 

123 else: 

124 self._startup.set_result(None) 

125 

126 async for msg in stream: 

127 try: 

128 await self.parse_event(msg.data) 

129 except ValidationError: 

130 logger.exception("Event validation failed")