Async multirunner racer with auto task spawning for python by runie

multirunner_racer.py

from asyncio import Event, create_task, run, sleep, Task
from random import randint

class MultiRunner:
    def __init__(self, coro, n):
        self._coro = coro
        self._n = n

        self._result = None
        self.args = []
        self.kwargs = {}
        self._tasks = []
        self._event = Event()

    async def __call__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs

        for n in range(self._n):
            self.spawn()

        await self._event.wait()
        return self._result

    def spawn(self):
        print('Spawning task...')
        task = create_task(self._coro(*self.args, **self.kwargs))
        task.add_done_callback(self.callback)
        self._tasks.append(task)

    def callback(self, task: Task):
        # don't do anything if result is already set
        if self._event.is_set():
            return

        # get result, if it's an exception, spawn a new task
        try:
            result = task.result()
        except:
            self.spawn()
            return

        # if we got a result, cancel all other tasks
        for _task in self._tasks:
            _task.cancel()

        # set result and set event flag so .result() can finish awaiting and return the result
        self._result = result
        self._event.set()


async def some_task():
    to_wait = randint(1, 100) / 10
    should_fail = randint(0, 40) > 1
    await sleep(to_wait)
    if should_fail:
        print('Failing on purpose...')
        raise Exception('generic-error')
    return to_wait


async def main():
    runner = MultiRunner(some_task, 8)
    result = await runner()
    print(result)


run(main())