取消python asyncio中的嵌套协同程序

Question

在我的应用程序中，我有一个协程，可以等待其他几个协同程序，如果这个协同程序，可以等待另一个协同程序，等等。 如果其中一个协同程序失败，则无需执行尚未执行的所有其他协同程序。 （在我的情况下，这甚至是有害的，我想启动几个回滚协同程序）。 那么，如何取消所有嵌套协同程序的执行？这就是我现在所拥有的：

import asyncio

async def foo():
    for i in range(5):
        print('Foo', i)
        await asyncio.sleep(0.5)
    print('Foo2 done')

async def bar():
    await asyncio.gather(bar1(), bar2())


async def bar1():
    await asyncio.sleep(1)
    raise Exception('Boom!')


async def bar2():
    for i in range(5):
        print('Bar2', i)
        await asyncio.sleep(0.5)
    print('Bar2 done')


async def baz():
    for i in range(5):
        print('Baz', i)
        await asyncio.sleep(0.5)

async def main():
    task_foo = asyncio.Task(foo())
    task_bar = asyncio.Task(bar())
    try:
        await asyncio.gather(task_foo, task_bar)
    except Exception:
        print('One task failed. Canceling all')
        task_foo.cancel()
        task_bar.cancel()
    print('Now we want baz')
    await baz()

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    try:
        loop.run_until_complete(main())
    finally:
        loop.close()

这显然不起作用。正如您所看到的，foo coroutine已根据需要取消，但bar2仍在运行：

Foo 0
Bar2 0
Foo 1
Bar2 1
Foo 2
Bar2 2
One task failed. Canceling all
Now we want baz
Baz 0
Bar2 3
Baz 1
Bar2 4
Baz 2
Bar2 done
Baz 3
Baz 4

所以，我肯定做错了什么。这里的正确方法是什么？

Answer 1

当你致电task_bar.cancel()时，任务已经完成，所以没有任何效果。正如gather docs state：

  

如果return_exceptions为true，则任务中的异常将被视为与成功结果相同，并在结果列表中收集; 否则，第一个引发的异常将立即传播到返回的未来。

这正是发生的事情，您对task_bar协程稍作修改即：

async def bar():
    try:
        await asyncio.gather(bar1(), bar2())
    except Exception:
        print("Got a generic exception on bar")
        raise

输出：

Foo 0
Bar2 0
Foo 1
Bar2 1
Foo 2
Bar2 2
Got a generic exception on bar
One task failed. Canceling all
<Task finished coro=<bar() done, defined at cancel_nested_coroutines.py:11> exception=Exception('Boom!',)>
Now we want baz
Baz 0
Bar2 3
Baz 1
Bar2 4
Baz 2
Bar2 done
Baz 3
Baz 4

我也会在task_bar来电之前打印task_bar.cancel()，请注意它已完成，因此调用cancel无效。

就解决方案而言，我认为产生协程需要处理它所安排的协同程序的取消，因为在协程完成后我无法找到检索它们的方法（除了滥用Task.all_tasks之外听起来不对劲。）

说过我必须使用wait代替gather并返回第一个例外，这里有一个完整的例子：

import asyncio


async def foo():
    for i in range(5):
        print('Foo', i)
        await asyncio.sleep(0.5)
    print('Foo done')


async def bar():
    done, pending = await asyncio.wait(
        [bar1(), bar2()], return_when=asyncio.FIRST_EXCEPTION)

    for task in pending:
        task.cancel()

    for task in done:
        task.result()  # needed to raise the exception if it happened


async def bar1():
    await asyncio.sleep(1)
    raise Exception('Boom!')


async def bar2():
    for i in range(5):
        print('Bar2', i)
        await asyncio.sleep(0.5)
    print('Bar2 done')


async def baz():
    for i in range(5):
        print('Baz', i)
        await asyncio.sleep(0.5)


async def main():
    task_foo = asyncio.Task(foo())
    task_bar = asyncio.Task(bar())
    try:
        await asyncio.gather(task_foo, task_bar)
    except Exception:
        print('One task failed. Canceling all')
        print(task_bar)
        task_foo.cancel()
        task_bar.cancel()

    print('Now we want baz')
    await baz()

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    try:
        loop.run_until_complete(main())
    finally:
        loop.close()

哪个输出：

Foo 0
Bar2 0
Foo 1
Bar2 1
Foo 2
Bar2 2
One task failed. Canceling all
<Task finished coro=<bar() done, defined at cancel_nested_coroutines_2.py:11> exception=Exception('Boom!',)>
Now we want baz
Baz 0
Baz 1
Baz 2
Baz 3
Baz 4

它不是很好，但它有效。

Answer 2

据我所知，取消协程本身时，无法自动取消协同程序的所有子任务。所以你必须手动清理子任务。 在等待asyncio.gather future时抛出异常时，您可以通过Gathering_future对象的_children属性访问剩余的任务。 你的例子工作：

import asyncio

async def foo():
    for i in range(5):
        print('Foo', i)
        await asyncio.sleep(0.5)
    print('Foo2 done')

async def bar():
    gathering = asyncio.gather(bar1(), bar2())
    try:
        await gathering
    except Exception:
        # cancel all subtasks of this coroutine
        [task.cancel() for task in gathering._children]
        raise

async def bar1():
    await asyncio.sleep(1)
    raise Exception('Boom!')

async def bar2():
    for i in range(5):
        print('Bar2', i)
        try:
            await asyncio.sleep(0.5)
        except asyncio.CancelledError:
            # you can cleanup here
            print("Bar2 cancelled")
            break
    else:
        print('Bar2 done')

async def baz():
    for i in range(5):
        print('Baz', i)
        await asyncio.sleep(0.5)

async def main():
    task_foo = asyncio.Task(foo())
    task_bar = asyncio.Task(bar())
    try:
        task = asyncio.gather(task_foo, task_bar)
        await task
    except Exception:
        print('One task failed. Canceling all')
        task_foo.cancel()
        task_bar.cancel()
    print('Now we want baz')
    await baz()

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    try:
        loop.run_until_complete(main())
    finally:
        loop.close()

返回

Foo 0
Bar2 0
Foo 1
Bar2 1
Foo 2
Bar2 2
Bar2 cancelled
One task failed. Canceling all
Now we want baz
Baz 0
Baz 1
Baz 2
Baz 3
Baz 4