确保仅在取消等待的协程时才取消包装协程

Question

我需要包装一个返回数据的协程。如果返回数据，则不再可用。如果协程被取消，数据在下次调用时可用。我需要包装协程具有相同的行为，但有时它会在包装协程已经完成时被取消。

我可以使用以下代码重现此行为。

import asyncio
loop = asyncio.get_event_loop()

fut = asyncio.Future()

async def wait():
    return await fut

task = asyncio.ensure_future(wait())

async def test():
    await asyncio.sleep(0.1)
    fut.set_result('data')
    print ('fut', fut)
    print ('task', task)
    task.cancel()
    await asyncio.sleep(0.1)
    print ('fut', fut)
    print ('task', task)

loop.run_until_complete(test())

输出清楚地表明，在协程完成后，包装协程被取消，这意味着数据永远丢失了。我不能屏蔽这两个呼叫，因为如果我被取消，我无论如何都没有数据可以返回。

fut <Future finished result='data'>
task <Task pending coro=<wait() running at <ipython-input-8-6d115ded09c6>:7> wait_for=<Future finished result='data'>>
fut <Future finished result='data'>
task <Task cancelled coro=<wait() done, defined at <ipython-input-8-6d115ded09c6>:6>>

在我的情况下，这是由于有两个期货，一个验证包装协程，一个取消包装协程，有时一起验证。我可能会选择延迟取消（通过asyncio.sleep(0)），但我确定它永远不会发生意外吗？

---

问题在任务中更有意义：

import asyncio
loop = asyncio.get_event_loop()

data = []
fut_data = asyncio.Future()

async def get_data():
    while not data:
        await asyncio.shield(fut_data)
    return data.pop()

fut_wapper = asyncio.Future()

async def wrapper_data():
    task = asyncio.ensure_future(get_data())
    return await task

async def test():
    task = asyncio.ensure_future(wrapper_data())
    await asyncio.sleep(0)
    data.append('data')
    fut_data.set_result(None)
    await asyncio.sleep(0)
    print ('wrapper_data', task)
    task.cancel()
    await asyncio.sleep(0)
    print ('wrapper_data', task)
    print ('data', data)

loop.run_until_complete(test())

task <Task cancelled coro=<wrapper_data() done, defined at <ipython-input-2-93645b78e9f7>:16>>
data []

数据已被消费，但任务已取消，因此无法检索数据。直接等待get_data() 可以，但不能取消。

Answer 1

我认为您需要首先保护等待的未来不被取消，然后检测您自己的取消。如果未来尚未完成，则将取消传播到其中（有效地撤消shield()）并传播出去。如果未来已经完成，忽略取消并返回数据。

代码看起来像这样，也更改为避免使用全局变量并使用asyncio.run()（如果您使用的是 Python 3.6，则可以转到run_until_complete()）：

import asyncio

async def wait(fut):
    try:
        return await asyncio.shield(fut)
    except asyncio.CancelledError:
        if fut.done():
            # we've been canceled, but we have the data - ignore the
            # cancel request
            return fut.result()
        # otherwise, propagate the cancellation into the future
        fut.cancel()
        # ...and to the caller
        raise

async def test():
    loop = asyncio.get_event_loop()
    fut = loop.create_future()
    task = asyncio.create_task(wait(fut))
    await asyncio.sleep(0.1)
    fut.set_result('data')
    print ('fut', fut)
    print ('task', task)
    task.cancel()
    await asyncio.sleep(0.1)
    print ('fut', fut)
    print ('task', task)

asyncio.run(test())

请注意，忽略取消请求可被视为滥用取消机制。但是，如果知道任务会在之后继续进行（理想情况下是立即完成），那么在您的情况下这可能是正确的事情。建议谨慎。