Coverage for core / src / sensorkit / common / aio.py: 96%

1# SPDX-License-Identifier: Apache-2.0 

2"""Async utilities: scoped futures, periodic loops, value latches, and observer queues.""" 

3 

4from __future__ import annotations 

5 

6import asyncio 

7import contextlib 

8from collections.abc import AsyncGenerator, Awaitable, Callable 

9from typing import AsyncContextManager, ClassVar, Self, overload 

10 

11from sensorkit.common.logging import limited_logger 

12 

13 

14def scoped_waiter[T](aw: Awaitable[T]) -> AsyncContextManager[asyncio.Future[T]]: 

15 """Create an async context manager wrapping an awaitable in a scoped `asyncio.Future`. 

16 

17 The future is cancelled if it has not completed before the `async with` block 

18 exits or if an exception propagates out of the block. 

19 

20 Important: 

21 Any exceptions raised by the awaitable are silently discarded. This makes 

22 `scoped_waiter` suitable only for "waiter" type tasks where exceptions either 

23 won't occur or are not meaningful to handle (e.g. background monitoring tasks, 

24 optional notifications, or advisory operations). 

25 

26 Args: 

27 aw: The awaitable object to be wrapped in an `asyncio.Future`. 

28 

29 Returns: 

30 An async context manager that yields the `asyncio.Future` for the awaitable. 

31 """ 

32 return _scoped_waiter(aw) 

33 

34 

35@overload 

36def cleanup_future(fut: asyncio.Task): ... 

37 

38@overload 

39def cleanup_future(fut: asyncio.Future): ... 

40 

41def cleanup_future(fut: asyncio.Future | asyncio.Task): 

42 """Safely clean up a Future or an already-completed Task. 

43 

44 Args: 

45 fut: The asyncio.Future to clean up. 

46 """ 

47 if not fut.done(): 

48 fut.cancel() 

49 elif not fut.cancelled(): 

50 fut.exception() 

51 

52 

53@contextlib.asynccontextmanager 

54async def _scoped_waiter(aw: Awaitable) -> AsyncGenerator[asyncio.Task]: 

55 fut = asyncio.ensure_future(aw) 

56 

57 try: 

58 yield fut 

59 finally: 

60 if not fut.done(): 

61 fut.cancel() 

62 with contextlib.suppress(asyncio.CancelledError): 

63 await fut 

64 else: 

65 with contextlib.suppress(asyncio.CancelledError): 

66 fut.exception() 

67 

68 

69class AsyncValueLatch[T]: 

70 """Stores a value and stages pending changes.""" 

71 

72 def __init__(self, initial_value: T): 

73 self.value: T = initial_value 

74 self.pending_value: T = initial_value 

75 self._pending = asyncio.Event() 

76 

77 def update(self, value: T, only_if_different=True): 

78 """Flag a pending value change and set that pending value.""" 

79 if only_if_different and value == self.pending_value: 

80 return False 

81 

82 self.pending_value = value 

83 self._pending.set() 

84 return True 

85 

86 def pending_change(self): 

87 """Return True if there is a pending value change.""" 

88 return self._pending.is_set() 

89 

90 async def wait_until_pending(self): 

91 """Wait until there is a pending value change.""" 

92 await self._pending.wait() 

93 

94 def apply(self): 

95 """Apply the pending value if there is one and return the current value.""" 

96 if self._pending.is_set(): 

97 self._pending.clear() 

98 self.value = self.pending_value 

99 

100 return self.value 

101 

102 

103class AsyncObserver[T]: 

104 """Fan out notified values to independent, bounded subscriber queues. 

105 

106 Each subscriber receives values notified after it subscribes, in order. With 

107 `initial_value=True`, it first receives the current value, if one has been set. 

108 When a subscriber's queue is full, its oldest pending value is discarded. 

109 """ 

110 

111 NOT_SET: ClassVar[object] = object() 

112 DEFAULT_MAXSIZE: ClassVar[int] = 1024 

113 

114 def __init__(self, initial_value: T = NOT_SET): 

115 self._observers: set[asyncio.Queue[T]] = set() 

116 self._value = initial_value 

117 self._dropped = 0 

118 

119 def subscribe(self, *, initial_value: bool = False, maxsize: int = DEFAULT_MAXSIZE): 

120 """Create and return a new observer queue. 

121 

122 Args: 

123 initial_value: Whether to seed the queue with the current value, if one is set. 

124 maxsize: Bound on values queued for this subscriber. 

125 

126 Raises: 

127 ValueError: If `maxsize` is not positive, which would leave the queue unbounded. 

128 """ 

129 if maxsize < 1: 

130 raise ValueError(f"maxsize must be positive, got {maxsize}") 

131 

132 queue: asyncio.Queue[T] = asyncio.Queue(maxsize) 

133 self._observers.add(queue) 

134 

135 if initial_value and self._value is not self.NOT_SET: 

136 queue.put_nowait(self._value) 

137 

138 return queue 

139 

140 def unsubscribe(self, queue: asyncio.Queue[T]): 

141 """Remove an observer queue.""" 

142 self._observers.discard(queue) 

143 queue.shutdown() 

144 

145 @contextlib.contextmanager 

146 def subscription(self, *, initial_value: bool = False, maxsize: int = DEFAULT_MAXSIZE): 

147 """Context manager that provides a subscriber queue and unsubscribes it on exit.""" 

148 queue = self.subscribe(initial_value=initial_value, maxsize=maxsize) 

149 

150 try: 

151 yield queue 

152 finally: 

153 self.unsubscribe(queue) 

154 

155 def notify(self, value: T): 

156 """Update the current value and notify all observers. 

157 

158 Delivery never blocks on a subscriber. A subscriber whose queue is full loses its 

159 oldest pending value to make room. 

160 """ 

161 self._value = value 

162 

163 for queue in self._observers: 

164 try: 

165 queue.put_nowait(value) 

166 except asyncio.QueueFull: 

167 queue.get_nowait() 

168 queue.task_done() 

169 queue.put_nowait(value) 

170 self._dropped += 1 

171 

172 limited_logger().warning( 

173 f"observer queue full at {queue.maxsize}; dropping the oldest value. " 

174 "A subscriber is not keeping up, or leaked without unsubscribing." 

175 ) 

176 

177 @property 

178 def value(self) -> T: 

179 """Get the current value.""" 

180 if self._value is self.NOT_SET: 

181 raise RuntimeError("Observed value must be set prior to access") 

182 

183 return self._value 

184 

185 @property 

186 def subscriber_count(self) -> int: 

187 """The number of queues currently subscribed.""" 

188 return len(self._observers) 

189 

190 @property 

191 def dropped(self) -> int: 

192 """The number of values dropped because a subscriber's queue was full.""" 

193 return self._dropped 

194 

195 async def consume(self, *, initial_value: bool = False, maxsize: int = DEFAULT_MAXSIZE): 

196 """Yield successive values as an async generator, unsubscribing automatically on exit.""" 

197 queue = self.subscribe(initial_value=initial_value, maxsize=maxsize) 

198 

199 try: 

200 while True: 

201 yield await queue.get() 

202 queue.task_done() 

203 except GeneratorExit: 

204 queue.task_done() 

205 except asyncio.QueueShutDown as e: 

206 raise StopAsyncIteration() from e 

207 finally: 

208 self.unsubscribe(queue) 

209 

210 

211class AsyncLoop: 

212 """Runs a function repeatedly on a fixed interval. 

213 

214 The interval is a plain attribute, so assigning to it retunes a running loop from 

215 its next sleep onward. 

216 """ 

217 

218 def __init__( 

219 self, 

220 func: Callable[[], Awaitable[None] | None], 

221 *, 

222 interval: float, 

223 log: bool = False, 

224 label: str | None = None, 

225 ): 

226 """Initialize the loop without starting it. 

227 

228 Args: 

229 func: Called once per iteration. Takes no arguments. 

230 interval: Seconds to sleep between iterations. 

231 log: Whether to log exceptions raised by `func`. 

232 label: Optional label for logging. 

233 """ 

234 self.func = func 

235 self.interval = interval 

236 self._log = log 

237 self._label = label or getattr(func, "__qualname__", str(func)) 

238 self._task: asyncio.Task | None = None 

239 

240 @property 

241 def active(self) -> bool: 

242 """Report whether the loop is currently running.""" 

243 return self._task is not None and not self._task.done() 

244 

245 def start(self) -> Self: 

246 """Start the loop, returning self. 

247 

248 Does nothing if the loop is already running, so callers that cannot easily 

249 tell may start it unconditionally. 

250 """ 

251 if not self.active: 

252 self._task = asyncio.create_task(self._run()) 

253 

254 return self 

255 

256 async def stop(self): 

257 """Cancel the loop and wait for it to unwind. 

258 

259 Safe to call on a loop that was never started, and on one already stopped. 

260 """ 

261 if self._task is None: 

262 return 

263 

264 task = self._task 

265 task.cancel() 

266 self._task = None 

267 

268 with contextlib.suppress(asyncio.CancelledError): 

269 await task 

270 

271 async def _run(self): 

272 while True: 

273 try: 

274 if coro := self.func(): 

275 await coro 

276 except Exception as e: 

277 if self._log: 

278 from loguru import logger 

279 

280 msg = f": {e}" if str(e) else "" 

281 logger.exception(f"{type(e).__name__} in {self._label} loop{msg}") 

282 

283 await asyncio.sleep(self.interval)