在 Python asyncio 中创建并发任务之间的依赖关系

Question

我在消费者/生产者关系中有两个任务，由asyncio.Queue 。 如果生产者任务失败，我希望消费者任务也尽快失败，而不是无限期地在队列中等待。 消费者任务可以独立于生产者任务创建（产生）。

一般而言，我想在两个任务之间实现依赖关系，这样一个任务的失败也是另一个任务的失败，同时保持这两个任务并发（即一个不会直接等待另一个）。

这里可以使用什么样的解决方案（例如模式）？

基本上，我在考虑erlang 的“链接” 。

我认为可以使用回调实现类似的东西，即asyncio.Task.add_done_callback

谢谢！

Answer 1

一种方法是通过队列传播异常，结合工作处理的委托：

class ValidWorkLoad:
    async def do_work(self, handler):
        await handler(self)


class HellBrokeLoose:
    def __init__(self, exception):
        self._exception = exception

    async def do_work(self, handler):
        raise self._exception


async def worker(name, queue):
    async def handler(work_load):
        print(f'{name} handled')

    while True:
        next_work = await queue.get()
        try:
            await next_work.do_work(handler)
        except Exception as e:
            print(f'{name} caught exception: {type(e)}: {e}')
            break
        finally:
            queue.task_done()


async def producer(name, queue):
    i = 0
    while True:
        try:
            # Produce some work, or fail while trying
            new_work = ValidWorkLoad()
            i += 1
            if i % 3 == 0:
                raise ValueError(i)
            await queue.put(new_work)
            print(f'{name} produced')
            await asyncio.sleep(0)  # Preempt just for the sake of the example
        except Exception as e:
            print('Exception occurred')
            await queue.put(HellBrokeLoose(e))
            break


loop = asyncio.get_event_loop()
queue = asyncio.Queue(loop=loop)
producer_coro = producer('Producer', queue)
consumer_coro = worker('Consumer', queue)
loop.run_until_complete(asyncio.gather(producer_coro, consumer_coro))
loop.close()

哪些输出：

生产者生产

消费者处理

生产者生产

消费者处理

发生异常

消费者捕获异常：<class 'ValueError'>：3

或者，您可以跳过委托，并指定一个指示工作人员停止的项目。 在生产者中捕获异常时，您将该指定项目放入队列中。

Answer 2

从评论：

我试图避免的行为是消费者无视生产者的死亡并无限期地等待队列。 我希望消费者收到生产者死亡的通知，并有机会做出反应。 或者只是失败，即使它也在等待队列。

除了Yigal 提出的答案之外，另一种方法是设置第三个任务来监视这两个任务并在另一个任务完成时取消一个任务。 这可以推广到任何两个任务：

async def cancel_when_done(source, target):
    assert isinstance(source, asyncio.Task)
    assert isinstance(target, asyncio.Task)
    try:
        await source
    except:
        # SOURCE is a task which we expect to be awaited by someone else
        pass
    target.cancel()

现在在设置生产者和消费者时，您可以将它们与上述功能联系起来。 例如：

async def producer(q):
    for i in itertools.count():
        await q.put(i)
        await asyncio.sleep(.2)
        if i == 7:
            1/0

async def consumer(q):
    while True:
        val = await q.get()
        print('got', val)

async def main():
    loop = asyncio.get_event_loop()
    queue = asyncio.Queue()
    p = loop.create_task(producer(queue))
    c = loop.create_task(consumer(queue))
    loop.create_task(cancel_when_done(p, c))
    await asyncio.gather(p, c)

asyncio.get_event_loop().run_until_complete(main())

Answer 3

另一种可能的解决方案：

import asyncio
def link_tasks(t1: Union[asyncio.Task, asyncio.Future], t2: Union[asyncio.Task, asyncio.Future]):
    """
    Link the fate of two asyncio tasks,
    such that the failure or cancellation of one
    triggers the cancellation of the other
    """
    def done_callback(other: asyncio.Task, t: asyncio.Task):
        # TODO: log cancellation due to link propagation
        if t.cancelled():
            other.cancel()
        elif t.exception():
            other.cancel()
    t1.add_done_callback(functools.partial(done_callback, t2))
    t2.add_done_callback(functools.partial(done_callback, t1))

这使用asyncio.Task.add_done_callback来注册回调，如果其中一个任务失败或被取消，该回调将取消另一个任务。