asyncio: prevent task from being cancelled twice
<p>Sometimes, my coroutine cleanup code includes some blocking parts (in the <code>asyncio</code> sense, i.e. they may yield).</p>

<p>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.</p>

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

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

<p>Is there a Pythonic way to do so?</p>

<pre><code>#!/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()
</code></pre>

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

<p>Unlike <code>asyncio.shield</code>, which protects the callee, but raises <code>CancelledError</code> in its caller, this function suppresses <code>CancelledError</code> altogether.</p>

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

<pre><code>@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
</code></pre>

python python-asyncio coroutine