Cancellation and Cleanup

Cancellation is the part of asyncio that decides whether your library is safe to use in a real program. Timeouts cancel. A failing sibling in a TaskGroup cancels. A user closing a browser tab cancels. If your code leaks a connection each time, it leaks under exactly the conditions that already mean something is wrong.

What cancelling actually does

task.cancel() doesn’t stop anything immediately. It arranges for CancelledError to be raised inside the task, at its current await point, the next time the loop runs it.

import asyncio

log = []


async def work():
    log.append("started")
    try:
        await asyncio.sleep(1)
        log.append("finished normally")
    except asyncio.CancelledError:
        log.append("cancelled here")
        raise
    finally:
        log.append("cleanup ran")


async def main():
    task = asyncio.create_task(work())
    await asyncio.sleep(0.01)              # let it reach the sleep
    log.append("calling cancel()")
    task.cancel()
    log.append("cancel() returned, nothing has happened yet")
    try:
        await task
    except asyncio.CancelledError:
        log.append("await task raised CancelledError")
    print("\n".join(log))
    print("task.cancelled():", task.cancelled())


asyncio.run(main())

Note the order: cancel() returns before the task has noticed. Cancellation is a request, delivered at the next opportunity, and the task gets to run its except and finally blocks on the way out.

CancelledError is not an Exception

This is the detail that causes the bugs. Since Python 3.8:

class CancelledError(BaseException): ...

So except Exception does not catch it — which is deliberate, because cancellation should propagate. The damage is in cleanup code written as:

        try:
            value = await source.read()
        except Exception:
            await source.close()      # never runs when cancelled
            raise
        await source.close()
        return value

import asyncio

open_count = 0


class Source:
    async def open(self):
        global open_count
        open_count += 1

    async def close(self):
        global open_count
        open_count -= 1

    async def read(self):
        await asyncio.sleep(1)
        return "data"


async def careless():
    source = Source()
    await source.open()
    try:
        value = await source.read()
    except Exception:                  # CancelledError is not an Exception
        await source.close()
        raise
    await source.close()
    return value


async def careful():
    source = Source()
    await source.open()
    try:
        return await source.read()
    finally:                           # runs whatever happens
        await source.close()


async def cancel_after_start(coroutine):
    global open_count
    open_count = 0
    task = asyncio.create_task(coroutine)
    await asyncio.sleep(0.01)
    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        pass
    return open_count


async def main():
    print("careless leaves open:", await cancel_after_start(careless()))
    print("careful  leaves open:", await cancel_after_start(careful()))


asyncio.run(main())

One source left open per cancelled read, in code that looks like it handles errors. The fix is finally, or the async with that wraps it up for you — which is why today’s library does the work inside async with Source(...).

    except asyncio.CancelledError:
        pass            # almost always wrong

Cleanup that awaits

finally runs during cancellation, but an await inside it can itself be cancelled — by a second cancellation, or by a timeout you wrap the cleanup in. (The CancelledError that triggered the finally isn’t re-raised into it; the risk is a fresh interruption arriving while the cleanup is awaiting.) For a quick local close that’s fine — source.close() doesn’t await anything that would be interrupted. For a “tell the server we’re going” that involves I/O, you have to decide what “must complete” actually means, and the tempting recipe gets it wrong:

    finally:
        with contextlib.suppress(TimeoutError):
            async with asyncio.timeout(2):
                await asyncio.shield(source.goodbye())   # does NOT do what it looks like

shield protects goodbye() from the timeout’s cancellation — which is the problem. When the deadline fires, the cancellation hits the await, so control leaves the block, but the shielded goodbye() keeps running, now detached, with nobody awaiting it. You didn’t bound the cleanup; you abandoned it, and left a task running past the finally.

There are two honest options, and neither is the recipe above:

• Bound it and mean it. Drop the shield. asyncio.timeout then cancels goodbye() itself at the deadline, so nothing leaks — at the cost that a slow goodbye is cut off. That’s usually the right trade in cleanup.

      finally:
          with contextlib.suppress(TimeoutError):
              async with asyncio.timeout(2):
                  await source.goodbye()

• Let it outlive this scope on purpose. If the goodbye genuinely must finish even though this task is being cancelled, hand it to something that lives longer — a task owned by the pipeline or the application — and let that own its completion and cancellation. Don’t pretend the finally waited for it.

Shutting down what you started

Anything created with create_task outside a TaskGroup is yours to stop. Say you started a background worker that loops forever — the code that starts it has to end it:

    worker = asyncio.create_task(run_forever())
    try:
        ...                                  # use it
    finally:
        worker.cancel()
        await asyncio.gather(worker, return_exceptions=True)

Both lines are needed. cancel() requests; gather(..., return_exceptions=True) waits for the task to actually finish unwinding and absorbs the CancelledError that comes back. Skip the wait and the program can exit while the task is still in its finally; skip the absorb and shutdown raises the cancellation you asked for.

import asyncio
import unittest

started, finished, cleaned = [], [], []


class WorkerPool:
    def __init__(self, workers=2):
        self.queue = asyncio.Queue()
        self.tasks = [asyncio.create_task(self.work(index)) for index in range(workers)]

    async def work(self, index):
        try:
            while True:
                job = await self.queue.get()
                try:
                    started.append(job)
                    await asyncio.sleep(0.05)
                    finished.append(job)
                finally:
                    self.queue.task_done()
        finally:
            cleaned.append(index)

    def submit(self, job):
        self.queue.put_nowait(job)

    async def aclose(self, *, drain=False):
        """Stop the workers.

        drain=True: let the queued jobs finish first, then stop.
        drain=False: stop now, even mid-job.

        Either way, no task may be left running when this returns.
        TODO: write it.
        """
        raise NotImplementedError("aclose")


class Pool(unittest.TestCase):
    def setUp(self):
        started.clear(); finished.clear(); cleaned.clear()

    def test_draining_finishes_the_queued_jobs(self):
        async def scenario():
            pool = WorkerPool(workers=2)
            for job in range(4):
                pool.submit(job)
            await pool.aclose(drain=True)
            return [task for task in pool.tasks if not task.done()]

        self.assertEqual(asyncio.run(scenario()), [])
        self.assertEqual(sorted(finished), [0, 1, 2, 3])

    def test_stopping_now_abandons_work_in_progress(self):
        async def scenario():
            pool = WorkerPool(workers=2)
            for job in range(4):
                pool.submit(job)
            await asyncio.sleep(0.01)          # jobs 0 and 1 are in flight
            await pool.aclose(drain=False)
            return [task for task in pool.tasks if not task.done()]

        self.assertEqual(asyncio.run(scenario()), [])
        self.assertEqual(finished, [])
        self.assertEqual(sorted(cleaned), [0, 1], "both workers should have unwound")

    def test_closing_twice_is_harmless(self):
        async def scenario():
            pool = WorkerPool(workers=1)
            await pool.aclose()
            await pool.aclose()

        asyncio.run(scenario())


unittest.TextTestRunner(verbosity=2).run(
    unittest.defaultTestLoader.loadTestsFromTestCase(Pool)
)

    async def aclose(self, *, drain=False):
        if drain:
            await self.queue.join()            # wait for the backlog
        for task in self.tasks:
            task.cancel()
        await asyncio.gather(*self.tasks, return_exceptions=True)
        self.tasks = []

What you learned

• cancel() requests cancellation; it’s delivered at the task’s next await, and except/finally still run.
• CancelledError inherits from BaseException, so except Exception misses it — use finally or async with for cleanup, and always re-raise.
• asyncio.shield protects an inner step from an outer cancel but leaves the await cancellable; to bound cleanup, use a timeout without a shield, or own the shielded task’s lifetime yourself.
• Anything you create_task you must cancel and await, absorbing the CancelledError.