asyncio：防止任务被取消两次

Question

有时，我的协程清理代码包含一些阻塞部分（在asyncio 意义上，即它们可能会产生）。

我尝试仔细设计它们，以免它们无限期地阻塞。因此，“按照约定”，协程一旦进入其清理片段就绝不能被中断。

不幸的是，我找不到防止这种情况的方法，并且当它发生时会发生不好的事情（无论是由实际的双重cancel 调用引起的；或者当它几乎自己完成时，正在清理，并且恰好是从其他地方取消）。

理论上，我可以将清理工作委托给其他函数，用shield 保护它，然后用try-except 循环包围它，但这很丑。

有没有 Pythonic 的方式来做到这一点？

#!/usr/bin/env python3

import asyncio

@asyncio.coroutine
def foo():
    """
    This is the function in question,
    with blocking cleanup fragment.
    """

    try:
        yield from asyncio.sleep(1)
    except asyncio.CancelledError:
        print("Interrupted during work")
        raise
    finally:
        print("I need just a couple more seconds to cleanup!")
        try:
            # upload results to the database, whatever
            yield from asyncio.sleep(1)
        except asyncio.CancelledError:
            print("Interrupted during cleanup :(")
        else:
            print("All cleaned up!")

@asyncio.coroutine
def interrupt_during_work():
    # this is a good example, all cleanup
    # finishes successfully

    t = asyncio.async(foo())

    try:
        yield from asyncio.wait_for(t, 0.5)
    except asyncio.TimeoutError:
        pass
    else:
        assert False, "should've been timed out"

    t.cancel()

    # wait for finish
    try:
        yield from t
    except asyncio.CancelledError:
        pass

@asyncio.coroutine
def interrupt_during_cleanup():
    # here, cleanup is interrupted

    t = asyncio.async(foo())

    try:
        yield from asyncio.wait_for(t, 1.5)
    except asyncio.TimeoutError:
        pass
    else:
        assert False, "should've been timed out"

    t.cancel()

    # wait for finish
    try:
        yield from t
    except asyncio.CancelledError:
        pass

@asyncio.coroutine
def double_cancel():
    # cleanup is interrupted here as well
    t = asyncio.async(foo())

    try:
        yield from asyncio.wait_for(t, 0.5)
    except asyncio.TimeoutError:
        pass
    else:
        assert False, "should've been timed out"

    t.cancel()

    try:
        yield from asyncio.wait_for(t, 0.5)
    except asyncio.TimeoutError:
        pass
    else:
        assert False, "should've been timed out"

    # although double cancel is easy to avoid in 
    # this particular example, it might not be so obvious
    # in more complex code
    t.cancel()

    # wait for finish
    try:
        yield from t
    except asyncio.CancelledError:
        pass

@asyncio.coroutine
def comain():
    print("1. Interrupt during work")
    yield from interrupt_during_work()

    print("2. Interrupt during cleanup")
    yield from interrupt_during_cleanup()

    print("3. Double cancel")
    yield from double_cancel()

def main():
    loop = asyncio.get_event_loop()
    task = loop.create_task(comain())
    loop.run_until_complete(task)

if __name__ == "__main__":
    main()

Answer 1

我最终编写了一个简单的函数，可以提供更强大的屏蔽，可以这么说。

与保护被调用者的asyncio.shield 不同，但在其调用者中引发CancelledError，此函数完全抑制CancelledError。

缺点是这个函数不允许你稍后处理CancelledError。你不会看到它是否曾经发生过。这样做需要稍微更复杂的东西。

@asyncio.coroutine
def super_shield(arg, *, loop=None):
    arg = asyncio.async(arg)
    while True:
        try:
            return (yield from asyncio.shield(arg, loop=loop))
        except asyncio.CancelledError:
            continue

Answer 2

我在遇到类似问题时找到了 WGH 的解决方案。我想等待一个线程，但是常规的 asyncio 取消（有或没有屏蔽）只会取消等待者并使线程四处浮动，不受控制。这是对super_shield 的修改，它可以选择性地允许对取消请求做出反应，也可以在等待内处理取消：

await protected(aw, lambda: print("Cancel request"))

这保证等待对象已经完成或从内部引发CancelledError。如果您的任务可以通过其他方式取消（例如设置线程观察到的标志），您可以使用可选的取消回调来启用取消。

实施：

async def protect(aw, cancel_cb: typing.Callable = None):
    """
    A variant of `asyncio.shield` that protects awaitable as well
    as the awaiter from being cancelled.

    Cancellation events from the awaiter are turned into callbacks
    for handling cancellation requests manually.

    :param aw: Awaitable.
    :param cancel_cb: Optional cancellation callback.
    :return: Result of awaitable.
    """
    task = asyncio.ensure_future(aw)
    while True:
        try:
            return await asyncio.shield(task)
        except asyncio.CancelledError:
            if task.done():
                raise
            if cancel_cb is not None:
                cancel_cb()