Applying Backpressure in an Asyncio OSM Pipeline Jump to heading

A PBF reader can produce features far faster than a database can accept them, and an unbounded queue between the two converts that difference into memory consumption until the process dies.

Prerequisites Jump to heading

Conceptual minimum Jump to heading

Backpressure is not a mechanism you add; it is what you get when you remove unbounded buffering. asyncio.Queue(maxsize=N) makes await queue.put() suspend once the queue is full, and that suspension propagates backwards through every stage until the reader itself stops reading. The queue bound is the entire implementation.

Three details decide whether it works.

Every queue needs a bound. One unbounded queue anywhere in the chain absorbs the whole difference in rate, and the bounds on the others become decoration.

The bound is in items, but memory is in bytes. A queue of 1000 batches of 50,000 features is not a bound anybody intended. Size the queue against the batch size, and prefer small queues of large batches to large queues of small ones, because the per-item overhead dominates otherwise.

Failure has to travel both directions. A consumer that dies leaves the producer blocked on a full queue forever, and a producer that dies leaves the consumer waiting on an empty one. TaskGroup cancels siblings on an exception, which handles the first; a sentinel per consumer handles the second.

CPU-bound work is the case where this reasoning breaks down. Parsing a PBF block is not I/O, so an async def that does it blocks the event loop and no amount of queue bounding helps. That work belongs in a ProcessPoolExecutor reached through run_in_executor, with the bound applied to the number of outstanding submissions.

What an unbounded queue does to a rate mismatch Two panels. With an unbounded queue, a reader producing eighty thousand features a second feeds a writer accepting twelve thousand, and the difference of sixty-eight thousand a second accumulates in memory until the process is killed, with throughput unchanged by the buffering. With a bounded queue, the reader blocks on put once the bound is reached, memory stays flat at the bound times the batch size, and the pipeline runs at the writer's rate, which is the rate it was always going to run at. Same rates, two outcomes Unbounded queue Reader: 80k features/s Writer: 12k features/s 68k/s accumulates in RAM Killed, not slowed Throughput unchanged Bounded queue Reader blocks on put Memory flat at bound Runs at the writer's rate The rate it always was No surprise at 3am Buffering never made the pipeline faster; it only postponed the moment the slow stage set the pace, and paid memory for the delay.
The bounded version is not slower. It is the same speed, without the failure.

Runnable solution Jump to heading

python
from __future__ import annotations

import asyncio
import logging
import os
from collections.abc import AsyncIterator, Iterable
from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger("osm.async.backpressure")

BATCH = 5_000          # features per batch: big, so per-item overhead is small
QUEUE_DEPTH = 4        # batches in flight: small, so memory is bounded tightly
WRITERS = 4

SENTINEL: object = object()


@dataclass(frozen=True)
class Batch:
    features: tuple[dict, ...]
    source_block: int


def parse_block(raw: bytes, index: int) -> Batch:
    """CPU-bound. Runs in a worker process, never on the event loop."""
    # Stand-in for the real decode; the point is that it does not await.
    features = tuple({"id": index * 1000 + i, "raw_len": len(raw)}
                     for i in range(BATCH))
    return Batch(features, index)


async def reader(blocks: Iterable[tuple[int, bytes]],
                 queue: asyncio.Queue,
                 pool: ProcessPoolExecutor) -> None:
    """Decode blocks off-loop, and block on put once the queue is full.

    The await on put IS the backpressure. Nothing else throttles this.
    """
    loop = asyncio.get_running_loop()
    for index, raw in blocks:
        batch = await loop.run_in_executor(pool, parse_block, raw, index)
        await queue.put(batch)          # suspends when the queue is full
    for _ in range(WRITERS):
        await queue.put(SENTINEL)       # one per consumer, never one shared
    logger.info("reader finished")


async def writer(name: str, queue: asyncio.Queue, sink) -> int:
    written = 0
    while True:
        item = await queue.get()
        try:
            if item is SENTINEL:
                return written
            await sink.write(item.features)
            written += len(item.features)
        finally:
            queue.task_done()


class SlowSink:
    """Stand-in for a database: the stage that sets the pipeline's real rate."""

    def __init__(self, rows_per_second: int) -> None:
        self.rate = rows_per_second

    async def write(self, features: tuple[dict, ...]) -> None:
        await asyncio.sleep(len(features) / self.rate)


async def run(blocks: Iterable[tuple[int, bytes]]) -> int:
    queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_DEPTH)
    sink = SlowSink(rows_per_second=12_000)
    total = 0

    with ProcessPoolExecutor(max_workers=os.cpu_count()) as pool:
        # TaskGroup cancels every sibling when one raises: a writer that dies
        # cannot leave the reader blocked on a queue nobody will drain.
        async with asyncio.TaskGroup() as tg:
            tg.create_task(reader(blocks, queue, pool))
            writers = [tg.create_task(writer(f"w{i}", queue, sink))
                       for i in range(WRITERS)]
        total = sum(w.result() for w in writers)

    logger.info("wrote %d features", total)
    return total


async def monitor(queue: asyncio.Queue, period_s: float = 5.0) -> None:
    """A persistently full queue means the sink is the bottleneck.

    A persistently empty one means the reader is. Either is useful to know;
    not knowing which is the usual state of an async pipeline.
    """
    while True:
        await asyncio.sleep(period_s)
        depth = queue.qsize()
        logger.info("queue depth %d/%d (%s-bound)", depth, QUEUE_DEPTH,
                    "sink" if depth >= QUEUE_DEPTH - 1 else "reader")


if __name__ == "__main__":
    fake = [(i, b"\x00" * 1024) for i in range(20)]
    asyncio.run(run(fake))

Step-by-step walkthrough Jump to heading

  1. Bound the queue, in batches. maxsize=4 with 5,000-feature batches caps in-flight work at 20,000 features, which is a number you can multiply by a feature’s size and reason about.
  2. Let await put do the throttling. No rate limiter, no sleep, no token bucket. The suspension is the mechanism, and adding anything on top of it fights the thing that already works.
  3. Keep CPU work off the loop. run_in_executor with a process pool means block decoding does not stall every other coroutine, which is the failure that looks like backpressure not working.
  4. Send one sentinel per consumer. A single sentinel stops one writer and leaves the others waiting forever, and it is the most common way this shape deadlocks.
  5. Use TaskGroup. An exception in any task cancels the rest, which is what stops a dead writer from leaving the reader blocked on a full queue.
  6. Measure the queue depth. Persistently full means the sink is the bottleneck; persistently empty means the reader is. Without that signal you are guessing at which half to optimise.
  7. Size batches upward before queues. Per-item overhead dominates at small batch sizes, so a queue of four large batches outperforms a queue of four hundred small ones at the same memory.
How a stall propagates backwards through bounded stages Four stages. The sink slows, perhaps because the database is checkpointing. Its writers stop taking from the queue, so the queue reaches its bound. The reader's await on put suspends, so it stops submitting blocks to the process pool. Memory stays flat at the bound rather than growing, and when the sink recovers every stage resumes in order with nothing lost. The chain only works if every queue in it is bounded, since one unbounded queue absorbs the entire mismatch. A stall travelling upstream sink slows database checkpoints writes take longer queue fills consumers stop taking reaches its bound reader suspends await put blocks stops submitting work memory flat nothing accumulates resumes in order One unbounded queue anywhere in this chain absorbs the whole mismatch and makes every other bound decorative.
Nothing here is a rate limiter. The suspension on a bounded put is the entire implementation.
Five ways this shape deadlocks or leaks, and the signal each gives A grid of five failure modes against their observable signal and the fix. An unbounded queue shows memory climbing while throughput stays constant, fixed by setting an explicit maxsize. A single sentinel for several consumers shows the run hanging at the end with writers idle, fixed by emitting one per consumer. CPU work on the event loop shows an empty queue with the reader slow, fixed by run_in_executor. A dead consumer shows the reader blocked on put with no progress, fixed by a task group. An ungated process pool shows memory climbing despite a bounded queue, fixed by a semaphore around submission. Failure modes and their signals Observable signal Fix Unbounded queue memory climbs, rate flat set an explicit maxsize One shared sentinel hangs at the end, idle one sentinel per consumer CPU on the loop queue empty, reader slow run_in_executor Dead consumer reader blocked on put use a TaskGroup Ungated pool memory climbs anyway semaphore on submit The last row catches people who bounded every queue carefully and then handed unlimited work to an executor that buffers it just as happily.
Each signal is cheap to observe and none of them is visible without logging the queue depth.

Verification Jump to heading

  • Memory stays flat. Run against a deliberately slow sink and watch resident memory; it should plateau rather than climb.
  • Throughput matches the sink. Total features divided by wall-clock should land near the sink’s rate, not the reader’s.
  • The queue depth reports usefully. Confirm the monitor distinguishes a full queue from an empty one under both fast and slow sinks.
  • A dying writer stops the run. Raise inside one writer and confirm the whole group cancels rather than hanging.
  • No deadlock on completion. Confirm every writer exits, which proves the sentinel count matches the consumer count.

Common errors and fixes Jump to heading

Symptom Root cause One-line fix
Memory climbs until the process dies Unbounded queue somewhere in the chain Give every queue an explicit maxsize
Bound set but memory still grows Bound counts items, batches are huge Size the bound against the batch, in bytes
Pipeline hangs at the end One sentinel for several consumers Send one sentinel per consumer
Everything stalls, queue stays empty CPU work running on the event loop Move decoding to run_in_executor
Reader blocked forever A consumer died silently Use TaskGroup so siblings are cancelled
Throughput far below the sink’s rate Too few consumers for a latency-bound sink Add writers; concurrency hides per-write latency
Cannot tell which stage is slow Queue depth never observed Log qsize() against maxsize periodically

Specification reference Jump to heading

asyncio.Queue(maxsize=0) creates a queue of infinite size. If maxsize is greater than zero, put() blocks when the queue reaches maxsize until an item is removed by get(). asyncio.TaskGroup provides a context manager holding a group of tasks; if any task fails with an exception other than CancelledError, the remaining tasks in the group are cancelled. See the Python asyncio queue and task-group documentation.

Frequently Asked Questions Jump to heading

Why not use a rate limiter instead of a bounded queue?

Because a rate limiter needs a rate, and the correct rate is whatever the sink happens to manage right now — which changes with load, with checkpoints, with the time of day. A bounded queue discovers that rate continuously without being told it. Configuring a limiter means either setting it below the sink’s capacity, wasting throughput, or above it, which is the unbounded case with extra code.

How large should the queue bound be?

Small. The queue exists to absorb short-term jitter, not to buffer work, and two to eight batches covers nearly all jitter. Going larger trades memory for a benefit that stops accruing almost immediately, because once the queue is deep enough to bridge a momentary stall, further depth only delays the point at which the reader learns the sink is slow.

Does this apply when the sink is a file rather than a database?

Yes, and more subtly, because the operating system’s page cache provides its own unbounded buffer. Writes appear instant until dirty-page limits are reached, at which point the throughput collapses to the disk’s actual rate. The queue bound still works, but the observed sink rate during the first minute is not the real one, so measure over a long enough window to get past the cache.

What if one stage is CPU-bound and the rest are I/O-bound?

Put the CPU-bound stage in a process pool and bound the number of outstanding submissions rather than a queue. A ProcessPoolExecutor will happily accept unlimited work, serialising every argument into memory as it queues, so the executor is itself an unbounded buffer unless you gate it — a semaphore acquired before submitting and released on completion is the usual shape.

Up one level: Async PBF Parsing with Pyrosm.