Coverage for async-durable-execution/src/async_durable_execution/runner/local/scheduler.py: 97%

1"""A single-threaded asyncio scheduler for local runner callbacks.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import logging 

7from typing import TYPE_CHECKING, Any 

8 

9if TYPE_CHECKING: 

10 from collections.abc import Awaitable, Callable 

11 

12logger = logging.getLogger(__name__) 

13 

14 

15class Event: 

16 """An event created by Scheduler that will block on wait until it's set.""" 

17 

18 def __init__(self, scheduler: Scheduler, event: asyncio.Event) -> None: 

19 self._scheduler: Scheduler = scheduler 

20 self._event: asyncio.Event = event 

21 self._exception: Exception | None = None 

22 

23 def set(self): 

24 """Set the event with this to unblock wait.""" 

25 self._scheduler.set_event(self._event) 

26 

27 def set_exception(self, exception: Exception): 

28 """Set exception and unblock waiters.""" 

29 self._exception = exception 

30 self._scheduler.set_event(self._event) 

31 

32 async def wait_async( 

33 self, timeout: float | None = None, *, clear_on_set: bool = True 

34 ) -> bool: 

35 """Wait until the event is set.""" 

36 result = await self._scheduler.wait_for_event(self._event, timeout) 

37 if clear_on_set: 

38 self._scheduler.remove_event(self._event) 

39 if result and self._exception: 

40 raise self._exception 

41 return result 

42 

43 def wait(self, timeout: float | None = None, *, clear_on_set: bool = True) -> bool: 

44 """Synchronously wait for compatibility with scheduler unit tests.""" 

45 loop = self._scheduler.get_loop() 

46 if loop.is_running(): 

47 msg = "Event.wait() cannot block a running event loop; use wait_async()." 

48 raise RuntimeError(msg) 

49 return loop.run_until_complete( 

50 self.wait_async(timeout=timeout, clear_on_set=clear_on_set) 

51 ) 

52 

53 def remove(self): 

54 """Remove the event from the Scheduler.""" 

55 self._scheduler.remove_event(self._event) 

56 

57 

58class Scheduler: 

59 """A Scheduler to run callables later on one asyncio event loop.""" 

60 

61 def __init__(self) -> None: 

62 self._running: bool = False 

63 self._stopping: bool = False 

64 self._events: set[asyncio.Event] = set() 

65 self._tasks: set[asyncio.Future[Any]] = set() 

66 self._running_tasks: dict[asyncio.Future[Any], asyncio.Task[Any]] = {} 

67 self._loop: asyncio.AbstractEventLoop | None = None 

68 

69 def __enter__(self): 

70 self.start() 

71 return self 

72 

73 def __exit__(self, exc_type, exc_val, exc_tb): 

74 self.stop() 

75 

76 def start(self): 

77 """Start the scheduler. Not thread-safe.""" 

78 if self._running: 

79 return 

80 self._loop = self._get_or_create_loop() 

81 self._running = True 

82 

83 def stop(self): 

84 """Stop the scheduler, releasing resources. Not thread-safe.""" 

85 if not self._running: 

86 return 

87 

88 self._running = False 

89 self._stopping = True 

90 

91 running_tasks = list(self._running_tasks.values()) 

92 futures = list(self._tasks) 

93 self._events.clear() 

94 self._running_tasks.clear() 

95 self._tasks.clear() 

96 

97 for task in running_tasks: 

98 task.cancel() 

99 for future in futures: 

100 future.cancel() 

101 

102 self._stopping = False 

103 

104 def get_loop(self) -> asyncio.AbstractEventLoop: 

105 """Return the scheduler event loop, creating one for sync compatibility.""" 

106 if self._loop is None or self._loop.is_closed(): 

107 self._loop = self._get_or_create_loop() 

108 return self._loop 

109 

110 def is_started(self) -> bool: 

111 """Return True if the scheduler is started.""" 

112 return self._running 

113 

114 def event_count(self) -> int: 

115 """Return the number of events in the scheduler.""" 

116 return len(self._events) 

117 

118 def task_count(self) -> int: 

119 """Return the number of scheduled tasks that are not done.""" 

120 return sum(1 for task in self._tasks if not task.done()) 

121 

122 def call_later( 

123 self, 

124 func: Callable[[], Awaitable[Any]], 

125 delay: float = 0, 

126 count: int | None = 1, # noqa: ARG002 

127 completion_event: Event | None = None, 

128 ) -> asyncio.Future[Any]: 

129 """Call func after the delay.""" 

130 loop = self.get_loop() 

131 if not self._running or self._stopping: 

132 cancelled_future: asyncio.Future[Any] = loop.create_future() 

133 cancelled_future.cancel() 

134 return cancelled_future 

135 

136 future: asyncio.Future[Any] = loop.create_future() 

137 if count == 0: 

138 future.set_result(None) 

139 return future 

140 

141 def cleanup(_future: asyncio.Future[Any]) -> None: 

142 self._tasks.discard(_future) 

143 if task := self._running_tasks.pop(_future, None): 

144 task.cancel() 

145 if _future.done() and not _future.cancelled(): 

146 _future.exception() 

147 

148 async def execute() -> None: 

149 if future.cancelled(): 

150 return 

151 try: 

152 await asyncio.sleep(delay) 

153 if not self._running or self._stopping or future.cancelled(): 

154 return 

155 result = await func() 

156 except Exception as err: 

157 if completion_event: 

158 completion_event.set_exception(err) 

159 else: 

160 msg: str = "error in scheduled task" 

161 logger.exception(msg) 

162 if not future.done(): 

163 future.set_exception(err) 

164 else: 

165 if not future.done(): 

166 future.set_result(result) 

167 

168 task = loop.create_task(execute()) 

169 self._running_tasks[future] = task 

170 

171 def discard_task(_task: asyncio.Task[Any]) -> None: 

172 self._running_tasks.pop(future, None) 

173 

174 task.add_done_callback(discard_task) 

175 future.add_done_callback(cleanup) 

176 self._tasks.add(future) 

177 return future 

178 

179 def create_event(self) -> Event: 

180 """Create an event controlled by the Scheduler.""" 

181 event = asyncio.Event() 

182 self._events.add(event) 

183 return Event(self, event) 

184 

185 async def wait_for_event( 

186 self, event: asyncio.Event, timeout: float | None = None 

187 ) -> bool: 

188 """Wait for an event if it is still tracked by the Scheduler.""" 

189 if event not in self._events: 

190 return False 

191 try: 

192 await asyncio.wait_for(event.wait(), timeout=timeout) 

193 except asyncio.TimeoutError: 

194 return False 

195 return event in self._events 

196 

197 def set_event(self, event: asyncio.Event): 

198 """Set event if it is still tracked by the Scheduler.""" 

199 should_set = event in self._events 

200 if should_set: 

201 event.set() 

202 

203 def remove_event(self, event: asyncio.Event): 

204 """Remove event from Scheduler.""" 

205 self._events.discard(event) 

206 event.set() 

207 

208 @staticmethod 

209 def _get_or_create_loop() -> asyncio.AbstractEventLoop: 

210 try: 

211 return asyncio.get_running_loop() 

212 except RuntimeError: 

213 pass 

214 

215 try: 

216 loop = asyncio.get_event_loop() 

217 except RuntimeError: 

218 loop = asyncio.new_event_loop() 

219 asyncio.set_event_loop(loop) 

220 return loop