Race Async Tasks and Return the First Success

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Iterable
from typing import TypeVar

T = TypeVar("T")


async def first_success(awaitables: Iterable[Awaitable[T]], *, timeout: float) -> T:
    tasks = [asyncio.create_task(awaitable) for awaitable in awaitables]
    if not tasks:
        raise ValueError("at least one awaitable is required")

    loop = asyncio.get_running_loop()
    deadline = loop.time() + timeout
    pending = set(tasks)
    failures: list[BaseException] = []
    try:
        while pending:
            remaining = deadline - loop.time()
            if remaining <= 0:
                raise TimeoutError("no replica succeeded before the deadline")
            done, pending = await asyncio.wait(
                pending,
                timeout=remaining,
                return_when=asyncio.FIRST_COMPLETED,
            )
            if not done:
                raise TimeoutError("no replica succeeded before the deadline")
            for task in done:
                error = task.exception()
                if error is None:
                    return task.result()
                failures.append(error)
        messages = "; ".join(str(error) for error in failures)
        raise RuntimeError(f"all replicas failed: {messages}")
    finally:
        for task in pending:
            task.cancel()
        if pending:
            await asyncio.gather(*pending, return_exceptions=True)


async def fetch(name: str, delay: float, *, fail: bool = False) -> str:
    await asyncio.sleep(delay)
    if fail:
        raise ConnectionError(f"{name} unavailable")
    return name


async def main() -> None:
    result = await first_success(
        [fetch("slow", 0.20), fetch("fast", 0.05), fetch("broken", 0.01, fail=True)],
        timeout=1.0,
    )
    print(f"winner: {result}")
    assert result == "fast"

    cancelled = asyncio.Event()

    async def cancellable() -> str:
        try:
            await asyncio.sleep(30)
        except asyncio.CancelledError:
            cancelled.set()
            raise
        return "never"

    assert await first_success(
        [fetch("quick", 0.01), cancellable()],
        timeout=1.0,
    ) == "quick"
    await asyncio.sleep(0)
    assert cancelled.is_set()

    try:
        await first_success(
            [fetch("a", 0.01, fail=True), fetch("b", 0.02, fail=True)],
            timeout=1.0,
        )
    except RuntimeError as exc:
        print(f"all failed: {exc}")


if __name__ == "__main__":
    asyncio.run(main())