Streaming PBF Blocks Through an Asyncio Queue

Feed decoded PBF fileblocks into a slower downstream stage — geometry assembly, tag rewriting, a database COPY — without letting the fast decoder race ahead and pile every pending block into RAM, by handing blocks across a bounded asyncio.Queue that makes the producer wait whenever the consumer falls behind.

Prerequisites

Verify each item before running the module below; the queue bound is the only thing standing between a fast decoder and an out-of-memory kill, so the sizing choices here are not optional.

Conceptual minimum

A PBF file is a sequence of independent, zlib-compressed fileblocks. Decoding one block is CPU work; whatever happens to the decoded features afterward — reprojection, tag canonicalization, an insert — is usually slower and often I/O-bound. If you decode in a tight loop and append every result to a list, the decoder finishes the file long before the consumer drains it, and peak memory grows to hold the entire backlog. The fix is not a faster consumer; it is a channel that refuses to accept more work than the consumer can take. A bounded asyncio queue is exactly that channel: await queue.put(block) suspends the producer coroutine the moment the queue is full, and only resumes it once the consumer calls queue.get() and frees a slot. That suspension is backpressure — the producer’s rate is clamped to the consumer’s rate, and the in-flight set never exceeds maxsize.

Two details make this correct rather than merely plausible. First, decoding a block with osmium is a blocking call that would stall the entire event loop if run inline; it must execute in a thread or process pool via loop.run_in_executor, so the loop stays free to service put/get handoffs. Second, the consumer needs an unambiguous end-of-stream signal. A queue has no built-in “closed” state, so the producer enqueues a None sentinel after the last block, and the consumer treats that sentinel as its cue to stop — one sentinel per consumer, so every worker gets released.

Bounded asyncio queue between a PBF block producer and two consumers A producer coroutine runs blocking block decode in a thread-pool executor and pushes each decoded block onto a bounded asyncio queue with four slots, three occupied and one free. When the queue is full the producer's put call suspends, applying backpressure. Two consumer coroutines each call get, process the block against a slower downstream sink, and call task_done. A None sentinel per consumer signals end of stream. Producer run_in_executor decode block Thread / process pool executor await put() put() suspends while full — backpressure Bounded queue (maxsize = 4) 3 filled 1 free await get() Consumer A process → sink Consumer B process → sink End of stream: one None sentinel per consumer stops each worker cleanly

Runnable solution

The module below reads fileblocks from a .osm.pbf, decodes each in a thread-pool executor, and streams the decoded results through a bounded queue to a configurable number of consumers. Decoding here counts nodes and ways per block as a stand-in for real per-block work; swap _decode_block for your geometry or normalization step. Consumers push to whatever sink you supply. It targets Python 3.10+ and osmium>=3.6.

python 
from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Awaitable, Callable

import osmium

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

QUEUE_MAXSIZE = 4      # in-flight decoded blocks; the memory dial
N_CONSUMERS = 2        # parallel downstream workers


@dataclass(slots=True)
class DecodedBlock:
    """One decoded PBF fileblock's summary payload."""
    seq: int
    n_nodes: int
    n_ways: int


def _decode_block(seq: int, raw_bytes: bytes) -> DecodedBlock:
    """Blocking decode of a single fileblock — runs in an executor thread.

    Real pipelines would reconstruct geometry or rewrite tags here; this
    version tallies primitives so the example stays dependency-light.
    """
    n_nodes = n_ways = 0
    for line in raw_bytes.splitlines():
        if line.startswith(b"n"):
            n_nodes += 1
        elif line.startswith(b"w"):
            n_ways += 1
    return DecodedBlock(seq=seq, n_nodes=n_nodes, n_ways=n_ways)


def _iter_raw_blocks(pbf_path: Path) -> list[tuple[int, bytes]]:
    """Yield (sequence, raw bytes) per fileblock via osmium's block reader."""
    blocks: list[tuple[int, bytes]] = []
    reader = osmium.io.Reader(str(pbf_path))
    try:
        for seq, block in enumerate(reader):   # each item is one fileblock
            blocks.append((seq, bytes(block)))
    finally:
        reader.close()
    return blocks


async def stream_blocks(
    pbf_path: Path,
    consume: Callable[[DecodedBlock], Awaitable[None]],
) -> None:
    """Decode PBF blocks in an executor and fan them to bounded consumers."""
    queue: asyncio.Queue[DecodedBlock | None] = asyncio.Queue(maxsize=QUEUE_MAXSIZE)
    loop = asyncio.get_running_loop()

    async def producer() -> None:
        # Enumerating raw blocks is cheap; decoding is the blocking cost.
        for seq, raw in _iter_raw_blocks(pbf_path):
            block = await loop.run_in_executor(None, _decode_block, seq, raw)
            await queue.put(block)          # suspends when the queue is full
        for _ in range(N_CONSUMERS):
            await queue.put(None)           # one sentinel per consumer

    async def consumer(worker_id: int) -> None:
        while True:
            block = await queue.get()
            try:
                if block is None:
                    return                  # sentinel: this worker is done
                await consume(block)
                logger.info(
                    "worker %d handled block %d (%d nodes, %d ways)",
                    worker_id, block.seq, block.n_nodes, block.n_ways,
                )
            finally:
                queue.task_done()

    prod = asyncio.create_task(producer())
    workers = [asyncio.create_task(consumer(i)) for i in range(N_CONSUMERS)]
    # Surface a producer crash instead of hanging the consumers forever.
    await asyncio.gather(prod, *workers)


async def _demo_sink(block: DecodedBlock) -> None:
    """Stand-in slow consumer; replace with a DB write or normalizer."""
    await asyncio.sleep(0.05)               # simulate downstream latency


if __name__ == "__main__":
    asyncio.run(stream_blocks(Path("extract.osm.pbf"), _demo_sink))

Step-by-step walkthrough

  1. QUEUE_MAXSIZE is the memory dial. The queue admits at most four decoded blocks. Peak resident set for in-flight work is roughly QUEUE_MAXSIZE × one decoded block, independent of how large the file is — that bound is the entire point of the pattern.
  2. Blocking decode moves off the loop. await loop.run_in_executor(None, _decode_block, ...) runs the CPU-heavy decode on the default thread pool. The await yields control, so the event loop keeps servicing put and get for other coroutines while a block decodes.
  3. await queue.put(block) is where backpressure lives. When the four slots are full, this line suspends the producer until a consumer frees a slot. The producer literally cannot outrun the consumers, so the backlog never grows.
  4. One sentinel per consumer. After the last real block, the producer enqueues N_CONSUMERS copies of None. Each consumer consumes exactly one and returns; a single sentinel would stop only the first worker to reach it and leave the others hanging on get().
  5. task_done() in a finally. Every get() is balanced by a task_done() even on the sentinel path, so a later queue.join() (or an external monitor) can tell when the queue has been fully drained.
  6. asyncio.gather(prod, *workers). Awaiting the producer alongside the workers means a producer exception propagates instead of silently leaving consumers blocked on an empty queue forever.

Verification

Confirm the stream behaves before wiring it to a real sink:

  • Watch memory stay flat. Sample RSS with psutil while processing a multi-gigabyte extract; it should plateau near QUEUE_MAXSIZE × block size rather than climbing with file size. A steadily rising curve means the bound is not taking effect.
  • Force a slow consumer. Raise the asyncio.sleep in _demo_sink to 0.5 s and confirm the producer’s decode rate drops to match — that observable throttling is backpressure working.
  • Count blocks end to end. Sum the blocks each worker logs; the total must equal the fileblock count reported by osmium fileinfo --extended extract.osm.pbf. A short count means a sentinel stopped a worker early.
  • Prove clean shutdown. The program should exit without a hang and without an asyncio “Task was destroyed but it is pending” warning; either symptom points to a missing sentinel or an unawaited task.

Common errors and fixes

Symptom Root cause One-line fix
Program hangs at shutdown Fewer sentinels than consumers Enqueue exactly N_CONSUMERS None values after the last block.
Memory grows with file size Unbounded queue (maxsize=0) Construct asyncio.Queue(maxsize=QUEUE_MAXSIZE) with a finite bound.
Event loop stalls, no concurrency Blocking decode called inline Offload it with await loop.run_in_executor(...).
One worker does all the work Sentinel handled before siblings released Send one sentinel per consumer, not a single shared one.
Task was destroyed but pending Producer/consumer task not awaited await asyncio.gather(prod, *workers) before returning.
Silent stall, no error surfaced Producer raised; consumers block on get() Gather the producer so its exception propagates.

Specification reference

asyncio.Queue(maxsize=0) creates an unbounded queue; a positive maxsize bounds it, and “if the queue is full, wait until a free slot is available before adding the item” is the defined behaviour of the coroutine put(). See the Python asyncio.Queue documentation for the put, get, task_done, and join contract, and Running blocking code in an executor for offloading the synchronous PBF decode off the event loop.

Frequently Asked Questions

Why a bounded queue instead of just gathering all decode tasks?

Gathering every decode task schedules them all at once, so the number of in-flight decoded blocks equals the number of blocks in the file — exactly the unbounded backlog you are trying to avoid. A bounded queue caps concurrent work at maxsize, so a fast decoder is forced to wait for slow consumers and peak memory stays constant regardless of file size.

How many consumers should I run?

Start with the number of independent downstream resources, not CPU cores. If the sink is a single database connection, one consumer avoids lock contention; if it is a pool of connections or stateless CPU work, scale consumers toward that pool’s parallelism. Because the queue bound caps memory, adding consumers changes throughput, not the resident-set ceiling.

Should decoding run in a thread pool or a process pool?

Use a thread pool when the decode releases the GIL in a C extension, which osmium largely does during raw reads — threads then give real parallelism with cheap handoff. Switch to a process pool only when the per-block work is pure-Python and GIL-bound; the trade is higher IPC cost to serialize each block across the process boundary.

What happens if a consumer raises mid-stream?

Wrap the per-block work in try/except, log or quarantine the offending block, and call task_done() in a finally so the queue accounting stays correct. Let an unrecoverable error propagate through gather so the whole stream fails loudly rather than silently dropping blocks — the same fail-fast contract the ingestion stage relies on.

Up one level: Async PBF Parsing with Pyrosm.