Coverage for async_durable_execution/_runner/local/scheduler.py: 96%
131 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +0000
1"""A single-threaded asyncio scheduler for local runner callbacks."""
3from __future__ import annotations
5import asyncio
6import logging
7from typing import TYPE_CHECKING, Any
9if TYPE_CHECKING:
10 from collections.abc import Awaitable, Callable
12logger = logging.getLogger(__name__)
15class Event:
16 """An event created by Scheduler that will block on wait until it's set."""
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
23 def set(self) -> None:
24 """Set the event with this to unblock wait."""
25 self._scheduler.set_event(self._event)
27 def set_exception(self, exception: Exception) -> None:
28 """Set exception and unblock waiters."""
29 self._exception = exception
30 self._scheduler.set_event(self._event)
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
44class Scheduler:
45 """A Scheduler to run callables later on one asyncio event loop."""
47 def __init__(self) -> None:
48 self._running: bool = False
49 self._stopping: bool = False
50 self._events: set[asyncio.Event] = set()
51 self._tasks: set[asyncio.Future[Any]] = set()
52 self._running_tasks: dict[asyncio.Future[Any], asyncio.Task[Any]] = {}
53 self._loop: asyncio.AbstractEventLoop | None = None
55 def __enter__(self) -> Scheduler:
56 self.start()
57 return self
59 def __exit__(self, exc_type, exc_val, exc_tb) -> None:
60 self.stop()
62 def start(self) -> None:
63 """Start the scheduler. Not thread-safe."""
64 if self._running:
65 return
66 self._loop = self._get_or_create_loop()
67 self._running = True
69 def stop(self) -> None:
70 """Stop the scheduler, releasing resources. Not thread-safe."""
71 if not self._running:
72 return
74 self._running = False
75 self._stopping = True
77 running_tasks = list(self._running_tasks.values())
78 futures = list(self._tasks)
79 self._events.clear()
80 self._running_tasks.clear()
81 self._tasks.clear()
83 for task in running_tasks:
84 task.cancel()
85 for future in futures:
86 future.cancel()
88 self._stopping = False
90 def get_loop(self) -> asyncio.AbstractEventLoop:
91 """Return the scheduler event loop, creating one for sync compatibility."""
92 if self._loop is None or self._loop.is_closed(): 92 ↛ 93line 92 didn't jump to line 93 because the condition on line 92 was never true
93 self._loop = self._get_or_create_loop()
94 return self._loop
96 def call_later(
97 self,
98 func: Callable[[], Awaitable[Any]],
99 delay: float = 0,
100 count: int | None = 1,
101 completion_event: Event | None = None,
102 ) -> asyncio.Future[Any]:
103 """Call func after the delay."""
104 loop = self.get_loop()
105 if not self._running or self._stopping:
106 cancelled_future: asyncio.Future[Any] = loop.create_future()
107 cancelled_future.cancel()
108 return cancelled_future
110 future: asyncio.Future[Any] = loop.create_future()
111 if count == 0:
112 future.set_result(None)
113 return future
115 def cleanup(_future: asyncio.Future[Any]) -> None:
116 self._tasks.discard(_future)
117 if task := self._running_tasks.pop(_future, None):
118 task.cancel()
119 if _future.done() and not _future.cancelled():
120 _future.exception()
122 async def execute() -> None:
123 if future.cancelled():
124 return
125 try:
126 await asyncio.sleep(delay)
127 if not self._running or self._stopping or future.cancelled(): 127 ↛ 128line 127 didn't jump to line 128 because the condition on line 127 was never true
128 return
129 result = await func()
130 except Exception as err:
131 if completion_event:
132 completion_event.set_exception(err)
133 else:
134 msg: str = "error in scheduled task"
135 logger.exception(msg)
136 if not future.done(): 136 ↛ exitline 136 didn't return from function 'execute' because the condition on line 136 was always true
137 future.set_exception(err)
138 else:
139 if not future.done(): 139 ↛ exitline 139 didn't return from function 'execute' because the condition on line 139 was always true
140 future.set_result(result)
142 task = loop.create_task(execute())
143 self._running_tasks[future] = task
145 def discard_task(_task: asyncio.Task[Any]) -> None:
146 self._running_tasks.pop(future, None)
148 task.add_done_callback(discard_task)
149 future.add_done_callback(cleanup)
150 self._tasks.add(future)
151 return future
153 def create_event(self) -> Event:
154 """Create an event controlled by the Scheduler."""
155 event = asyncio.Event()
156 self._events.add(event)
157 return Event(self, event)
159 async def wait_for_event(
160 self, event: asyncio.Event, timeout: float | None = None
161 ) -> bool:
162 """Wait for an event if it is still tracked by the Scheduler."""
163 if event not in self._events:
164 return False
165 try:
166 await asyncio.wait_for(event.wait(), timeout=timeout)
167 except asyncio.TimeoutError:
168 return False
169 return event in self._events
171 def set_event(self, event: asyncio.Event) -> None:
172 """Set event if it is still tracked by the Scheduler."""
173 should_set = event in self._events
174 if should_set:
175 event.set()
177 def remove_event(self, event: asyncio.Event) -> None:
178 """Remove event from Scheduler."""
179 self._events.discard(event)
180 event.set()
182 @staticmethod
183 def _get_or_create_loop() -> asyncio.AbstractEventLoop:
184 try:
185 return asyncio.get_running_loop()
186 except RuntimeError:
187 pass
189 try:
190 loop = asyncio.get_event_loop()
191 except RuntimeError:
192 loop = asyncio.new_event_loop()
193 asyncio.set_event_loop(loop)
194 return loop