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

Question

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

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

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

基本上，我在想erlang's "links"。

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

谢谢！

Answer 1

来自评论：

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

除了answer presented by 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 2

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

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()

哪些输出：

制片人制作

消费者处理

制片人制作

消费者处理

发生异常

消费者捕获异常：: 3

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

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 注册回调，如果其中一个任务失败或被取消，将取消另一个任务。