Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/events/services/event_persister.py: 93%

1""" 

2The event persister moves event messages from the event bus to storage 

3storage as fast as it can. Never gets tired. 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9from contextlib import asynccontextmanager 

10from datetime import timedelta 

11from typing import TYPE_CHECKING, AsyncGenerator, List, NoReturn 

12 

13from prefect.logging import get_logger 

14from prefect.server.database import provide_database_interface 

15from prefect.server.events.schemas.events import ReceivedEvent 

16from prefect.server.events.storage.database import write_events 

17from prefect.server.services.base import RunInEphemeralServers, Service 

18from prefect.server.utilities.messaging import ( 

19 Consumer, 

20 Message, 

21 MessageHandler, 

22 create_consumer, 

23) 

24from prefect.server.utilities.messaging._consumer_names import ( 

25 generate_unique_consumer_name, 

26) 

27from prefect.settings.context import get_current_settings 

28from prefect.settings.models.server.services import ServerServicesEventPersisterSettings 

29 

30if TYPE_CHECKING: 30 ↛ 31line 30 didn't jump to line 31 because the condition on line 30 was never true

31 import logging 

32 

33logger: "logging.Logger" = get_logger(__name__) 

34 

35 

36class EventPersister(RunInEphemeralServers, Service): 

37 """A service that persists events to the database as they arrive.""" 

38 

39 consumer_task: asyncio.Task[None] | None = None 

40 

41 @classmethod 

42 def service_settings(cls) -> ServerServicesEventPersisterSettings: 

43 return get_current_settings().server.services.event_persister 

44 

45 def __init__(self): 

46 super().__init__() 

47 self._started_event: asyncio.Event | None = None 

48 

49 @property 

50 def started_event(self) -> asyncio.Event: 

51 if self._started_event is None: 

52 self._started_event = asyncio.Event() 

53 return self._started_event 

54 

55 @started_event.setter 

56 def started_event(self, value: asyncio.Event) -> None: 

57 self._started_event = value 

58 

59 async def start(self) -> NoReturn: 

60 assert self.consumer_task is None, "Event persister already started" 

61 self.consumer: Consumer = create_consumer( 

62 "events", 

63 group="event-persister", 

64 name=generate_unique_consumer_name("event-persister"), 

65 read_batch_size=self.service_settings().read_batch_size, 

66 ) 

67 

68 settings = self.service_settings() 

69 async with create_handler( 

70 batch_size=settings.batch_size, 

71 flush_every=timedelta(seconds=settings.flush_interval), 

72 queue_max_size=settings.queue_max_size, 

73 max_flush_retries=settings.max_flush_retries, 

74 ) as handler: 

75 self.consumer_task = asyncio.create_task(self.consumer.run(handler)) 

76 logger.debug("Event persister started") 

77 self.started_event.set() 

78 

79 try: 

80 await self.consumer_task 

81 except asyncio.CancelledError: 

82 pass 

83 

84 async def stop(self) -> None: 

85 assert self.consumer_task is not None, "Event persister not started" 

86 self.consumer_task.cancel() 

87 try: 

88 await self.consumer_task 

89 except asyncio.CancelledError: 

90 pass 

91 finally: 

92 await self.consumer.cleanup() 

93 self.consumer_task = None 

94 if self.started_event: 94 ↛ 96line 94 didn't jump to line 96 because the condition on line 94 was always true

95 self.started_event.clear() 

96 logger.debug("Event persister stopped") 

97 

98 

99@asynccontextmanager 

100async def create_handler( 

101 batch_size: int = 20, 

102 flush_every: timedelta = timedelta(seconds=5), 

103 queue_max_size: int = 50_000, 

104 max_flush_retries: int = 5, 

105) -> AsyncGenerator[MessageHandler, None]: 

106 """ 

107 Set up a message handler that will accumulate and send events to 

108 the database every `batch_size` messages, or every `flush_every` interval to flush 

109 any remaining messages. 

110 

111 Event trimming/retention is handled by the db_vacuum service 

112 (vacuum_old_events and vacuum_events_with_retention_overrides tasks). 

113 

114 Args: 

115 batch_size: Number of events to accumulate before flushing 

116 flush_every: Maximum time between flushes 

117 queue_max_size: Maximum events in queue before dropping new events 

118 max_flush_retries: Consecutive flush failures before dropping events 

119 """ 

120 db = provide_database_interface() 

121 

122 queue: asyncio.Queue[ReceivedEvent] = asyncio.Queue(maxsize=queue_max_size) 

123 flush_lock = asyncio.Lock() 

124 consecutive_failures = 0 

125 

126 async def flush() -> None: 

127 nonlocal consecutive_failures 

128 

129 async with flush_lock: 

130 if queue.qsize() == 0: 130 ↛ 131line 130 didn't jump to line 131 because the condition on line 130 was never true

131 return 

132 

133 # Log warning when queue reaches 80% capacity 

134 if queue_max_size > 0 and queue.qsize() > queue_max_size * 0.8: 134 ↛ 135line 134 didn't jump to line 135 because the condition on line 134 was never true

135 logger.warning( 

136 "Event queue at %d%% capacity (%d/%d)", 

137 int(queue.qsize() / queue_max_size * 100), 

138 queue.qsize(), 

139 queue_max_size, 

140 ) 

141 

142 logger.debug("Persisting %d events...", queue.qsize()) 

143 

144 batch: List[ReceivedEvent] = [] 

145 

146 while queue.qsize() > 0: 

147 batch.append(await queue.get()) 

148 

149 try: 

150 async with db.session_context() as session: 

151 await write_events(session=session, events=batch) 

152 await session.commit() 

153 logger.debug("Finished persisting events.") 

154 consecutive_failures = 0 # Reset on success 

155 except Exception: 

156 consecutive_failures += 1 

157 if consecutive_failures >= max_flush_retries: 

158 logger.error( 

159 "Max flush retries (%d) reached, dropping %d events", 

160 max_flush_retries, 

161 len(batch), 

162 exc_info=True, 

163 ) 

164 consecutive_failures = 0 

165 else: 

166 logger.debug( 

167 "Error flushing events (attempt %d/%d), restoring to queue", 

168 consecutive_failures, 

169 max_flush_retries, 

170 exc_info=True, 

171 ) 

172 for event in batch: 

173 queue.put_nowait(event) 

174 

175 async def flush_periodically(): 

176 try: 

177 while True: 

178 await asyncio.sleep(flush_every.total_seconds()) 

179 if queue.qsize(): 

180 await flush() 

181 except asyncio.CancelledError: 

182 return 

183 

184 async def message_handler(message: Message): 

185 if not message.data: 

186 return 

187 

188 event = ReceivedEvent.model_validate_json(message.data) 

189 

190 logger.debug( 

191 "Received event: %s with id: %s for resource: %s", 

192 event.event, 

193 event.id, 

194 event.resource.get("prefect.resource.id"), 

195 ) 

196 

197 try: 

198 queue.put_nowait(event) 

199 except asyncio.QueueFull: 

200 logger.warning( 

201 "Event queue full (%d/%d), dropping event id=%s", 

202 queue.qsize(), 

203 queue_max_size, 

204 event.id, 

205 ) 

206 return 

207 

208 if queue.qsize() >= batch_size: 

209 await flush() 

210 

211 periodic_flush = asyncio.create_task(flush_periodically()) 

212 

213 try: 

214 yield message_handler 

215 finally: 

216 periodic_flush.cancel() 

217 if queue.qsize(): 217 ↛ exitline 217 didn't return from function 'create_handler' because the condition on line 217 was always true

218 await flush()