Coverage for core / src / sensorkit / auto / constraint.py: 95%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5import json 

6from abc import ABC, abstractmethod 

7from collections.abc import Sequence 

8from dataclasses import dataclass 

9from datetime import UTC, datetime, timedelta 

10from typing import Any, Literal, override 

11 

12from loguru import logger 

13from pydantic import model_validator 

14 

15from sensorkit.common.aio import cleanup_future 

16from sensorkit.common.condition import AnyCondition, resolve_field 

17from sensorkit.common.model import ModelRegistry, RegistryBaseModel 

18from sensorkit.core.client import SensorKit 

19 

20_constraint_registry = ModelRegistry(discriminator="kind") 

21_details_registry = ModelRegistry(discriminator="kind") 

22 

23 

24class Constraint(RegistryBaseModel, ABC): 

25 """Abstract operating constraint that runs a background monitoring task.""" 

26 

27 kind: str 

28 ttl: float = 30.0 

29 hold: float = 0.0 

30 optional: bool = False 

31 

32 @classmethod 

33 def model_registry(cls): 

34 return _constraint_registry 

35 

36 # FIXME: Use Context or similar instead of **kwargs 

37 @abstractmethod 

38 async def check_task(self, evaluator: ConstraintEvaluator, /, **kwargs) -> None: 

39 """Long-running coroutine that monitors the constraint and updates the evaluator.""" 

40 

41 @model_validator(mode="before") 

42 @classmethod 

43 def _compat(cls, data: Any) -> Any: 

44 if isinstance(data, dict): 

45 if "time_to_live" in data: 

46 data["ttl"] = data.pop("time_to_live") 

47 if "hold_duration" in data: 

48 data["hold"] = data.pop("hold_duration") 

49 if "activate_on_timeout" in data: 

50 data["optional"] = not data.pop("activate_on_timeout") 

51 

52 return data 

53 

54 

55class ConstraintDetails(RegistryBaseModel): 

56 """Base for structured constraint state details. Subclasses auto-register.""" 

57 

58 kind: str 

59 

60 @classmethod 

61 def model_registry(cls): 

62 return _details_registry 

63 

64 

65@dataclass 

66class ConstraintState: 

67 """Per-constraint state.""" 

68 

69 kind: str 

70 state: Literal["not_ready", "clear", "holding", "active"] = "not_ready" 

71 reason: str = "" 

72 details: ConstraintDetails | None = None 

73 

74 

75class ConstraintEvaluator: 

76 """State interface for one check_task run. Created fresh by ConstraintManager on each restart.""" 

77 

78 _QUEUE_MAXSIZE = 10 

79 

80 def __init__( 

81 self, 

82 constraint: Constraint, 

83 *, 

84 timeout: asyncio.Timeout, 

85 initial_hold: float = 0.0, 

86 ): 

87 self.constraint = constraint 

88 self.holding_until: float | None = None 

89 """Event loop deadline of the hold in effect, if there is one. 

90 

91 A hold runs on the monotonic event loop clock so that it cannot be cut short by a 

92 wall clock adjustment. Use `holding_until_dt` to express it as an instant. 

93 """ 

94 

95 self._timeout = timeout 

96 self._ready = asyncio.Event() 

97 self._hold_task: asyncio.Task | None = None 

98 

99 if self.constraint.optional: 

100 self.intended_state = ConstraintState( 

101 kind=constraint.kind, state="clear", reason="optional constraint" 

102 ) 

103 self._ready.set() 

104 else: 

105 self.intended_state = ConstraintState( 

106 kind=constraint.kind, state="not_ready", reason="awaiting data" 

107 ) 

108 

109 self.actual_state = self.intended_state 

110 self._state_updated = asyncio.Event() 

111 

112 if initial_hold > 0: 

113 self._begin_hold(initial_hold) 

114 

115 @property 

116 def is_active(self) -> bool: 

117 return self.actual_state.state != "clear" 

118 

119 @property 

120 def is_ready(self) -> bool: 

121 return self._ready.is_set() 

122 

123 async def wait_ready(self): 

124 await self._ready.wait() 

125 

126 async def next_update(self) -> ConstraintState: 

127 await self._state_updated.wait() 

128 state = self.actual_state 

129 self._state_updated.clear() 

130 return state 

131 

132 def ready(self): 

133 """Mark constraint ready.""" 

134 self._timeout.reschedule(asyncio.get_event_loop().time() + self.constraint.ttl) 

135 

136 if not self._ready.is_set(): 

137 self._ready.set() 

138 self._update_state() 

139 

140 def constrain(self, reason: str, *, details: ConstraintDetails | None = None): 

141 """Set constraint active. 

142 

143 Returns: 

144 bool: True if state changed (clear → active) 

145 

146 Raises: 

147 QueueFull: if the status queue is full 

148 """ 

149 self._timeout.reschedule(asyncio.get_event_loop().time() + self.constraint.ttl) 

150 self._cancel_hold() 

151 self.intended_state = ConstraintState( 

152 kind=self.constraint.kind, 

153 state="active", 

154 reason=reason, 

155 details=details, 

156 ) 

157 self._update_state() 

158 

159 def clear(self, reason: str = "", *, details: ConstraintDetails | None = None): 

160 """Set constraint inactive. If hold > 0, defers the clear until the hold expires. 

161 

162 Raises: 

163 QueueFull: if the status queue is full 

164 """ 

165 self._timeout.reschedule(asyncio.get_event_loop().time() + self.constraint.ttl) 

166 

167 if self.constraint.hold > 0 and self.intended_state.state == "active": 

168 self._begin_hold(self.constraint.hold) 

169 

170 self.intended_state = ConstraintState( 

171 kind=self.constraint.kind, 

172 state="clear", 

173 reason=reason, 

174 details=details, 

175 ) 

176 self._update_state() 

177 

178 async def cleanup(self): 

179 """Stop the hold task, leaving its deadline readable by `hold_remaining`.""" 

180 self._stop_hold_task() 

181 

182 def hold_remaining(self) -> float: 

183 """Hold a replacement evaluator should start with, or zero if there is nothing to resume. 

184 

185 A hold delays releasing a constraint, so it has to outlive the evaluator that started it. 

186 Without this the replacement starts with nothing pending and releases on its first clear 

187 reading, cutting the delay short. 

188 """ 

189 if self.holding_until is not None: 

190 return max(0.0, self.holding_until - asyncio.get_event_loop().time()) 

191 

192 return self.constraint.hold if self.intended_state.state == "active" else 0.0 

193 

194 @property 

195 def holding_until_dt(self) -> datetime | None: 

196 """`holding_until` as a UTC instant, or None if no hold is in effect. 

197 

198 Projected from the current wall clock on each access rather than stored, so that a clock 

199 adjustment cannot move a hold that is already running. A deadline that has already passed 

200 projects into the past. 

201 """ 

202 if self.holding_until is None: 

203 return None 

204 

205 remaining = self.holding_until - asyncio.get_event_loop().time() 

206 

207 return datetime.now(UTC) + timedelta(seconds=remaining) 

208 

209 def _update_state(self): 

210 if not self._ready.is_set(): 

211 state = ConstraintState(kind=self.constraint.kind, state="not_ready") 

212 elif self.holding_until is not None: 

213 state = ConstraintState(kind=self.constraint.kind, state="holding") 

214 else: 

215 state = self.intended_state 

216 

217 self.actual_state = state 

218 self._state_updated.set() 

219 

220 async def _do_hold(self, wait_secs: float): 

221 logger.debug(f"holding {self.constraint.kind} constraint for {wait_secs:.1f}s") 

222 await asyncio.sleep(wait_secs) 

223 

224 # It's safe to clear the hold early here since there are no further awaits. 

225 self._hold_task = None 

226 self.holding_until = None 

227 self._update_state() 

228 

229 def _begin_hold(self, wait_secs: float): 

230 """Start a hold running for the given duration. 

231 

232 Raises: 

233 RuntimeError: if a hold is already in effect. 

234 """ 

235 if self.holding_until is not None: 

236 raise RuntimeError(f"{self.constraint.kind} constraint is already holding") 

237 

238 self.holding_until = asyncio.get_event_loop().time() + wait_secs 

239 self._hold_task = asyncio.create_task(self._do_hold(wait_secs)) 

240 self._hold_task.add_done_callback(cleanup_future) 

241 

242 def _cancel_hold(self): 

243 """Abandon any hold, discarding its deadline so it cannot be carried over.""" 

244 self._stop_hold_task() 

245 self.holding_until = None 

246 

247 def _stop_hold_task(self): 

248 if self._hold_task is not None and not self._hold_task.done(): 

249 self._hold_task.cancel() 

250 

251 self._hold_task = None 

252 

253 

254class ConstraintManager: 

255 """Manages the lifecycle of all constraints.""" 

256 

257 CONSTRAINT_RESTART_GRACE = 5.0 

258 CONSTRAINT_RESTART_DELAY = 5.0 

259 

260 def __init__(self, constraints: list[Constraint]): 

261 self._constraints = constraints 

262 self._entries: list[ConstraintState] = [ConstraintState(kind=c.kind) for c in constraints] 

263 self._constrained_set: set[int] = set() 

264 self._ready_set: set[int] = set() 

265 self._ready_event = asyncio.Event() 

266 

267 if not constraints: 

268 self._ready_event.set() 

269 

270 @property 

271 def entries(self) -> Sequence[ConstraintState]: 

272 """Current constraint state in config order.""" 

273 return self._entries 

274 

275 def is_constrained(self) -> bool: 

276 return len(self._constrained_set) > 0 

277 

278 async def start( 

279 self, 

280 *, 

281 task_group: asyncio.TaskGroup, 

282 ready_timeout: float | None = None, 

283 **kwargs, 

284 ) -> bool: 

285 """Start all constraint monitoring tasks. 

286 

287 Creates supervisor tasks for each constraint that will monitor and update 

288 their states continuously. Waits for all constraints to become ready or 

289 until the ready_timeout expires. 

290 

291 Args: 

292 task_group: The asyncio.TaskGroup to create constraint supervisor tasks in. 

293 ready_timeout: Maximum time in seconds to wait for all constraints to become ready. 

294 If None, waits indefinitely. 

295 **kwargs: Additional keyword arguments passed to each constraint's check_task method. 

296 

297 Returns: 

298 bool: True if all constraints became ready within the timeout, False if one or more 

299 constraints remain unready. 

300 """ 

301 logger.debug(f"ConstraintManager starting with {len(self._constraints)} constraints") 

302 

303 # Start a supervisor task to manage each constraint's evaluation loop. 

304 for i, constraint in enumerate(self._constraints): 

305 task_group.create_task(self._constraint_supervisor(i, constraint, **kwargs)) 

306 

307 try: 

308 async with asyncio.timeout(ready_timeout): 

309 await self._ready_event.wait() 

310 except asyncio.TimeoutError: 

311 return False 

312 

313 return True 

314 

315 def _set_state(self, idx: int, current: ConstraintState): 

316 prev = self._entries[idx] 

317 self._entries[idx] = current 

318 

319 if current.state != prev.state: 

320 state_str = current.state.replace("_", " ") 

321 reason = f"(reason: {current.reason})" if current.reason else "" 

322 logger.info(f"{current.kind.capitalize()} constraint is {state_str} {reason}") 

323 

324 if current.state == "not_ready": 

325 self._ready_set.discard(idx) 

326 self._constrained_set.add(idx) 

327 else: 

328 self._ready_set.add(idx) 

329 

330 if current.state == "clear": 

331 self._constrained_set.discard(idx) 

332 else: 

333 self._constrained_set.add(idx) 

334 

335 if len(self._ready_set) == len(self._entries): 

336 self._ready_event.set() 

337 else: 

338 self._ready_event.clear() 

339 

340 async def _constraint_supervisor(self, idx: int, constraint: Constraint, **kwargs): 

341 restarting = False 

342 initial_hold = 0.0 

343 

344 while True: 

345 evaluator: ConstraintEvaluator | None = None 

346 

347 try: 

348 initial_timeout = ( 

349 None 

350 if constraint.ttl is None 

351 else self.CONSTRAINT_RESTART_GRACE + constraint.ttl 

352 ) 

353 

354 async with asyncio.timeout(initial_timeout) as timeout: 

355 # Create a constraint evaluator and set the initial state. 

356 evaluator = ConstraintEvaluator( 

357 constraint, timeout=timeout, initial_hold=initial_hold 

358 ) 

359 self._set_state(idx, evaluator.actual_state) 

360 

361 # Sleep if this was a restart. 

362 if restarting: 

363 timeout.reschedule( 

364 None 

365 if (when := timeout.when()) is None 

366 else when + self.CONSTRAINT_RESTART_DELAY 

367 ) 

368 await asyncio.sleep(self.CONSTRAINT_RESTART_DELAY) 

369 

370 # Run the constraint and process updates. 

371 await self._constraint_task(idx, constraint, evaluator, **kwargs) 

372 except asyncio.CancelledError: 

373 # Propagate only direct cancellation. 

374 if asyncio.current_task().cancelling(): 

375 raise 

376 

377 logger.error(f"{constraint.kind.capitalize()} constraint cancelled unexpectedly") 

378 except asyncio.TimeoutError: 

379 # Only log timeouts for non-optional constraints. 

380 if not constraint.optional: 

381 logger.error(f"{constraint.kind.capitalize()} constraint timed out") 

382 except Exception: 

383 logger.exception(f"{constraint.kind.capitalize()} constraint error") 

384 

385 # A hold outlives the evaluator that started it, so the replacement resumes it. 

386 initial_hold = evaluator.hold_remaining() if evaluator is not None else 0.0 

387 restarting = True 

388 

389 async def _constraint_task( 

390 self, 

391 idx: int, 

392 constraint: Constraint, 

393 evaluator: ConstraintEvaluator, 

394 **kwargs, 

395 ): 

396 if not constraint.optional: 

397 logger.debug(f"starting {constraint.kind} constraint check") 

398 

399 # Start the constraint check task. 

400 check_task = asyncio.create_task(constraint.check_task(evaluator, **kwargs)) 

401 update_task: asyncio.Task | None = None 

402 

403 try: 

404 # Wait for the check task to signal that it is ready. 

405 await evaluator.wait_ready() 

406 

407 # Read events emitted by the check task. 

408 while True: 

409 update_task: asyncio.Task = asyncio.create_task(evaluator.next_update()) 

410 done, _ = await asyncio.wait( 

411 {update_task, check_task}, 

412 return_when=asyncio.FIRST_COMPLETED, 

413 ) 

414 

415 if update_task in done: 

416 state = update_task.result() 

417 self._set_state(idx, state) 

418 

419 if check_task in done: 

420 # Re-raise check task exception, if any. 

421 check_task.result() 

422 break 

423 finally: 

424 aws = [evaluator.cleanup(), check_task] 

425 check_task.cancel() 

426 

427 if update_task: 

428 update_task.cancel() 

429 aws.append(update_task) 

430 

431 await asyncio.gather(*aws, return_exceptions=True) 

432 

433 

434class GenericConstraint(Constraint): 

435 """Generic constraint driven by a Condition evaluated against any entity keyword.""" 

436 

437 kind: Literal["conditional"] = "conditional" 

438 entity: str 

439 keyword: str 

440 field: str | None = None 

441 condition: AnyCondition 

442 

443 @override 

444 async def check_task(self, evaluator: ConstraintEvaluator, kit: SensorKit): 

445 label = f"{self.entity}.{self.keyword}" 

446 if self.field: 

447 label += f".{self.field}" 

448 logger.debug(f"monitoring conditional constraint on {label}") 

449 

450 _NONE_YET = object() 

451 previous = _NONE_YET 

452 was_active = False 

453 

454 client = kit.entity(self.entity) 

455 consumer = await client._stream.consume(self.keyword) 

456 

457 async for msg in consumer: 

458 try: 

459 data = json.loads(msg.data) 

460 except Exception: 

461 continue 

462 

463 current = resolve_field(data, self.field) if self.field else data 

464 

465 if previous is not _NONE_YET: 

466 _, is_active = self.condition.evaluate(current, previous, was_active) 

467 reason = f"{label} = {current}" 

468 

469 if is_active: 

470 evaluator.constrain(reason) 

471 else: 

472 evaluator.clear(reason) 

473 

474 evaluator.ready() 

475 was_active = is_active 

476 

477 previous = current