import asyncio
|
import contextlib
|
|
# PY3.9: Import Callable from typing until we drop Python 3.9 support
|
# https://github.com/python/cpython/issues/87131
|
from typing import (
|
TYPE_CHECKING,
|
Any,
|
Awaitable,
|
Callable,
|
Iterable,
|
List,
|
Optional,
|
Set,
|
Tuple,
|
TypeVar,
|
Union,
|
)
|
|
_T = TypeVar("_T")
|
|
RE_RAISE_EXCEPTIONS = (SystemExit, KeyboardInterrupt)
|
|
|
def _set_result(wait_next: "asyncio.Future[None]") -> None:
|
"""Set the result of a future if it is not already done."""
|
if not wait_next.done():
|
wait_next.set_result(None)
|
|
|
async def _wait_one(
|
futures: "Iterable[asyncio.Future[Any]]",
|
loop: asyncio.AbstractEventLoop,
|
) -> _T:
|
"""Wait for the first future to complete."""
|
wait_next = loop.create_future()
|
|
def _on_completion(fut: "asyncio.Future[Any]") -> None:
|
if not wait_next.done():
|
wait_next.set_result(fut)
|
|
for f in futures:
|
f.add_done_callback(_on_completion)
|
|
try:
|
return await wait_next
|
finally:
|
for f in futures:
|
f.remove_done_callback(_on_completion)
|
|
|
async def staggered_race(
|
coro_fns: Iterable[Callable[[], Awaitable[_T]]],
|
delay: Optional[float],
|
*,
|
loop: Optional[asyncio.AbstractEventLoop] = None,
|
) -> Tuple[Optional[_T], Optional[int], List[Optional[BaseException]]]:
|
"""
|
Run coroutines with staggered start times and take the first to finish.
|
|
This method takes an iterable of coroutine functions. The first one is
|
started immediately. From then on, whenever the immediately preceding one
|
fails (raises an exception), or when *delay* seconds has passed, the next
|
coroutine is started. This continues until one of the coroutines complete
|
successfully, in which case all others are cancelled, or until all
|
coroutines fail.
|
|
The coroutines provided should be well-behaved in the following way:
|
|
* They should only ``return`` if completed successfully.
|
|
* They should always raise an exception if they did not complete
|
successfully. In particular, if they handle cancellation, they should
|
probably reraise, like this::
|
|
try:
|
# do work
|
except asyncio.CancelledError:
|
# undo partially completed work
|
raise
|
|
Args:
|
----
|
coro_fns: an iterable of coroutine functions, i.e. callables that
|
return a coroutine object when called. Use ``functools.partial`` or
|
lambdas to pass arguments.
|
|
delay: amount of time, in seconds, between starting coroutines. If
|
``None``, the coroutines will run sequentially.
|
|
loop: the event loop to use. If ``None``, the running loop is used.
|
|
Returns:
|
-------
|
tuple *(winner_result, winner_index, exceptions)* where
|
|
- *winner_result*: the result of the winning coroutine, or ``None``
|
if no coroutines won.
|
|
- *winner_index*: the index of the winning coroutine in
|
``coro_fns``, or ``None`` if no coroutines won. If the winning
|
coroutine may return None on success, *winner_index* can be used
|
to definitively determine whether any coroutine won.
|
|
- *exceptions*: list of exceptions returned by the coroutines.
|
``len(exceptions)`` is equal to the number of coroutines actually
|
started, and the order is the same as in ``coro_fns``. The winning
|
coroutine's entry is ``None``.
|
|
"""
|
loop = loop or asyncio.get_running_loop()
|
exceptions: List[Optional[BaseException]] = []
|
tasks: Set[asyncio.Task[Optional[Tuple[_T, int]]]] = set()
|
|
async def run_one_coro(
|
coro_fn: Callable[[], Awaitable[_T]],
|
this_index: int,
|
start_next: "asyncio.Future[None]",
|
) -> Optional[Tuple[_T, int]]:
|
"""
|
Run a single coroutine.
|
|
If the coroutine fails, set the exception in the exceptions list and
|
start the next coroutine by setting the result of the start_next.
|
|
If the coroutine succeeds, return the result and the index of the
|
coroutine in the coro_fns list.
|
|
If SystemExit or KeyboardInterrupt is raised, re-raise it.
|
"""
|
try:
|
result = await coro_fn()
|
except RE_RAISE_EXCEPTIONS:
|
raise
|
except BaseException as e:
|
exceptions[this_index] = e
|
_set_result(start_next) # Kickstart the next coroutine
|
return None
|
|
return result, this_index
|
|
start_next_timer: Optional[asyncio.TimerHandle] = None
|
start_next: Optional[asyncio.Future[None]]
|
task: asyncio.Task[Optional[Tuple[_T, int]]]
|
done: Union[asyncio.Future[None], asyncio.Task[Optional[Tuple[_T, int]]]]
|
coro_iter = iter(coro_fns)
|
this_index = -1
|
try:
|
while True:
|
if coro_fn := next(coro_iter, None):
|
this_index += 1
|
exceptions.append(None)
|
start_next = loop.create_future()
|
task = loop.create_task(run_one_coro(coro_fn, this_index, start_next))
|
tasks.add(task)
|
start_next_timer = (
|
loop.call_later(delay, _set_result, start_next) if delay else None
|
)
|
elif not tasks:
|
# We exhausted the coro_fns list and no tasks are running
|
# so we have no winner and all coroutines failed.
|
break
|
|
while tasks or start_next:
|
done = await _wait_one(
|
(*tasks, start_next) if start_next else tasks, loop
|
)
|
if done is start_next:
|
# The current task has failed or the timer has expired
|
# so we need to start the next task.
|
start_next = None
|
if start_next_timer:
|
start_next_timer.cancel()
|
start_next_timer = None
|
|
# Break out of the task waiting loop to start the next
|
# task.
|
break
|
|
if TYPE_CHECKING:
|
assert isinstance(done, asyncio.Task)
|
|
tasks.remove(done)
|
if winner := done.result():
|
return *winner, exceptions
|
finally:
|
# We either have:
|
# - a winner
|
# - all tasks failed
|
# - a KeyboardInterrupt or SystemExit.
|
|
#
|
# If the timer is still running, cancel it.
|
#
|
if start_next_timer:
|
start_next_timer.cancel()
|
|
#
|
# If there are any tasks left, cancel them and than
|
# wait them so they fill the exceptions list.
|
#
|
for task in tasks:
|
task.cancel()
|
with contextlib.suppress(asyncio.CancelledError):
|
await task
|
|
return None, None, exceptions
|
