Zachary Turner
Zachary Turner

Reputation: 772

asyncio wait on multiple tasks with timeout and cancellation

I have some code that runs multiple tasks in a loop like this:

    done, running = await asyncio.wait(running, timeout=timeout_seconds,
                                       return_when=asyncio.FIRST_COMPLETED)

I need to be able to determine which of these timed out. According to the documentation:

Note that this function does not raise asyncio.TimeoutError. Futures or Tasks that aren’t done when the timeout occurs are simply returned in the second set.

I could use wait_for() instead, but that function only accepts a single awaitable, whereas I need to specify multiple. Is there any way to determine which one from the set of awaitables I passed to wait() was responsible for the timeout?

Alternatively, is there a way to use wait_for() with multiple awaitables?

Upvotes: 3

Views: 3129

Answers (2)

Break
Break

Reputation: 342

Here is how I do it:

done, pending = await asyncio.wait({
    asyncio.create_task(task, name=index)
        for index, task in enumerate([
            my_coroutine(),
            my_coroutine(),
            my_coroutine(),
        ])
    },
    return_when=asyncio.FIRST_COMPLETED
)

num = next(t.get_name() for t in done)

if num == 2:
    pass

Use enumerate to name the tasks as they are created.

Upvotes: 1

Ivan Sushkov
Ivan Sushkov

Reputation: 76

Your can try that tricks, probably it is not good solution:

import asyncio

async def foo():
    return 42

async def need_some_sleep():
    await asyncio.sleep(1000)
    return 42


async def coro_wrapper(coro):
    result = await asyncio.wait_for(coro(), timeout=10)
    return result

loop = asyncio.get_event_loop()

done, running = loop.run_until_complete(asyncio.wait(
    [coro_wrapper(foo), coro_wrapper(need_some_sleep)],
    return_when=asyncio.FIRST_COMPLETED
    )
)

for item in done:
    print(item.result())

print(done, running)

Upvotes: 1

Related Questions