asyncio: prevent task from being cancelled twice

Question

Sometimes, my coroutine cleanup code includes some blocking parts (in the asyncio sense, i.e. they may yield).

I try to design them carefully, so they don't block indefinitely. So "by contract", coroutine must never be interrupted once it's inside its cleanup fragment.

Unfortunately, I can't find a way to prevent this, and bad things occur when it happens (whether it's caused by actual double cancel call; or when it's almost finished by itself, doing cleanup, and happens to be cancelled from elsewhere).

Theoretically, I can delegate cleanup to some other function, protect it with a shield, and surround it with try-except loop, but it's just ugly.

Is there a Pythonic way to do so?

#!/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

I found WGH's solution when encountering a similar problem. I'd like to await a thread, but regular asyncio cancellation (with or without shield) will just cancel the awaiter and leave the thread floating around, uncontrolled. Here is a modification of super_shield that optionally allows reacting on cancel requests and also handles cancellation from within the awaitable:

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

This guarantees that the awaitable has finished or raised CancelledError from within. If your task could be cancelled by other means (e.g. setting a flag observed by a thread), you can use the optional cancel callback to enable cancellation.

Implementation:

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

Answer 2

I ended up writing a simple function that provides a stronger shield, so to speak.

Unlike asyncio.shield, which protects the callee, but raises CancelledError in its caller, this function suppresses CancelledError altogether.

The drawback is that this function doesn't allow you to handle CancelledError later. You won't see whether it has ever happened. Something slightly more complex would be required to do so.

@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