-
Notifications
You must be signed in to change notification settings - Fork 15
/
impl.py
457 lines (382 loc) · 15.4 KB
/
impl.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
# Copyright 2021 The Duet Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Internal implementation details for duet."""
import enum
import functools
import heapq
import itertools
import signal
import threading
import time
from concurrent.futures import Future
from contextvars import ContextVar
from typing import (
Any,
Awaitable,
Callable,
cast,
Coroutine,
Generic,
Iterator,
List,
Optional,
Set,
TypeVar,
)
import duet.futuretools as futuretools
T = TypeVar("T")
class Interrupt(BaseException):
def __init__(self, task, error):
self.task = task
self.error = error
class TaskState(enum.Enum):
WAITING = 0
SUCCEEDED = 1
FAILED = 2
class TaskStateError(Exception):
def __init__(self, state: TaskState, expected_state: TaskState) -> None:
self.state = state
self.expected_state = expected_state
super().__init__(f"state: {state}, expected: {expected_state}")
# Sentinel local variable name that we insert into coroutines.
# This allows us to detect whether a task is running when we get Ctrl-C.
LOCALS_TASK_SCHEDULER = "__duet_task_scheduler__"
class Task(Generic[T]):
def __init__(
self,
awaitable: Awaitable[T],
scheduler: "Scheduler",
main_task: Optional["Task"],
) -> None:
self.scheduler = scheduler
self.main_task = main_task
self._state = TaskState.WAITING
self._future: Optional[Future] = None
self._ready_future = futuretools.AwaitableFuture[None]()
self._ready_future.set_result(None) # Ready to advance.
self.interruptible = True
self._interrupt: Optional[Interrupt] = None
self._result: Optional[T] = None
self._error: Optional[Exception] = None
self._deadlines: List[DeadlineEntry] = []
if main_task and main_task.deadline is not None:
self.push_deadline(main_task.deadline)
self._generator = awaitable.__await__() # Returns coroutine generator.
if isinstance(awaitable, Coroutine):
awaitable.cr_frame.f_locals.setdefault(LOCALS_TASK_SCHEDULER, scheduler)
def _check_state(self, expected_state: TaskState) -> None:
if self._state != expected_state:
raise TaskStateError(self._state, expected_state)
@property
def future(self) -> Optional[Future]:
self._check_state(TaskState.WAITING)
return self._future
@property
def result(self) -> T:
self._check_state(TaskState.SUCCEEDED)
return cast(T, self._result)
@property
def done(self) -> bool:
return self._state == TaskState.SUCCEEDED or self._state == TaskState.FAILED
def add_ready_callback(self, callback: Callable[["Task"], Any]) -> None:
self._check_state(TaskState.WAITING)
self._ready_future.add_done_callback(lambda _: callback(self))
def advance(self):
if self.done:
return
if self._state == TaskState.WAITING:
self._ready_future.result()
token = _current_task.set(self)
try:
if self._interrupt:
interrupt = self._interrupt
self._interrupt = None
if interrupt.task is self:
error = interrupt.error
else:
error = interrupt
f = self._generator.throw(error)
else:
f = next(self._generator)
except StopIteration as e:
self._result = e.value
self._state = TaskState.SUCCEEDED
return
except (Interrupt, Exception) as error:
self._error = error
self._state = TaskState.FAILED
if self.main_task:
self.main_task.interrupt(self, error)
return
else:
raise
else:
if not isinstance(f, Future):
raise TypeError(f"expected Future, got {type(f)}: {f}")
ready_future = futuretools.AwaitableFuture()
f.add_done_callback(lambda _: ready_future.try_set_result(None))
self._future = f
self._ready_future = ready_future
self._state = TaskState.WAITING
finally:
_current_task.reset(token)
def push_deadline(self, deadline: float) -> None:
if self._deadlines:
deadline = min(self._deadlines[-1].deadline, deadline)
entry = self.scheduler.add_deadline(self, deadline)
self._deadlines.append(entry)
def pop_deadline(self) -> None:
entry = self._deadlines.pop(-1)
entry.valid = False
@property
def deadline(self) -> Optional[float]:
return self._deadlines[-1].deadline if self._deadlines else None
def interrupt(self, task, error):
if self.done or not self.interruptible or self._interrupt:
return
self._interrupt = Interrupt(task, error)
self._ready_future.try_set_result(None)
if self._future:
self._future.cancel()
def close(self):
self._generator.close()
self.scheduler = None
self.main_task = None
_current_task: ContextVar[Task] = ContextVar("current_task")
def current_task() -> Task:
"""Gets the currently-running duet task.
This must be called from within a running async function, or else it will
raise a RuntimeError.
"""
try:
return _current_task.get()
except LookupError:
raise RuntimeError("Can only be called from an async function.")
def current_scheduler() -> "Scheduler":
"""Gets the currently-running duet scheduler.
This must be called from within a running async function, or else it will
raise a RuntimeError.
"""
return current_task().scheduler
def any_ready(tasks: Set[Task]) -> futuretools.AwaitableFuture[None]:
"""Returns a Future that will fire when any of the given tasks is ready."""
if not tasks or any(task.done for task in tasks):
return futuretools.completed_future(None)
f = futuretools.AwaitableFuture[None]()
for task in tasks:
task.add_ready_callback(lambda _: f.try_set_result(None))
return f
class ReadySet:
"""Container for an ordered set of tasks that are ready to advance."""
def __init__(self):
self._cond = threading.Condition()
self._buffer = futuretools.BufferGroup()
self._tasks: List[Task] = []
self._task_set: Set[Task] = set()
def register(self, task: Task) -> None:
"""Registers task to be added to this set when it is ready."""
self._buffer.add(task.future)
task.add_ready_callback(self._add)
def _add(self, task: Task) -> None:
"""Adds the given task to the ready set, if it is not already there."""
with self._cond:
if task not in self._task_set:
self._task_set.add(task)
self._tasks.append(task)
self._cond.notify()
def get_all(self, timeout: Optional[float] = None) -> List[Task]:
"""Gets all ready tasks and clears the ready set.
If no tasks are ready yet, we flush buffered futures to notify them
that they should proceed, and then block until one or more tasks become
ready.
"""
with self._cond:
if self._tasks:
return self._pop_tasks()
# Flush buffered futures to ensure we make progress. Note that we must
# release the condition lock before flushing to avoid a deadlock if
# buffered futures complete and trigger a call to self._add.
self._buffer.flush()
with self._cond:
if not self._tasks:
if not self._cond.wait(timeout):
raise TimeoutError()
return self._pop_tasks()
def _pop_tasks(self) -> List[Task]:
tasks = self._tasks
self._tasks = []
self._task_set.clear()
return tasks
def interrupt(self) -> None:
with self._cond:
self._cond.notify()
@functools.total_ordering
class DeadlineEntry:
"""A entry for one Deadline in the Scheduler's priority queue.
This follows the implementation notes in the stdlib heapq docs:
https://docs.python.org/3/library/heapq.html#priority-queue-implementation-notes
Attributes:
task: The task associated with this deadline.
deadline: Absolute time when the deadline will elapse.
count: Monotonically-increasing counter to preserve creation order when
comparing entries with the same deadline.
valid: Flag indicating whether the deadline is still valid. If the task
exits its scope before the deadline elapses, we mark the deadline as
invalid but leave it in the scheduler's priority queue since removal
would require an O(n) scan. The scheduler ignores invalid deadlines
when they elapse.
"""
_counter = itertools.count()
def __init__(self, task: Task, deadline: float):
self.task = task
self.deadline = deadline
self.count = next(self._counter)
self._cmp_val = (deadline, self.count)
self.valid = True
def __eq__(self, other: Any) -> bool:
if not isinstance(other, DeadlineEntry):
return NotImplemented
return self._cmp_val == other._cmp_val
def __lt__(self, other: Any) -> bool:
if not isinstance(other, DeadlineEntry):
return NotImplemented
return self._cmp_val < other._cmp_val
def __repr__(self) -> str:
return f"DeadlineEntry({self.task}, {self.deadline}, {self.count})"
class Scheduler:
def __init__(self) -> None:
self.active_tasks: Set[Task] = set()
self._ready_tasks = ReadySet()
self._prev_signal: Optional[Callable] = None
self._interrupted = False
self._deadlines: List[DeadlineEntry] = []
def spawn(self, awaitable: Awaitable[Any], main_task: Optional[Task] = None) -> Task:
"""Spawns a new Task to run an awaitable in this Scheduler.
Note that the task will not be advanced until the next scheduler tick.
Also, note that this function is safe to call from sync code (such as
duet.run) or async code (such as within a scope).
Args:
func: The async function to run.
*args: Args for func.
**kwds: Keyword args for func.
Returns:
A Task to run the given awaitable.
"""
task = Task(awaitable, scheduler=self, main_task=main_task)
self.active_tasks.add(task)
self._ready_tasks.register(task)
return task
def time(self) -> float:
return time.time()
def add_deadline(self, task: Task, deadline: float) -> DeadlineEntry:
entry = DeadlineEntry(task, deadline=deadline)
heapq.heappush(self._deadlines, entry)
return entry
def get_next_deadline(self) -> Optional[float]:
while self._deadlines:
if not self._deadlines[0].valid:
heapq.heappop(self._deadlines)
continue
return self._deadlines[0].deadline
return None
def get_deadline_tasks(self, deadline: float) -> Iterator[Task]:
while self._deadlines and self._deadlines[0].deadline <= deadline:
entry = heapq.heappop(self._deadlines)
if entry.valid:
yield entry.task
def tick(self):
"""Runs the scheduler ahead by one tick.
This waits for at least one active task to complete, then advances all
ready tasks and sets up a new future to be notified later by tasks that
are still active (or yet to be spawned). Raises a RuntimeError if there
are no currently active tasks.
"""
if not self.active_tasks:
raise RuntimeError("tick called with no active tasks")
if self._interrupted:
task = next(iter(self.active_tasks))
task.interrupt(task, KeyboardInterrupt)
self._interrupted = False
deadline = self.get_next_deadline()
timeout = None if deadline is None else max(deadline - self.time(), 0)
try:
ready_tasks = self._ready_tasks.get_all(timeout)
except TimeoutError:
for task in self.get_deadline_tasks(deadline):
task.interrupt(task, TimeoutError())
ready_tasks = self._ready_tasks.get_all(None)
for task in ready_tasks:
try:
task.advance()
finally:
if task.done:
task.close()
self.active_tasks.discard(task)
else:
self._ready_tasks.register(task)
def _interrupt(self, signum: int, frame: Optional[Any]) -> None:
"""Interrupt signal handler used while this scheduler is running.
This is inspired by trio's interrupt handling, described here:
https://vorpus.org/blog/control-c-handling-in-python-and-trio/
If the interrupted frame is inside a running task, which we detect by
looking for a special local variable inserted into the task coroutine,
we simply raise a KeyboardInterrupt as usual. Otherwise we set a flag
which will get checked on the next tick() and cause a task to be
interrupted.
One important difference from trio is that duet is reentrant, so when
detecting whether we are in a task we have to check whether the task's
scheduler is self. If the interrupted frame is running in a task of a
different scheduler, that should not raise KeyboardInterrupt directly.
"""
if self._in_task(frame):
raise KeyboardInterrupt
else:
self._interrupted = True
self._ready_tasks.interrupt()
def _in_task(self, frame) -> bool:
while frame is not None:
if frame.f_locals.get(LOCALS_TASK_SCHEDULER, None) is self:
return True
frame = frame.f_back
return False
def __enter__(self):
if (
threading.current_thread() == threading.main_thread()
and signal.getsignal(signal.SIGINT) == signal.default_int_handler
):
self._prev_signal = signal.signal(signal.SIGINT, self._interrupt)
return self
def __exit__(self, exc_type, exc, tb):
def finish_tasks(error=None):
if error:
for task in self.active_tasks:
task.interrupt(None, error)
while self.active_tasks:
try:
self.tick()
except Exception:
if not error:
raise
try:
if exc:
finish_tasks(exc)
else:
try:
finish_tasks()
except Exception as exc:
finish_tasks(exc)
raise
finally:
if self._prev_signal:
signal.signal(signal.SIGINT, self._prev_signal)