Skip to content

Making Async Iterators Cancellation-Safe

A consumer reads events from an async iterator and wraps each anext() in a five-second timeout, so a quiet stream does not block its heartbeat. The first quiet period passes, the timeout fires, the consumer loops around to read again — and the iterator reports that it is exhausted, although the stream is still producing. Another consumer is cancelled during a deploy, restarts, and a message that was fetched but never processed is simply gone. Both are cancellation-safety bugs in the iterator. Cancellation arrives as an exception at whatever await is in progress, and an iterator that holds state across that await — a popped item, an acknowledged message, a suspended generator frame — can be left broken or lossy. This guide shows the two failure modes in running code, fixes each with a design that keeps state consistent at every await point, and adds a test that cancels at every possible point.

Prerequisites

Two ways cancellation breaks an iterator 2 columns contrasting broken iterator, lossy iterator. Two ways cancellation breaks an iterator broken iterator async generator frame timeout cancels inner await error thrown into frame generator finishes next read: exhausted lossy iterator state held across await item popped from buffer await for enrichment cancelled here item never returned One fails loudly on the next read; the other fails silently forever.

1. See a generator die after one timeout

A timeout around anext() cancels whatever the iterator is awaiting. For an async generator, the CancelledError is thrown into the generator's frame; unless the generator catches it, the generator finishes with that exception and is permanently exhausted. A class-based iterator has no frame to kill, so it keeps working.

import asyncio


async def events_generator(queue: asyncio.Queue):
    while True:
        yield await queue.get()                    # a timeout cancels this await, in this frame


class EventsIterator:
    def __init__(self, queue: asyncio.Queue) -> None:
        self.queue = queue

    def __aiter__(self) -> "EventsIterator":
        return self

    async def __anext__(self):
        return await self.queue.get()               # asyncio.Queue.get is cancellation-safe


async def main() -> None:
    for kind in ("generator", "class"):
        queue: asyncio.Queue[str] = asyncio.Queue()
        it = events_generator(queue) if kind == "generator" else EventsIterator(queue)
        try:
            async with asyncio.timeout(0.05):      # a quiet period
                await anext(it)
        except TimeoutError:
            pass
        queue.put_nowait("late event")
        try:
            print(kind, "->", await asyncio.wait_for(anext(it), 0.5))
        except StopAsyncIteration:
            print(kind, "-> exhausted after one timeout")


asyncio.run(main())

The generator reports exhausted after one timeout; the class returns late event. The fix is not to catch CancelledError inside the generator and carry on — that suppresses real cancellation. Either apply timeouts to something that can be cancelled without killing the iterator, as in step 4, or use a class-based iterator whose state lives in attributes rather than in a suspended frame.

Verify: the output shows the generator exhausted and the class-based iterator still delivering after the same timeout.

2. Do not hold an item across an await

The second failure loses data. An iterator that removes an item from its buffer and then awaits something — decoding, enrichment, a lookup — drops that item if cancellation lands in the await: the item is no longer in the buffer and was never returned.

import asyncio


async def fetch_batch() -> list[int]:
    await asyncio.sleep(0)
    return [1, 2, 3]


class LossyIterator:
    def __init__(self) -> None:
        self.buffer: list[int] = []

    async def __anext__(self) -> int:
        if not self.buffer:
            self.buffer.extend(await fetch_batch())
        item = self.buffer.pop(0)                  # state changed...
        await asyncio.sleep(0.05)                  # ...then an await: cancellation loses `item`
        return item


class SafeIterator:
    def __init__(self) -> None:
        self.buffer: list[int] = []

    async def __anext__(self) -> int:
        if not self.buffer:
            batch = await fetch_batch()            # await first...
            self.buffer.extend(batch)              # ...then mutate, with no await in between
        await asyncio.sleep(0.05)                  # any slow work happens before the pop
        return self.buffer.pop(0)                  # commit and return in one synchronous step


async def cancel_once_then_read(it) -> int:
    task = asyncio.create_task(it.__anext__())
    await asyncio.sleep(0.01)
    task.cancel()
    await asyncio.gather(task, return_exceptions=True)
    return await it.__anext__()


async def main() -> None:
    print("lossy:", await cancel_once_then_read(LossyIterator()))   # 2: item 1 is gone
    print("safe: ", await cancel_once_then_read(SafeIterator()))    # 1: nothing lost


asyncio.run(main())

The rule generalises: between the last await and the return, perform all state changes in one synchronous stretch. If an item must be transformed by awaited work, transform a copy while it is still in the buffer (or in a pending slot), and remove it only when the result is ready to return.

Verify: after a cancellation, the lossy iterator returns 2 and the safe one returns 1.

Where the state change happens 3 lanes over time. Where the state change happens lossy pop item await enrichment return safe await enrichment pop item return cancellation lands here time → Cancellation during the await is harmless only when nothing was committed yet.

3. Acknowledge after processing, not after receiving

For message sources with acknowledgements — brokers, work queues, database-backed outboxes — the iterator must not acknowledge when it hands out a message. If the consumer is cancelled while processing, an early acknowledgement turns a restart into data loss. Yield messages with an explicit ack(), and let the consumer acknowledge after its work completes.

import asyncio
from dataclasses import dataclass, field


@dataclass
class Delivery:
    body: str
    tag: int
    _broker: "FakeBroker" = field(repr=False)

    async def ack(self) -> None:
        await self._broker.ack(self.tag)


class FakeBroker:
    def __init__(self, bodies: list[str]) -> None:
        self.pending = list(enumerate(bodies))    # unacknowledged messages are redelivered
        self.acked: set[int] = set()

    async def receive(self) -> tuple[int, str]:
        await asyncio.sleep(0)
        for tag, body in self.pending:
            if tag not in self.acked:
                return tag, body
        raise StopAsyncIteration

    async def ack(self, tag: int) -> None:
        await asyncio.sleep(0)
        self.acked.add(tag)


class Deliveries:
    def __init__(self, broker: FakeBroker) -> None:
        self.broker = broker

    def __aiter__(self) -> "Deliveries":
        return self

    async def __anext__(self) -> Delivery:
        tag, body = await self.broker.receive()
        return Delivery(body, tag, self.broker)   # no ack here


async def consume(broker: FakeBroker, cancel_during: str | None = None) -> list[str]:
    done = []
    async for delivery in Deliveries(broker):
        if delivery.body == cancel_during:
            raise asyncio.CancelledError           # simulate a deploy mid-processing
        await asyncio.sleep(0)                     # process
        done.append(delivery.body)
        await delivery.ack()                       # only now is it safe to forget
    return done


async def main() -> None:
    broker = FakeBroker(["a", "b", "c"])
    try:
        await consume(broker, cancel_during="b")
    except asyncio.CancelledError:
        pass
    print("after restart:", await consume(broker))   # ['b', 'c']: 'b' was redelivered


asyncio.run(main())

This gives at-least-once delivery: a message interrupted mid-processing is delivered again after restart, so processing must be idempotent. The same discipline applies to real clients — manual acknowledgement in processing RabbitMQ messages with aio-pika and committing offsets after processing in consuming Kafka topics with aiokafka.

Verify: the restarted consumer processes b and c; nothing acknowledged before processing finished.

4. Time out waiting without breaking the iterator

Consumers often need to wake up periodically even when the source is quiet — to send heartbeats, flush batches or check for shutdown. Instead of cancelling anext() on the iterator itself, move the receive into a long-lived task or a queue owned by the iterator, and apply the timeout to waiting for the result.

import asyncio


class TimeoutFriendly:
    """Wraps any async iterator; per-call timeouts never cancel the underlying read."""

    def __init__(self, source) -> None:
        self._source = source
        self._pending: asyncio.Task | None = None

    async def next(self, timeout: float):
        if self._pending is None:
            self._pending = asyncio.ensure_future(anext(self._source))
        done, _ = await asyncio.wait({self._pending}, timeout=timeout)
        if not done:
            raise TimeoutError                        # the read keeps going in the background
        task, self._pending = self._pending, None
        return task.result()

    async def aclose(self) -> None:
        if self._pending is not None:
            self._pending.cancel()
            await asyncio.gather(self._pending, return_exceptions=True)


async def main() -> None:
    queue: asyncio.Queue[str] = asyncio.Queue()
    reader = TimeoutFriendly(events_generator(queue))
    heartbeats = 0
    for _ in range(3):
        try:
            print("event:", await reader.next(timeout=0.05))
        except TimeoutError:
            heartbeats += 1                           # quiet: do periodic work, keep reading
            if heartbeats == 2:
                queue.put_nowait("finally something")
    await reader.aclose()


asyncio.run(main())

asyncio.wait() with a timeout does not cancel the task it waits on, so the generator's await queue.get() is never interrupted and the generator stays alive across quiet periods. Here two heartbeats pass and the generator — the same one that died in step 1 — then delivers the event.

Bounding a read without breaking it A grid of 3 rows by 2 columns. Bounding a read without breaking it approach read cancelled? generator survives? asyncio.timeout around anext yes no asyncio.wait_for(anext(it)) yes no asyncio.wait on pending read no yes Time out the wait, not the read.

Verify: the loop prints event: finally something on its third pass, after two silent timeouts, using the generator-based source that step 1 showed dying under a direct timeout.

5. Test cancellation at every await point

Cancellation bugs hide in the one await nobody thought about. A small harness can cancel an iterator call after each possible number of loop iterations and assert the invariant after every one: nothing lost, nothing duplicated, iterator still usable.

import asyncio


async def cancel_at_every_point(make_iterator, expected: list, max_steps: int = 20) -> None:
    for step in range(max_steps):
        it = make_iterator()
        task = asyncio.create_task(it.__anext__())
        for _ in range(step):
            await asyncio.sleep(0)                  # let the call advance `step` iterations
        cancelled = task.cancel()
        results = await asyncio.gather(task, return_exceptions=True)
        got = [] if cancelled and isinstance(results[0], asyncio.CancelledError) else [results[0]]
        while len(got) < len(expected):
            got.append(await it.__anext__())
        assert got == expected, f"step {step}: got {got}"


async def main() -> None:
    await cancel_at_every_point(SafeIterator, [1, 2, 3])
    print("SafeIterator: consistent at every cancellation point")
    try:
        await cancel_at_every_point(LossyIterator, [1, 2, 3])
    except AssertionError as exc:
        print("LossyIterator:", exc)


asyncio.run(main())

The harness drives the call forward one scheduler step at a time before cancelling, which reaches every await inside __anext__ without knowing the implementation. It is the iterator-specific version of the techniques in testing cancellation and cleanup paths.

Verify: the safe iterator passes at every step, and the lossy iterator fails with a message naming the step where an item disappeared.

Verification

An async iterator is cancellation-safe when:

  • Timeouts do not end iteration: after a timed-out wait, the next read still returns data.
  • No item is held across an await: state is committed in one synchronous step immediately before returning.
  • Acknowledgement follows processing: interrupted messages are redelivered, and processing is idempotent.
  • Waiting can be bounded without cancelling reads: periodic work uses a wrapper that times out the wait, not the read.
  • A harness proves it: cancelling at every scheduler step never loses, duplicates or breaks.

Pitfalls & edge cases

  • Catching CancelledError inside a generator to keep it alive. That also swallows real cancellation from shutdown. Restructure the timeout instead.
  • Assuming library reads are cancel-safe. asyncio.Queue.get() and StreamReader.readuntil() keep data when cancelled, but a custom read() that pops from a buffer before awaiting may not. Check each source.
  • Wrapping with wait_for on the iterator. asyncio.wait_for(anext(gen), t) cancels the read on timeout, exactly like step 1. Use the non-cancelling wrapper.
  • Background reads that outlive the consumer. The wrapper's pending task must be cancelled on close, or it keeps consuming from the source.
  • Batching iterators. A partially filled batch is state; when cancelled mid-batch, either return the partial batch on the next call or requeue its items. Batching is covered in batching queue items by size and time.

Frequently Asked Questions

Why does my async generator stop after asyncio.timeout fires?

The timeout cancels the await currently running inside the generator, so CancelledError is thrown into the generator's frame. The generator finishes with that exception and is exhausted, so the next anext raises StopAsyncIteration. Use a class-based iterator or time out waiting for the result without cancelling the read itself.

What makes an async iterator cancellation-safe?

Cancellation at any await inside anext must leave the iterator consistent: no item removed from its source or buffer without being returned, no acknowledgement sent before processing, and the iterator still usable afterwards. Perform state changes in one synchronous step right before returning, after the last await.

Is asyncio.Queue.get cancellation-safe?

Yes. If a task awaiting queue.get is cancelled, no item is removed, so the item remains available for the next reader. Custom iterators built on top of it can still lose items if they remove an item and then await something else before returning it.

How do I add a timeout to async for without losing messages?

Keep the read running in a task owned by a wrapper, wait for that task with asyncio.wait and a timeout, and do periodic work when the wait times out while leaving the read in progress. Acknowledge messages only after processing so an interrupted consumer receives them again.