Python asyncio 超时/重试设计模式

Question

请看下面的代码（为了简单起见，我没有使用 pydantic 来对 corutine、重试、超时进行分组）：

import asyncio
import typing as tp
import random

async def my_func(wait_time: int) -> str:
    random_number = random.random()
    random_time = wait_time - random_number if random.random() < 0.5 else wait_time + random_number
    print(f"waiting for {wait_time}{random_time:+} seconds")
    await asyncio.sleep(wait_time)
    return f"waited for {wait_time}{random_time:+} seconds"

async def main() -> None:

    task1 = asyncio.create_task(my_func(wait_time=1), name='task1')
    task2 = asyncio.create_task(my_func(wait_time=2), name='task2')
    task3 = asyncio.create_task(my_func(wait_time=3), name='task3')

    task1_timeout = 1.2
    task2_timeout = 2.2
    task3_timeout = 3.2

    task1_retry = 4
    task2_retry = 3
    task3_retry = 2

    total_timeout = 5

    <what to put here?>

    return task1_result, task2_result, task3_result

asyncio.run(main())

如您所见，我有函数 my_func （在现实生活中我将有多个不同的函数）。 在 main() 中，我定义了 3 个任务。每个任务都有其超时和重试。 例如，task1 超时 2 秒，重试 3 次。

此外，我还有另一个（全局）超时，total_timeout，它表示 main() 必须完成的时间。

例如，如果task1 开始运行并且在 1.2 秒内没有得到结果，我们应该最多重试 4 次，所以在我们根本无法得到结果的情况下，我们仍然低于 timeout_total 的 5秒。

对于task2，在2.2秒内超时，可以重复3次，在4.4秒第二次重复完成后，如果我们再次重试，它将在第5秒被total_timeout截断。

对于task3，如果我们第一次尝试没有完成，我们没有足够的时间进行第二次尝试（total_timeout）。

我想同时执行所有三个任务，尊重他们各自的超时和重试，以及total_timeout。最后，最多 5 秒后，我将得到三个元素的元组，它们将是 str（my_func 的输出）或 None（以防所有重复失败，或任务已被 total_timeout 切断）。 所以输出可以是(str, str, str)、(str, None, str) 或(None, None, None)。

有人可以提供一些示例代码来完成我所描述的吗？

Answer 1

我认为这是一个很好的问题。我提出了这个解决方案，它结合了asyncio.gather()和asyncio.wait_for().

这里要求第三个任务等待5秒与3.2 秒超时（重试 2 次），并将返回 None，如asyncio.TimeoutError将被提升（并被抓住）。

import asyncio
import random
import sys


total_timeout = float(sys.argv[1]) if len(sys.argv) > 1 else 5.0


async def work_coro(wait_time: int) -> str:
    random_number = random.random()
    random_time = wait_time - random_number if 
        random.random() < 0.5 else wait_time + random_number

    if random_number > 0.7:
        raise RuntimeError('Random sleep time too high')

    print(f"waiting for {wait_time}{random_time:+} seconds")

    await asyncio.sleep(random_time)

    return f"waited for {wait_time}{random_time:+} seconds"


async def coro_trunner(wait_time: int,
                       retry: int,
                       timeout: float) -> str:
    """
    Run work_coro in a controlled timing environment

    :param int wait_time: How long the coroutine will sleep on each run
    :param int retry: Retry count (if the coroutine times out, retry x times)
    :param float timeout: Timeout for the coroutine
    """

    for attempt in range(0, retry):
        try:
            start_time = loop.time()
            print(f'{work_coro}: ({wait_time}, {retry}, {timeout}): '
                  'spawning')

            return await asyncio.wait_for(work_coro(wait_time),
                                          timeout)
        except asyncio.TimeoutError:
            diff_time = loop.time() - start_time
            print(f'{work_coro}: ({wait_time}, {retry}, {timeout}): '
                  f'timeout (diff_time: {diff_time}')
            continue
        except asyncio.CancelledError:
            print(f'{work_coro}: ({wait_time}, {retry}, {timeout}): '
                  'cancelled')
            break
        except Exception as err:
            # Unknown error raised in the worker: give it another chance

            print(f'{work_coro}: ({wait_time}, {retry}, {timeout}): '
                  f'error in worker: {err}')
            continue


async def main() -> list:
    tasks = [
        asyncio.create_task(coro_trunner(1, 2, 1.2)),
        asyncio.create_task(coro_trunner(2, 3, 2.2)),
        asyncio.create_task(coro_trunner(5, 5, 5.2))
    ]

    try:
        gaf = asyncio.gather(*tasks)
        results = await asyncio.wait_for(gaf,
                                         total_timeout)
    except (asyncio.TimeoutError,
            asyncio.CancelledError,
            Exception):
        # Total timeout reached: get the results that are ready

        # Consume the gather exception
        exc = gaf.exception()  # noqa

        results = []

        for task in tasks:
            if task.done():
                results.append(task.result())
            else:
                # We want to know when a task yields nothing
                results.append(None)
                task.cancel()

    return results


print(f'Total timeout: {total_timeout}')


loop = asyncio.get_event_loop()
start_time = loop.time()

results = asyncio.run(main())
end_time = loop.time()

print(f'{end_time - start_time} --> {results}')