Coverage for core / src / sensorkit / core / state.py: 94%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5import inspect 

6import threading 

7from typing import ClassVar, Self, get_type_hints 

8 

9from pydantic import BaseModel 

10 

11from sensorkit.backend.base import KVError 

12from sensorkit.backend.event import Event 

13from sensorkit.core.entity import EntityBase, EntityClient 

14from sensorkit.core.impl.entity import EntityImpl 

15 

16 

17class EventSourcedState(BaseModel): 

18 """State model that can recover values of Event-typed fields from an event stream.""" 

19 

20 _event_fields: ClassVar[dict[type[Event], str]] = {} 

21 _event_fields_lock: ClassVar[threading.Lock] = threading.Lock() 

22 

23 def model_post_init(self, __context): 

24 """Introspect fields at init time to identify Event-typed fields.""" 

25 super().model_post_init(__context) 

26 self._set_event_fields() 

27 self._update_lock = asyncio.Lock() 

28 

29 @classmethod 

30 def _set_event_fields(cls): 

31 if cls._event_fields: 

32 return 

33 

34 with cls._event_fields_lock: 

35 # Double-check pattern to avoid redundant work if another thread already initialized 

36 if cls._event_fields: 

37 return 

38 

39 # Store at class level. 

40 cls._event_fields = cls._introspect_event_fields() 

41 

42 @classmethod 

43 def _introspect_event_fields(cls): 

44 type_hints = get_type_hints(cls) 

45 event_fields = {} 

46 

47 for field, field_type in type_hints.items(): 

48 # Check if field type is a subclass of Event 

49 if inspect.isclass(field_type) and issubclass(field_type, Event): 

50 if field_type in event_fields: 

51 raise RuntimeError( 

52 f"Duplicate event type {field_type.model_tag()}" 

53 f" found at {cls.__name__}.{field}" 

54 ) 

55 

56 event_fields[field_type] = field 

57 

58 return event_fields 

59 

60 async def update( 

61 self, 

62 entity: EntityImpl, 

63 *events: Event, 

64 publish_state=True, 

65 return_snapshot=False, 

66 ): 

67 """Apply events to the state, emit them on the entity's stream, and optionally publish the state to KV. 

68 

69 Returns a deep copy of the updated state if return_snapshot is True, else None. 

70 """ 

71 async with self._update_lock: 

72 for event in events: 

73 field = self._event_fields.get(type(event)) 

74 

75 if field is None: 

76 raise KeyError(f"No {event.event_model} event field exists") 

77 

78 setattr(self, field, event) 

79 await entity.emit_event(event) 

80 

81 # Ordering is crucial here: events must be emitted first since they are the source of 

82 # truth. 

83 if publish_state: 

84 await entity.kv_put_model(self) 

85 

86 # Make a snapshot while we hold the lock, 

87 snapshot = self.model_copy(deep=True) if return_snapshot else None 

88 

89 return snapshot 

90 

91 @classmethod 

92 async def event_stream[T: Event](cls, entity: EntityClient, event_type: type[T]): 

93 """Yield a continuous stream of events of the given type, starting from the current stored state.""" 

94 field = cls._event_fields[event_type] 

95 stream = await entity.monitor_event(event_type) 

96 state = await entity.kv_get_model(cls) 

97 original: T | None = getattr(state, field) 

98 original_id = original.event_id.int if original is not None else 0 

99 

100 yield original 

101 

102 # FIXME: For this to work properly, we need the backend to be enhanced to expose a method 

103 # of starting the stream at a given timestamp. NATS supports this so this is just a 

104 # much needed backend iteration. Below is a poor-man's substitute that kind of works 

105 # only because we currently always publish a state update subsequent to every event. 

106 # 

107 # Skip any stream events already reflected in `original` (or that produced it), then yield 

108 # everything after. We order by the full uuid7 event_id rather than timestamp(): the latter 

109 # is only millisecond-resolution, so two events emitted in the same millisecond compare 

110 # equal and a genuinely newer event would be silently dropped, hanging the consumer. 

111 async for event in stream: 

112 if event.event_id.int > original_id: 

113 yield event 

114 break 

115 

116 async for event in stream: 

117 yield event 

118 

119 raise RuntimeError("unexpected end of stream") 

120 

121 @classmethod 

122 async def recover(cls, entity: EntityBase) -> Self: 

123 """Recover state from the KV store and validate Event fields are up to date.""" 

124 # Retrieve the stored data from KV. 

125 obj = await entity.kv_get_model(cls) 

126 

127 # Verify that each cached event is the latest event of that type. If it isn't, retrieve 

128 # the latest event. 

129 for field in cls._event_fields.values(): 

130 event: Event = getattr(obj, field) 

131 newer = await cls._get_latest_event_if_newer(entity, event) 

132 

133 if newer is not None: 

134 setattr(obj, field, newer) 

135 

136 return obj 

137 

138 @classmethod 

139 async def recover_or_init(cls, entity: EntityBase, **kwargs): 

140 """Recover state from KV, or create and publish a new instance using kwargs if none exists.""" 

141 try: 

142 return await cls.recover(entity) 

143 except KVError: 

144 new = cls(**kwargs) 

145 await entity.kv_put_model(new) 

146 return new 

147 

148 @staticmethod 

149 async def _get_latest_event_if_newer(entity: EntityBase, event: Event): 

150 # TODO: Implement. 

151 return event or None